From acb3d671be3647b59881a28da3cbdd62e00d4ec8 Mon Sep 17 00:00:00 2001 From: JackM323 Date: Mon, 3 Aug 2026 22:27:26 +0200 Subject: [PATCH] Reworked training new reward/penalty system learning phases with curriculum learning new training parameters cleanup of old code better logging while training multiple environments instead of robots (they could bumb into each other) --- README.md | 266 +++++++++++++++++++++++++++--------- Robot.py | 4 +- ml/SimManager.py | 33 +++-- ml/env.py | 341 +++++++++++++++++++++++++++++++++++++---------- ml/evaluate.py | 5 +- ml/run_eval.py | 4 +- ml/run_train.py | 6 +- ml/train.py | 89 +++++++++++-- 8 files changed, 585 insertions(+), 163 deletions(-) diff --git a/README.md b/README.md index cdb10b7..b8a59fc 100644 --- a/README.md +++ b/README.md @@ -4,12 +4,45 @@ JackBot is a modular 3D hexapod robot control and machine learning framework bui --- -## Key Features +## What JackBot Is -* **Unified Robot Abstraction (`Robot.py`):** Virtual backends (`RobotBackend` protocol) allow seamless switching between 3D PyBullet simulation and physical hardware (ESP32 / Arduino) without changing high-level logic. -* **Flexible Input Pipeline:** Pluggable input handlers supporting Pygame gamepad controllers, manual GUI sliders, or randomized direction vectors. -* **Parallel Multi-Robot Training:** Vectorized Gymnasium environment (`JackBotEnv`) capable of simulating and training $N$ parallel hexapods simultaneously in PyBullet for PPO reinforcement learning. -* **Live Telemetry & Visual Tracking:** Built-in PyBullet overlay features including real-time performance HUDs, floating leader crown tracking ($\text{👑}$) for top-reward robots, and visual failure feedback (failed robots turn semi-transparent dark gray). +JackBot is a Python-based hexapod robot project that combines: + +* a robot control stack for a six-legged walking robot, +* a physics simulator for testing in software before using real hardware, and +* a reinforcement learning pipeline that teaches the robot how to move through trial and error. + +The project is not just one script or one model. It is a full control system that can: + +* run the robot in a PyBullet simulation, +* accept human commands from a GUI or gamepad, +* stream target joint positions to physical hardware, and +* train a policy using PPO so the robot can learn locomotion automatically. + +## Why the Project Exists + +A hexapod is hard to control manually because each leg has several joints and the robot has to maintain balance while moving. Instead of hard-coding every motion rule, this project uses the robot model and physics simulation as a testbed to learn walking behavior. + +In practice, the workflow looks like this: + +1. The robot starts from a standing pose. +2. The policy receives sensory information from the simulation. +3. The policy chooses new joint target motions. +4. The simulation updates the physics. +5. The policy receives a reward for surviving and moving in the right direction. +6. Over many iterations, PPO improves the controller. + +--- + +## Key Components + +* **Robot abstraction (`Robot.py`)**: wraps the robot body, robot state, and joint control API. +* **Kinematics (`kinematics.py`)**: turns target leg poses into actual joint angles through inverse kinematics. +* **Simulation (`simulation.py`)**: creates and updates the PyBullet world where the robot can be tested safely. +* **Inputs (`inputs/`)**: lets the robot be controlled from either a GUI, a gamepad, or a generated target command. +* **ML environment (`ml/env.py`)**: provides the Gymnasium environment that the policy interacts with. +* **PPO training (`ml/train.py`)**: trains a policy using stable-baselines3. +* **Evaluation (`ml/evaluate.py`)**: loads a saved model and runs policy rollouts to inspect performance. --- @@ -17,36 +50,40 @@ JackBot is a modular 3D hexapod robot control and machine learning framework bui ```text JackBot/ -├── main.py # Primary application entry point for manual & hardware control -├── Robot.py # Core Robot class, kinematics wrapper, and Backend protocols -├── config.py # Global settings (backend selection, URDF path, communication specs) -├── kinematics.py # Forward and Inverse Kinematics (IKPy) -├── simulation.py # Base PyBullet GUI wrapper for single-robot interactive simulation -├── robot_init.py # Default stance angles and neutral leg positions -├── DataTypes.py # Strongly typed arrays (PosArray, RadArray, DegArray) & structs +├── main.py # manual control and hardware streaming entry point +├── Robot.py # robot wrapper and backend abstraction +├── config.py # global settings and communication configuration +├── kinematics.py # inverse kinematics used to map motion commands to joints +├── simulation.py # PyBullet simulation shell for the robot +├── robot_init.py # neutral standing pose and initial joint values +├── DataTypes.py # typed data structures for positions and joint angles │ -├── states/ # Finite State Machine (FSM) gait states -│ ├── State.py # Base State class -│ ├── idle.py # Neutral stance state -│ └── walking.py # Inverse-kinematics tripod gait state +├── states/ # gait/state-machine behavior logic +│ ├── State.py +│ ├── IdleState.py +│ ├── WalkingState.py +│ └── WaveState.py │ -├── inputs/ # Input providers -│ ├── InputProvider.py # Base input abstraction -│ └── PygameController.py # Asynchronous gamepad loop (process-isolated) +├── inputs/ # command sources for the robot +│ ├── InputProvider.py +│ ├── PygameController.py +│ └── RandomeInputProvider.py │ -├── gui/ # Control interface -│ └── MainWindow.py # Pygame / parameter GUI layout and command resolver +├── gui/ # visual control frontend +│ └── MainWindow.py │ -├── EspCommunication.py # WiFi socket sender for ESP32 hardware -├── ArduinoCommunication.py # Serial communication wrapper for Arduino hardware -├── JackBotUrdf.urdf # Kinematic 3D model definition (18 active joints) +├── EspCommunication.py # ESP32 communication layer +├── ArduinoCommunication.py # Arduino serial communication layer +├── JackBotUrdf.urdf # robot mesh and joint definition │ -└── ml/ # Machine Learning Subsystem - ├── env.py # JackBotEnv (Gymnasium multi-robot vector environment) - ├── SimManager.py # Physics server initialization and scene loading - ├── MetricsOverlay.py # PyBullet HUD (MetricsHUD) & leader crown tracking (LeaderCrown) - ├── run_train.py # PPO training execution script - └── run_eval.py # Model evaluation script +└── ml/ # reinforcement learning subsystem + ├── env.py # Gymnasium environment used by PPO + ├── SimManager.py # PyBullet scene setup and stepping + ├── MetricsOverlay.py # HUD and visual overlays + ├── train.py # PPO training entry point + ├── run_train.py # command-line training launcher + ├── evaluate.py # policy evaluation loop + └── run_eval.py # CLI wrapper for evaluation ``` --- @@ -102,40 +139,144 @@ Edit `config.py` prior to launching `main.py` to configure execution mode and co --- -### 2. Machine Learning: PPO Training & Evaluation (`ml/`) +### 2. How the Machine Learning System Works -The `ml/` directory contains tools to train RL policies using **Proximal Policy Optimization (PPO)**. The agent receives observations ($18\text{ joint angles} + 4\text{ velocity/turning commands}$) and outputs continuous joint delta actions in $[-1, 1]$. +The machine learning subsystem teaches the robot to move by interacting with a simulation environment instead of relying on a manually written gait controller. -#### A. Training a Model (`run_train.py`) +At a high level: -Train a single robot policy: +* The policy receives a vector of observations from the simulator. +* The policy outputs a set of continuous joint deltas. +* Those outputs are applied to the robot in the physics simulation. +* The environment computes a reward based on survival, movement, stability, and command-following quality. +* PPO updates the policy over many training iterations to improve the controller. + +#### What the policy sees and does + +The policy input contains: + +* 18 joint angles from the robot, and +* 4 command dimensions describing the desired motion direction and yaw rate. + +The policy output is: + +* 18 continuous values, one per joint, describing how much each joint should change. + +This means the learning method is not controlling a discrete gait state directly. Instead, it learns a continuous control policy for the whole body. + +#### Methods implemented + +The current ML stack uses: + +* **Gymnasium** as the environment API. +* **PyBullet** as the physics engine. +* **Stable-Baselines3 PPO** as the learning algorithm. +* **Curriculum learning** to gradually unlock harder command regimes. + +The key ideas are: + +1. **Standing first**: the robot must stay upright and stable. +2. **Forward motion second**: the robot must learn to move in a commanded direction. +3. **Turning next**: after forward command-following is stable, the policy gets a yaw/turning challenge. +4. **Omni-direction later**: the full command space is gradually exposed only once the simpler skills are reliable. + +This staged approach is important because learning to balance and walk all at once is difficult for a real hexapod. + +#### Training flow + +Training is launched through `ml/run_train.py` and hands off to `ml/train.py`. + +Typical training starts as: ```bash python ml/run_train.py --timesteps 100000 --model ml/checkpoints/ppo_joint_command ``` -Train using multi-robot parallel vectorization with visual GUI enabled: -```bash -python ml/run_train.py --timesteps 500000 --model ml/checkpoints/ppo_joint_command --num-robots 16 --robot-spacing 0.75 --gui -``` +Training uses the simulator as the environment, runs PPO updates, and periodically saves checkpoint models. -**Training CLI Arguments:** -* `--timesteps`: Total training timesteps. -* `--num-robots`: Number of parallel robot instances spawned in a grid layout (e.g., 16 to 64). -* `--robot-spacing`: Distance in meters between robot spawn origins. -* `--start-pose`: Stance pose at environment reset (`init_deg` or `init90_deg`). -* `--gui`: Renders the live PyBullet GUI with metrics HUD, leader crown, and failure graying. +#### Evaluation flow -#### B. Evaluating a Model (`run_eval.py`) +Evaluation uses `ml/run_eval.py` and `ml/evaluate.py`. -Run an evaluation loop using a saved model checkpoint: +A saved PPO model is loaded and run in deterministic evaluation mode. This is used to answer questions like: + +* Does the policy actually move? +* Does it survive longer than before? +* Does it remain stable under the commanded motion? + +Example: ```bash python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 5 --gui ``` -Multi-robot evaluation with custom stance pose: -```bash -python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3 --gui --num-robots 4 --robot-spacing 0.8 --start-pose init_deg -``` +--- + +## What the Reward Is Trying to Teach + +The reward function is not a single number with one purpose. It combines several terms so the policy learns to: + +* survive in the environment, +* move in the commanded direction, +* avoid unwanted sideways drift, +* stay upright without excessive tilt, +* avoid giant control jumps, +* keep the robot away from a low collapsed posture. + +The main ideas are: + +* a small alive bonus for staying upright, +* directional movement rewards, +* penalties for drifting or standing still when a command is active, +* penalties for excessive tilt or collapse, +* a penalty for large abrupt control changes. + +This style of reward shaping encourages the policy to learn real locomotion rather than just freezing in place. + +### Curriculum / phase progression + +The environment uses a staged curriculum to expose command complexity gradually. + +* **Phase 0**: basic standing / forward-focused behavior +* **Phase 1**: turning and directional regularization +* **Phase 2**: omni-directional movement +* **Phase 3**: full-command challenge + +The goal is to prevent the learner from being asked to master all difficult motion goals at once. + +--- + +## How the Control Loop Works in Practice + +A simple, human-readable pipeline is: + +1. `main.py` or the training/evaluation wrapper creates the robot environment. +2. The simulation runs the robot in PyBullet. +3. The policy observes the latest joint states and command vectors. +4. PPO predicts a new action. +5. The action is applied to the robot joints. +6. The simulator steps forward one frame. +7. The reward is computed from the new state. +8. The policy improves based on that reward. + +That loop is the essence of the machine learning part of the project. + +--- + +## What Makes this Project Useful + +This repository is useful because it combines several layers that are often separate: + +* robot control and kinematics, +* physical simulation, +* command input sources, +* RL environment construction, +* PPO training and evaluation. + +For a beginner, the easiest way to think about the repo is: + +- `main.py` is for direct manual control, +- `ml/env.py` is the simulator-to-policy interface, +- `ml/train.py` is where training happens, +- `ml/evaluate.py` is where you check whether the learned policy is actually good. --- @@ -143,15 +284,18 @@ python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 3 ### 1. Reward Function Formulation -The per-robot step reward ($R_{\text{step}}$) incentivizes tracking directional velocity commands while maintaining body stability and smooth joint actuation. The total environment step reward is the sum of all individual active robot rewards: +The per-robot step reward ($R_{\text{step}}$) is shaped to reward tracked motion and survival while strongly discouraging stillness once a command is present. The current implementation uses the following structure: -$$R_{\text{step}} = R_{\text{alive}} + R_{\text{tracking}} + R_{\text{rotation}} - P_{\text{stability}} - P_{\text{action}}$$ +$$R_{\text{step}} = R_{\text{alive}} + R_{\text{tracking}} + R_{\text{rotation}} - P_{\text{drift}} - P_{\text{still}}(t) - P_{\text{height}} - P_{\text{stability}} - P_{\text{large-delta}}$$ -* **Alive Bonus ($R_{\text{alive}} = +0.1$):** A constant positive baseline awarded every timestep the robot remains upright. -* **Velocity Tracking ($R_{\text{tracking}} = v_{x,\text{cmd}} \cdot v_x + v_{y,\text{cmd}} \cdot v_y$):** Rewards linear movement in the target command direction ($v_x, v_y$). -* **Yaw Rotation Tracking ($R_{\text{rotation}} = \omega_{\text{cmd}} \cdot \omega_z$):** Rewards turning along the vertical yaw axis according to angular command $\omega_{\text{cmd}}$. -* **Body Stability Penalty ($P_{\text{stability}} = 0.2 \cdot (|\text{roll}| + |\text{pitch}|)$):** Penalizes tilting away from a level horizontal posture. -* **Action Energy Penalty ($P_{\text{action}} = 0.01 \cdot \sum a_i^2$):** Penalizes excessive joint delta actions to encourage smooth, energy-efficient leg movements and reduce jitter. +* **Alive Bonus ($R_{\text{alive}} \approx +0.02$):** A small positive baseline awarded every timestep the robot remains upright. +* **Velocity Tracking ($R_{\text{tracking}}$):** Rewards linear movement aligned with the commanded direction vector. +* **Yaw Rotation Tracking ($R_{\text{rotation}}$):** Rewards turning in the commanded yaw direction. +* **Drift Penalty ($P_{\text{drift}}$):** Penalizes lateral motion that does not align with the commanded direction. +* **Age-Ramped Stillness Penalty ($P_{\text{still}}(t)$):** Penalizes remaining stationary when a command is active; the penalty ramps up as the episode gets older so the policy cannot settle into a frozen local optimum. +* **Height Penalty ($P_{\text{height}}$):** Penalizes falling below the measured post-settle standing height. +* **Body Stability Penalty ($P_{\text{stability}}$):** Penalizes roll and pitch tilt. +* **Large-Delta Penalty ($P_{\text{large-delta}}$):** Penalizes only large joint-control jumps, not normal continuous command usage, so the robot is free to use its joints continuously. --- @@ -160,11 +304,11 @@ $$R_{\text{step}} = R_{\text{alive}} + R_{\text{tracking}} + R_{\text{rotation}} In multi-robot vectorized training (`JackBotEnv`), individual robot failures are handled independently to allow maximum simulation efficiency: * **Individual Failure Masking:** A robot is flagged as failed (`failed_robots_mask[idx] = True`) if either condition is met: - * **Severe Tilt:** Base roll or pitch orientation exceeds $0.7\text{ radians}$ ($\approx 40^\circ$). - * **Base Collapse:** Base height drops below $0.05\text{ meters}$ above the ground plane. + * **Severe Tilt:** Base roll or pitch orientation exceeds $0.9\text{ radians}$. + * **Base Collapse:** Base height drops below a relative threshold derived from the measured settled standing height. * **Visual Failure Feedback:** When a robot fails during GUI execution (`--gui`), its 3D URDF mesh immediately updates to a **semi-transparent dark gray** visual state (`COLOR_FAILED = [0.3, 0.3, 0.3, 0.6]`) to distinguish it from active learners. -* **Environment Termination (`terminated=True`):** The entire environment step resets when any Robot fails. -* **Environment Truncation (`truncated=True`):** Occurs when the episode reaches the maximum allowable step budget (`max_episode_steps = 3000`). +* **Environment Termination (`terminated=True`):** The environment as a whole is marked as terminated when the failure ratio reaches the configured threshold. +* **Environment Truncation (`truncated=True`):** Occurs when the episode reaches the maximum allowable step budget (`max_episode_steps`). --- diff --git a/Robot.py b/Robot.py index d06f164..b0aa496 100644 --- a/Robot.py +++ b/Robot.py @@ -84,6 +84,7 @@ class Robot: urdf_path: str = cfg.urdf_path ): self.urdf_path = urdf_path + self.start_pose = start_pose # --- BACKEND FACTORY CREATION --- if isinstance(backend_type, BackendType): @@ -142,7 +143,8 @@ class Robot: self.backend.step_simulation() def reset_to_init(self) -> None: - self.current_rad = ri.init_deg.to_rad() + pose_deg = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg + self.current_rad = pose_deg.to_rad() self.current_pos = kin.ikpyForward(self.current_rad) self.set_joint_angles(self.current_rad) self.step_sim() diff --git a/ml/SimManager.py b/ml/SimManager.py index d712ce7..c13e52e 100644 --- a/ml/SimManager.py +++ b/ml/SimManager.py @@ -34,30 +34,29 @@ class SimManager: p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client) def load_scene( - self, urdf_path: str, num_robots: int, robot_spacing: float, base_pos_fn + self, urdf_path: str, robot_spacing: float, base_pos_fn ) -> Tuple[int, List[int], List[List[int]]]: - """Loads plane and hexapod bodies into the simulation scene.""" + """Loads the plane and a single hexapod body into the simulation scene.""" plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client) robots = [] robot_joint_indices = [] self.robot_joints.clear() - for r_id in range(num_robots): - base_pos = base_pos_fn(r_id, num_robots, robot_spacing) - robot = p.loadURDF( - urdf_path, - basePosition=base_pos, - useFixedBase=False, - physicsClientId=self.physics_client - ) - robots.append(robot) + base_pos = base_pos_fn(0, robot_spacing) + robot = p.loadURDF( + urdf_path, + basePosition=base_pos, + useFixedBase=False, + physicsClientId=self.physics_client + ) + robots.append(robot) - joint_indices = [ - i for i in range(p.getNumJoints(robot, physicsClientId=self.physics_client)) - if p.getJointInfo(robot, i, physicsClientId=self.physics_client)[2] == p.JOINT_REVOLUTE - ] - robot_joint_indices.append(joint_indices) - self.robot_joints[robot] = joint_indices + joint_indices = [ + i for i in range(p.getNumJoints(robot, physicsClientId=self.physics_client)) + if p.getJointInfo(robot, i, physicsClientId=self.physics_client)[2] == p.JOINT_REVOLUTE + ] + robot_joint_indices.append(joint_indices) + self.robot_joints[robot] = joint_indices return plane_id, robots, robot_joint_indices diff --git a/ml/env.py b/ml/env.py index 4c1e6f2..525669a 100644 --- a/ml/env.py +++ b/ml/env.py @@ -26,29 +26,25 @@ class JackBotEnv(gym.Env): self, use_gui: bool = True, random_command: bool = True, - num_robots: int = 1, robot_spacing: float = 0.5, start_pose: str = "init_deg", - max_episode_steps: int = 3000, + max_episode_steps: int = 5000, urdf_path: str = cfg.urdf_path, - termination_threshold: float = 0.001, # Termination ratio threshold ): super().__init__() self.use_gui = use_gui self.random_command = random_command - self.num_robots = num_robots self.robot_spacing = robot_spacing self.start_pose = start_pose self.max_episode_steps = max_episode_steps self.urdf_path = urdf_path - self.termination_threshold = termination_threshold self.episode_count = 0 self.step_count = 0 self.total_steps = 0 self.cumulative_reward = 0.0 - self.robot_rewards = [0.0 for _ in range(self.num_robots)] - self.failed_robots_mask = [False for _ in range(self.num_robots)] + self.robot_rewards = [0.0] + self.failed_robots_mask = [False] self._first_reset = True # Initialize Simulation Manager @@ -57,7 +53,7 @@ class JackBotEnv(gym.Env): # Connect physics world self.plane, self.pb_robots, self.robot_joint_indices = self.sim_manager.load_scene( - self.urdf_path, self.num_robots, self.robot_spacing, self._robot_base_position + self.urdf_path, self.robot_spacing, self._robot_base_position ) # Instantiate Robot Python wrappers per PyBullet body ID @@ -71,14 +67,27 @@ class JackBotEnv(gym.Env): ] # Action (18 joint deltas per robot) & Observation (18 angles + 4 command dims per robot) - action_dim = self.num_robots * 18 - obs_dim = self.num_robots * (18 + 4) + 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((self.num_robots, 4), dtype=np.float32) + self.commands = np.zeros((1, 4), dtype=np.float32) self.last_action = np.zeros(action_dim, dtype=np.float32) + self.target_height = 0.14 + self.collapse_height_fraction = 0.55 + self.tilt_failure_rad = 0.9 + self.start_positions = [[0.0, 0.0, 0.0]] + self.max_distance_from_start = [0.0] + self.max_survival_steps = 0 + self.curriculum_phase = 0 + self.curriculum_episode_limit = 150 + self.curriculum_stage_requirements = { + 1: {"survival_steps": 120, "distance": 0.04, "stability_roll_pitch": 0.35}, + 2: {"survival_steps": 250, "distance": 0.08, "stability_roll_pitch": 0.30}, + 3: {"survival_steps": 450, "distance": 0.12, "stability_roll_pitch": 0.25}, + } # Floating HUD & Leader Crown Visualizers self.hud = MetricsHUD(physics_client_id=self.sim_manager.physics_client) @@ -92,19 +101,32 @@ class JackBotEnv(gym.Env): for j in range(num_joints): p.changeVisualShape(pb_id, j, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client) - def _robot_base_position(self, robot_id: int, num_robots: int = 1, spacing: float = 0.5) -> list[float]: - cols = int(math.sqrt(num_robots - 1)) + 1 - row = robot_id // cols - col = robot_id % cols - x = (col - (cols - 1) / 2.0) * spacing - y = (row - (cols - 1) / 2.0) * spacing - return [x, y, 0.2] + def _robot_base_position(self, robot_id: int, spacing: float = 0.5) -> list[float]: + return [0.0, 0.0, 0.14] def sample_command(self) -> np.ndarray: - vx = np.random.uniform(-1.0, 1.0) - vy = np.random.uniform(-0.5, 0.5) - vz = 0.0 - omega = np.random.uniform(-1.0, 1.0) + """Curriculum command sampler with survival-gated difficulty progression.""" + phase = self.curriculum_phase + + if phase == 0: + # Phase 1: Forward Walking Focus + vx = np.random.uniform(0.5, 1.0) + vy = 0.0 + vz = 0.0 + omega = 0.0 + elif phase == 1: + # Phase 2: Forward/Backward + Turning + vx = np.random.uniform(-1.0, 1.0) + vy = 0.0 + vz = 0.0 + omega = np.random.uniform(-0.8, 0.8) + else: + # Phase 3: Full Omnidirectional Movement + vx = np.random.uniform(-1.0, 1.0) + vy = np.random.uniform(-0.5, 0.5) + vz = 0.0 + omega = np.random.uniform(-1.0, 1.0) + return np.array([vx, vy, vz, omega], dtype=np.float32) def reset(self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None): @@ -112,15 +134,13 @@ class JackBotEnv(gym.Env): self.episode_count += 1 self.step_count = 0 self.cumulative_reward = 0.0 - self.robot_rewards = [0.0 for _ in range(self.num_robots)] - self.failed_robots_mask = [False for _ in range(self.num_robots)] + self.robot_rewards = [0.0] + self.failed_robots_mask = [False] for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)): - # Get default spawn position - spawn_pos = self._robot_base_position(idx, self.num_robots, self.robot_spacing) + spawn_pos = self._robot_base_position(idx, self.robot_spacing) spawn_orn = [0, 0, 0, 1] - # Teleport base back to start p.resetBasePositionAndOrientation( pb_id, spawn_pos, spawn_orn, physicsClientId=self.sim_manager.physics_client ) @@ -129,23 +149,38 @@ class JackBotEnv(gym.Env): physicsClientId=self.sim_manager.physics_client ) - # Reset joint angles directly without reloading URDF robot_obj.reset_to_init() - # Restore original default visual color (clears failure dark gray) if self.use_gui: self._set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0]) self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32) + self.max_distance_from_start = [0.0] + self.max_survival_steps = 0 + + if self._first_reset: + self.curriculum_phase = 0 + self.curriculum_episode_limit = min(self.max_episode_steps, 150) + self._first_reset = False if self.random_command: - self.commands = np.stack([self.sample_command() for _ in range(self.num_robots)]) + self.commands = np.stack([self.sample_command() for _ in range(1)]) else: - self.commands = np.zeros((self.num_robots, 4), dtype=np.float32) + self.commands = np.zeros((1, 4), dtype=np.float32) - for _ in range(100): + for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)): + pos, _ = p.getBasePositionAndOrientation( + pb_id, physicsClientId=self.sim_manager.physics_client + ) + self.start_positions[idx] = [float(pos[0]), float(pos[1]), float(pos[2])] + + for _ in range(200): self.sim_manager.step() + self.target_height = self._measure_settled_height() + if self.target_height <= 0.0: + self.target_height = 0.14 + if self.use_gui: self.hud.reset() self.leader_crown.reset() @@ -153,6 +188,61 @@ class JackBotEnv(gym.Env): return self._get_obs(), {} + def get_robot_velocities(self) -> list: + """Exposes velocities for the SB3 metrics callback.""" + vels = [] + for pb_id in self.pb_robots: + lin_v, ang_v = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client) + vels.append((lin_v, ang_v)) + return vels + + def get_robot_distance_metrics(self) -> list: + """Exposes distance-from-start metrics for logging without affecting reward.""" + metrics = [] + for idx, pb_id in enumerate(self.pb_robots): + pos, _ = p.getBasePositionAndOrientation( + pb_id, physicsClientId=self.sim_manager.physics_client + ) + start_x, start_y, _ = self.start_positions[idx] + dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32))) + self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist) + metrics.append((dist, self.max_distance_from_start[idx])) + + self.max_survival_steps = max(self.max_survival_steps, self.step_count) + return metrics + + def get_current_robot_metrics(self) -> list: + """Returns current per-robot reward, distance, and velocity summaries for alive robots only.""" + metrics = [] + for idx, pb_id in enumerate(self.pb_robots): + if self.failed_robots_mask[idx]: + continue + + pos, _ = p.getBasePositionAndOrientation( + pb_id, physicsClientId=self.sim_manager.physics_client + ) + linear_vel, angular_vel = p.getBaseVelocity( + pb_id, physicsClientId=self.sim_manager.physics_client + ) + start_x, start_y, _ = self.start_positions[idx] + dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32))) + speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32))) + yaw_rate = float(abs(angular_vel[2])) + metrics.append({ + "reward": float(self.robot_rewards[idx]), + "distance_from_start": dist, + "speed": speed, + "yaw_rate": yaw_rate, + "alive": True, + "survival_steps": int(self.step_count), + }) + + return metrics + + def get_survival_steps(self) -> int: + """Returns the current survival length for the environment's current episode.""" + return int(self.step_count) + def _get_obs(self) -> np.ndarray: obs_list = [] for idx, (pb_id, joint_indices) in enumerate(zip(self.pb_robots, self.robot_joint_indices)): @@ -167,73 +257,191 @@ class JackBotEnv(gym.Env): return np.concatenate(obs_list).astype(np.float32) + def _measure_settled_height(self) -> float: + heights = [] + for pb_id in self.pb_robots: + pos, _ = p.getBasePositionAndOrientation( + pb_id, physicsClientId=self.sim_manager.physics_client + ) + heights.append(pos[2]) + return float(np.mean(heights)) if heights else 0.14 + def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]: self.step_count += 1 self.total_steps += 1 + previous_action = self.last_action.copy() self.last_action = action.copy() + self._update_curriculum() - action_per_robot = action.reshape(self.num_robots, 18) + # Resample commands every 300 steps during long episodes + if self.random_command and (self.step_count % 300 == 0 or self._curriculum_advanced): + self.commands = np.stack([self.sample_command() for _ in range(1)]) + + action_per_robot = action.reshape(1, 18) for robot, act in zip(self.robots, action_per_robot): robot.apply_rl_action(act) self.sim_manager.step() - # Update failure status & gray coloring self._update_robot_failures() + self.get_robot_distance_metrics() obs = self._get_obs() - reward, per_robot_step_rewards = self._compute_reward() + reward, per_robot_step_rewards = self._compute_reward(action, previous_action) self.cumulative_reward += reward for idx, r_step in enumerate(per_robot_step_rewards): self.robot_rewards[idx] += r_step terminated = self._is_done() - truncated = self.step_count >= self.max_episode_steps + truncated = self.step_count >= self.curriculum_episode_limit self._update_hud() self._update_leader_visuals() return obs, reward, terminated, truncated, {} - def _compute_reward(self) -> Tuple[float, list[float]]: - rewards = [] - for idx, pb_id in enumerate(self.pb_robots): - # Stop rewarding robots that have already collapsed or flipped - if self.failed_robots_mask[idx]: - rewards.append(-0.5) # Penalty per step while collapsed - continue + def _phase_progress_ready(self, phase: int) -> bool: + if phase not in self.curriculum_stage_requirements: + return False - linear_vel, angular_vel = p.getBaseVelocity( + if not self.pb_robots or not self.max_distance_from_start: + return False + + req = self.curriculum_stage_requirements[phase] + survival_ok = self.max_survival_steps >= req["survival_steps"] + distance_ok = self.max_distance_from_start[0] >= req["distance"] + + position, orientation = p.getBasePositionAndOrientation( + self.pb_robots[0], physicsClientId=self.sim_manager.physics_client + ) + roll, pitch, _ = p.getEulerFromQuaternion(orientation) + stability_ok = abs(roll) <= req["stability_roll_pitch"] and abs(pitch) <= req["stability_roll_pitch"] + height_ok = position[2] >= max(0.09, self.target_height * 0.85) + + return survival_ok and distance_ok and stability_ok and height_ok + + def _update_curriculum(self): + self._curriculum_advanced = False + + phase_labels = { + 0: "stand-and-forward", + 1: "turn-and-direction", + 2: "omni-direction", + 3: "full-command", + } + + if self.curriculum_phase < 1 and self._phase_progress_ready(1): + self.curriculum_phase = 1 + self.curriculum_episode_limit = min(self.max_episode_steps, 400) + self._curriculum_advanced = True + print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}") + elif self.curriculum_phase < 2 and self._phase_progress_ready(2): + self.curriculum_phase = 2 + self.curriculum_episode_limit = min(self.max_episode_steps, 700) + self._curriculum_advanced = True + print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}") + elif self.curriculum_phase < 3 and self._phase_progress_ready(3): + self.curriculum_phase = 3 + self.curriculum_episode_limit = min(self.max_episode_steps, self.max_episode_steps) + self._curriculum_advanced = True + print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}") + + def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> Tuple[float, list[float]]: + rewards = [] + current_actions = action.reshape(1, 18) + previous_actions = previous_action.reshape(1, 18) + + for idx, pb_id in enumerate(self.pb_robots): + pos, orientation = p.getBasePositionAndOrientation( pb_id, physicsClientId=self.sim_manager.physics_client ) - _, orientation = p.getBasePositionAndOrientation( + linear_vel, angular_vel = p.getBaseVelocity( pb_id, physicsClientId=self.sim_manager.physics_client ) roll, pitch, _ = p.getEulerFromQuaternion(orientation) command = self.commands[idx] - forward_reward = command[0] * linear_vel[0] + command[1] * linear_vel[1] - rotation_reward = command[3] * angular_vel[2] - stability_penalty = abs(roll) + abs(pitch) - - # Calculate per-robot action penalty - action_dim_per_robot = 18 - start_idx = idx * action_dim_per_robot - end_idx = start_idx + action_dim_per_robot - robot_action = self.last_action[start_idx:end_idx] - action_penalty = float(np.sum(np.square(robot_action))) * 0.01 + cmd_vx = command[0] + cmd_vy = command[1] + cmd_yaw = command[3] - r_step = 0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty + # ------------------------------------------------------------- + # 1. LINEAR VECTOR SPEED MAXIMIZATION (Magnitude + Direction) + # ------------------------------------------------------------- + cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32) + cmd_norm = np.linalg.norm(cmd_dir) + + if cmd_norm > 0.05: + # Normalize target direction vector + unit_cmd_dir = cmd_dir / cmd_norm + actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32) + + # Speed aligned with target direction (m/s) + aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir)) + + # Moderate movement reward to encourage directional motion + linear_speed_reward = 2.5 * aligned_speed + + # Penalize sideways drift (perpendicular velocity to commanded direction) + perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir + drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel)) + else: + # If no linear command given, penalize all horizontal movement + linear_speed_reward = 0.0 + drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2) + + # ------------------------------------------------------------- + # 2. TURNING SPEED MAXIMIZATION + # ------------------------------------------------------------- + actual_yaw_rate = angular_vel[2] # rad/s in PyBullet Z-axis + + if abs(cmd_yaw) > 0.05: + # Reward turning in the commanded direction, but less aggressively + turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw) + else: + # Penalize unwanted rotation when joystick turn is centered + turning_reward = -0.6 * (actual_yaw_rate ** 2) + + # ------------------------------------------------------------- + # 3. SMALL ALIVE BONUS + AGE-RAMPED STILLNESS PENALTY + # ------------------------------------------------------------- + alive_reward = 0.02 + + age_ratio = min(1.0, self.step_count / max(1, self.curriculum_episode_limit)) + still_penalty = 0.0 + if (cmd_norm > 0.1 or abs(cmd_yaw) > 0.1) and (abs(linear_vel[0]) < 0.02 and abs(actual_yaw_rate) < 0.05): + still_penalty = 0.35 + 0.85 * age_ratio + + # ------------------------------------------------------------- + # 4. POSTURE & STABILITY PENALTIES + # ------------------------------------------------------------- + height_penalty = 4.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0 + stability_penalty = 1.5 * (roll**2 + pitch**2) + + # Penalize only large control jumps; allow continuous low-amplitude motion + control_delta = np.abs(current_actions[idx] - previous_actions[idx]) + large_delta_mask = control_delta > 0.12 + large_delta_penalty = 0.002 * float(np.sum(np.square(control_delta[large_delta_mask]))) if np.any(large_delta_mask) else 0.0 + + # Combined Step Reward + r_step = ( + alive_reward + + linear_speed_reward + + turning_reward + - drift_penalty + - still_penalty + - height_penalty + - stability_penalty + - large_delta_penalty + ) rewards.append(r_step) return float(np.sum(rewards)), rewards def _is_done(self) -> bool: - """Returns True only when the percentage of failed robots exceeds the threshold.""" - failed_count = sum(self.failed_robots_mask) - failure_ratio = failed_count / self.num_robots - return failure_ratio >= self.termination_threshold + """Returns True when the robot has entered a failed state.""" + return bool(self.failed_robots_mask[0]) if self.failed_robots_mask else False def _update_hud(self): if not self.use_gui or not self.pb_robots: @@ -267,8 +475,7 @@ class JackBotEnv(gym.Env): ) def _update_leader_visuals(self): - """Positions floating crown above top robot without altering active materials.""" - if not self.use_gui or self.num_robots <= 1: + if not self.use_gui: return best_idx = int(np.argmax(self.robot_rewards)) @@ -279,23 +486,23 @@ class JackBotEnv(gym.Env): self.leader_crown.update(leader_pos) def _update_robot_failures(self): - """Checks failure condition for each robot and turns failed ones gray.""" + """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 # Already marked failed + continue position, orientation = p.getBasePositionAndOrientation( pb_id, physicsClientId=self.sim_manager.physics_client ) roll, pitch, _ = p.getEulerFromQuaternion(orientation) - is_tilted = abs(roll) > 0.7 or abs(pitch) > 0.7 - is_collapsed = position[2] < 0.05 + 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: - # Turn failed robot semi-transparent dark gray self._set_robot_color(pb_id, COLOR_FAILED) def close(self): diff --git a/ml/evaluate.py b/ml/evaluate.py index 8630ade..0392b17 100644 --- a/ml/evaluate.py +++ b/ml/evaluate.py @@ -13,10 +13,9 @@ def evaluate( model_path: str, episodes: int = 5, use_gui: bool = True, - num_robots: int = 16, # Default to 16 to match your trained (352,) observation space robot_spacing: float = 0.5, start_pose: str = "init_deg", - random_command: bool = True, + random_command: bool = False, save_json: Optional[str] = None, ) -> Dict[str, Any]: try: @@ -36,7 +35,6 @@ def evaluate( env = JackBotEnv( use_gui=use_gui, random_command=random_command, - num_robots=num_robots, robot_spacing=robot_spacing, start_pose=start_pose, ) @@ -72,7 +70,6 @@ def evaluate( metrics = { "model_path": str(model_path), "episodes_evaluated": episodes, - "num_robots": num_robots, "mean_reward": float(np.mean(episode_rewards)), "std_reward": float(np.std(episode_rewards)), "mean_episode_length": float(np.mean(episode_lengths)), diff --git a/ml/run_eval.py b/ml/run_eval.py index b5dc09f..b02ac82 100644 --- a/ml/run_eval.py +++ b/ml/run_eval.py @@ -20,9 +20,9 @@ def main(): parser.add_argument("--model", type=str, required=True, help="Path to the trained model file (.zip)") parser.add_argument("--episodes", type=int, default=3, help="Number of evaluation episodes to run.") parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during evaluation") - parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the evaluation environment") 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("--random-command", action="store_true", help="Randomize command samples during evaluation") parser.add_argument("--save-metrics", type=str, default=None, help="Optional path to save JSON metrics report") args = parser.parse_args() @@ -30,9 +30,9 @@ def main(): model_path=args.model, episodes=args.episodes, use_gui=args.gui, - num_robots=args.num_robots, robot_spacing=args.robot_spacing, start_pose=args.start_pose, + random_command=args.random_command, save_json=args.save_metrics, ) diff --git a/ml/run_train.py b/ml/run_train.py index 173de21..c325097 100644 --- a/ml/run_train.py +++ b/ml/run_train.py @@ -1,7 +1,7 @@ """Minimal training launcher for quick experiments. Usage: - python ml/run_train.py --timesteps 50000 --model ml/checkpoints/ppo_joint_command + 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. @@ -20,12 +20,12 @@ from ml.train import train def main(): parser = argparse.ArgumentParser() - parser.add_argument("--timesteps", type=int, default=50000, help="Total number of 'practice steps'.") + 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 with one robot each") + 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") args = parser.parse_args() diff --git a/ml/train.py b/ml/train.py index 9b98fb2..72a14ac 100644 --- a/ml/train.py +++ b/ml/train.py @@ -1,6 +1,8 @@ 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 @@ -37,6 +39,67 @@ class MilestoneCheckpointCallback(BaseCallback): 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: + """Required by SB3 BaseCallback; no-op here because the rollout summary is emitted at rollout end.""" + return True + + def _on_rollout_end(self) -> bool: + """Executed right before PPO outputs the log table to console.""" + 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) + + # Backward-compatible aliases so old dashboards keep a stable field name. + 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.") @@ -44,19 +107,18 @@ def parse_args(): 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("--use-gui", action="store_true", help="Enable PyBullet GUI during training") + 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(num_robots, robot_spacing, start_pose, use_gui, rank, seed=0): +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, # Only rank 0 gets GUI if requested random_command=True, - num_robots=num_robots, robot_spacing=robot_spacing, start_pose=start_pose, ) @@ -71,7 +133,6 @@ def train( seed: int = 0, device: str = "auto", use_gui: bool = False, - num_robots: int = 1, num_workers: int = 8, robot_spacing: float = 0.5, start_pose: str = "init_deg", @@ -84,13 +145,13 @@ def train( # Create multi-process vector environment if num_workers > 1: env_fns = [ - make_env(num_robots, robot_spacing, start_pose, use_gui, rank=i, seed=seed) + 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(num_robots, robot_spacing, start_pose, use_gui, rank=0, seed=seed) + make_env(robot_spacing, start_pose, use_gui, rank=0, seed=seed) ]) def resolve_device(requested_device: str) -> str: @@ -122,6 +183,17 @@ def train( env, verbose=1, seed=seed, + learning_rate=3.5e-4, # Cut LR in half (from 3e-4) to smooth out updates + n_steps=2048, # Larger rollout buffer per env for stable gradients + batch_size=128, # Larger minibatches reduce noise + n_epochs=10, # Number of epoch updates per rollout + gamma=0.99, # Discount factor + gae_lambda=0.95, # GAE smoothing + clip_range=0.2, # Standard PPO clipping + target_kl=0.03, # EARLY STOPPING: Halts policy update if KL > 0.015! + ent_coef=0.03, # Entropy coefficient to encourage exploration + vf_coef=0.5, + max_grad_norm=0.5, device=device, tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"), ) @@ -135,7 +207,9 @@ def train( step_interval=100_000 ) - model.learn(total_timesteps=total_timesteps, callback=milestone_cb) + metrics_callback = JackBotMetricsCallback() + + model.learn(total_timesteps=total_timesteps, callback=[milestone_cb, metrics_callback]) Path(model_path).parent.mkdir(parents=True, exist_ok=True) model.save(model_path) @@ -150,7 +224,6 @@ if __name__ == "__main__": seed=args.seed, device=args.device, use_gui=args.use_gui, - num_robots=args.num_robots, num_workers=args.num_workers, robot_spacing=args.robot_spacing, start_pose=args.start_pose,