reward ajustments

robot starts balancing without changing phase and farms alive bonus
fix ->external force and bigger height penalty

Robot.py

@@ -172,9 +172,7 @@ class Robot:
     # --- RL METHODS ---
 
     def apply_rl_action(self, action: np.ndarray) -> None:
-        """
-        Applies continuous RL action deltas [-1, 1] to current joint angles.
-        """
+        """Applies continuous RL action deltas [-1, 1] to current joint angles."""
         action = np.asarray(action, dtype=np.float32)
         scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale
 
@@ -191,22 +189,21 @@ class Robot:
     def get_observation(self, command: Optional[np.ndarray] = None) -> np.ndarray:
         """
         Returns observation vector [18 joint angles] + [optional 4 command dimensions].
-        Queries PyBullet if backend is PyBulletBackend; otherwise falls back to internal state.
+        Queries SimManager helper if PyBulletBackend is used; falls back to internal state otherwise.
         """
-        if isinstance(self.backend, PyBulletBackend) and self.backend.sim and hasattr(self.backend.sim, 'physics_client'):
-            physics_client = self.backend.sim.physics_client
+        if isinstance(self.backend, PyBulletBackend) and self.backend.sim:
             body_id = self.backend.body_id if self.backend.body_id is not None else 0
-            
-            # Retrieve joint mapping from SimManager/Simulation if available
-            if hasattr(self.backend.sim, 'robot_joints') and body_id in self.backend.sim.robot_joints:
-                joint_indices = self.backend.sim.robot_joints[body_id]
-            else:
-                joint_indices = list(range(18))
-
+            if hasattr(self.backend.sim, 'get_robot_joint_angles'):
+                joint_angles = self.backend.sim.get_robot_joint_angles(body_id)
+            elif hasattr(self.backend.sim, 'physics_client'):
+                physics_client = self.backend.sim.physics_client
+                joint_indices = self.backend.sim.robot_joints.get(body_id, list(range(18))) if hasattr(self.backend.sim, 'robot_joints') else list(range(18))
                 joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=physics_client)
                 joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32)
             else:
                 joint_angles = self.current_rad.data.flatten().astype(np.float32)
+        else:
+            joint_angles = self.current_rad.data.flatten().astype(np.float32)
 
         if command is not None:
             cmd = np.asarray(command, dtype=np.float32).flatten()

ml/SimManager.py

@@ -1,9 +1,10 @@
 """
 ml/SimManager.py - PyBullet Simulation & Multi-Body Manager
 """
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Optional
 import pybullet as p
 import pybullet_data
+import numpy as np
 import DataTypes as dt
 
 
@@ -85,3 +86,107 @@ class SimManager:
         if self.physics_client is not None and p.isConnected(self.physics_client):
             p.disconnect(self.physics_client)
             self.physics_client = None
