pretrain logic and small fixes

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
+        stride_len = cfg.step_length
-        step_height = 0.035
+        step_height = cfg.step_height
 
         center_data = (
             self.center_points.data if hasattr(self.center_points, 'data') else np.array(self.center_points)

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 = (
+        lines = [
-            f"=== JACKBOT METRICS ===\n"
+            "=== JACKBOT TELEMETRY ===",
-            f"Episode:     {episode}\n"
+            f"Mode:        {mode.upper()}",
-            f"Global Step: {step}\n"
+            f"Curriculum:  {phase}",
-            f"FPS:         {fps:.1f}\n"
+            f"Status:      {status}",
-            f"----------------------\n"
+            f"Episode:     {episode} (Step {ep_step})",
-            f"Top Rewards: [{top1}, {top2}, {top3}]\n"
+            f"Global Step: {step}",
-            f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]\n"
+            f"FPS:         {fps:.1f}",
-            f"Height:      {height:.3f} m\n"
+            "-------------------------",
-            f"Roll/Pitch:  {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.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
         )

ml/env.py

@@ -16,8 +16,6 @@ from simulation import Simulation
 from Robot import Robot, PyBulletBackend
 from ml.MetricsOverlay import MetricsHUD
 
-COLOR_FAILED = [0.3, 0.3, 0.3, 0.6]
-
 
 class CurriculumPhase(IntEnum):
     STAND_ONLY = 0
@@ -154,9 +152,6 @@ class JackBotEnv(gym.Env):
         # 2. Reset internal kinematics & hard reset joints in PyBullet
         self.robot.reset_to_init()
 
-        if self.use_gui:
-            self.sim.set_robot_color([1.0, 1.0, 1.0, 1.0])
-
         action_dim = self.action_space.shape[0]
         self.last_action = np.zeros(action_dim, dtype=np.float32)
         self.last_last_action = np.zeros(action_dim, dtype=np.float32)
@@ -212,12 +207,11 @@ 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
+        joint_range = np.minimum(
-        target_angles = np.where(
+            self.default_joint_angles - self.min_joint_limits,
-            action < 0.0,
+            self.max_joint_limits - self.default_joint_angles
-            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)
         )
+        target_angles = self.default_joint_angles + action * joint_range
         self.robot.tick(action=target_angles)
 
         if self.robot_mode != "kinematics" and self.step_count % 60 == 0:
@@ -421,8 +415,6 @@ class JackBotEnv(gym.Env):
 
         if is_tilted or is_collapsed:
             self.is_failed = True
-            if self.use_gui:
-                self.sim.set_robot_color(COLOR_FAILED)
 
     def close(self):
         self.sim.disconnect()

ml/model.py

@@ -1,49 +0,0 @@
-import torch
-import torch.nn as nn
-from torch.distributions import Normal
-
-
-class ActorCritic(nn.Module):
-    def __init__(self, obs_dim: int, action_dim: int, hidden_sizes=(256, 256)):
-        super().__init__()
-        self.backbone = nn.Sequential(
-            nn.Linear(obs_dim, hidden_sizes[0]),
-            nn.ReLU(),
-            nn.Linear(hidden_sizes[0], hidden_sizes[1]),
-            nn.ReLU(),
-        )
-
-        self.mean_head = nn.Linear(hidden_sizes[1], action_dim)
-        self.value_head = nn.Linear(hidden_sizes[1], 1)
-        self.log_std = nn.Parameter(torch.zeros(action_dim, dtype=torch.float32))
-
-    def forward(self, obs: torch.Tensor):
-        x = self.backbone(obs)
-        mean = self.mean_head(x)
-        std = self.log_std.exp()
-        value = self.value_head(x).squeeze(-1)
-        return mean, std, value
-
-    def get_action(self, obs: torch.Tensor):
-        mean, std, value = self.forward(obs)
-        dist = Normal(mean, std)
-        action = dist.sample()
-        log_prob = dist.log_prob(action).sum(-1)
-        return action, log_prob, value
-
-    def evaluate_actions(self, obs: torch.Tensor, actions: torch.Tensor):
-        mean, std, value = self.forward(obs)
-        dist = Normal(mean, std)
-        log_prob = dist.log_prob(actions).sum(-1)
-        entropy = dist.entropy().sum(-1)
-        return value, log_prob, entropy
-
-    def save(self, path: str):
-        torch.save(self.state_dict(), path)
-
-    @classmethod
-    def load(cls, path: str, obs_dim: int, action_dim: int, hidden_sizes=(256, 256)):
-        model = cls(obs_dim, action_dim, hidden_sizes)
-        model.load_state_dict(torch.load(path, map_location=torch.device("cpu")))
-        model.eval()
-        return model

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
+    )

ml/run_eval.py

