train_humanoid_gpu.py

"""
Humanoid SAC — GPU Training (Vast.ai / A6000)
===============================================
Custom reward shaping for UPRIGHT walking:
  - Big reward for keeping torso vertical
  - Capped forward velocity (walking, not sprinting)
  - Penalties for wobble, flailing, jerky motion
  - Smooth action rewards for natural gait

Expected time: ~1-2 hours on A6000 for 50M steps.

Usage:
  bash setup.sh          # install deps
  python train_humanoid_gpu.py   # train

Author: Milan Narula
Project: ARCSA AIR Scholarship Application
"""

import os
import json
import time
import numpy as np
import gymnasium as gym
from stable_baselines3 import SAC
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.vec_env import SubprocVecEnv

# ── Config ───────────────────────────────────────────────────────────────
TIMESTEPS = 10_000_000
SAVE_DIR = "models"
RESULTS_DIR = "results"

import torch
if torch.cuda.is_available():
    DEVICE = "cuda"
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
else:
    DEVICE = "cpu"
    print("WARNING: No GPU detected — this will be very slow")


class UprightHumanoidWrapper(gym.Wrapper):
    """
    Custom reward for UPRIGHT walking.

    The default MuJoCo reward produces "controlled falling" — the robot
    leans forward and catches itself. This wrapper instead rewards:

    1. Torso staying vertical (uprightness from quaternion)
    2. Maintaining standing height (~1.25m)
    3. Moderate forward velocity (capped — no sprinting)
    4. Low energy usage (no flailing)
    5. Smooth actions (no jerky movements)
    6. Walking straight (minimal lateral drift)
    """

    def __init__(self, env):
        super().__init__(env)
        self._prev_action = None

    def step(self, action):
        obs, _, terminated, truncated, info = self.env.step(action)

        qpos = self.env.unwrapped.data.qpos
        qvel = self.env.unwrapped.data.qvel

        z_height = qpos[2]
        w, qx, qy, qz = qpos[3:7]
        forward_vel = qvel[0]
        lateral_vel = qvel[1]

        # Uprightness: 1.0 = perfectly vertical, 0.0 = horizontal
        upright = 1.0 - 2.0 * (qx * qx + qy * qy)

        # ── Reward components ──

        # Forward velocity (capped at 1.5 m/s for walking speed)
        forward_reward = min(forward_vel, 1.5) * 1.0

        # Upright bonus — THE most important signal
        upright_reward = 3.0 * max(0, upright)

        # Height maintenance
        height_reward = 2.0 * max(0, 1.0 - abs(z_height - 1.25))

        # Lean penalty
        lean_penalty = -2.0 * max(0, 1.0 - upright)

        # Energy penalty
        energy_penalty = -0.05 * np.sum(np.square(action))

        # Lateral penalty
        lateral_penalty = -0.5 * abs(lateral_vel)

        # Smoothness penalty
        smooth_penalty = 0.0
        if self._prev_action is not None:
            smooth_penalty = -0.02 * np.sum(np.square(action - self._prev_action))
        self._prev_action = action.copy()

        # Alive bonus
        alive_bonus = 2.0

        shaped_reward = (
            forward_reward + upright_reward + height_reward +
            lean_penalty + energy_penalty + lateral_penalty +
            smooth_penalty + alive_bonus
        )

        # Only terminate if truly fallen
        fallen = z_height < 0.7 or upright < 0.3
        terminated = fallen

        info["upright"] = upright
        info["z_height"] = z_height
        info["forward_vel"] = forward_vel

        return obs, shaped_reward, terminated, truncated, info

    def reset(self, **kwargs):
        self._prev_action = None
        return self.env.reset(**kwargs)


