""" 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()