""" ml/pretrain_bc.py - Behavioral Cloning teacher-data pipeline. This script collects observation/action pairs from the kinematics solver, then uses those pairs as training data for a PPO policy. In practice it serves as a teacher- student pretraining step: the kinematics controller generates sample trajectories, and this file teaches the policy to imitate those behavior patterns before PPO training. """ 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})...") # Disable random command resampling inside env.step so manual command locks persist env = JackBotEnv(use_gui=use_gui, robot_mode="kinematics", random_command=False) 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. Update command and vector targets every 120 steps if i % 120 == 0: env.command = env.sample_command() 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)] # 2. Capture observation BEFORE stepping environment current_obs = env._get_obs() # 3. Step environment ONCE (updates kinematics solver, PyBullet physics, and computes IK) obs, _, terminated, truncated, _ = env.step(np.zeros(18, dtype=np.float32)) # 4. Extract procedural IK joint targets computed during this step target_ik_rad = env.robot.current_rad.data.flatten().copy() # Step 5: Convert target radians directly to [-1, 1] relative to joint limits min_lim = env.min_joint_limits.flatten() max_lim = env.max_joint_limits.flatten() normalized_action = 2.0 * (target_ik_rad - min_lim) / (max_lim - min_lim) - 1.0 normalized_action = np.clip(normalized_action, -1.0, 1.0) # Step 6: Store matching input (obs) and target ground truth (normalized_action) observations.append(current_obs.copy()) actions.append(normalized_action.copy()) if use_gui: time.sleep(1.0 / 60.0) # 7. Handle episode boundaries using terminated and truncated 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 )