class TrainingCallback(BaseCallback):
    def __init__(self, check_freq=50_000):
        super().__init__()
        self.check_freq = check_freq
        self.metrics = {
            "timesteps": [], "mean_reward": [], "mean_survival": [],
            "mean_upright": [], "best_reward": float("-inf"),
        }
        self.start_time = time.time()

    def _on_step(self) -> bool:
        if self.n_calls % self.check_freq == 0:
            results = self._evaluate()
            mean_r = float(np.mean(results["rewards"]))
            mean_s = float(np.mean(results["survival"]))
            mean_u = float(np.mean(results["upright"]))

            self.metrics["timesteps"].append(int(self.num_timesteps))
            self.metrics["mean_reward"].append(round(mean_r, 1))
            self.metrics["mean_survival"].append(round(mean_s, 0))
            self.metrics["mean_upright"].append(round(mean_u, 3))

            if mean_r > self.metrics["best_reward"]:
                self.metrics["best_reward"] = mean_r
                self.model.save(os.path.join(SAVE_DIR, "best_humanoid_sac"))

            elapsed = (time.time() - self.start_time) / 3600
            rate = self.num_timesteps / (time.time() - self.start_time)
            eta = (TIMESTEPS - self.num_timesteps) / rate / 3600 if rate > 0 else 0

            print(
                f"  {self.num_timesteps:>10,} | "
                f"R: {mean_r:>7.1f} | "
                f"Surv: {mean_s:>5.0f} | "
                f"Up: {mean_u:.2f} | "
                f"Best: {self.metrics['best_reward']:.1f} | "
                f"{elapsed:.1f}h / ETA {eta:.1f}h"
            )

            os.makedirs(RESULTS_DIR, exist_ok=True)
            with open(os.path.join(RESULTS_DIR, "humanoid_sac_metrics.json"), "w") as f:
                json.dump(self.metrics, f, indent=2)

        return True

    def _evaluate(self, n_episodes=3):
        results = {"rewards": [], "survival": [], "upright": []}
        for ep in range(n_episodes):
            env = UprightHumanoidWrapper(gym.make("Humanoid-v5"))
            obs, _ = env.reset(seed=ep + 9000)
            total_r, steps, uprights = 0, 0, []
            done = False
            while not done and steps < 1000:
                action, _ = self.model.predict(obs, deterministic=True)
                obs, r, terminated, truncated, info = env.step(action)
                total_r += r
                steps += 1
                uprights.append(info.get("upright", 0))
                done = terminated or truncated
            results["rewards"].append(total_r)
            results["survival"].append(steps)
            results["upright"].append(np.mean(uprights))
            env.close()
        return results


def main():
    print("=" * 60)
    print("HUMANOID SAC — Upright Walking (GPU)")
    print(f"Device: {DEVICE} | Steps: {TIMESTEPS:,}")
    print("=" * 60)

    os.makedirs(SAVE_DIR, exist_ok=True)
    os.makedirs(RESULTS_DIR, exist_ok=True)

    N_ENVS = 16

    def make_env(rank):
        def _init():
            env = UprightHumanoidWrapper(gym.make("Humanoid-v5"))
            return env
        return _init

    env = SubprocVecEnv([make_env(i) for i in range(N_ENVS)])

    model = SAC(
        policy="MlpPolicy",
        env=env,
        learning_rate=3e-4,
        buffer_size=1_000_000,
        batch_size=512,
        tau=0.005,
        gamma=0.99,
        learning_starts=25_000,
        train_freq=1,
        gradient_steps=2,
        verbose=0,
        device=DEVICE,
        seed=42,
        policy_kwargs=dict(net_arch=[512, 256, 256]),
    )

    start = time.time()
    callback = TrainingCallback(check_freq=50_000)
    model.learn(total_timesteps=TIMESTEPS, callback=callback)
    elapsed = time.time() - start

    model.save(os.path.join(SAVE_DIR, "final_humanoid_sac"))
    env.close()

    print(f"\n  Done in {elapsed / 3600:.1f} hours")
    print(f"  Best reward: {callback.metrics['best_reward']:.1f}")
    print(f"\n  Download models/best_humanoid_sac.zip back to your laptop")


if __name__ == "__main__":
    main()