JackM323
acb3d671be
new reward/penalty system learning phases with curriculum learning new training parameters cleanup of old code better logging while training multiple environments instead of robots (they could bumb into each other)
509 lines
21 KiB
Python
509 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
|
|
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": 120, "distance": 0.04, "stability_roll_pitch": 0.35},
|
|
2: {"survival_steps": 250, "distance": 0.08, "stability_roll_pitch": 0.30},
|
|
3: {"survival_steps": 450, "distance": 0.12, "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
|
|
|
|
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 = 4.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() |