reward ajustments

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

ml/SimManager.py

@@ -164,3 +164,29 @@ class SimManager:
         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,6 +3,7 @@ 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
@@ -18,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."""
 
@@ -44,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
@@ -87,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": 1000, "distance": 0.00, "stability_roll_pitch": 0.35},
+            CurriculumPhase.FORWARD: {
-            2: {"survival_steps": 350, "distance": 10.0, "stability_roll_pitch": 0.30},
+                "survival_steps": 400,          # Must survive ~8 seconds
-            3: {"survival_steps": 550, "distance": 15.00, "stability_roll_pitch": 0.25},
+                "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
@@ -107,24 +143,36 @@ class JackBotEnv(gym.Env):
         """Curriculum command sampler with survival-gated difficulty progression."""
         phase = self.curriculum_phase
 
-        if phase == 0:
+        if phase == CurriculumPhase.STAND_ONLY:
-            # Phase 1: Forward Walking Focus
+            # 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)
 
@@ -135,6 +183,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
 
         for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
             spawn_pos = self._robot_base_position(idx, self.robot_spacing)
@@ -152,8 +201,8 @@ class JackBotEnv(gym.Env):
         self.exploration_bonus_active = False
 
         if self._first_reset:
-            self.curriculum_phase = 0
+            self.curriculum_phase = CurriculumPhase.STAND_ONLY
-            self.curriculum_episode_limit = min(self.max_episode_steps, 150)
+            self.curriculum_episode_limit = min(self.max_episode_steps, 500)
             self._first_reset = False
 
         if self.random_command:
@@ -190,6 +239,9 @@ class JackBotEnv(gym.Env):
         metrics = []
         for idx, pb_id in enumerate(self.pb_robots):
             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)
@@ -251,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()
@@ -270,48 +329,72 @@ class JackBotEnv(gym.Env):
         self._update_leader_visuals()
         return obs, reward, terminated, truncated, {}
 
-    def _phase_progress_ready(self, phase: int) -> bool:
+    def _phase_progress_ready(self, next_phase: CurriculumPhase) -> bool:
-        if phase not in self.curriculum_stage_requirements:
+        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:
+        req = self.curriculum_stage_requirements[next_phase]
-            return False
         
-        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"]
 
+        # 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 = {
+        if self.curriculum_phase < CurriculumPhase.FORWARD and self._phase_progress_ready(CurriculumPhase.FORWARD):
-            0: "stand-and-forward",
+            self.curriculum_phase = CurriculumPhase.FORWARD
-            1: "turn-and-direction",
+            self.curriculum_episode_limit = min(self.max_episode_steps, 600)
-            2: "omni-direction",
+            self._curriculum_advanced = True
-            3: "full-command",
+            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):
+        elif self.curriculum_phase < CurriculumPhase.TURN_AND_DIRECTION and self._phase_progress_ready(CurriculumPhase.TURN_AND_DIRECTION):
-            self.curriculum_phase = 1
+            self.curriculum_phase = CurriculumPhase.TURN_AND_DIRECTION
-            self.curriculum_episode_limit = min(self.max_episode_steps, 400)
+            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}")
+            print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
-        elif self.curriculum_phase < 2 and self._phase_progress_ready(2):
+
-            self.curriculum_phase = 2
+        elif self.curriculum_phase < CurriculumPhase.OMNI_DIRECTION and self._phase_progress_ready(CurriculumPhase.OMNI_DIRECTION):
-            self.curriculum_episode_limit = min(self.max_episode_steps, 700)
+            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}")
+            print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
-        elif self.curriculum_phase < 3 and self._phase_progress_ready(3):
+
-            self.curriculum_phase = 3
+        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 = []
@@ -361,9 +444,19 @@ class JackBotEnv(gym.Env):
                 still_penalty = 0.35 + 0.85 * age_ratio
 
             # 4. POSTURE & STABILITY PENALTIES
-            height_penalty = 10.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)
-            print(-height_penalty)
+
             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

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

@@ -189,8 +189,8 @@ def train(
         verbose=1,
         seed=seed,
         learning_rate=1.5e-4,      # Cut LR in half (from 3e-4) to smooth out updates
-        n_steps=1024,              # Larger rollout buffer per env for stable gradients
+        n_steps=256,               # Larger rollout buffer per env for stable gradients
-        batch_size=128,            # Larger minibatches reduce noise
+        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