code reduction for training and eval
reward and everything else changed again wont be the last time
This commit is contained in:
+18
-62
@@ -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,69 +76,25 @@ 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:
|
||||
self._text_id = p.addUserDebugText(
|
||||
text=hud_text,
|
||||
textPosition=text_position,
|
||||
textColorRGB=text_color,
|
||||
textSize=0.1,
|
||||
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
|
||||
)
|
||||
# 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,
|
||||
textColorRGB=text_color,
|
||||
textSize=0.1,
|
||||
textOrientation=text_orientation,
|
||||
physicsClientId=self.client_id
|
||||
)
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Removes the active debug text item so a new episode starts with a clean overlay."""
|
||||
if self._text_id is not None:
|
||||
try:
|
||||
p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
|
||||
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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
+2
-4
@@ -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"]
|
||||
|
||||
@@ -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 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
|
||||
)
|
||||
|
||||
# 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 (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
|
||||
vx = np.random.uniform(0.15, 0.50)
|
||||
|
||||
vy, vz, omega = 0.0, 0.0, 0.0
|
||||
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.robot_reward = 0.0
|
||||
self.is_failed = False
|
||||
|
||||
self.episode_height_sum = 0.0
|
||||
self.episode_roll_sum = 0.0
|
||||
self.episode_pitch_sum = 0.0
|
||||
|
||||
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.0, 0.0, 1.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])
|
||||
if self.use_gui:
|
||||
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]
|
||||
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)
|
||||
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]))
|
||||
|
||||
return [{
|
||||
"reward": float(self.robot_reward),
|
||||
"distance_from_start": dist,
|
||||
"speed": speed,
|
||||
"yaw_rate": yaw_rate,
|
||||
"alive": True,
|
||||
"survival_steps": 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()
|
||||
|
||||
if self.use_gui and self.step_count % render_freq == 0:
|
||||
self.sim_manager.set_rendering(True)
|
||||
|
||||
self._update_robot_failures()
|
||||
self.get_robot_distance_metrics()
|
||||
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
|
||||
|
||||
self._update_hud()
|
||||
self._update_leader_visuals()
|
||||
if self.step_count % 120 == 0 and self.use_gui:
|
||||
self._update_hud()
|
||||
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)
|
||||
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
|
||||
)
|
||||
|
||||
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, cmd_vy, _, cmd_yaw = self.command
|
||||
cmd_norm = math.hypot(cmd_vx, cmd_vy)
|
||||
|
||||
cmd_vx = command[0] # Normalized -1.0 to 1.0
|
||||
cmd_vy = command[1] # Normalized -1.0 to 1.0
|
||||
cmd_yaw = command[3]
|
||||
# 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
|
||||
|
||||
cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32)
|
||||
cmd_norm = np.linalg.norm(cmd_dir)
|
||||
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
|
||||
|
||||
actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)
|
||||
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]
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# 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))
|
||||
# 3. Posture & Stability Sub-Rewards
|
||||
height_error = pos[2] - self.target_height
|
||||
r_height = math.exp(-150.0 * (height_error ** 2))
|
||||
|
||||
target_speed = cmd_norm * self.max_robot_speed # e.g., 0.35 * 0.8 = 0.28 m/s
|
||||
orientation_error = roll**2 + pitch**2
|
||||
r_stability = math.exp(-25.0 * orientation_error)
|
||||
|
||||
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)
|
||||
joint_error = np.mean(np.square(current_joints - self.default_joint_angles))
|
||||
r_pose = math.exp(-2.0 * joint_error)
|
||||
|
||||
# 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)
|
||||
action_delta = np.mean(np.square(action - previous_action))
|
||||
r_smoothness = math.exp(-0.1 * action_delta)
|
||||
|
||||
# -----------------------------------------------------------------
|
||||
# 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
|
||||
# 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
|
||||
|
||||
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
|
||||
total_reward = (
|
||||
(w_height * r_height)
|
||||
+ (w_stability * r_stability)
|
||||
+ (w_pose * r_pose)
|
||||
+ (w_smoothness * r_smoothness)
|
||||
)
|
||||
rewards.append(r_step)
|
||||
else:
|
||||
# WALKING / TURNING MODE
|
||||
is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0)
|
||||
|
||||
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
|
||||
# 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(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
|
||||
|
||||
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)
|
||||
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.sim_manager.set_robot_color(pb_id, COLOR_FAILED)
|
||||
if is_tilted or is_collapsed:
|
||||
self.is_failed = True
|
||||
if self.use_gui:
|
||||
self.sim_manager.set_robot_color(self.pb_robot, COLOR_FAILED)
|
||||
|
||||
def close(self):
|
||||
self.sim_manager.disconnect()
|
||||
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
|
||||
-114
@@ -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
|
||||
+92
-21
@@ -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()
|
||||
+113
-28
@@ -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()
|
||||
main()
|
||||
-257
@@ -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,
|
||||
)
|
||||
Reference in New Issue
Block a user