+
+    # --- ROBOT GETTERS AND SETTERS ---
+
+    def reset_robot_base(
+        self,
+        body_id: int,
+        position: List[float],
+        orientation: Optional[List[float]] = None,
+        linear_velocity: Optional[List[float]] = None,
+        angular_velocity: Optional[List[float]] = None
+    ) -> None:
+        """Resets a robot body's base position, orientation, and velocities."""
+        if orientation is None:
+            orientation = [0.0, 0.0, 0.0, 1.0]
+        if linear_velocity is None:
+            linear_velocity = [0.0, 0.0, 0.0]
+        if angular_velocity is None:
+            angular_velocity = [0.0, 0.0, 0.0]
+
+        p.resetBasePositionAndOrientation(
+            body_id, position, orientation, physicsClientId=self.physics_client
+        )
+        p.resetBaseVelocity(
+            body_id, linearVelocity=linear_velocity, angularVelocity=angular_velocity,
+            physicsClientId=self.physics_client
+        )
+
+    def get_robot_pose(self, body_id: int) -> Tuple[List[float], List[float]]:
+        """Returns base position (x, y, z) and orientation quaternion (x, y, z, w)."""
+        pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+        return list(pos), list(orn)
+
+    def get_robot_rpy(self, body_id: int) -> Tuple[float, float, float]:
+        """Returns roll, pitch, yaw angles in radians for the given robot body."""
+        _, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+        roll, pitch, yaw = p.getEulerFromQuaternion(orn)
+        return float(roll), float(pitch), float(yaw)
+
+    def get_robot_pose_and_rpy(self, body_id: int) -> Tuple[List[float], Tuple[float, float, float]]:
+        """Returns base position and (roll, pitch, yaw) tuple in radians."""
+        pos, orn = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+        roll, pitch, yaw = p.getEulerFromQuaternion(orn)
+        return list(pos), (float(roll), float(pitch), float(yaw))
+
+    def get_robot_velocity(self, body_id: int) -> Tuple[List[float], List[float]]:
+        """Returns linear velocity (vx, vy, vz) and angular velocity (wx, wy, wz)."""
+        lin_v, ang_v = p.getBaseVelocity(body_id, physicsClientId=self.physics_client)
+        return list(lin_v), list(ang_v)
+
+    def get_robot_joint_angles(self, body_id: int, joint_indices: Optional[List[int]] = None) -> np.ndarray:
+        """Returns joint angles as a 1D numpy array float32 for specified or registered joint indices."""
+        if joint_indices is None:
+            joint_indices = self.robot_joints.get(body_id, list(range(18)))
+        joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=self.physics_client)
+        return np.array([state[0] for state in joint_states], dtype=np.float32)
+
+    def set_robot_color(self, body_id: int, rgba: List[float]) -> None:
+        """Changes visual color RGBA of base link and all joints of the specified robot body."""
+        num_joints = p.getNumJoints(body_id, physicsClientId=self.physics_client)
+        p.changeVisualShape(body_id, -1, rgbaColor=rgba, physicsClientId=self.physics_client)
+        for j in range(num_joints):
+            p.changeVisualShape(body_id, j, rgbaColor=rgba, physicsClientId=self.physics_client)
+
+    def measure_robot_heights(self, robot_ids: List[int]) -> List[float]:
+        """Gets current Z height for all specified robot body IDs."""
+        heights = []
+        for body_id in robot_ids:
+            pos, _ = p.getBasePositionAndOrientation(body_id, physicsClientId=self.physics_client)
+            heights.append(pos[2])
+        return heights
+
+    def settle_and_measure_height(self, robot_ids: List[int], steps: int = 200, fallback_height: float = 0.14) -> float:
+        """Steps simulation for designated steps so robot settles, then calculates target standing height."""
+        for _ in range(steps):
+            self.step()
+        heights = self.measure_robot_heights(robot_ids)
+        mean_height = float(np.mean(heights)) if heights else fallback_height
+        return mean_height if mean_height > 0.0 else fallback_height
+
+    def apply_external_force(
+        self,
+        body_id: int,
+        force: list[float] | np.ndarray,
+        link_index: int = -1,
+        position: list[float] | np.ndarray = (0.0, 0.0, 0.0),
+        frame: int = p.WORLD_FRAME,
+    ):
+        """
+        Applies a 3D force vector (in Newtons) to a robot link.
+        
+        :param body_id: PyBullet body ID.
+        :param force: [fx, fy, fz] force vector in Newtons.
+        :param link_index: Target link index (-1 refers to the base/torso).
+        :param position: Offset [x, y, z] relative to link center where force is applied.
+        :param frame: p.WORLD_FRAME (global axes) or p.LINK_FRAME (robot's body axes).
+        """
+        p.applyExternalForce(
+            objectUniqueId=body_id,
+            linkIndex=link_index,
+            forceObj=list(force),
+            posObj=list(position),
+            flags=frame,
+            physicsClientId=self.physics_client,
+        )

ml/env.py

