JackM323/JackBot
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects 2 Releases Wiki Activity

code reduction for training and eval

Browse Source 
reward and everything else changed
again
wont be the last time
This commit is contained in:
JackM323
2026-08-05 22:44:05 +02:00
parent 15e0206739
commit c93c524a10
10 changed files with 606 additions and 814 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1 -1
View File
@@ -4,7 +4,7 @@ node_modules/
.vscode/
__pycache__/
ml/checkpoints/
ml/tensorboard/
ml/logs/
# Ignore environment files with private passwords/keys
.env
ml/MetricsOverlay.py
+10 -54
View File
@@ -44,7 +44,7 @@ class MetricsHUD:
robot_rewards: List[float],
cmd_vel: np.ndarray,
fps: float = 0.0,
avg_height: float = 0.0,
height: float = 0.0,
roll_pitch: Tuple[float, float] = (0.0, 0.0)
) -> None:
"""Updates floating black text block in 3D space with billboarding."""
@@ -65,7 +65,7 @@ class MetricsHUD:
f"----------------------\n"
f"Top Rewards: [{top1}, {top2}, {top3}]\n"
f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]\n"
f"Height: {avg_height:.3f} m\n"
f"Height: {height:.3f} m\n"
f"Roll/Pitch: {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°"
)
@@ -76,7 +76,14 @@ class MetricsHUD:
# Calculate dynamic orientation to align text flat against camera plane
text_orientation = self._get_camera_facing_orientation()
if self._text_id is None:
# Safely remove the old text to prevent PyBullet ghosting/overlapping
if self._text_id is not None:
try:
p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
except Exception:
pass
# Draw fresh text
self._text_id = p.addUserDebugText(
text=hud_text,
textPosition=text_position,
@@ -85,16 +92,6 @@ class MetricsHUD:
textOrientation=text_orientation,
physicsClientId=self.client_id
)
else:
self._text_id = p.addUserDebugText(
text=hud_text,
textPosition=text_position,
textColorRGB=text_color,
textSize=0.1,
textOrientation=text_orientation,
replaceItemUniqueId=self._text_id,
physicsClientId=self.client_id
)
def reset(self) -> None:
"""Removes the active debug text item so a new episode starts with a clean overlay."""
@@ -104,44 +101,3 @@ class MetricsHUD:
except Exception:
pass
self._text_id = None
class LeaderCrown:
"""Renders a floating crown or star emoji above the leading robot in PyBullet."""
def __init__(self, physics_client_id: int = 0):
self.client_id = physics_client_id
self._text_id = None
def update(self, leader_pos: list[float]):
"""Positions a floating crown ~0.35m directly above the lead robot's base."""
crown_pos = [leader_pos[0], leader_pos[1], leader_pos[2] + 0.35]
# You can use "👑 CROWN", "⭐ LEADER", or "★ TOP1"
crown_text = "👑"
if self._text_id is None:
self._text_id = p.addUserDebugText(
text=crown_text,
textPosition=crown_pos,
textColorRGB=[1.0, 0.84, 0.0],
textSize=2.0,
physicsClientId=self.client_id
)
else:
self._text_id = p.addUserDebugText(
text=crown_text,
textPosition=crown_pos,
textColorRGB=[1.0, 0.84, 0.0],
textSize=2.0,
replaceItemUniqueId=self._text_id,
physicsClientId=self.client_id
)
def reset(self):
if self._text_id is not None:
try:
p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
except Exception:
pass
self._text_id = None
ml/SimManager.py
+40
View File
@@ -82,11 +82,40 @@ class SimManager:
def step(self):
p.stepSimulation(physicsClientId=self.physics_client)
def set_rendering(self, enabled: bool) -> None:
"""Toggles PyBullet 3D rendering to speed up simulation."""
if self.physics_client is not None and p.isConnected(self.physics_client):
p.configureDebugVisualizer(
p.COV_ENABLE_RENDERING,
1 if enabled else 0,
physicsClientId=self.physics_client
)
def disconnect(self):
if self.physics_client is not None and p.isConnected(self.physics_client):
p.disconnect(self.physics_client)
self.physics_client = None
def get_contact_points(
self,
bodyA: int = -1,
bodyB: int = -1,
linkIndexA: int = -1,
linkIndexB: int = -1,
):
"""Wrapper around pybullet.getContactPoints bound to this simulation client."""
kwargs = {"physicsClientId": self.physics_client}
if bodyA != -1:
kwargs["bodyA"] = bodyA
if bodyB != -1:
kwargs["bodyB"] = bodyB
if linkIndexA != -1:
kwargs["linkIndexA"] = linkIndexA
if linkIndexB != -1:
kwargs["linkIndexB"] = linkIndexB
return p.getContactPoints(**kwargs)
# --- ROBOT GETTERS AND SETTERS ---
def reset_robot_base(
@@ -124,6 +153,17 @@ class SimManager:
roll, pitch, yaw = p.getEulerFromQuaternion(orn)
return float(roll), float(pitch), float(yaw)
def get_foot_link_indices(self, body_id: int) -> list[int]:
"""Inspects URDF joint structure to extract link IDs for leg tips and tibias."""
foot_indices = []
num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client)
for j_idx in range(num_joints):
info = p.getJointInfo(body_id, j_idx, physicsClientId=self.physics_client)
link_name = info[12].decode("utf-8")
if "tip" in link_name or "tibia" in link_name:
foot_indices.append(j_idx)
return foot_indices
def get_robot_pose_and_rpy(self, body_id: int) -> Tuple[List[float], Tuple[float, float, float]]:
"""Returns base position and (roll, pitch, yaw) tuple in radians."""
pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
ml/__init__.py
+2 -4
View File
@@ -1,6 +1,4 @@
from .env import JackBotEnv
from .env import JackBotEnv, CurriculumCallback
from .model import ActorCritic
from .train import train
from .evaluate import evaluate
__all__ = ["JackBotEnv", "ActorCritic", "train", "evaluate"]
__all__ = ["JackBotEnv", "CurriculumCallback", "ActorCritic"]
ml/env.py
+326 -314
View File
@@ -9,14 +9,15 @@ from typing import Optional, Tuple, Dict, Any, List
import gymnasium as gym
from gymnasium import spaces
import numpy as np
from stable_baselines3.common.callbacks import BaseCallback
from config import cfg
from Robot import Robot, PyBulletBackend
from ml.SimManager import SimManager
from ml.MetricsOverlay import MetricsHUD, LeaderCrown
from ml.MetricsOverlay import MetricsHUD
# Color Palette RGBA for Terminated/Failed Robots
COLOR_FAILED = [0.3, 0.3, 0.3, 0.6] # Collapsed / Tilted Robot (Dark Semi-Transparent Gray)
COLOR_FAILED = [0.3, 0.3, 0.3, 0.6]
class CurriculumPhase(IntEnum):
@@ -28,93 +29,92 @@ class CurriculumPhase(IntEnum):
class JackBotEnv(gym.Env):
"""Gymnasium environment wrapping JackBot hexapods with floating text & crown overlay."""
"""Gymnasium environment wrapping a single JackBot hexapod."""
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,
max_episode_steps: int = 3000,
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.max_robot_speed = 0.8
self.max_robot_speed = 0.6
self.episode_count = 0
self.step_count = 0
self.total_steps = 0
self.cumulative_reward = 0.0
self.robot_rewards = [0.0]
self.consecutive_still_steps = [0]
self.failed_robots_mask = [False]
self.robot_reward = 0.0
self.consecutive_still_steps = 0
self.is_failed = False
self.episode_count = 0
self.episode_height_sum = 0.0
self.episode_roll_sum = 0.0
self.episode_pitch_sum = 0.0
self._curriculum_advanced = False
self._first_reset = True
# Dynamic Command Resampling Timing (60 Hz control loop)
self.control_freq = 60
self.min_cmd_hold_steps = int(2.0 * self.control_freq) # 120 steps (2s)
self.max_cmd_hold_steps = int(6.0 * self.control_freq) # 360 steps (6s)
self.next_cmd_resample_step = 0
self.initial_stand_steps = 120 # Mandatory 2s standing window at episode reset
# 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
# Connect physics world & load single robot
self.plane, pb_robots, robot_joint_indices = self.sim_manager.load_scene(
self.urdf_path, 0.0, self._robot_base_position
)
self.pb_robot = pb_robots[0]
self.joint_indices = robot_joint_indices[0]
# 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,
# Instantiate Robot Python wrapper (start_pose is managed inside Robot.py)
self.robot = Robot(
backend_type=PyBulletBackend(self.sim_manager, body_id=self.pb_robot),
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 (18 joint deltas) & Observation (18 angles + 4 command dims)
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.command = np.zeros(4, dtype=np.float32)
self.last_action = np.zeros(action_dim, dtype=np.float32)
self.target_height = 0.14
self.target_height = 0.122
self.collapse_height_fraction = 0.55
self.tilt_failure_rad = 0.9
# Small random exploration pulse settings
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.foot_link_indices = self._find_foot_link_indices()
self.start_position = [0.0, 0.0, 0.0]
self.max_distance_from_start = 0.0
self.max_survival_steps = 0
self.default_joint_angles = np.zeros(18, dtype=np.float32)
# Curriculum Initialization via Enum
self.curriculum_phase = CurriculumPhase.STAND_ONLY
self.curriculum_episode_limit = 300 # 300 steps limit gives headroom for 400-step requirement
# Gates required to unlock each target phase
self.curriculum_stage_requirements = {
CurriculumPhase.FORWARD: {
"survival_steps": 400, # Must survive ~8 seconds
"min_avg_height_ratio": 0.90, # Average height >= 90% of target
"stability_roll_pitch": 0.18, # Max ~10 degrees tilt
"survival_steps": 300,
"min_avg_height_ratio": 0.88,
"max_avg_roll_pitch": 0.18, # ~10 degrees average
},
CurriculumPhase.TURN_AND_DIRECTION: {
"survival_steps": 500,
"min_forward_distance": 2.5, # Must walk +2.5m forward (+X)
"max_lateral_drift": 0.8, # Max 0.8m drift on Y axis
"min_forward_distance": 2.5,
"max_lateral_drift": 0.8,
"min_avg_height_ratio": 0.85,
"stability_roll_pitch": 0.25,
},
@@ -132,51 +132,42 @@ class JackBotEnv(gym.Env):
},
}
# 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 _robot_base_position(self, robot_id: int, spacing: float = 0.5) -> list[float]:
return [0.0, 0.0, 0.14]
def _robot_base_position(self, robot_id: int, spacing: float = 0.0) -> list[float]:
return [0.0, 0.0, 0.13]
def _find_foot_link_indices(self) -> list:
return self.sim_manager.get_foot_link_indices(self.pb_robot)
def sample_command(self) -> np.ndarray:
"""Curriculum command sampler with survival-gated difficulty progression."""
phase = self.curriculum_phase
if phase == CurriculumPhase.STAND_ONLY:
# Phase 0: Pure Standing (Zero commands)
vx = 0.0
vy = 0.0
vz = 0.0
omega = 0.0
elif phase == CurriculumPhase.FORWARD:
# Phase 1: Straight Forward Walking + Standing Checks
if np.random.random() < 0.30:
# 30% chance: Stand Still Command (vx = 0)
vx = 0.0
else:
# 70% chance: Forward Walk Command
stand_probabilities = {
CurriculumPhase.STAND_ONLY: 1.0,
CurriculumPhase.FORWARD: 0.25,
CurriculumPhase.TURN_AND_DIRECTION: 0.20,
CurriculumPhase.OMNI_DIRECTION: 0.15,
CurriculumPhase.FULL_COMMAND: 0.15,
}
if np.random.random() < stand_probabilities.get(phase, 0.15):
return np.zeros(4, dtype=np.float32)
if phase == CurriculumPhase.FORWARD:
vx = np.random.uniform(0.15, 0.50)
vy, vz, omega = 0.0, 0.0, 0.0
elif phase == CurriculumPhase.TURN_AND_DIRECTION:
# Phase 2: Forward/Backward + Turning + Standing
if np.random.random() < 0.20:
vx, omega = 0.0, 0.0
else:
vx = np.random.uniform(-1.0, 1.0)
omega = np.random.uniform(-0.8, 0.8)
vx = np.random.uniform(-0.8, 0.8)
vy, vz = 0.0, 0.0
omega = np.random.uniform(-0.8, 0.8)
elif phase == CurriculumPhase.OMNI_DIRECTION:
# Phase 3: Full Omnidirectional Movement
vx = np.random.uniform(-1.0, 1.0)
vx = np.random.uniform(-0.8, 0.8)
vy = np.random.uniform(-0.5, 0.5)
vz = 0.0
omega = np.random.uniform(-1.0, 1.0)
omega = np.random.uniform(-0.8, 0.8)
else:
# Phase 4: Full Unconstrained Commands
vx = np.random.uniform(-1.0, 1.0)
vy = np.random.uniform(-1.0, 1.0)
vz = 0.0
@@ -189,177 +180,158 @@ class JackBotEnv(gym.Env):
self.episode_count += 1
self.step_count = 0
self.cumulative_reward = 0.0
self.robot_rewards = [0.0]
self.failed_robots_mask = [False]
self.episode_height_sum = 0.0
self.robot_reward = 0.0
self.is_failed = 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)
self.episode_height_sum = 0.0
self.episode_roll_sum = 0.0
self.episode_pitch_sum = 0.0
spawn_pos = self._robot_base_position(0)
spawn_orn = [0.0, 0.0, 0.0, 1.0]
self.sim_manager.reset_robot_base(pb_id, spawn_pos, spawn_orn)
robot_obj.reset_to_init()
self.sim_manager.reset_robot_base(self.pb_robot, spawn_pos, spawn_orn)
self.robot.reset_to_init()
if self.use_gui:
self.sim_manager.set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
self.sim_manager.set_robot_color(self.pb_robot, [1.0, 1.0, 1.0, 1.0])
self.consecutive_still_steps = [0 for _ in range(len(self.pb_robots))]
self.consecutive_still_steps = 0
self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32)
self.max_distance_from_start = [0.0]
self.max_distance_from_start = 0.0
self.max_survival_steps = 0
self.exploration_bonus_active = False
if self._first_reset:
self.curriculum_phase = CurriculumPhase.STAND_ONLY
self.curriculum_episode_limit = min(self.max_episode_steps, 500)
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)
self.command = np.zeros(4, dtype=np.float32)
self.next_cmd_resample_step = self.initial_stand_steps
for idx, pb_id in enumerate(self.pb_robots):
pos, _ = self.sim_manager.get_robot_pose(pb_id)
self.start_positions[idx] = [float(pos[0]), float(pos[1]), float(pos[2])]
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
self.start_position = [float(pos[0]), float(pos[1]), float(pos[2])]
# Settle for 200 steps after dropping in a standing position, then measure target height
self.target_height = self.sim_manager.settle_and_measure_height(
self.pb_robots, steps=200, fallback_height=0.14
[self.pb_robot], steps=200, fallback_height=0.122
)
self.default_joint_angles = np.array(
self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices),
dtype=np.float32
)
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 = self.sim_manager.get_robot_velocity(pb_id)
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, _ = self.sim_manager.get_robot_pose(pb_id)
if idx == 0:
self.episode_height_sum += float(pos[2])
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
"""Returns metric summary for the callback."""
if self.is_failed:
return []
pos, _ = self.sim_manager.get_robot_pose(pb_id)
linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id)
start_x, start_y, _ = self.start_positions[idx]
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
linear_vel, angular_vel = self.sim_manager.get_robot_velocity(self.pb_robot)
start_x, start_y, _ = self.start_position
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]),
return [{
"reward": float(self.robot_reward),
"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)
"phase_name": self.curriculum_phase.name,
}]
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_angles = self.sim_manager.get_robot_joint_angles(pb_id, joint_indices)
robot_obs = np.concatenate([joint_angles, self.commands[idx]])
obs_list.append(robot_obs)
return np.concatenate(obs_list).astype(np.float32)
joint_angles = self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices)
return np.concatenate([joint_angles, self.command]).astype(np.float32)
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)])
if self.random_command and (self.step_count >= self.next_cmd_resample_step or self._curriculum_advanced):
self.command = self.sample_command()
random_interval = np.random.randint(self.min_cmd_hold_steps, self.max_cmd_hold_steps + 1)
self.next_cmd_resample_step = self.step_count + random_interval
action_per_robot = action.reshape(1, 18)
for robot, act in zip(self.robots, action_per_robot):
robot.apply_rl_action(act)
self.robot.apply_rl_action(action)
if self.step_count % 60 == 0:
random_force = np.random.uniform(-2.0, 2.0, size=2) # X and Y push (Newtons)
random_force = np.random.uniform(-2.0, 2.0, size=2)
self.sim_manager.apply_external_force(
body_id=self.pb_robots[0],
body_id=self.pb_robot,
force=[random_force[0], random_force[1], 0.0]
)
render_freq = 10 # Only draw 1 in every 10 frames
if self.use_gui and self.step_count % render_freq != 0:
self.sim_manager.set_rendering(False)
self.sim_manager.step()
self._update_robot_failures()
self.get_robot_distance_metrics()
if self.use_gui and self.step_count % render_freq == 0:
self.sim_manager.set_rendering(True)
self._update_robot_failure()
self._update_distance_metrics()
self._update_curriculum()
obs = self._get_obs()
reward, per_robot_step_rewards = self._compute_reward(action, previous_action)
reward = 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
self.robot_reward += reward
terminated = self._is_done()
truncated = self.step_count >= self.curriculum_episode_limit
terminated = self.is_failed
truncated = self.step_count >= self.max_episode_steps
if self.step_count % 120 == 0 and self.use_gui:
self._update_hud()
self._update_leader_visuals()
return obs, reward, terminated, truncated, {}
def _update_distance_metrics(self):
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
self.episode_height_sum += float(pos[2])
start_x, start_y, _ = self.start_position
dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
self.max_distance_from_start = max(self.max_distance_from_start, dist)
self.max_survival_steps = max(self.max_survival_steps, self.step_count)
def _phase_progress_ready(self, next_phase: CurriculumPhase) -> bool:
if next_phase not in self.curriculum_stage_requirements or not self.pb_robots:
if next_phase not in self.curriculum_stage_requirements:
return False
req = self.curriculum_stage_requirements[next_phase]
# 1. Survival Step Check
# 1. Survival Check
survival_ok = self.max_survival_steps >= req["survival_steps"]
# 2. Body Posture & Tilt Check
position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robots[0])
stability_ok = abs(roll) <= req["stability_roll_pitch"] and abs(pitch) <= req["stability_roll_pitch"]
# 2. Smooth Average Stability Checks (Prevents 1-frame spikes from failing curriculum)
avg_roll = self.episode_roll_sum / max(1, self.step_count)
avg_pitch = self.episode_pitch_sum / max(1, self.step_count)
max_allowed_angle = req.get("max_avg_roll_pitch", 0.20)
stability_ok = (avg_roll <= max_allowed_angle) and (avg_pitch <= max_allowed_angle)
# 3. Average Height Ratio Check Across the Episode
# 3. Average Height Check
avg_height = self.episode_height_sum / max(1, self.step_count)
required_min_avg_height = self.target_height * req.get("min_avg_height_ratio", 0.85)
height_ok = avg_height >= required_min_avg_height
# 4. Distance and Displacement Checks
start_x, start_y, _ = self.start_positions[0]
dx = position[0] - start_x
dy = position[1] - start_y
# 4. Distance and Drift Checks
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
start_x, start_y, _ = self.start_position
dx = pos[0] - start_x
dy = pos[1] - start_y
dist_2d = math.hypot(dx, dy)
distance_ok = True
@@ -368,203 +340,243 @@ class JackBotEnv(gym.Env):
elif "min_distance" in req:
distance_ok = dist_2d >= req["min_distance"]
displacement_ok = True
if "max_displacement" in req:
displacement_ok = dist_2d <= req["max_displacement"]
drift_ok = True
if "max_lateral_drift" in req:
drift_ok = abs(dy) <= req["max_lateral_drift"]
return survival_ok and height_ok and stability_ok and displacement_ok and distance_ok and drift_ok
return survival_ok and height_ok and stability_ok and distance_ok and drift_ok
def _update_curriculum(self):
self._curriculum_advanced = False
forced_forward = (self.curriculum_phase < CurriculumPhase.FORWARD) and (self.step_count >= 2500)
if self.curriculum_phase < CurriculumPhase.FORWARD and self._phase_progress_ready(CurriculumPhase.FORWARD):
if self.curriculum_phase < CurriculumPhase.FORWARD and (self._phase_progress_ready(CurriculumPhase.FORWARD) or forced_forward):
self.curriculum_phase = CurriculumPhase.FORWARD
self.curriculum_episode_limit = min(self.max_episode_steps, 600)
self._curriculum_advanced = True
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
reason = "FORCED (2500 steps)" if forced_forward else "MET"
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked [{reason}] at step {self.step_count}")
elif self.curriculum_phase < CurriculumPhase.TURN_AND_DIRECTION and self._phase_progress_ready(CurriculumPhase.TURN_AND_DIRECTION):
self.curriculum_phase = CurriculumPhase.TURN_AND_DIRECTION
self.curriculum_episode_limit = min(self.max_episode_steps, 800)
self._curriculum_advanced = True
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
elif self.curriculum_phase < CurriculumPhase.OMNI_DIRECTION and self._phase_progress_ready(CurriculumPhase.OMNI_DIRECTION):
self.curriculum_phase = CurriculumPhase.OMNI_DIRECTION
self.curriculum_episode_limit = min(self.max_episode_steps, 1000)
self._curriculum_advanced = True
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
elif self.curriculum_phase < CurriculumPhase.FULL_COMMAND and self._phase_progress_ready(CurriculumPhase.FULL_COMMAND):
self.curriculum_phase = CurriculumPhase.FULL_COMMAND
self.curriculum_episode_limit = min(self.max_episode_steps, self.max_episode_steps)
self._curriculum_advanced = True
print(f"[Curriculum] Phase {self.curriculum_phase.name} 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, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id)
linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id)
command = self.commands[idx]
cmd_vx = command[0] # Normalized -1.0 to 1.0
cmd_vy = command[1] # Normalized -1.0 to 1.0
cmd_yaw = command[3]
cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32)
cmd_norm = np.linalg.norm(cmd_dir)
actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)
# -----------------------------------------------------------------
# 1. NORMALIZED SPEED TRACKING (1.0 = v_max)
# -----------------------------------------------------------------
if cmd_norm > 0.05:
unit_cmd_dir = cmd_dir / cmd_norm
aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir))
target_speed = cmd_norm * self.max_robot_speed # e.g., 0.35 * 0.8 = 0.28 m/s
if aligned_speed <= 0.0:
# Zero or backward movement gets ZERO reward
linear_speed_reward = 0.0
elif aligned_speed <= target_speed:
# Smoothly scales from 0.0 to +5.0 as speed approaches target
linear_speed_reward = 5.0 * (aligned_speed / target_speed)
else:
# Gently penalize overspeeding beyond target
overspeed_ratio = (aligned_speed - target_speed) / target_speed
linear_speed_reward = max(0.0, 5.0 - 2.5 * overspeed_ratio)
# Drift penalty for sideways sliding
perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir
drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel))
else:
aligned_speed = 0.0
linear_speed_reward = 0.0
drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2)
# -----------------------------------------------------------------
# 2. TIME-ACCUMULATING STILLNESS PENALTY
# -----------------------------------------------------------------
actual_yaw_rate = angular_vel[2]
is_moving = abs(aligned_speed) > 0.03 or abs(actual_yaw_rate) > 0.08
if (cmd_norm > 0.05 or abs(cmd_yaw) > 0.05) and not is_moving:
# Increment counter every step the robot stalls under command
self.consecutive_still_steps[idx] += 1
else:
# Reset counter as soon as robot makes a valid move!
self.consecutive_still_steps[idx] = 0
# Penalty grows by 0.05 every step frozen, capped at -3.5 per step
still_penalty = min(3.5, 0.05 * self.consecutive_still_steps[idx])
# -----------------------------------------------------------------
# 3. TURNING & POSTURE PENALTIES
# -----------------------------------------------------------------
if abs(cmd_yaw) > 0.05:
turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw)
else:
turning_reward = -0.6 * (actual_yaw_rate ** 2)
alive_reward = 0.20
# Height drop penalty
min_valid_h = self.target_height * 0.90
if pos[2] < min_valid_h:
drop = (min_valid_h - pos[2]) / self.target_height
height_penalty = 4.0 * drop + 20.0 * (drop ** 2)
else:
height_penalty = 0.0
stability_penalty = 8 * (roll**2 + pitch**2)
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
r_step = (
alive_reward
+ linear_speed_reward
+ turning_reward
- drift_penalty
- still_penalty
- height_penalty
- stability_penalty
- large_delta_penalty
def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> float:
"""
Calculates task rewards using normalized Exponential Kernels.
Includes a deadband filter for jittering and zero-reward gating when stationary.
"""
# 1. Fetch Robot State
pos, (roll, pitch, yaw) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot)
linear_vel, angular_vel = self.sim_manager.get_robot_velocity(self.pb_robot)
current_joints = np.array(
self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices),
dtype=np.float32
)
rewards.append(r_step)
return float(np.sum(rewards)), rewards
cmd_vx, cmd_vy, _, cmd_yaw = self.command
cmd_norm = math.hypot(cmd_vx, cmd_vy)
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
# 2. Velocity Deadband Filtering (Ignores jittering & micro-movements)
VEL_DEADBAND = 0.04 # 4 cm/s threshold
YAW_DEADBAND = 0.05 # 0.05 rad/s threshold
raw_speed = math.hypot(linear_vel[0], linear_vel[1])
if raw_speed < VEL_DEADBAND:
filtered_vx, filtered_vy = 0.0, 0.0
filtered_speed = 0.0
else:
filtered_vx, filtered_vy = linear_vel[0], linear_vel[1]
filtered_speed = raw_speed
raw_yaw_rate = abs(angular_vel[2])
if raw_yaw_rate < YAW_DEADBAND:
filtered_yaw_rate = 0.0
else:
filtered_yaw_rate = angular_vel[2]
# 3. Posture & Stability Sub-Rewards
height_error = pos[2] - self.target_height
r_height = math.exp(-150.0 * (height_error ** 2))
orientation_error = roll**2 + pitch**2
r_stability = math.exp(-25.0 * orientation_error)
joint_error = np.mean(np.square(current_joints - self.default_joint_angles))
r_pose = math.exp(-2.0 * joint_error)
action_delta = np.mean(np.square(action - previous_action))
r_smoothness = math.exp(-0.1 * action_delta)
# 4. Mode Logic
if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05:
# STANDING MODE: Reward clean posture, height, and stability
w_height = 0.35
w_stability = 0.35
w_pose = 0.20
w_smoothness = 0.10
total_reward = (
(w_height * r_height)
+ (w_stability * r_stability)
+ (w_pose * r_pose)
+ (w_smoothness * r_smoothness)
)
else:
# WALKING / TURNING MODE
is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0)
# HARD GATE: If commanded to move but standing still/jittering, reward is strictly 0.0
if not is_moving:
return 0.0
target_vx = cmd_vx * self.max_robot_speed
target_vy = cmd_vy * self.max_robot_speed
lin_vel_error = (filtered_vx - target_vx)**2 + (filtered_vy - target_vy)**2
r_lin_vel = math.exp(-25.0 * lin_vel_error)
ang_vel_error = (filtered_yaw_rate - cmd_yaw)**2
r_ang_vel = math.exp(-15.0 * ang_vel_error)
# Stillness Check: Commanded to move, but staying virtually still
stillness_penalty = 0.0
if math.hypot(target_vx, target_vy) > 0.08 and math.hypot(linear_vel[0], linear_vel[1]) < 0.03:
r_lin_vel = 0.0 # Strip velocity credit completely
stillness_penalty = -0.25
w_lin_vel = 0.55
w_ang_vel = 0.15
w_height = 0.10
w_stability = 0.12
w_smoothness = 0.08
total_reward = (
(w_lin_vel * r_lin_vel)
+ (w_ang_vel * r_ang_vel)
+ (w_height * r_height)
+ (w_stability * r_stability)
+ (w_smoothness * r_smoothness)
+ stillness_penalty
)
# Scaled reward for policy stability
return float(total_reward / 10.0)
def _update_hud(self):
if not self.use_gui or not self.pb_robots:
if not self.use_gui:
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, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id)
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)))
pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot)
self.hud.update(
episode=self.episode_count,
step=self.total_steps,
robot_rewards=self.robot_rewards,
cmd_vel=self.commands[0],
robot_rewards=[self.robot_reward],
cmd_vel=self.command,
fps=fps,
avg_height=avg_height,
roll_pitch=avg_roll_pitch
height=pos[2],
roll_pitch=(math.degrees(roll), math.degrees(pitch))
)
def _update_leader_visuals(self):
if not self.use_gui:
def _update_robot_failure(self):
if self.is_failed:
return
best_idx = int(np.argmax(self.robot_rewards))
leader_pb_id = self.pb_robots[best_idx]
leader_pos, _ = self.sim_manager.get_robot_pose(leader_pb_id)
self.leader_crown.update(leader_pos)
if self.step_count < 15:
return
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, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id)
collapse_threshold = max(0.06, self.collapse_height_fraction * self.target_height)
position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot)
collapse_threshold = max(0.04, 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
self.is_failed = True
if self.use_gui:
self.sim_manager.set_robot_color(pb_id, COLOR_FAILED)
self.sim_manager.set_robot_color(self.pb_robot, COLOR_FAILED)
def close(self):
self.sim_manager.disconnect()
class CurriculumCallback(BaseCallback):
"""Logs curriculum phase breakdown and best performance metrics to TensorBoard."""
def __init__(self, verbose=0):
super().__init__(verbose)
self.best_speed = 0.0
self.best_yaw_rate = 0.0
self.best_distance = 0.0
self.best_survival_steps = 0.0
self.best_reward = -float('inf')
def _on_step(self) -> bool:
return True
def _on_rollout_end(self) -> bool:
try:
vec_env = self.training_env
alive_metrics = vec_env.env_method("get_current_robot_metrics")
self.best_speed = 0.0
self.best_yaw_rate = 0.0
self.best_distance = 0.0
self.best_survival_steps = 0.0
self.best_reward = -float('inf')
phase_counts = {
"stand_only": 0,
"forward": 0,
"turn_and_direction": 0,
"omni_direction": 0,
"full_command": 0,
}
for worker_res in alive_metrics:
for metrics in worker_res:
if not metrics.get("alive", False):
continue
phase_key = metrics.get("phase_name", "STAND_ONLY").lower()
if phase_key in phase_counts:
phase_counts[phase_key] += 1
if metrics["reward"] > self.best_reward:
self.best_reward = float(metrics["reward"])
if metrics["speed"] > self.best_speed:
self.best_speed = float(metrics["speed"])
if metrics["yaw_rate"] > self.best_yaw_rate:
self.best_yaw_rate = float(metrics["yaw_rate"])
if metrics["distance_from_start"] > self.best_distance:
self.best_distance = float(metrics["distance_from_start"])
if metrics["survival_steps"] > self.best_survival_steps:
self.best_survival_steps = float(metrics["survival_steps"])
for phase_name, count in phase_counts.items():
self.logger.record(f"phase/{phase_name}", count)
self.logger.record("custom/best_reward", float(self.best_reward) if np.isfinite(self.best_reward) else 0.0)
self.logger.record("custom/best_survival_steps", float(self.best_survival_steps))
self.logger.record("custom/best_distance_from_start_m", float(self.best_distance))
self.logger.record("custom/best_speed_mps", float(self.best_speed))
self.logger.record("custom/best_yaw_rate_rads", float(self.best_yaw_rate))
except Exception:
pass
return True
ml/evaluate.py
-114
View File
@@ -1,114 +0,0 @@
"""
ml/evaluate.py - Curriculum-aware evaluation routine for trained JackBot PPO policies.
"""
import argparse
import json
import time
from pathlib import Path
from typing import Optional, Dict, List, Any
import numpy as np
def evaluate(
model_path: str,
episodes: int = 5,
use_gui: bool = True,
robot_spacing: float = 0.5,
start_pose: str = "init_deg",
random_command: bool = True, # Default to True so curriculum commands are sampled
save_json: Optional[str] = None,
) -> Dict[str, Any]:
try:
from stable_baselines3 import PPO
except ImportError as exc:
raise ImportError(
"stable-baselines3 is required for evaluation. Install with: pip install stable-baselines3"
) from exc
from .env import JackBotEnv
print(f"[Eval] Loading policy model from: {model_path}")
# Force device="cpu" to prevent AMD ROCm/hipBLASLt matrix multiplication crashes
model = PPO.load(model_path, device="cpu")
# Initialize environment with random commands enabled for curriculum progression
env = JackBotEnv(
use_gui=use_gui,
random_command=random_command,
robot_spacing=robot_spacing,
start_pose=start_pose,
)
episode_rewards: List[float] = []
episode_lengths: List[int] = []
episode_phases: List[str] = []
for ep in range(episodes):
obs, _ = env.reset()
done = False
total_reward = 0.0
steps = 0
initial_phase = env.curriculum_phase.name
print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} [Phase: {initial_phase}] ---")
while not done:
action, _ = model.predict(obs, deterministic=True)
# Store phase before step to detect live phase transitions
prev_phase = env.curriculum_phase
obs, reward, terminated, truncated, _ = env.step(action)
# Done on either physical failure (terminated) or phase step limit (truncated)
done = terminated or truncated
total_reward += float(reward)
steps += 1
# Log live phase transition if unlocked during this step
if env.curriculum_phase != prev_phase:
print(f" └─ [Eval Milestone] Curriculum advanced to {env.curriculum_phase.name} at episode step {steps}!")
if use_gui:
time.sleep(1.0 / 240.0)
status_str = "FAILED (Terminated)" if terminated else "COMPLETED (Truncated)"
final_phase = env.curriculum_phase.name
episode_rewards.append(total_reward)
episode_lengths.append(steps)
episode_phases.append(final_phase)
print(
f"Episode {ep + 1} Finished [{status_str}]: "
f"Phase = {final_phase} | Total Reward = {total_reward:.2f} | Steps = {steps}"
)
env.close()
metrics = {
"model_path": str(model_path),
"episodes_evaluated": episodes,
"final_curriculum_phase": env.curriculum_phase.name,
"mean_reward": float(np.mean(episode_rewards)),
"std_reward": float(np.std(episode_rewards)),
"mean_episode_length": float(np.mean(episode_lengths)),
"raw_rewards": episode_rewards,
"episode_phases": episode_phases,
}
print("\n" + "=" * 60)
print(f"EVALUATION COMPLETE ({episodes} Episodes)")
print(f"Final Reached Phase: {metrics['final_curriculum_phase']}")
print(f"Mean Reward: {metrics['mean_reward']:.2f} ± {metrics['std_reward']:.2f}")
print(f"Mean Episode Length: {metrics['mean_episode_length']:.1f} steps")
print("=" * 60)
if save_json:
out_path = Path(save_json)
out_path.parent.mkdir(parents=True, exist_ok=True)
with open(out_path, "w") as f:
json.dump(metrics, f, indent=4)
print(f"[Eval] Saved evaluation metrics to: {out_path.resolve()}")
return metrics
ml/run_eval.py
+92 -21
View File
@@ -5,38 +5,109 @@ Usage:
"""
import argparse
import json
import time
import sys
from pathlib import Path
ROOT_DIR = Path(__file__).resolve().parent.parent
if str(ROOT_DIR) not in sys.path:
sys.path.insert(0, str(ROOT_DIR))
from ml.evaluate import evaluate
import numpy as np
from stable_baselines3 import PPO
sys.path.append(str(Path(__file__).resolve().parent.parent))
from ml.env import JackBotEnv
def main():
parser = argparse.ArgumentParser(description="JackBot Curriculum Policy Evaluator Wrapper")
parser.add_argument("--model", type=str, required=True, help="Path to the trained model file (.zip)")
parser.add_argument("--episodes", type=int, default=5, help="Number of evaluation episodes to run.")
parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during evaluation")
parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
parser.add_argument("--no-random-command", dest="random_command", action="store_false", help="Disable command sampling and lock to zero commands")
parser = argparse.ArgumentParser(description="JackBot Policy Evaluator")
parser.add_argument("--model", type=str, required=True, help="Path to trained model checkpoint (.zip)")
parser.add_argument("--episodes", type=int, default=5, help="Number of evaluation episodes to run")
parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
parser.add_argument("--no-random-command", dest="random_command", action="store_false", help="Lock commands to zero (disable random command sampling)")
parser.set_defaults(random_command=True)
parser.add_argument("--save-metrics", type=str, default=None, help="Optional path to save JSON metrics report")
parser.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to output evaluation summary")
args = parser.parse_args()
evaluate(
model_path=args.model,
episodes=args.episodes,
print(f"[Eval] Loading policy model from: {args.model}")
model = PPO.load(args.model, device="cpu")
# Instantiate single evaluation environment matching train setup
env = JackBotEnv(
use_gui=args.gui,
robot_spacing=args.robot_spacing,
start_pose=args.start_pose,
random_command=args.random_command,
save_json=args.save_metrics,
)
episode_rewards = []
episode_lengths = []
episode_phases = []
try:
for ep in range(args.episodes):
obs, _ = env.reset()
done = False
total_reward = 0.0
steps = 0
initial_phase = env.curriculum_phase.name
print(f"\n--- Starting Evaluation Episode {ep + 1}/{args.episodes} [Phase: {initial_phase}] ---")
while not done:
# Deterministic prediction matches evaluation standards
action, _ = model.predict(obs, deterministic=True)
prev_phase = env.curriculum_phase
obs, reward, terminated, truncated, _ = env.step(action)
done = terminated or truncated
total_reward += float(reward)
steps += 1
if env.curriculum_phase != prev_phase:
print(f" └─ [Eval Milestone] Reached {env.curriculum_phase.name} at step {steps}!")
if args.gui:
time.sleep(1.0 / 240.0)
status_str = "FAILED (Terminated)" if terminated else "COMPLETED (Truncated)"
final_phase = env.curriculum_phase.name
episode_rewards.append(total_reward)
episode_lengths.append(steps)
episode_phases.append(final_phase)
print(
f"Episode {ep + 1} Finished [{status_str}]: "
f"Phase = {final_phase} | Total Reward = {total_reward:.2f} | Steps = {steps}"
)
finally:
env.close()
# Calculate metrics report
mean_reward = float(np.mean(episode_rewards))
std_reward = float(np.std(episode_rewards))
mean_length = float(np.mean(episode_lengths))
print("\n" + "=" * 60)
print(f"EVALUATION COMPLETE ({args.episodes} Episodes)")
print(f"Final Reached Phase: {env.curriculum_phase.name}")
print(f"Mean Reward: {mean_reward:.2f} ± {std_reward:.2f}")
print(f"Mean Episode Length: {mean_length:.1f} steps")
print("=" * 60)
if args.save_metrics:
metrics = {
"model_path": str(args.model),
"episodes_evaluated": args.episodes,
"final_curriculum_phase": env.curriculum_phase.name,
"mean_reward": mean_reward,
"std_reward": std_reward,
"mean_episode_length": mean_length,
"raw_rewards": episode_rewards,
"episode_phases": episode_phases,
}
out_path = Path(args.save_metrics)
out_path.parent.mkdir(parents=True, exist_ok=True)
with open(out_path, "w") as f:
json.dump(metrics, f, indent=4)
print(f"[Eval] Saved evaluation report to: {out_path.resolve()}")
if __name__ == "__main__":
main()
ml/run_train.py
+112 -27
View File
@@ -1,47 +1,132 @@
"""Minimal training launcher for quick experiments.
Usage:
python ml/run_train.py --timesteps 500000 --model ml/checkpoints/ppo_joint_command
This is a convenience wrapper around `ml.train.train` with friendly defaults
for interactive experimentation.
python ml/run_train.py --total-timesteps 1500000 --gui
"""
import argparse
import os
import re
import sys
import warnings
from pathlib import Path
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
ROOT_DIR = Path(__file__).resolve().parent.parent
if str(ROOT_DIR) not in sys.path:
sys.path.insert(0, str(ROOT_DIR))
sys.path.append(str(Path(__file__).resolve().parent.parent))
from ml.env import JackBotEnv, CurriculumCallback
from ml.train import train
# Silence SB3's UserWarning about SubprocVecEnv vs DummyVecEnv
warnings.filterwarnings("ignore", category=UserWarning, module="stable_baselines3")
def get_next_run_number(save_dir: str) -> int:
"""Scans the save directory for existing ppo<number> patterns and returns the next integer."""
if not os.path.exists(save_dir):
return 1
existing_numbers = []
for item in os.listdir(save_dir):
# Match pattern ppo followed by numbers (e.g., ppo1, ppo_1, jackbot_ppo12)
matches = re.findall(r"ppo_?(\d+)", item, re.IGNORECASE)
for m in matches:
existing_numbers.append(int(m))
return max(existing_numbers, default=0) + 1
def make_env(rank: int, use_gui: bool = False, seed: int = 0):
"""Utility helper to instantiate parallel JackBot environments."""
def _init():
env = JackBotEnv(
use_gui=use_gui,
random_command=True,
)
env.reset(seed=seed + rank)
return env
return _init
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--timesteps", type=int, default=500000, help="Total number of 'practice steps'.")
parser.add_argument("--model", type=str, default="ml/checkpoints/ppo_joint_command", help="Path to save the trained model")
parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducibility")
parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto'")
parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during training")
parser.add_argument("--num-workers", type=int, default=1, help="Number of training environment workers")
parser.add_argument("--robot-spacing", type=float, default=1.5, help="Spacing between robots in meters")
parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
parser = argparse.ArgumentParser(description="JackBot PPO Curriculum Trainer")
parser.add_argument("--num-workers", type=int, default=16, help="Number of parallel sub-process environments")
parser.add_argument("--total-timesteps", type=int, default=1_500_000, help="Total training timesteps")
parser.add_argument("--log-dir", type=str, default="ml/logs", help="Directory for TensorBoard logs")
parser.add_argument("--save-dir", type=str, default="ml/checkpoints", help="Directory for model checkpoints")
parser.add_argument("--save-freq", type=int, default=50_000, help="Checkpoint save frequency (steps)")
parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
args = parser.parse_args()
Path(args.model).parent.mkdir(parents=True, exist_ok=True)
train(
total_timesteps=args.timesteps,
model_path=args.model,
seed=args.seed,
device=args.device,
use_gui=args.gui,
num_workers=args.num_workers,
robot_spacing=args.robot_spacing,
start_pose=args.start_pose,
os.makedirs(args.log_dir, exist_ok=True)
os.makedirs(args.save_dir, exist_ok=True)
# Automatically determine next run number (e.g. ppo1, ppo2, ppo3...)
run_num = get_next_run_number(args.save_dir)
ppo_name = f"ppo{run_num}"
print(f"[Train] Initializing Run #{run_num} ('{ppo_name}') with {args.num_workers} parallel workers...")
env_fns = [
make_env(rank=i, use_gui=(args.gui if i == 0 else False))
for i in range(args.num_workers)
]
vec_env = SubprocVecEnv(env_fns)
# Initialize PPO Policy Hyperparameters
model = PPO(
policy="MlpPolicy",
env=vec_env,
learning_rate=3e-4,
n_steps=256,
batch_size=256,
n_epochs=10,
gamma=0.99,
gae_lambda=0.95,
clip_range=0.2,
ent_coef=0.03,
target_kl=0.05,
vf_coef=0.5,
max_grad_norm=0.5,
verbose=1,
tensorboard_log=args.log_dir,
device="cpu",
)
# Setup Callbacks with ppo<number> naming
checkpoint_callback = CheckpointCallback(
save_freq=max(1, args.save_freq // args.num_workers),
save_path=args.save_dir,
name_prefix=f"jackbot_{ppo_name}",
)
curriculum_callback = CurriculumCallback()
eval_env = DummyVecEnv([lambda: JackBotEnv(use_gui=False, random_command=True)])
best_model_path = os.path.join(args.save_dir, f"best_model_{ppo_name}")
eval_callback = EvalCallback(
eval_env,
best_model_save_path=best_model_path,
log_path="ml/logs/results",
eval_freq=max(1, 10_000 // args.num_workers),
deterministic=True,
render=False,
)
print(f"[Train] Starting training for {args.total_timesteps} timesteps...")
try:
model.learn(
total_timesteps=args.total_timesteps,
callback=[checkpoint_callback, curriculum_callback, eval_callback],
progress_bar=True,
)
final_model_path = os.path.join(args.save_dir, f"jackbot_{ppo_name}_final.zip")
model.save(final_model_path)
print(f"[Train] Training complete! Saved final model to {final_model_path}")
finally:
vec_env.close()
eval_env.close()
if __name__ == "__main__":
main()
ml/train.py
-257
View File
@@ -1,257 +0,0 @@
import os
import argparse
from pathlib import Path
import numpy as np
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv
from .env import JackBotEnv
class MilestoneCheckpointCallback(BaseCallback):
"""
Saves a model checkpoint the FIRST time total_timesteps
crosses every multiple of step_interval (e.g., 100,000).
Saves inside the matching PPO_X subfolder as created by TensorBoard.
"""
def __init__(self, save_path: str, name_prefix: str = "ppo_jackbot", step_interval: int = 100_000, verbose: int = 1):
super().__init__(verbose)
self.base_save_path = save_path
self.run_save_path = save_path
self.name_prefix = name_prefix
self.step_interval = step_interval
self.last_milestone = 0
def _on_training_start(self) -> None:
"""Executed right before training loop starts. Resolves TensorBoard's run folder name (e.g. PPO_1)."""
if self.logger and self.logger.dir:
run_folder_name = Path(self.logger.dir).name # Extracts "PPO_1", "PPO_2", etc.
self.run_save_path = os.path.join(self.base_save_path, run_folder_name)
else:
self.run_save_path = self.base_save_path
os.makedirs(self.run_save_path, exist_ok=True)
def _on_step(self) -> bool:
current_milestone = self.num_timesteps // self.step_interval
if current_milestone > self.last_milestone:
self.last_milestone = current_milestone
milestone_step = current_milestone * self.step_interval
save_file = os.path.join(
self.run_save_path,
f"{self.name_prefix}_{milestone_step}_steps.zip"
)
self.model.save(save_file)
if self.verbose > 0:
print(f"\n[Checkpoint] Saved milestone model at {self.num_timesteps} steps -> {save_file}\n")
return True
class JackBotMetricsCallback(BaseCallback):
"""
Tracks the best current alive-robot performance for the most recent rollout,
instead of logging lifetime maxima from the entire training run.
"""
def __init__(self, verbose=0):
super().__init__(verbose)
self.best_alive_speed = 0.0
self.best_alive_yaw_rate = 0.0
self.best_alive_distance = 0.0
self.best_alive_survival_steps = 0.0
self.best_alive_reward = -float('inf')
def _on_step(self) -> bool:
return True
def _on_rollout_end(self) -> bool:
try:
vec_env = self.training_env
alive_metrics = vec_env.env_method("get_current_robot_metrics")
self.best_alive_speed = 0.0
self.best_alive_yaw_rate = 0.0
self.best_alive_distance = 0.0
self.best_alive_survival_steps = 0.0
self.best_alive_reward = -float('inf')
for worker_res in alive_metrics:
for metrics in worker_res:
if not metrics.get("alive", False):
continue
if metrics["reward"] > self.best_alive_reward:
self.best_alive_reward = float(metrics["reward"])
if metrics["speed"] > self.best_alive_speed:
self.best_alive_speed = float(metrics["speed"])
if metrics["yaw_rate"] > self.best_alive_yaw_rate:
self.best_alive_yaw_rate = float(metrics["yaw_rate"])
if metrics["distance_from_start"] > self.best_alive_distance:
self.best_alive_distance = float(metrics["distance_from_start"])
if metrics["survival_steps"] > self.best_alive_survival_steps:
self.best_alive_survival_steps = float(metrics["survival_steps"])
self.logger.record("custom/best_alive_speed_mps", float(self.best_alive_speed))
self.logger.record("custom/best_alive_yaw_rate_rads", float(self.best_alive_yaw_rate))
self.logger.record("custom/best_alive_distance_from_start_m", float(self.best_alive_distance))
self.logger.record("custom/best_alive_survival_steps", float(self.best_alive_survival_steps))
self.logger.record("custom/best_alive_reward", float(self.best_alive_reward) if np.isfinite(self.best_alive_reward) else 0.0)
self.logger.record("custom/max_speed_mps", float(self.best_alive_speed))
self.logger.record("custom/max_yaw_rate_rads", float(self.best_alive_yaw_rate))
self.logger.record("custom/max_distance_from_start_m", float(self.best_alive_distance))
self.logger.record("custom/max_survival_steps", float(self.best_alive_survival_steps))
self.logger.record("custom/best_episode_reward", float(self.best_alive_reward) if np.isfinite(self.best_alive_reward) else 0.0)
except Exception:
pass
return True
def parse_args():
parser = argparse.ArgumentParser(description="Train a joint-command policy for JackBot.")
parser.add_argument("--timesteps", type=int, default=500_000, help="Total training timesteps")
parser.add_argument("--model-path", type=str, default="ml/checkpoints/ppo_joint_command", help="Where to save the trained model")
parser.add_argument("--seed", type=int, default=0, help="Random seed")
parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto' to autodetect")
parser.add_argument("--gui", "--use-gui", dest="use_gui", action="store_true", help="Enable PyBullet GUI during training")
parser.add_argument("--num-workers", type=int, default=8, help="Number of parallel CPU worker processes")
parser.add_argument("--robot-spacing", type=float, default=3.0, help="Spacing between robots in meters")
parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
return parser.parse_args()
def make_env(robot_spacing, start_pose, use_gui, rank, seed=0):
def _init():
env = JackBotEnv(
use_gui=use_gui if rank == 0 else False,
random_command=True,
robot_spacing=robot_spacing,
start_pose=start_pose,
)
env.reset(seed=seed + rank)
return env
return _init
def train(
total_timesteps: int,
model_path: str,
seed: int = 0,
device: str = "auto",
use_gui: bool = False,
num_workers: int = 8,
robot_spacing: float = 0.5,
start_pose: str = "init_deg",
):
try:
from stable_baselines3 import PPO
except ImportError as exc:
raise ImportError("stable-baselines3 is required. Install with: pip install stable-baselines3") from exc
if num_workers > 1:
env_fns = [
make_env(robot_spacing, start_pose, use_gui, rank=i, seed=seed)
for i in range(num_workers)
]
env = SubprocVecEnv(env_fns)
else:
env = DummyVecEnv([
make_env(robot_spacing, start_pose, use_gui, rank=0, seed=seed)
])
def resolve_device(requested_device: str) -> str:
try:
import torch
except ImportError:
if requested_device != "cpu":
raise RuntimeError("PyTorch is not installed.")
return "cpu"
hip_supported = getattr(torch.version, "hip", None) is not None
cuda_available = torch.cuda.is_available()
hip_available = hip_supported and getattr(torch.backends, "hip", None) is not None and torch.backends.hip.is_available()
if requested_device == "auto":
return "cuda" if (hip_available or cuda_available) else "cpu"
if requested_device in {"cuda", "gpu", "hip"}:
if hip_available or cuda_available:
return "cuda"
raise RuntimeError(f"GPU requested ({requested_device}) but not available.")
return "cpu"
device = resolve_device(device)
policy_kwargs = dict(
log_std_init=-1.5,
net_arch=dict(pi=[256, 256], vf=[256, 256])
)
model = PPO(
"MlpPolicy",
env,
verbose=1,
seed=seed,
learning_rate=1.5e-4,
n_steps=256,
batch_size=256,
n_epochs=10,
gamma=0.99,
gae_lambda=0.95,
clip_range=0.2,
target_kl=0.03,
ent_coef=0.03,
vf_coef=0.5,
max_grad_norm=0.5,
device=device,
policy_kwargs=policy_kwargs,
tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"),
)
# Base folder where model runs will be stored
save_dir = str(Path(model_path).parent)
model_prefix = Path(model_path).stem
milestone_cb = MilestoneCheckpointCallback(
save_path=save_dir,
name_prefix=model_prefix,
step_interval=100_000
)
metrics_callback = JackBotMetricsCallback()
model.learn(total_timesteps=total_timesteps, callback=[milestone_cb, metrics_callback])
# Save the final model inside the matching PPO_X directory as well
if model.logger and model.logger.dir:
run_folder_name = Path(model.logger.dir).name
final_dir = Path(model_path).parent / run_folder_name
else:
final_dir = Path(model_path).parent
final_dir.mkdir(parents=True, exist_ok=True)
final_save_path = final_dir / f"{model_prefix}_final.zip"
model.save(str(final_save_path))
if model.verbose > 0:
print(f"[Training Complete] Saved final model to -> {final_save_path}")
env.close()
if __name__ == "__main__":
args = parse_args()
train(
args.timesteps,
args.model_path,
seed=args.seed,
device=args.device,
use_gui=args.use_gui,
num_workers=args.num_workers,
robot_spacing=args.robot_spacing,
start_pose=args.start_pose,
)
requirements.txt
+2 -1
View File
@@ -8,6 +8,7 @@ matplotlib
# Deep learning / RL
torch
stable-baselines3
gymnasium[box2d]
stable-baselines3[extra]
gymnasium
shimmy
tensorboard