diff --git a/ml/SimManager.py b/ml/SimManager.py index 3a96701..3f01a0e 100644 --- a/ml/SimManager.py +++ b/ml/SimManager.py @@ -163,4 +163,30 @@ class SimManager: 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 \ No newline at end of file + 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, + ) \ No newline at end of file diff --git a/ml/env.py b/ml/env.py index 62078ba..2fa971d 100644 --- a/ml/env.py +++ b/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}, - 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 @@ -107,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) @@ -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_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: @@ -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: - 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[phase] + req = self.curriculum_stage_requirements[next_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 = { - 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 = [] @@ -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 diff --git a/ml/evaluate.py b/ml/evaluate.py index 0392b17..2303859 100644 --- a/ml/evaluate.py +++ b/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 diff --git a/ml/train.py b/ml/train.py index bbddf2e..87da066 100644 --- a/ml/train.py +++ b/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 - batch_size=128, # Larger minibatches reduce noise + 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