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

@@ -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.
+JackBot is a Python-based hexapod robot project that combines:
-* **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.
+* a robot control stack for a six-legged walking robot,
-* **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).
+* 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
+├── main.py                  # manual control and hardware streaming entry point
-├── Robot.py                 # Core Robot class, kinematics wrapper, and Backend protocols
+├── Robot.py                 # robot wrapper and backend abstraction
-├── config.py                # Global settings (backend selection, URDF path, communication specs)
+├── config.py                # global settings and communication configuration
-├── kinematics.py            # Forward and Inverse Kinematics (IKPy)
+├── kinematics.py            # inverse kinematics used to map motion commands to joints
-├── simulation.py            # Base PyBullet GUI wrapper for single-robot interactive simulation
+├── simulation.py            # PyBullet simulation shell for the robot
-├── robot_init.py            # Default stance angles and neutral leg positions
+├── robot_init.py            # neutral standing pose and initial joint values
-├── DataTypes.py             # Strongly typed arrays (PosArray, RadArray, DegArray) & structs
+├── DataTypes.py             # typed data structures for positions and joint angles
 │
-├── states/                  # Finite State Machine (FSM) gait states
+├── states/                  # gait/state-machine behavior logic
-│   ├── State.py             # Base State class
+│   ├── State.py
-│   ├── idle.py              # Neutral stance state
+│   ├── IdleState.py
-│   └── walking.py           # Inverse-kinematics tripod gait state
+│   ├── WalkingState.py
+│   └── WaveState.py
 │
-├── inputs/                  # Input providers
+├── inputs/                  # command sources for the robot
-│   ├── InputProvider.py     # Base input abstraction
+│   ├── InputProvider.py
-│   └── PygameController.py  # Asynchronous gamepad loop (process-isolated)
+│   ├── PygameController.py
+│   └── RandomeInputProvider.py
 │
-├── gui/                     # Control interface
+├── gui/                     # visual control frontend
-│   └── MainWindow.py        # Pygame / parameter GUI layout and command resolver
+│   └── MainWindow.py
 │
-├── EspCommunication.py      # WiFi socket sender for ESP32 hardware
+├── EspCommunication.py      # ESP32 communication layer
-├── ArduinoCommunication.py  # Serial communication wrapper for Arduino hardware
+├── ArduinoCommunication.py  # Arduino serial communication layer
-├── JackBotUrdf.urdf         # Kinematic 3D model definition (18 active joints)
+├── JackBotUrdf.urdf         # robot mesh and joint definition
 │
-└── ml/                      # Machine Learning Subsystem
+└── ml/                      # reinforcement learning subsystem
-    ├── env.py               # JackBotEnv (Gymnasium multi-robot vector environment)
+    ├── env.py               # Gymnasium environment used by PPO
-    ├── SimManager.py        # Physics server initialization and scene loading
+    ├── SimManager.py        # PyBullet scene setup and stepping
-    ├── MetricsOverlay.py    # PyBullet HUD (MetricsHUD) & leader crown tracking (LeaderCrown)
+    ├── MetricsOverlay.py    # HUD and visual overlays
-    ├── run_train.py         # PPO training execution script
+    ├── train.py             # PPO training entry point
-    └── run_eval.py          # Model evaluation script
+    ├── 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:
+Training uses the simulator as the environment, runs PPO updates, and periodically saves checkpoint models.
-```bash
-python ml/run_train.py --timesteps 500000 --model ml/checkpoints/ppo_joint_command --num-robots 16 --robot-spacing 0.75 --gui
-```
 