@@ -1,7 +1,6 @@
-"""Run a trained policy in the PyBullet sim with full curriculum progression.
+"""
-
+ml/run_eval.py - Phase-by-Phase Policy Evaluator for JackBot
-Usage:
+Evaluates a trained model across all curriculum phases with fixed command vectors.
-    python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 5 --gui
 """
 
 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.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to save evaluation summary")
-    parser.set_defaults(random_command=True)
-    parser.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to output 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 = []
+    # Multi-Phase Configurations Suite
-    episode_lengths = []
+    phase_configs = [
-    episode_phases = []
+        (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):
+        print("\n" + "=" * 75)
-            obs, _ = env.reset()
+        print(" STARTING MULTI-PHASE EVALUATION SUITE")
-            done = False
+        print("=" * 75)
-            total_reward = 0.0
-            steps = 0
-            initial_phase = env.curriculum_phase.name
 
-            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:
+            ep_rewards = []
-                # Deterministic prediction matches evaluation standards
+            ep_steps = []
-                action, _ = model.predict(obs, deterministic=True)
+            ep_distances = []
-                
-                prev_phase = env.curriculum_phase
-                obs, reward, terminated, truncated, _ = env.step(action)
-                done = terminated or truncated
 
-                total_reward += float(reward)
+            for ep in range(args.episodes_per_phase):
-                steps += 1
+                obs, _ = env.reset()
 
-                if env.curriculum_phase != prev_phase:
+                # Force environment into active curriculum phase and lock command
-                    print(f"  └─ [Eval Milestone] Reached {env.curriculum_phase.name} at step {steps}!")
+                env.curriculum_phase = phase_enum
+                env.command = test_cmd.copy()
+                obs = env._get_obs()
 
-                if args.gui:
+                done = False
-                    time.sleep(1.0 / 240.0)
+                total_reward = 0.0
+                steps = 0
 
-            status_str = "FAILED (Terminated)" if terminated else "COMPLETED (Truncated)"
+                while not done:
-            final_phase = env.curriculum_phase.name
+                    # Enforce locked command each step
-            
+                    env.command = test_cmd.copy()
-            episode_rewards.append(total_reward)
-            episode_lengths.append(steps)
-            episode_phases.append(final_phase)
 
-            print(
+                    # Predict deterministic action from policy
-                f"Episode {ep + 1} Finished [{status_str}]: "
+                    action, _ = model.predict(obs, deterministic=True)
-                f"Phase = {final_phase} | Total Reward = {total_reward:.2f} | Steps = {steps}"
+                    
-            )
+                    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
+    # Print Summary Table
-    mean_reward = float(np.mean(episode_rewards))
+    print("\n" + "=" * 80)
-    std_reward = float(np.std(episode_rewards))
+    print(" EVALUATION SUMMARY REPORT")
-    mean_length = float(np.mean(episode_lengths))
+    print("=" * 80)
-
+    print(f"{'Phase ID & Name':<25} | {'Label':<22} | {'Reward':<8} | {'Steps':<6} | {'Max Dist':<8}")
-    print("\n" + "=" * 60)
+    print("-" * 80)
-    print(f"EVALUATION COMPLETE ({args.episodes} Episodes)")
+    for res in phase_summary:
-    print(f"Final Reached Phase:  {env.curriculum_phase.name}")
+        phase_str = f"[{res['phase_id']}] {res['phase_name']}"
-    print(f"Mean Reward:          {mean_reward:.2f} ± {std_reward:.2f}")
+        print(
-    print(f"Mean Episode Length:  {mean_length:.1f} steps")
+            f"{phase_str:<25} | "
-    print("=" * 60)
+            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)
+            json.dump({"model_path": str(args.model), "summary": phase_summary}, f, indent=4)
-        print(f"[Eval] Saved evaluation report to: {out_path.resolve()}")
+        print(f"\n[Eval] Saved report to: {out_path.resolve()}")
 
 
 if __name__ == "__main__":

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)),

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
+    # Initialize or Load PPO Policy Model
-    model = PPO(
+    if args.pretrained_model and os.path.exists(args.pretrained_model):
-        policy="MlpPolicy",
+        print(f"[Train] Loading pre-trained base knowledge from: {args.pretrained_model}")
-        env=vec_env,
+        model = PPO.load(
-        learning_rate=1e-4,
+            args.pretrained_model,
-        n_steps=256,
+            env=vec_env,
-        batch_size=256,
+            learning_rate=1e-4,  # Lower learning rate so RL fine-tunes without destroying base gait
-        n_epochs=10,
+            tensorboard_log=args.log_dir,
-        gamma=0.99,
+            device="cpu",
-        gae_lambda=0.95,
+        )
-        clip_range=0.2,
+    else:
-        ent_coef=0.03,
+        print("[Train] No base model provided. Starting training from scratch...")
-        target_kl=0.05,
+        model = PPO(
-        vf_coef=0.5,
+            policy="MlpPolicy",
-        max_grad_norm=0.5,
+            env=vec_env,
-        verbose=1,
+            learning_rate=1e-4,
-        tensorboard_log=args.log_dir,
+            n_steps=256,
-        device="cpu", 
+            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<number> naming
     checkpoint_callback = CheckpointCallback(

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: