pretrain logic and small fixes

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