519 lines
21 KiB
Python
519 lines
21 KiB
Python
"""
|
|
ml/env.py - Gymnasium Environment for JackBot Hexapod RL Training
|
|
"""
|
|
import time
|
|
import math
|
|
from typing import Optional, Tuple, Dict, Any, List
|
|
|
|
import gymnasium as gym
|
|
from gymnasium import spaces
|
|
import numpy as np
|
|
import pybullet as p
|
|
|
|
from config import cfg
|
|
from Robot import Robot, PyBulletBackend
|
|
from ml.SimManager import SimManager
|
|
from ml.MetricsOverlay import MetricsHUD, LeaderCrown
|
|
|
|
# Color Palette RGBA for Terminated/Failed Robots
|
|
COLOR_FAILED = [0.3, 0.3, 0.3, 0.6] # Collapsed / Tilted Robot (Dark Semi-Transparent Gray)
|
|
|
|
|
|
class JackBotEnv(gym.Env):
|
|
"""Gymnasium environment wrapping JackBot hexapods with floating text & crown overlay."""
|
|
|
|
def __init__(
|
|
self,
|
|
use_gui: bool = True,
|
|
random_command: bool = True,
|
|
robot_spacing: float = 0.5,
|
|
start_pose: str = "init_deg",
|
|
max_episode_steps: int = 5000,
|
|
urdf_path: str = cfg.urdf_path,
|
|
):
|
|
super().__init__()
|
|
self.use_gui = use_gui
|
|
self.random_command = random_command
|
|
self.robot_spacing = robot_spacing
|
|
self.start_pose = start_pose
|
|
self.max_episode_steps = max_episode_steps
|
|
self.urdf_path = urdf_path
|
|
|
|
self.episode_count = 0
|
|
self.step_count = 0
|
|
self.total_steps = 0
|
|
self.cumulative_reward = 0.0
|
|
self.robot_rewards = [0.0]
|
|
self.failed_robots_mask = [False]
|
|
self._first_reset = True
|
|
|
|
# Initialize Simulation Manager
|
|
self.sim_manager = SimManager(use_gui=self.use_gui)
|
|
self.sim_manager.connect()
|
|
|
|
# Connect physics world
|
|
self.plane, self.pb_robots, self.robot_joint_indices = self.sim_manager.load_scene(
|
|
self.urdf_path, self.robot_spacing, self._robot_base_position
|
|
)
|
|
|
|
# Instantiate Robot Python wrappers per PyBullet body ID
|
|
self.robots = [
|
|
Robot(
|
|
backend_type=PyBulletBackend(self.sim_manager, body_id=pb_id),
|
|
start_pose=self.start_pose,
|
|
urdf_path=self.urdf_path
|
|
)
|
|
for pb_id in self.pb_robots
|
|
]
|
|
|
|
# Action (18 joint deltas per robot) & Observation (18 angles + 4 command dims per robot)
|
|
action_dim = 18
|
|
obs_dim = 18 + 4
|
|
|
|
self.action_space = spaces.Box(-1.0, 1.0, shape=(action_dim,), dtype=np.float32)
|
|
self.observation_space = spaces.Box(-np.inf, np.inf, shape=(obs_dim,), dtype=np.float32)
|
|
|
|
self.commands = np.zeros((1, 4), dtype=np.float32)
|
|
self.last_action = np.zeros(action_dim, dtype=np.float32)
|
|
self.target_height = 0.14
|
|
self.collapse_height_fraction = 0.55
|
|
self.tilt_failure_rad = 0.9
|
|
|
|
# Small random exploration pulse to break local optima.
|
|
# The bonus is only granted on a rare step and only when the robot
|
|
# is still stable enough that a bold move is likely to remain safe.
|
|
self.exploration_bonus_prob = 0.03
|
|
self.exploration_bonus_interval = 120
|
|
self.exploration_bonus_scale = 0.08
|
|
self.exploration_bonus_active = False
|
|
|
|
self.start_positions = [[0.0, 0.0, 0.0]]
|
|
self.max_distance_from_start = [0.0]
|
|
self.max_survival_steps = 0
|
|
self.curriculum_phase = 0
|
|
self.curriculum_episode_limit = 150
|
|
self.curriculum_stage_requirements = {
|
|
1: {"survival_steps": 200, "distance": 0.00, "stability_roll_pitch": 0.35},
|
|
2: {"survival_steps": 350, "distance": 10.0, "stability_roll_pitch": 0.30},
|
|
3: {"survival_steps": 550, "distance": 15.00, "stability_roll_pitch": 0.25},
|
|
}
|
|
|
|
# Floating HUD & Leader Crown Visualizers
|
|
self.hud = MetricsHUD(physics_client_id=self.sim_manager.physics_client)
|
|
self.leader_crown = LeaderCrown(physics_client_id=self.sim_manager.physics_client)
|
|
self.last_time = time.time()
|
|
|
|
def _set_robot_color(self, pb_id: int, rgba: List[float]):
|
|
"""Helper to change the visual color of a robot body and all its links."""
|
|
num_joints = p.getNumJoints(pb_id, physicsClientId=self.sim_manager.physics_client)
|
|
p.changeVisualShape(pb_id, -1, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client)
|
|
for j in range(num_joints):
|
|
p.changeVisualShape(pb_id, j, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client)
|
|
|
|
def _robot_base_position(self, robot_id: int, spacing: float = 0.5) -> list[float]:
|
|
return [0.0, 0.0, 0.14]
|
|
|
|
def sample_command(self) -> np.ndarray:
|
|
"""Curriculum command sampler with survival-gated difficulty progression."""
|
|
phase = self.curriculum_phase
|
|
|
|
if phase == 0:
|
|
# Phase 1: Forward Walking Focus
|
|
vx = np.random.uniform(0.5, 1.0)
|
|
vy = 0.0
|
|
vz = 0.0
|
|
omega = 0.0
|
|
elif phase == 1:
|
|
# Phase 2: Forward/Backward + Turning
|
|
vx = np.random.uniform(-1.0, 1.0)
|
|
vy = 0.0
|
|
vz = 0.0
|
|
omega = np.random.uniform(-0.8, 0.8)
|
|
else:
|
|
# Phase 3: Full Omnidirectional Movement
|
|
vx = np.random.uniform(-1.0, 1.0)
|
|
vy = np.random.uniform(-0.5, 0.5)
|
|
vz = 0.0
|
|
omega = np.random.uniform(-1.0, 1.0)
|
|
|
|
return np.array([vx, vy, vz, omega], dtype=np.float32)
|
|
|
|
def reset(self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None):
|
|
super().reset(seed=seed)
|
|
self.episode_count += 1
|
|
self.step_count = 0
|
|
self.cumulative_reward = 0.0
|
|
self.robot_rewards = [0.0]
|
|
self.failed_robots_mask = [False]
|
|
|
|
for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
|
|
spawn_pos = self._robot_base_position(idx, self.robot_spacing)
|
|
spawn_orn = [0, 0, 0, 1]
|
|
|
|
p.resetBasePositionAndOrientation(
|
|
pb_id, spawn_pos, spawn_orn, physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
p.resetBaseVelocity(
|
|
pb_id, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 0],
|
|
physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
|
|
robot_obj.reset_to_init()
|
|
|
|
if self.use_gui:
|
|
self._set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
|
|
|
|
self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32)
|
|
self.max_distance_from_start = [0.0]
|
|
self.max_survival_steps = 0
|
|
self.exploration_bonus_active = False
|
|
|
|
if self._first_reset:
|
|
self.curriculum_phase = 0
|
|
self.curriculum_episode_limit = min(self.max_episode_steps, 150)
|
|
self._first_reset = False
|
|
|
|
if self.random_command:
|
|
self.commands = np.stack([self.sample_command() for _ in range(1)])
|
|
else:
|
|
self.commands = np.zeros((1, 4), dtype=np.float32)
|
|
|
|
for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
|
|
pos, _ = p.getBasePositionAndOrientation(
|
|
pb_id, physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
self.start_positions[idx] = [float(pos[0]), float(pos[1]), float(pos[2])]
|
|
|
|
for _ in range(200):
|
|
self.sim_manager.step()
|
|
|
|
self.target_height = self._measure_settled_height()
|
|
if self.target_height <= 0.0:
|
|
self.target_height = 0.14
|
|
|
|
if self.use_gui:
|
|
self.hud.reset()
|
|
self.leader_crown.reset()
|
|
self._update_hud()
|
|
|
|
return self._get_obs(), {}
|
|
|
|
def get_robot_velocities(self) -> list:
|
|
"""Exposes velocities for the SB3 metrics callback."""
|
|
vels = []
|
|
for pb_id in self.pb_robots:
|
|
lin_v, ang_v = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client)
|
|
vels.append((lin_v, ang_v))
|
|
return vels
|
|
|
|
def get_robot_distance_metrics(self) -> list:
|
|
"""Exposes distance-from-start metrics for logging without affecting reward."""
|
|
metrics = []
|
|
for idx, pb_id in enumerate(self.pb_robots):
|
|
pos, _ = p.getBasePositionAndOrientation(
|
|
pb_id, physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
start_x, start_y, _ = self.start_positions[idx]
|
|
dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
|
|
self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist)
|
|
metrics.append((dist, self.max_distance_from_start[idx]))
|
|
|
|
self.max_survival_steps = max(self.max_survival_steps, self.step_count)
|
|
return metrics
|
|
|
|
def get_current_robot_metrics(self) -> list:
|
|
"""Returns current per-robot reward, distance, and velocity summaries for alive robots only."""
|
|
metrics = []
|
|
for idx, pb_id in enumerate(self.pb_robots):
|
|
if self.failed_robots_mask[idx]:
|
|
continue
|
|
|
|
pos, _ = p.getBasePositionAndOrientation(
|
|
pb_id, physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
linear_vel, angular_vel = p.getBaseVelocity(
|
|
pb_id, physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
start_x, start_y, _ = self.start_positions[idx]
|
|
dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
|
|
speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)))
|
|
yaw_rate = float(abs(angular_vel[2]))
|
|
metrics.append({
|
|
"reward": float(self.robot_rewards[idx]),
|
|
"distance_from_start": dist,
|
|
"speed": speed,
|
|
"yaw_rate": yaw_rate,
|
|
"alive": True,
|
|
"survival_steps": int(self.step_count),
|
|
})
|
|
|
|
return metrics
|
|
|
|
def get_survival_steps(self) -> int:
|
|
"""Returns the current survival length for the environment's current episode."""
|
|
return int(self.step_count)
|
|
|
|
def _get_obs(self) -> np.ndarray:
|
|
obs_list = []
|
|
for idx, (pb_id, joint_indices) in enumerate(zip(self.pb_robots, self.robot_joint_indices)):
|
|
joint_states = p.getJointStates(
|
|
pb_id,
|
|
joint_indices,
|
|
physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32)
|
|
robot_obs = np.concatenate([joint_angles, self.commands[idx]])
|
|
obs_list.append(robot_obs)
|
|
|
|
return np.concatenate(obs_list).astype(np.float32)
|
|
|
|
def _measure_settled_height(self) -> float:
|
|
heights = []
|
|
for pb_id in self.pb_robots:
|
|
pos, _ = p.getBasePositionAndOrientation(
|
|
pb_id, physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
heights.append(pos[2])
|
|
return float(np.mean(heights)) if heights else 0.14
|
|
|
|
def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
|
|
self.step_count += 1
|
|
self.total_steps += 1
|
|
previous_action = self.last_action.copy()
|
|
self.last_action = action.copy()
|
|
self._update_curriculum()
|
|
|
|
# Resample commands every 300 steps during long episodes
|
|
if self.random_command and (self.step_count % 300 == 0 or self._curriculum_advanced):
|
|
self.commands = np.stack([self.sample_command() for _ in range(1)])
|
|
|
|
action_per_robot = action.reshape(1, 18)
|
|
|
|
for robot, act in zip(self.robots, action_per_robot):
|
|
robot.apply_rl_action(act)
|
|
|
|
self.sim_manager.step()
|
|
|
|
self._update_robot_failures()
|
|
self.get_robot_distance_metrics()
|
|
|
|
obs = self._get_obs()
|
|
reward, per_robot_step_rewards = self._compute_reward(action, previous_action)
|
|
|
|
self.cumulative_reward += reward
|
|
for idx, r_step in enumerate(per_robot_step_rewards):
|
|
self.robot_rewards[idx] += r_step
|
|
|
|
terminated = self._is_done()
|
|
truncated = self.step_count >= self.curriculum_episode_limit
|
|
|
|
self._update_hud()
|
|
self._update_leader_visuals()
|
|
return obs, reward, terminated, truncated, {}
|
|
|
|
def _phase_progress_ready(self, phase: int) -> bool:
|
|
if phase not in self.curriculum_stage_requirements:
|
|
return False
|
|
|
|
if not self.pb_robots or not self.max_distance_from_start:
|
|
return False
|
|
|
|
req = self.curriculum_stage_requirements[phase]
|
|
survival_ok = self.max_survival_steps >= req["survival_steps"]
|
|
distance_ok = self.max_distance_from_start[0] >= req["distance"]
|
|
|
|
position, orientation = p.getBasePositionAndOrientation(
|
|
self.pb_robots[0], physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
|
|
stability_ok = abs(roll) <= req["stability_roll_pitch"] and abs(pitch) <= req["stability_roll_pitch"]
|
|
height_ok = position[2] >= max(0.09, self.target_height * 0.85)
|
|
|
|
return survival_ok and distance_ok and stability_ok and height_ok
|
|
|
|
def _update_curriculum(self):
|
|
self._curriculum_advanced = False
|
|
|
|
phase_labels = {
|
|
0: "stand-and-forward",
|
|
1: "turn-and-direction",
|
|
2: "omni-direction",
|
|
3: "full-command",
|
|
}
|
|
|
|
if self.curriculum_phase < 1 and self._phase_progress_ready(1):
|
|
self.curriculum_phase = 1
|
|
self.curriculum_episode_limit = min(self.max_episode_steps, 400)
|
|
self._curriculum_advanced = True
|
|
print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
|
|
elif self.curriculum_phase < 2 and self._phase_progress_ready(2):
|
|
self.curriculum_phase = 2
|
|
self.curriculum_episode_limit = min(self.max_episode_steps, 700)
|
|
self._curriculum_advanced = True
|
|
print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
|
|
elif self.curriculum_phase < 3 and self._phase_progress_ready(3):
|
|
self.curriculum_phase = 3
|
|
self.curriculum_episode_limit = min(self.max_episode_steps, self.max_episode_steps)
|
|
self._curriculum_advanced = True
|
|
print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
|
|
|
|
def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> Tuple[float, list[float]]:
|
|
rewards = []
|
|
current_actions = action.reshape(1, 18)
|
|
previous_actions = previous_action.reshape(1, 18)
|
|
|
|
for idx, pb_id in enumerate(self.pb_robots):
|
|
pos, orientation = p.getBasePositionAndOrientation(
|
|
pb_id, physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
linear_vel, angular_vel = p.getBaseVelocity(
|
|
pb_id, physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
|
|
command = self.commands[idx]
|
|
|
|
cmd_vx = command[0]
|
|
cmd_vy = command[1]
|
|
cmd_yaw = command[3]
|
|
|
|
# -------------------------------------------------------------
|
|
# 1. LINEAR VECTOR SPEED MAXIMIZATION (Magnitude + Direction)
|
|
# -------------------------------------------------------------
|
|
cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32)
|
|
cmd_norm = np.linalg.norm(cmd_dir)
|
|
|
|
if cmd_norm > 0.05:
|
|
# Normalize target direction vector
|
|
unit_cmd_dir = cmd_dir / cmd_norm
|
|
actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)
|
|
|
|
# Speed aligned with target direction (m/s)
|
|
aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir))
|
|
|
|
# Moderate movement reward to encourage directional motion
|
|
linear_speed_reward = 2.5 * aligned_speed
|
|
|
|
# Penalize sideways drift (perpendicular velocity to commanded direction)
|
|
perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir
|
|
drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel))
|
|
else:
|
|
# If no linear command given, penalize all horizontal movement
|
|
linear_speed_reward = 0.0
|
|
drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2)
|
|
|
|
# -------------------------------------------------------------
|
|
# 2. TURNING SPEED MAXIMIZATION
|
|
# -------------------------------------------------------------
|
|
actual_yaw_rate = angular_vel[2] # rad/s in PyBullet Z-axis
|
|
|
|
if abs(cmd_yaw) > 0.05:
|
|
# Reward turning in the commanded direction, but less aggressively
|
|
turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw)
|
|
else:
|
|
# Penalize unwanted rotation when joystick turn is centered
|
|
turning_reward = -0.6 * (actual_yaw_rate ** 2)
|
|
|
|
# -------------------------------------------------------------
|
|
# 3. SMALL ALIVE BONUS + AGE-RAMPED STILLNESS PENALTY
|
|
# -------------------------------------------------------------
|
|
alive_reward = 0.02
|
|
|
|
age_ratio = min(1.0, self.step_count / max(1, self.curriculum_episode_limit))
|
|
still_penalty = 0.0
|
|
if (cmd_norm > 0.1 or abs(cmd_yaw) > 0.1) and (abs(linear_vel[0]) < 0.02 and abs(actual_yaw_rate) < 0.05):
|
|
still_penalty = 0.35 + 0.85 * age_ratio
|
|
|
|
# -------------------------------------------------------------
|
|
# 4. POSTURE & STABILITY PENALTIES
|
|
# -------------------------------------------------------------
|
|
height_penalty = 6.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0
|
|
stability_penalty = 1.5 * (roll**2 + pitch**2)
|
|
|
|
# Penalize only large control jumps; allow continuous low-amplitude motion
|
|
control_delta = np.abs(current_actions[idx] - previous_actions[idx])
|
|
large_delta_mask = control_delta > 0.12
|
|
large_delta_penalty = 0.002 * float(np.sum(np.square(control_delta[large_delta_mask]))) if np.any(large_delta_mask) else 0.0
|
|
|
|
# Combined Step Reward
|
|
r_step = (
|
|
alive_reward
|
|
+ linear_speed_reward
|
|
+ turning_reward
|
|
- drift_penalty
|
|
- still_penalty
|
|
- height_penalty
|
|
- stability_penalty
|
|
- large_delta_penalty
|
|
)
|
|
rewards.append(r_step)
|
|
|
|
return float(np.sum(rewards)), rewards
|
|
|
|
def _is_done(self) -> bool:
|
|
"""Returns True when the robot has entered a failed state."""
|
|
return bool(self.failed_robots_mask[0]) if self.failed_robots_mask else False
|
|
|
|
def _update_hud(self):
|
|
if not self.use_gui or not self.pb_robots:
|
|
return
|
|
|
|
now = time.time()
|
|
fps = 1.0 / max(now - self.last_time, 1e-5)
|
|
self.last_time = now
|
|
|
|
heights = []
|
|
rolls = []
|
|
pitches = []
|
|
for pb_id in self.pb_robots:
|
|
pos, orient = p.getBasePositionAndOrientation(pb_id, physicsClientId=self.sim_manager.physics_client)
|
|
roll, pitch, _ = p.getEulerFromQuaternion(orient)
|
|
heights.append(pos[2])
|
|
rolls.append(math.degrees(roll))
|
|
pitches.append(math.degrees(pitch))
|
|
|
|
avg_height = float(np.mean(heights))
|
|
avg_roll_pitch = (float(np.mean(rolls)), float(np.mean(pitches)))
|
|
|
|
self.hud.update(
|
|
episode=self.episode_count,
|
|
step=self.total_steps,
|
|
robot_rewards=self.robot_rewards,
|
|
cmd_vel=self.commands[0],
|
|
fps=fps,
|
|
avg_height=avg_height,
|
|
roll_pitch=avg_roll_pitch
|
|
)
|
|
|
|
def _update_leader_visuals(self):
|
|
if not self.use_gui:
|
|
return
|
|
|
|
best_idx = int(np.argmax(self.robot_rewards))
|
|
leader_pb_id = self.pb_robots[best_idx]
|
|
leader_pos, _ = p.getBasePositionAndOrientation(
|
|
leader_pb_id, physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
self.leader_crown.update(leader_pos)
|
|
|
|
def _update_robot_failures(self):
|
|
"""Checks failure condition and colors failed robots dark gray."""
|
|
for idx, pb_id in enumerate(self.pb_robots):
|
|
if self.failed_robots_mask[idx]:
|
|
continue
|
|
|
|
position, orientation = p.getBasePositionAndOrientation(
|
|
pb_id, physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
|
|
|
|
collapse_threshold = max(0.06, self.collapse_height_fraction * self.target_height)
|
|
is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad
|
|
is_collapsed = position[2] < collapse_threshold
|
|
|
|
if is_tilted or is_collapsed:
|
|
self.failed_robots_mask[idx] = True
|
|
if self.use_gui:
|
|
self._set_robot_color(pb_id, COLOR_FAILED)
|
|
|
|
def close(self):
|
|
self.sim_manager.disconnect() |