@@ -3,12 +3,12 @@ ml/env.py - Gymnasium Environment for JackBot Hexapod RL Training
 """
 import time
 import math
+from enum import IntEnum
 from typing import Optional, Tuple, Dict, Any, List
 
 import gymnasium as gym
 from gymnasium import spaces
 import numpy as np
-import pybullet as p
 
 from config import cfg
 from Robot import Robot, PyBulletBackend
@@ -19,6 +19,14 @@ from ml.MetricsOverlay import MetricsHUD, LeaderCrown
 COLOR_FAILED = [0.3, 0.3, 0.3, 0.6]   # Collapsed / Tilted Robot (Dark Semi-Transparent Gray)
 
 
+class CurriculumPhase(IntEnum):
+    STAND_ONLY = 0
+    FORWARD = 1
+    TURN_AND_DIRECTION = 2
+    OMNI_DIRECTION = 3
+    FULL_COMMAND = 4
+
+
 class JackBotEnv(gym.Env):
     """Gymnasium environment wrapping JackBot hexapods with floating text & crown overlay."""
 
@@ -45,6 +53,7 @@ class JackBotEnv(gym.Env):
         self.cumulative_reward = 0.0
         self.robot_rewards = [0.0]
         self.failed_robots_mask = [False]
+        self.episode_height_sum = 0.0
         self._first_reset = True
 
         # Initialize Simulation Manager
@@ -79,9 +88,7 @@ class JackBotEnv(gym.Env):
         self.collapse_height_fraction = 0.55
         self.tilt_failure_rad = 0.9
 
-        # Small random exploration pulse to break local optima.
-        # The bonus is only granted on a rare step and only when the robot
-        # is still stable enough that a bold move is likely to remain safe.
+        # Small random exploration pulse settings
         self.exploration_bonus_prob = 0.03
         self.exploration_bonus_interval = 120
         self.exploration_bonus_scale = 0.08
@@ -90,12 +97,38 @@ class JackBotEnv(gym.Env):
         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
+        
+        # Curriculum Initialization via Enum
+        self.curriculum_phase = CurriculumPhase.STAND_ONLY
+        self.curriculum_episode_limit = 500  # 500 steps limit gives headroom for 400-step requirement
+        
+        # Gates required to unlock each target phase
         self.curriculum_stage_requirements = {
-            1: {"survival_steps": 200, "distance": 0.00, "stability_roll_pitch": 0.35},
-            2: {"survival_steps": 350, "distance": 10.0, "stability_roll_pitch": 0.30},
-            3: {"survival_steps": 550, "distance": 15.00, "stability_roll_pitch": 0.25},
+            CurriculumPhase.FORWARD: {
+                "survival_steps": 400,          # Must survive ~8 seconds
+                "max_displacement": 0.25,       # Must remain within 0.25m radius
+                "min_avg_height_ratio": 0.90,   # Average height >= 90% of target
+                "stability_roll_pitch": 0.18,   # Max ~10 degrees tilt
+            },
+            CurriculumPhase.TURN_AND_DIRECTION: {
+                "survival_steps": 500,
+                "min_forward_distance": 2.5,    # Must walk +2.5m forward (+X)
+                "max_lateral_drift": 0.8,       # Max 0.8m drift on Y axis
+                "min_avg_height_ratio": 0.85,
+                "stability_roll_pitch": 0.25,
+            },
+            CurriculumPhase.OMNI_DIRECTION: {
+                "survival_steps": 600,
+                "min_distance": 5.0,
+                "min_avg_height_ratio": 0.85,
+                "stability_roll_pitch": 0.25,
+            },
+            CurriculumPhase.FULL_COMMAND: {
+                "survival_steps": 750,
+                "min_distance": 8.0,
+                "min_avg_height_ratio": 0.85,
+                "stability_roll_pitch": 0.20,
+            },
         }
 
         # Floating HUD & Leader Crown Visualizers
@@ -103,13 +136,6 @@ class JackBotEnv(gym.Env):
         self.leader_crown = LeaderCrown(physics_client_id=self.sim_manager.physics_client)
         self.last_time = time.time()
 
-    def _set_robot_color(self, pb_id: int, rgba: List[float]):
-        """Helper to change the visual color of a robot body and all its links."""
-        num_joints = p.getNumJoints(pb_id, physicsClientId=self.sim_manager.physics_client)
-        p.changeVisualShape(pb_id, -1, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client)
-        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, spacing: float = 0.5) -> list[float]:
         return [0.0, 0.0, 0.14]
 
@@ -117,24 +143,36 @@ class JackBotEnv(gym.Env):
         """Curriculum command sampler with survival-gated difficulty progression."""
         phase = self.curriculum_phase
 
-        if phase == 0:
-            # Phase 1: Forward Walking Focus
+        if phase == CurriculumPhase.STAND_ONLY:
+            # Phase 0: Pure Standing (Zero commands)
+            vx = 0.0
+            vy = 0.0
+            vz = 0.0
+            omega = 0.0
+        elif phase == CurriculumPhase.FORWARD:
+            # Phase 1: Straight Forward Walking
             vx = np.random.uniform(0.5, 1.0)
             vy = 0.0
             vz = 0.0
             omega = 0.0
-        elif phase == 1:
+        elif phase == CurriculumPhase.TURN_AND_DIRECTION:
             # 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:
+        elif phase == CurriculumPhase.OMNI_DIRECTION:
             # 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)
+        else:
+            # Phase 4: Full Unconstrained Commands
+            vx = np.random.uniform(-1.0, 1.0)
+            vy = np.random.uniform(-1.0, 1.0)
+            vz = 0.0
+            omega = np.random.uniform(-1.0, 1.0)
 
         return np.array([vx, vy, vz, omega], dtype=np.float32)
 
@@ -145,23 +183,17 @@ class JackBotEnv(gym.Env):
         self.cumulative_reward = 0.0
         self.robot_rewards = [0.0]
         self.failed_robots_mask = [False]
+        self.episode_height_sum = 0.0
 
         for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
             spawn_pos = self._robot_base_position(idx, self.robot_spacing)
-            spawn_orn = [0, 0, 0, 1]
-
-            p.resetBasePositionAndOrientation(
-                pb_id, spawn_pos, spawn_orn, physicsClientId=self.sim_manager.physics_client
-            )
-            p.resetBaseVelocity(
-                pb_id, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 0],
-                physicsClientId=self.sim_manager.physics_client
-            )
+            spawn_orn = [0.0, 0.0, 0.0, 1.0]
 
+            self.sim_manager.reset_robot_base(pb_id, spawn_pos, spawn_orn)
             robot_obj.reset_to_init()
 
             if self.use_gui:
-                self._set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
+                self.sim_manager.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]
@@ -169,8 +201,8 @@ class JackBotEnv(gym.Env):
         self.exploration_bonus_active = False
 
         if self._first_reset:
-            self.curriculum_phase = 0
-            self.curriculum_episode_limit = min(self.max_episode_steps, 150)
+            self.curriculum_phase = CurriculumPhase.STAND_ONLY
+            self.curriculum_episode_limit = min(self.max_episode_steps, 500)
             self._first_reset = False
 
         if self.random_command:
@@ -178,18 +210,14 @@ class JackBotEnv(gym.Env):
         else:
             self.commands = np.zeros((1, 4), dtype=np.float32)
 
-        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
-            )
+        for idx, pb_id in enumerate(self.pb_robots):
+            pos, _ = self.sim_manager.get_robot_pose(pb_id)
             self.start_positions[idx] = [float(pos[0]), float(pos[1]), float(pos[2])]
 
-        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
+        # Settle for 200 steps after dropping in a standing position, then measure target height
+        self.target_height = self.sim_manager.settle_and_measure_height(
+            self.pb_robots, steps=200, fallback_height=0.14
+        )
 
         if self.use_gui:
             self.hud.reset()
@@ -202,7 +230,7 @@ class JackBotEnv(gym.Env):
         """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)