-**Training CLI Arguments:**
+#### Evaluation flow
-* `--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.
 
-#### 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.
+* **Alive Bonus ($R_{\text{alive}} \approx +0.02$):** A small 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$).
+* **Velocity Tracking ($R_{\text{tracking}}$):** Rewards linear movement aligned with the commanded direction vector.
-* **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}}$.
+* **Yaw Rotation Tracking ($R_{\text{rotation}}$):** Rewards turning in the commanded yaw direction.
-* **Body Stability Penalty ($P_{\text{stability}} = 0.2 \cdot (|\text{roll}| + |\text{pitch}|)$):** Penalizes tilting away from a level horizontal posture.
+* **Drift Penalty ($P_{\text{drift}}$):** Penalizes lateral motion that does not align with the commanded direction.
-* **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.
+* **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$).
+  * **Severe Tilt:** Base roll or pitch orientation exceeds $0.9\text{ radians}$.
-  * **Base Collapse:** Base height drops below $0.05\text{ meters}$ above the ground plane.
+  * **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 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 = 3000`).
+* **Environment Truncation (`truncated=True`):** Occurs when the episode reaches the maximum allowable step budget (`max_episode_steps`).
 
 ---
 

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()

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(0, robot_spacing)
-            base_pos = base_pos_fn(r_id, num_robots, robot_spacing)
+        robot = p.loadURDF(
-            robot = p.loadURDF(
+            urdf_path,
-                urdf_path, 
+            basePosition=base_pos,
-                basePosition=base_pos, 
+            useFixedBase=False,
-                useFixedBase=False, 
+            physicsClientId=self.physics_client
-                physicsClientId=self.physics_client
+        )
-            )
+        robots.append(robot)
-            robots.append(robot)
 
-            joint_indices = [
+        joint_indices = [
-                i for i in range(p.getNumJoints(robot, physicsClientId=self.physics_client))
+            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
+            if p.getJointInfo(robot, i, physicsClientId=self.physics_client)[2] == p.JOINT_REVOLUTE
-            ]
+        ]
-            robot_joint_indices.append(joint_indices)
+        robot_joint_indices.append(joint_indices)
-            self.robot_joints[robot] = joint_indices
+        self.robot_joints[robot] = joint_indices
 
         return plane_id, robots, robot_joint_indices
 

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.robot_rewards = [0.0]
-        self.failed_robots_mask = [False for _ in range(self.num_robots)]
+        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
+        action_dim = 18
-        obs_dim = self.num_robots * (18 + 4)
+        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]:
+    def _robot_base_position(self, robot_id: int, spacing: float = 0.5) -> list[float]:
-        cols = int(math.sqrt(num_robots - 1)) + 1
+        return [0.0, 0.0, 0.14]
-        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 sample_command(self) -> np.ndarray:
-        vx = np.random.uniform(-1.0, 1.0)
+        """Curriculum command sampler with survival-gated difficulty progression."""
-        vy = np.random.uniform(-0.5, 0.5)
+        phase = self.curriculum_phase
-        vz = 0.0
+
-        omega = np.random.uniform(-1.0, 1.0)
+        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.robot_rewards = [0.0]
-        self.failed_robots_mask = [False for _ in range(self.num_robots)]
+        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.robot_spacing)
-            spawn_pos = self._robot_base_position(idx, self.num_robots, 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]]:
+    def _phase_progress_ready(self, phase: int) -> bool:
-        rewards = []
+        if phase not in self.curriculum_stage_requirements:
-        for idx, pb_id in enumerate(self.pb_robots):
+            return False
-            # 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
 
-            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]
+            cmd_vx = command[0]
-            rotation_reward = command[3] * angular_vel[2]
+            cmd_vy = command[1]
-            stability_penalty = abs(roll) + abs(pitch)
+            cmd_yaw = command[3]
-            
-            # 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
 
-            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."""
+        """Returns True when the robot has entered a failed state."""
-        failed_count = sum(self.failed_robots_mask)
+        return bool(self.failed_robots_mask[0]) if self.failed_robots_mask else False
-        failure_ratio = failed_count / self.num_robots
-        return failure_ratio >= self.termination_threshold
 
     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:
-        if not self.use_gui or self.num_robots <= 1:
             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
+            collapse_threshold = max(0.06, self.collapse_height_fraction * self.target_height)
-            is_collapsed = position[2] < 0.05
+            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):

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)),

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,
     )
 

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()

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,