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)

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
 

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

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,