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