diff --git a/Robot.py b/Robot.py index 5805571..a046289 100644 --- a/Robot.py +++ b/Robot.py @@ -220,8 +220,8 @@ class Robot: return self.gait_phase = (self.gait_phase + 0.22) % (2.0 * math.pi) - stride_len = 0.045 - step_height = 0.035 + stride_len = cfg.step_length + step_height = cfg.step_height center_data = ( self.center_points.data if hasattr(self.center_points, 'data') else np.array(self.center_points) diff --git a/ml/MetricsOverlay.py b/ml/MetricsOverlay.py index e7ba4ae..bb96b52 100644 --- a/ml/MetricsOverlay.py +++ b/ml/MetricsOverlay.py @@ -1,7 +1,7 @@ """ ml/MetricsOverlay.py - Camera-Facing (Billboard) 3D Floating Text Overlay """ -from typing import List, Tuple, Optional +from typing import List, Tuple, Optional, Dict import numpy as np import pybullet as p @@ -21,20 +21,15 @@ class MetricsHUD: yaw = cam_info[8] pitch = cam_info[9] - # Orient the text normal toward the camera view direction - # PyBullet text default faces local +Z/-Y depending on roll, - # converting visualizer yaw/pitch to Euler angles (roll, pitch, yaw in radians) - roll_rad = 0.0 pitch_rad = np.radians(pitch + 90.0) yaw_rad = np.radians(yaw) text_orientation = p.getQuaternionFromEuler( - [pitch_rad, roll_rad, yaw_rad], + [pitch_rad, 0.0, yaw_rad], physicsClientId=self.client_id ) return text_orientation except Exception: - # Fallback default orientation if camera info call fails return [0.0, 0.0, 0.0, 1.0] def update( @@ -45,29 +40,48 @@ class MetricsHUD: cmd_vel: np.ndarray, fps: float = 0.0, height: float = 0.0, - roll_pitch: Tuple[float, float] = (0.0, 0.0) + roll_pitch: Tuple[float, float] = (0.0, 0.0), + mode: str = "direct", + phase: str = "STAND_ONLY", + distance: float = 0.0, + status: str = "ALIVE", + reward_components: Optional[Dict[str, float]] = None, + ep_step: int = 0, ) -> None: - """Updates floating black text block in 3D space with billboarding.""" + """Updates floating text block in 3D space with expanded telemetry.""" sorted_rewards = sorted(robot_rewards, reverse=True) top1 = f"{sorted_rewards[0]:+.2f}" if len(sorted_rewards) > 0 else "0.00" - top2 = f"{sorted_rewards[1]:+.2f}" if len(sorted_rewards) > 1 else "0.00" - top3 = f"{sorted_rewards[2]:+.2f}" if len(sorted_rewards) > 2 else "0.00" vx = cmd_vel[0] if len(cmd_vel) > 0 else 0.0 vy = cmd_vel[1] if len(cmd_vel) > 1 else 0.0 - omega = cmd_vel[3] if len(cmd_vel) > 3 else 0.0 + omega = cmd_vel[2] if len(cmd_vel) > 2 else 0.0 - hud_text = ( - f"=== JACKBOT METRICS ===\n" - f"Episode: {episode}\n" - f"Global Step: {step}\n" - f"FPS: {fps:.1f}\n" - f"----------------------\n" - f"Top Rewards: [{top1}, {top2}, {top3}]\n" - f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]\n" - f"Height: {height:.3f} m\n" - f"Roll/Pitch: {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°" - ) + lines = [ + "=== JACKBOT TELEMETRY ===", + f"Mode: {mode.upper()}", + f"Curriculum: {phase}", + f"Status: {status}", + f"Episode: {episode} (Step {ep_step})", + f"Global Step: {step}", + f"FPS: {fps:.1f}", + "-------------------------", + f"Episode Rew: {top1}", + f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]", + f"Height: {height:.3f} m", + f"Roll/Pitch: {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°", + f"Max Dist: {distance:.2f} m", + ] + + if reward_components: + lin_v = reward_components.get("lin_vel", 0.0) + stab = reward_components.get("stability", 0.0) + h_rew = reward_components.get("height", 0.0) + jit = reward_components.get("jitter_penalty", 0.0) + lines.append("--- Reward Components ---") + lines.append(f"LinVel: {lin_v:.2f} | Stab: {stab:.2f}") + lines.append(f"Height: {h_rew:.2f} | Jitter: {jit:+.3f}") + + hud_text = "\n".join(lines) # Position above origin in simulation world text_position = [-0.8, -0.8, 1.2] @@ -76,7 +90,7 @@ class MetricsHUD: # Calculate dynamic orientation to align text flat against camera plane text_orientation = self._get_camera_facing_orientation() - # Safely remove the old text to prevent PyBullet ghosting/overlapping + # Safely remove old text to prevent PyBullet ghosting/overlapping if self._text_id is not None: try: p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id) @@ -88,7 +102,7 @@ class MetricsHUD: text=hud_text, textPosition=text_position, textColorRGB=text_color, - textSize=0.1, + textSize=0.085, textOrientation=text_orientation, physicsClientId=self.client_id ) diff --git a/ml/env.py b/ml/env.py index b62cb49..f5f3fd7 100644 --- a/ml/env.py +++ b/ml/env.py @@ -207,11 +207,9 @@ class JackBotEnv(gym.Env): self.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle" self.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)] - # Delegate execution tick to Robot instance - target_angles = np.where( - action < 0.0, - self.default_joint_angles + action * (self.default_joint_angles - self.min_joint_limits), - self.default_joint_angles + action * (self.max_joint_limits - self.default_joint_angles) + joint_range = np.minimum( + self.default_joint_angles - self.min_joint_limits, + self.max_joint_limits - self.default_joint_angles ) target_angles = self.default_joint_angles + action * joint_range self.robot.tick(action=target_angles) diff --git a/ml/pretrain_bc.py b/ml/pretrain_bc.py new file mode 100644 index 0000000..51c2dac --- /dev/null +++ b/ml/pretrain_bc.py @@ -0,0 +1,147 @@ +""" +ml/pretrain_bc.py - Behavioral Cloning from Kinematics Teacher +""" +import sys +import time +import argparse +from pathlib import Path +import numpy as np +import torch +import torch.nn as nn +from torch.utils.data import DataLoader, TensorDataset +from stable_baselines3 import PPO +from tqdm import tqdm # <--- Progress Bar Support + +# Ensure project root is in sys.path +sys.path.append(str(Path(__file__).resolve().parent.parent)) +from ml.env import JackBotEnv, CurriculumPhase + + +def collect_kinematics_dataset(num_samples: int = 100_000, use_gui: bool = False): + """Collects (Observation, Action) pairs directly from Kinematics Teacher.""" + print(f"\n[Pretrain] Collecting {num_samples} samples from Kinematics Teacher (GUI={use_gui})...") + + env = JackBotEnv(use_gui=use_gui, robot_mode="kinematics") + env.curriculum_phase = CurriculumPhase.FULL_COMMAND + + observations = [] + actions = [] + + obs, _ = env.reset() + + # --- PROGRESS BAR: Data Collection --- + pbar = tqdm(range(num_samples), desc=" Collecting Data", unit="step") + for i in pbar: + # 1. Sample random movement command + if i % 120 == 0: + env.command = env.sample_command() + + # 2. Let Robot compute Kinematics target angles + cmd_vx, cmd_vy, cmd_omega = env.command + env.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle" + env.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)] + + # Advance internal IK tick to calculate joint angles + env.robot.tick() + + # Extract .data and flatten (6, 3) matrix to 18-dim 1D array + target_ik_rad = env.robot.current_rad.data.flatten().copy() + + # 3. Convert target angles back to normalized [-1, 1] action space + normalized_action = np.where( + target_ik_rad >= env.default_joint_angles, + (target_ik_rad - env.default_joint_angles) / np.maximum(1e-5, env.max_joint_limits - env.default_joint_angles), + (target_ik_rad - env.default_joint_angles) / np.maximum(1e-5, env.default_joint_angles - env.min_joint_limits) + ) + normalized_action = np.clip(normalized_action, -1.0, 1.0) + + # 4. Save sample + observations.append(obs.copy()) + actions.append(normalized_action.copy()) + + # Step simulation environment + obs, _, terminated, truncated, _ = env.step(normalized_action) + + if use_gui: + time.sleep(1.0 / 60.0) + + if terminated or truncated: + obs, _ = env.reset() + + env.close() + print("[Pretrain] Data collection complete!\n") + return np.array(observations, dtype=np.float32), np.array(actions, dtype=np.float32) + + +def pretrain_policy( + save_path: str = "ml/checkpoints/jackbot_kinematics_base.zip", + epochs: int = 15, + batch_size: int = 256, + num_samples: int = 100_000, + use_gui: bool = False +): + # Collect dataset from Kinematics teacher + obs_data, action_data = collect_kinematics_dataset(num_samples=num_samples, use_gui=use_gui) + + # Initialize Dummy Env & Fresh SB3 PPO Model + dummy_env = JackBotEnv(use_gui=False, robot_mode="direct") + model = PPO("MlpPolicy", dummy_env, learning_rate=5e-4, verbose=0, device="cpu") + + # Extract PyTorch Policy Network & Optimizer + policy = model.policy + optimizer = torch.optim.Adam(policy.parameters(), lr=5e-4) + loss_fn = nn.MSELoss() + + # Convert to PyTorch Dataloader + dataset = TensorDataset(torch.tensor(obs_data), torch.tensor(action_data)) + loader = DataLoader(dataset, batch_size=batch_size, shuffle=True) + + print(f"[Pretrain] Pre-training Policy Network ({epochs} Epochs)...") + policy.train() + + # --- PROGRESS BAR: Epoch Training --- + for epoch in range(epochs): + epoch_loss = 0.0 + + batch_pbar = tqdm(loader, desc=f" Epoch {epoch + 1:02d}/{epochs:02d}", leave=True, unit="batch") + for batch_obs, batch_actions in batch_pbar: + optimizer.zero_grad() + + distribution = policy.get_distribution(batch_obs) + predicted_actions = distribution.distribution.mean + + loss = loss_fn(predicted_actions, batch_actions) + loss.backward() + optimizer.step() + + current_loss = loss.item() + epoch_loss += current_loss * len(batch_obs) + + # Dynamic loss update in progress bar tail + batch_pbar.set_postfix({"loss": f"{current_loss:.6f}"}) + + avg_loss = epoch_loss / len(dataset) + tqdm.write(f" └─ Epoch {epoch + 1:02d}/{epochs:02d} Complete | Mean MSE Loss: {avg_loss:.6f}") + + # Save SB3 Model Checkpoint + out_file = Path(save_path) + out_file.parent.mkdir(parents=True, exist_ok=True) + model.save(out_file) + dummy_env.close() + print(f"\n[Pretrain] Successfully saved pre-trained base model to: {out_file.resolve()}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="JackBot Behavioral Cloning Pre-trainer") + parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering during collection") + parser.add_argument("--num-samples", type=int, default=100_000, help="Number of dataset samples to collect") + parser.add_argument("--epochs", type=int, default=15, help="Number of BC training epochs") + parser.add_argument("--save-path", type=str, default="ml/checkpoints/jackbot_kinematics_base.zip", help="Output path for pre-trained model .zip") + args = parser.parse_args() + + pretrain_policy( + save_path=args.save_path, + epochs=args.epochs, + num_samples=args.num_samples, + use_gui=args.gui + ) \ No newline at end of file diff --git a/ml/run_eval.py b/ml/run_eval.py index 7e140e6..2188c56 100644 --- a/ml/run_eval.py +++ b/ml/run_eval.py @@ -1,7 +1,6 @@ -"""Run a trained policy in the PyBullet sim with full curriculum progression. - -Usage: - python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 5 --gui +""" +ml/run_eval.py - Phase-by-Phase Policy Evaluator for JackBot +Evaluates a trained model across all curriculum phases with fixed command vectors. """ import argparse @@ -11,102 +10,131 @@ import sys from pathlib import Path import numpy as np from stable_baselines3 import PPO + +# Ensure project root is in sys.path sys.path.append(str(Path(__file__).resolve().parent.parent)) -from ml.env import JackBotEnv +from ml.env import JackBotEnv, CurriculumPhase + def main(): - parser = argparse.ArgumentParser(description="JackBot Policy Evaluator") + parser = argparse.ArgumentParser(description="JackBot Phase-by-Phase Policy Evaluator") parser.add_argument("--model", type=str, required=True, help="Path to trained model checkpoint (.zip)") - parser.add_argument("--episodes", type=int, default=5, help="Number of evaluation episodes to run") + parser.add_argument("--episodes-per-phase", type=int, default=1, help="Number of test episodes per phase") + parser.add_argument("--max-steps-per-episode", type=int, default=600, help="Max simulation steps per episode") parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering") - parser.add_argument("--no-random-command", dest="random_command", action="store_false", help="Lock commands to zero (disable random command sampling)") - parser.set_defaults(random_command=True) - parser.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to output evaluation summary") + parser.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to save evaluation summary") args = parser.parse_args() print(f"[Eval] Loading policy model from: {args.model}") model = PPO.load(args.model, device="cpu") - # Instantiate single evaluation environment matching train setup + # Instantiate environment with random_command disabled so our test command stays locked env = JackBotEnv( use_gui=args.gui, - random_command=args.random_command, + random_command=False, + max_episode_steps=args.max_steps_per_episode, ) - episode_rewards = [] - episode_lengths = [] - episode_phases = [] + # Multi-Phase Configurations Suite + phase_configs = [ + (CurriculumPhase.STAND_ONLY, "STAND", np.array([0.0, 0.0, 0.0], dtype=np.float32)), + (CurriculumPhase.FORWARD, "FORWARD GAIT", np.array([1.0, 0.0, 0.0], dtype=np.float32)), + (CurriculumPhase.TURN_AND_DIRECTION, "FORWARD + YAW TURN", np.array([0.5, 0.0, 0.4], dtype=np.float32)), + (CurriculumPhase.OMNI_DIRECTION, "STRIDE LATERAL", np.array([0.5, 0.5, 0.0], dtype=np.float32)), + (CurriculumPhase.FULL_COMMAND, "FULL OMNI COMBINATION", np.array([0.3, 0.2, 0.3], dtype=np.float32)), + ] + + phase_summary = [] try: - for ep in range(args.episodes): - obs, _ = env.reset() - done = False - total_reward = 0.0 - steps = 0 - initial_phase = env.curriculum_phase.name + print("\n" + "=" * 75) + print(" STARTING MULTI-PHASE EVALUATION SUITE") + print("=" * 75) - print(f"\n--- Starting Evaluation Episode {ep + 1}/{args.episodes} [Phase: {initial_phase}] ---") + for phase_enum, phase_name, test_cmd in phase_configs: + print(f"\n▶ Testing Phase [{phase_enum.value}]: {phase_enum.name} ({phase_name})") + print(f" Target Command Vector [vx, vy, omega]: {test_cmd.tolist()}") - while not done: - # Deterministic prediction matches evaluation standards - action, _ = model.predict(obs, deterministic=True) - - prev_phase = env.curriculum_phase - obs, reward, terminated, truncated, _ = env.step(action) - done = terminated or truncated + ep_rewards = [] + ep_steps = [] + ep_distances = [] - total_reward += float(reward) - steps += 1 + for ep in range(args.episodes_per_phase): + obs, _ = env.reset() - if env.curriculum_phase != prev_phase: - print(f" └─ [Eval Milestone] Reached {env.curriculum_phase.name} at step {steps}!") + # Force environment into active curriculum phase and lock command + env.curriculum_phase = phase_enum + env.command = test_cmd.copy() + obs = env._get_obs() - if args.gui: - time.sleep(1.0 / 240.0) + done = False + total_reward = 0.0 + steps = 0 - status_str = "FAILED (Terminated)" if terminated else "COMPLETED (Truncated)" - final_phase = env.curriculum_phase.name - - episode_rewards.append(total_reward) - episode_lengths.append(steps) - episode_phases.append(final_phase) + while not done: + # Enforce locked command each step + env.command = test_cmd.copy() - print( - f"Episode {ep + 1} Finished [{status_str}]: " - f"Phase = {final_phase} | Total Reward = {total_reward:.2f} | Steps = {steps}" - ) + # Predict deterministic action from policy + action, _ = model.predict(obs, deterministic=True) + + obs, reward, terminated, truncated, _ = env.step(action) + done = terminated or truncated + + total_reward += float(reward) + steps += 1 + + if args.gui: + time.sleep(1.0 / 60.0) + + status_str = "FAILED (Collapsed)" if terminated else "SUCCESS (Completed)" + dist = float(env.max_distance_from_start) + print( + f" └─ Ep {ep + 1}/{args.episodes_per_phase}: {status_str:<19} | " + f"Steps: {steps:<4} | Reward: {total_reward:+.2f} | Max Dist: {dist:.2f}m" + ) + + ep_rewards.append(total_reward) + ep_steps.append(steps) + ep_distances.append(dist) + + phase_summary.append({ + "phase_id": phase_enum.value, + "phase_name": phase_enum.name, + "label": phase_name, + "command": test_cmd.tolist(), + "mean_reward": float(np.mean(ep_rewards)), + "mean_steps": float(np.mean(ep_steps)), + "mean_distance": float(np.mean(ep_distances)), + }) finally: env.close() - # Calculate metrics report - mean_reward = float(np.mean(episode_rewards)) - std_reward = float(np.std(episode_rewards)) - mean_length = float(np.mean(episode_lengths)) - - print("\n" + "=" * 60) - print(f"EVALUATION COMPLETE ({args.episodes} Episodes)") - print(f"Final Reached Phase: {env.curriculum_phase.name}") - print(f"Mean Reward: {mean_reward:.2f} ± {std_reward:.2f}") - print(f"Mean Episode Length: {mean_length:.1f} steps") - print("=" * 60) + # Print Summary Table + print("\n" + "=" * 80) + print(" EVALUATION SUMMARY REPORT") + print("=" * 80) + print(f"{'Phase ID & Name':<25} | {'Label':<22} | {'Reward':<8} | {'Steps':<6} | {'Max Dist':<8}") + print("-" * 80) + for res in phase_summary: + phase_str = f"[{res['phase_id']}] {res['phase_name']}" + print( + f"{phase_str:<25} | " + f"{res['label']:<22} | " + f"{res['mean_reward']:<+8.2f} | " + f"{res['mean_steps']:<6.0f} | " + f"{res['mean_distance']:<8.2f}m" + ) + print("=" * 80) + # Save JSON Report if requested if args.save_metrics: - metrics = { - "model_path": str(args.model), - "episodes_evaluated": args.episodes, - "final_curriculum_phase": env.curriculum_phase.name, - "mean_reward": mean_reward, - "std_reward": std_reward, - "mean_episode_length": mean_length, - "raw_rewards": episode_rewards, - "episode_phases": episode_phases, - } out_path = Path(args.save_metrics) out_path.parent.mkdir(parents=True, exist_ok=True) with open(out_path, "w") as f: - json.dump(metrics, f, indent=4) - print(f"[Eval] Saved evaluation report to: {out_path.resolve()}") + json.dump({"model_path": str(args.model), "summary": phase_summary}, f, indent=4) + print(f"\n[Eval] Saved report to: {out_path.resolve()}") if __name__ == "__main__": diff --git a/ml/run_eval_training.py b/ml/run_eval_training.py index f1990f2..3e84b25 100644 --- a/ml/run_eval_training.py +++ b/ml/run_eval_training.py @@ -21,6 +21,7 @@ def evaluate_kinematics(episode_length: int = 1000): print("=" * 70 + "\n") phase_configs = [ + (CurriculumPhase.STAND_ONLY, "STAND", np.array([1.0, 0.0, 0.0], dtype=np.float32)), (CurriculumPhase.FORWARD, "FORWARD GAIT", np.array([1.0, 0.0, 0.0], dtype=np.float32)), (CurriculumPhase.TURN_AND_DIRECTION, "FORWARD + YAW TURN", np.array([0.5, 0.0, 0.4], dtype=np.float32)), (CurriculumPhase.OMNI_DIRECTION, "STRIDE LATERAL", np.array([0.5, 0.5, 0.0], dtype=np.float32)), diff --git a/ml/run_train.py b/ml/run_train.py index 0c6e2ca..ddb3551 100644 --- a/ml/run_train.py +++ b/ml/run_train.py @@ -57,6 +57,12 @@ def main(): parser.add_argument("--save-dir", type=str, default="ml/checkpoints", help="Directory for model checkpoints") parser.add_argument("--save-freq", type=int, default=50_000, help="Checkpoint save frequency (steps)") parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering") + parser.add_argument( + "--pretrained-model", + type=str, + default=None, + help="Path to pre-trained base model checkpoint (.zip) to start PPO training from" + ) args = parser.parse_args() os.makedirs(args.log_dir, exist_ok=True) @@ -74,25 +80,36 @@ def main(): ] vec_env = SubprocVecEnv(env_fns) - # Initialize PPO Policy Hyperparameters - model = PPO( - policy="MlpPolicy", - env=vec_env, - learning_rate=1e-4, - n_steps=256, - batch_size=256, - n_epochs=10, - gamma=0.99, - gae_lambda=0.95, - clip_range=0.2, - ent_coef=0.03, - target_kl=0.05, - vf_coef=0.5, - max_grad_norm=0.5, - verbose=1, - tensorboard_log=args.log_dir, - device="cpu", - ) + # Initialize or Load PPO Policy Model + if args.pretrained_model and os.path.exists(args.pretrained_model): + print(f"[Train] Loading pre-trained base knowledge from: {args.pretrained_model}") + model = PPO.load( + args.pretrained_model, + env=vec_env, + learning_rate=1e-4, # Lower learning rate so RL fine-tunes without destroying base gait + tensorboard_log=args.log_dir, + device="cpu", + ) + else: + print("[Train] No base model provided. Starting training from scratch...") + model = PPO( + policy="MlpPolicy", + env=vec_env, + learning_rate=1e-4, + n_steps=256, + batch_size=256, + n_epochs=10, + gamma=0.99, + gae_lambda=0.95, + clip_range=0.2, + ent_coef=0.01, + target_kl=0.05, + vf_coef=0.5, + max_grad_norm=0.5, + verbose=1, + tensorboard_log=args.log_dir, + device="cpu", + ) # Setup Callbacks with ppo naming checkpoint_callback = CheckpointCallback( diff --git a/simulation.py b/simulation.py index 5ec2149..7c264f8 100644 --- a/simulation.py +++ b/simulation.py @@ -164,12 +164,6 @@ class Simulation: """Advances physics simulation by 1 time step.""" p.stepSimulation(physicsClientId=self.physics_client) - def set_robot_color(self, rgba: List[float]) -> None: - num_joints = p.getNumJoints(self.robot_id, physicsClientId=self.physics_client) - p.changeVisualShape(self.robot_id, -1, rgbaColor=rgba, physicsClientId=self.physics_client) - for j in range(num_joints): - p.changeVisualShape(self.robot_id, j, rgbaColor=rgba, physicsClientId=self.physics_client) - def settle_and_measure_height( self, target_angles: Optional[np.ndarray] = None, steps: int = 100, fallback_height: float = 0.122 ) -> float: