diff --git a/ml/env.py b/ml/env.py index 14053da..a1713e2 100644 --- a/ml/env.py +++ b/ml/env.py @@ -46,12 +46,14 @@ class JackBotEnv(gym.Env): self.start_pose = start_pose self.max_episode_steps = max_episode_steps self.urdf_path = urdf_path + self.max_robot_speed = 0.8 self.episode_count = 0 self.step_count = 0 self.total_steps = 0 self.cumulative_reward = 0.0 self.robot_rewards = [0.0] + self.consecutive_still_steps = [0] self.failed_robots_mask = [False] self.episode_height_sum = 0.0 self._first_reset = True @@ -106,7 +108,6 @@ class JackBotEnv(gym.Env): self.curriculum_stage_requirements = { 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 }, @@ -150,17 +151,24 @@ class JackBotEnv(gym.Env): 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 + # Phase 1: Straight Forward Walking + Standing Checks + if np.random.random() < 0.30: + # 30% chance: Stand Still Command (vx = 0) + vx = 0.0 + else: + # 70% chance: Forward Walk Command + vx = np.random.uniform(0.15, 0.50) + + vy, vz, omega = 0.0, 0.0, 0.0 + 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) + # Phase 2: Forward/Backward + Turning + Standing + if np.random.random() < 0.20: + vx, omega = 0.0, 0.0 + else: + vx = np.random.uniform(-1.0, 1.0) + omega = np.random.uniform(-0.8, 0.8) + vy, vz = 0.0, 0.0 elif phase == CurriculumPhase.OMNI_DIRECTION: # Phase 3: Full Omnidirectional Movement vx = np.random.uniform(-1.0, 1.0) @@ -195,6 +203,7 @@ class JackBotEnv(gym.Env): if self.use_gui: self.sim_manager.set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0]) + self.consecutive_still_steps = [0 for _ in range(len(self.pb_robots))] 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 @@ -406,56 +415,78 @@ class JackBotEnv(gym.Env): 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_vx = command[0] # Normalized -1.0 to 1.0 + cmd_vy = command[1] # Normalized -1.0 to 1.0 cmd_yaw = command[3] - # 1. LINEAR VECTOR SPEED MAXIMIZATION cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32) cmd_norm = np.linalg.norm(cmd_dir) + actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32) + + # ----------------------------------------------------------------- + # 1. NORMALIZED SPEED TRACKING (1.0 = v_max) + # ----------------------------------------------------------------- if cmd_norm > 0.05: unit_cmd_dir = cmd_dir / cmd_norm - actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32) - aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir)) - linear_speed_reward = 2.5 * aligned_speed + target_speed = cmd_norm * self.max_robot_speed # e.g., 0.35 * 0.8 = 0.28 m/s + + if aligned_speed <= 0.0: + # Zero or backward movement gets ZERO reward + linear_speed_reward = 0.0 + elif aligned_speed <= target_speed: + # Smoothly scales from 0.0 to +5.0 as speed approaches target + linear_speed_reward = 5.0 * (aligned_speed / target_speed) + else: + # Gently penalize overspeeding beyond target + overspeed_ratio = (aligned_speed - target_speed) / target_speed + linear_speed_reward = max(0.0, 5.0 - 2.5 * overspeed_ratio) + + # Drift penalty for sideways sliding perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel)) else: + aligned_speed = 0.0 linear_speed_reward = 0.0 drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2) - # 2. TURNING SPEED MAXIMIZATION + # ----------------------------------------------------------------- + # 2. TIME-ACCUMULATING STILLNESS PENALTY + # ----------------------------------------------------------------- actual_yaw_rate = angular_vel[2] + is_moving = abs(aligned_speed) > 0.03 or abs(actual_yaw_rate) > 0.08 + if (cmd_norm > 0.05 or abs(cmd_yaw) > 0.05) and not is_moving: + # Increment counter every step the robot stalls under command + self.consecutive_still_steps[idx] += 1 + else: + # Reset counter as soon as robot makes a valid move! + self.consecutive_still_steps[idx] = 0 + + # Penalty grows by 0.05 every step frozen, capped at -3.5 per step + still_penalty = min(3.5, 0.05 * self.consecutive_still_steps[idx]) + + # ----------------------------------------------------------------- + # 3. TURNING & POSTURE PENALTIES + # ----------------------------------------------------------------- if abs(cmd_yaw) > 0.05: turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw) else: turning_reward = -0.6 * (actual_yaw_rate ** 2) - # 3. SMALL ALIVE BONUS + AGE-RAMPED STILLNESS PENALTY - 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 : float - min_valid_h = self.target_height * 0.90 # 90% threshold + alive_reward = 0.20 + # Height drop penalty + min_valid_h = self.target_height * 0.90 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) + + stability_penalty = 8 * (roll**2 + pitch**2) control_delta = np.abs(current_actions[idx] - previous_actions[idx]) large_delta_mask = control_delta > 0.12 @@ -474,7 +505,7 @@ class JackBotEnv(gym.Env): rewards.append(r_step) return float(np.sum(rewards)), rewards - + def _is_done(self) -> bool: """Returns True when the robot has entered a failed state.""" return bool(self.failed_robots_mask[0]) if self.failed_robots_mask else False diff --git a/ml/train.py b/ml/train.py index 87da066..7361589 100644 --- a/ml/train.py +++ b/ml/train.py @@ -12,14 +12,25 @@ class MilestoneCheckpointCallback(BaseCallback): """ Saves a model checkpoint the FIRST time total_timesteps crosses every multiple of step_interval (e.g., 100,000). + Saves inside the matching PPO_X subfolder as created by TensorBoard. """ def __init__(self, save_path: str, name_prefix: str = "ppo_jackbot", step_interval: int = 100_000, verbose: int = 1): super().__init__(verbose) - self.save_path = save_path + self.base_save_path = save_path + self.run_save_path = save_path self.name_prefix = name_prefix self.step_interval = step_interval self.last_milestone = 0 - os.makedirs(self.save_path, exist_ok=True) + + def _on_training_start(self) -> None: + """Executed right before training loop starts. Resolves TensorBoard's run folder name (e.g. PPO_1).""" + if self.logger and self.logger.dir: + run_folder_name = Path(self.logger.dir).name # Extracts "PPO_1", "PPO_2", etc. + self.run_save_path = os.path.join(self.base_save_path, run_folder_name) + else: + self.run_save_path = self.base_save_path + + os.makedirs(self.run_save_path, exist_ok=True) def _on_step(self) -> bool: current_milestone = self.num_timesteps // self.step_interval @@ -29,7 +40,7 @@ class MilestoneCheckpointCallback(BaseCallback): milestone_step = current_milestone * self.step_interval save_file = os.path.join( - self.save_path, + self.run_save_path, f"{self.name_prefix}_{milestone_step}_steps.zip" ) self.model.save(save_file) @@ -39,6 +50,7 @@ class MilestoneCheckpointCallback(BaseCallback): return True + class JackBotMetricsCallback(BaseCallback): """ Tracks the best current alive-robot performance for the most recent rollout, @@ -53,11 +65,9 @@ class JackBotMetricsCallback(BaseCallback): 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") @@ -90,7 +100,6 @@ class JackBotMetricsCallback(BaseCallback): 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)) @@ -101,6 +110,7 @@ class JackBotMetricsCallback(BaseCallback): return True + def parse_args(): parser = argparse.ArgumentParser(description="Train a joint-command policy for JackBot.") parser.add_argument("--timesteps", type=int, default=500_000, help="Total training timesteps") @@ -117,7 +127,7 @@ def parse_args(): 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 + use_gui=use_gui if rank == 0 else False, random_command=True, robot_spacing=robot_spacing, start_pose=start_pose, @@ -142,7 +152,6 @@ def train( except ImportError as exc: raise ImportError("stable-baselines3 is required. Install with: pip install stable-baselines3") from exc - # Create multi-process vector environment if num_workers > 1: env_fns = [ make_env(robot_spacing, start_pose, use_gui, rank=i, seed=seed) @@ -179,24 +188,24 @@ 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]) -) + log_std_init=-1.5, + net_arch=dict(pi=[256, 256], vf=[256, 256]) + ) model = PPO( "MlpPolicy", env, verbose=1, seed=seed, - 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 - 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 + learning_rate=1.5e-4, + n_steps=256, + batch_size=256, + n_epochs=10, + gamma=0.99, + gae_lambda=0.95, + clip_range=0.2, + target_kl=0.03, + ent_coef=0.03, vf_coef=0.5, max_grad_norm=0.5, device=device, @@ -204,6 +213,7 @@ def train( tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"), ) + # Base folder where model runs will be stored save_dir = str(Path(model_path).parent) model_prefix = Path(model_path).stem @@ -217,8 +227,19 @@ def train( 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) + # Save the final model inside the matching PPO_X directory as well + if model.logger and model.logger.dir: + run_folder_name = Path(model.logger.dir).name + final_dir = Path(model_path).parent / run_folder_name + else: + final_dir = Path(model_path).parent + + final_dir.mkdir(parents=True, exist_ok=True) + final_save_path = final_dir / f"{model_prefix}_final.zip" + model.save(str(final_save_path)) + if model.verbose > 0: + print(f"[Training Complete] Saved final model to -> {final_save_path}") + env.close()