+            lin_v, ang_v = self.sim_manager.get_robot_velocity(pb_id)
             vels.append((lin_v, ang_v))
         return vels
 
@@ -210,9 +238,10 @@ class JackBotEnv(gym.Env):
         """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
-                )
+            pos, _ = self.sim_manager.get_robot_pose(pb_id)
+            if idx == 0:
+                self.episode_height_sum += float(pos[2])
+
             start_x, start_y, _ = self.start_positions[idx]
             dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
             self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist)
@@ -228,12 +257,8 @@ class JackBotEnv(gym.Env):
             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
-                )
+            pos, _ = self.sim_manager.get_robot_pose(pb_id)
+            linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id)
             start_x, start_y, _ = self.start_positions[idx]
             dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
             speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)))
@@ -256,26 +281,12 @@ class JackBotEnv(gym.Env):
     def _get_obs(self) -> np.ndarray:
         obs_list = []
         for idx, (pb_id, joint_indices) in enumerate(zip(self.pb_robots, self.robot_joint_indices)):
-            joint_states = p.getJointStates(
-                pb_id, 
-                joint_indices, 
-                physicsClientId=self.sim_manager.physics_client
-            )
-            joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32)
+            joint_angles = self.sim_manager.get_robot_joint_angles(pb_id, joint_indices)
             robot_obs = np.concatenate([joint_angles, self.commands[idx]])
             obs_list.append(robot_obs)
 
         return np.concatenate(obs_list).astype(np.float32)
 
-    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
@@ -292,6 +303,13 @@ class JackBotEnv(gym.Env):
         for robot, act in zip(self.robots, action_per_robot):
             robot.apply_rl_action(act)
 
+        if self.step_count % 60 == 0:
+            random_force = np.random.uniform(-2.0, 2.0, size=2)  # X and Y push (Newtons)
+            self.sim_manager.apply_external_force(
+                body_id=self.pb_robots[0], 
+                force=[random_force[0], random_force[1], 0.0]
+            )
+
         self.sim_manager.step()
 
         self._update_robot_failures()
@@ -311,51 +329,72 @@ class JackBotEnv(gym.Env):
         self._update_leader_visuals()
         return obs, reward, terminated, truncated, {}
 
-    def _phase_progress_ready(self, phase: int) -> bool:
-        if phase not in self.curriculum_stage_requirements:
+    def _phase_progress_ready(self, next_phase: CurriculumPhase) -> bool:
+        if next_phase not in self.curriculum_stage_requirements or not self.pb_robots:
             return False
 
-        if not self.pb_robots or not self.max_distance_from_start:
-            return False
+        req = self.curriculum_stage_requirements[next_phase]
         
-        req = self.curriculum_stage_requirements[phase]
+        # 1. Survival Step Check
         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)
+        # 2. Body Posture & Tilt Check
+        position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robots[0])
         stability_ok = abs(roll) <= req["stability_roll_pitch"] and abs(pitch) <= req["stability_roll_pitch"]
-        height_ok = position[2] >= max(0.09, self.target_height * 0.85)
 
-        return survival_ok and distance_ok and stability_ok and height_ok
+        # 3. Average Height Ratio Check Across the Episode
+        avg_height = self.episode_height_sum / max(1, self.step_count)
+        required_min_avg_height = self.target_height * req.get("min_avg_height_ratio", 0.85)
+        height_ok = avg_height >= required_min_avg_height
+
+        # 4. Distance and Displacement Checks
+        start_x, start_y, _ = self.start_positions[0]
+        dx = position[0] - start_x
+        dy = position[1] - start_y
+        dist_2d = math.hypot(dx, dy)
+
+        distance_ok = True
+        if "min_forward_distance" in req:
+            distance_ok = dx >= req["min_forward_distance"]
+        elif "min_distance" in req:
+            distance_ok = dist_2d >= req["min_distance"]
+
+        displacement_ok = True
+        if "max_displacement" in req:
+            displacement_ok = dist_2d <= req["max_displacement"]
+
+        drift_ok = True
+        if "max_lateral_drift" in req:
+            drift_ok = abs(dy) <= req["max_lateral_drift"]
+
+        return survival_ok and height_ok and stability_ok and displacement_ok and distance_ok and drift_ok
 
     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 < CurriculumPhase.FORWARD and self._phase_progress_ready(CurriculumPhase.FORWARD):
+            self.curriculum_phase = CurriculumPhase.FORWARD
+            self.curriculum_episode_limit = min(self.max_episode_steps, 600)
+            self._curriculum_advanced = True
+            print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
 
-        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)
+        elif self.curriculum_phase < CurriculumPhase.TURN_AND_DIRECTION and self._phase_progress_ready(CurriculumPhase.TURN_AND_DIRECTION):
+            self.curriculum_phase = CurriculumPhase.TURN_AND_DIRECTION
+            self.curriculum_episode_limit = min(self.max_episode_steps, 800)
             self._curriculum_advanced = True
-            print(f"[Curriculum] phase {self.curriculum_phase} ({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)
+            print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
+
+        elif self.curriculum_phase < CurriculumPhase.OMNI_DIRECTION and self._phase_progress_ready(CurriculumPhase.OMNI_DIRECTION):
+            self.curriculum_phase = CurriculumPhase.OMNI_DIRECTION
+            self.curriculum_episode_limit = min(self.max_episode_steps, 1000)
             self._curriculum_advanced = True
-            print(f"[Curriculum] phase {self.curriculum_phase} ({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
+            print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
+
+        elif self.curriculum_phase < CurriculumPhase.FULL_COMMAND and self._phase_progress_ready(CurriculumPhase.FULL_COMMAND):
+            self.curriculum_phase = CurriculumPhase.FULL_COMMAND
             self.curriculum_episode_limit = min(self.max_episode_steps, self.max_episode_steps)
             self._curriculum_advanced = True
-            print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
+            print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
 
     def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> Tuple[float, list[float]]:
         rewards = []
@@ -363,78 +402,65 @@ class JackBotEnv(gym.Env):
         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
-            )
-            linear_vel, angular_vel = p.getBaseVelocity(
-                pb_id, physicsClientId=self.sim_manager.physics_client
-            )
-            roll, pitch, _ = p.getEulerFromQuaternion(orientation)
+            pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id)
+            linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id)
             command = self.commands[idx]
 
             cmd_vx = command[0]
             cmd_vy = command[1]
             cmd_yaw = command[3]
 
-            # -------------------------------------------------------------
-            # 1. LINEAR VECTOR SPEED MAXIMIZATION (Magnitude + Direction)
-            # -------------------------------------------------------------
+            # 1. LINEAR VECTOR SPEED MAXIMIZATION
             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
+            actual_yaw_rate = angular_vel[2]
 
             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
+            alive_reward = 0.5
 
             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 = 6.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0
+            height_penalty : float
+            min_valid_h = self.target_height * 0.90  # 90% threshold
+
+            if pos[2] < min_valid_h:
+                # Normalized drop below the 90% mark
+                drop = (min_valid_h - pos[2]) / self.target_height
+                # Linear + quadratic penalty that rapidly outweighs the +0.50 alive bonus
+                height_penalty = 4.0 * drop + 20.0 * (drop ** 2)
+            else:
+                # Zero penalty inside the valid 90% - 110% zone!
+                height_penalty = 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
@@ -465,8 +491,7 @@ class JackBotEnv(gym.Env):
         rolls = []
         pitches = []
         for pb_id in self.pb_robots:
-            pos, orient = p.getBasePositionAndOrientation(pb_id, physicsClientId=self.sim_manager.physics_client)
-            roll, pitch, _ = p.getEulerFromQuaternion(orient)
+            pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id)
             heights.append(pos[2])
             rolls.append(math.degrees(roll))
             pitches.append(math.degrees(pitch))
@@ -490,9 +515,7 @@ class JackBotEnv(gym.Env):
 
         best_idx = int(np.argmax(self.robot_rewards))
         leader_pb_id = self.pb_robots[best_idx]
-        leader_pos, _ = p.getBasePositionAndOrientation(
-            leader_pb_id, physicsClientId=self.sim_manager.physics_client
-        )
+        leader_pos, _ = self.sim_manager.get_robot_pose(leader_pb_id)
         self.leader_crown.update(leader_pos)
 
     def _update_robot_failures(self):
@@ -501,10 +524,7 @@ class JackBotEnv(gym.Env):
             if self.failed_robots_mask[idx]:
                 continue
 
-            position, orientation = p.getBasePositionAndOrientation(
-                pb_id, physicsClientId=self.sim_manager.physics_client
-            )
-            roll, pitch, _ = p.getEulerFromQuaternion(orientation)
+            position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(pb_id)
 
             collapse_threshold = max(0.06, self.collapse_height_fraction * self.target_height)
             is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad
@@ -513,7 +533,7 @@ class JackBotEnv(gym.Env):
             if is_tilted or is_collapsed:
                 self.failed_robots_mask[idx] = True
                 if self.use_gui:
-                    self._set_robot_color(pb_id, COLOR_FAILED)
+                    self.sim_manager.set_robot_color(pb_id, COLOR_FAILED)
 
     def close(self):
         self.sim_manager.disconnect()

ml/evaluate.py

@@ -44,17 +44,16 @@ def evaluate(
 
     for ep in range(episodes):
         obs, _ = env.reset()
-        done = False
+        terminated = False
         total_reward = 0.0
         steps = 0
 
         print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} ---")
 
-        while not done:
+        while not terminated:
             action, _ = model.predict(obs, deterministic=True)
             obs, reward, terminated, truncated, _ = env.step(action)
             
-            done = terminated or truncated
             total_reward += float(reward)
             steps += 1
 

ml/train.py

@@ -178,14 +178,19 @@ def train(
 
     device = resolve_device(device)
 
+    policy_kwargs = dict(
+    log_std_init=-1.5,  # Sets initial std ~ 0.22 instead of 1.0
+    net_arch=dict(pi=[256, 256], vf=[256, 256])
+)
+
     model = PPO(
         "MlpPolicy",
         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
+        learning_rate=1.5e-4,      # Cut LR in half (from 3e-4) to smooth out updates
+        n_steps=256,               # Larger rollout buffer per env for stable gradients
+        batch_size=256,            # Larger minibatches reduce noise
         n_epochs=10,               # Number of epoch updates per rollout
         gamma=0.99,                # Discount factor
         gae_lambda=0.95,           # GAE smoothing
@@ -195,6 +200,7 @@ def train(
         vf_coef=0.5,
         max_grad_norm=0.5,
         device=device,
+        policy_kwargs=policy_kwargs,
         tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"),
     )