JackBot/ml/pretrain_bc.py

"""
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})...")
    
    # 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
    )