"""
Isaac Lab Humanoid — Upright Walking Training
===============================================
Uses NVIDIA Isaac Sim physics with Isaac Lab's humanoid environment.
Custom reward weights for upright, stable walking.

Runs 2048 parallel environments on GPU (vs 16 in MuJoCo).

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

import sys
sys.argv += ["--headless"]

from isaaclab.app import AppLauncher
launcher = AppLauncher({"headless": True})
sim_app = launcher.app

import gymnasium as gym
import torch
import isaaclab_tasks  # registers all Isaac Lab envs

from isaaclab_tasks.direct.humanoid.humanoid_env import HumanoidEnvCfg

# ── Custom config: heavy uprightness reward ──
cfg = HumanoidEnvCfg()
cfg.up_weight = 2.0              # default 0.1 — 20x increase
cfg.alive_reward_scale = 3.0     # reward for staying alive
cfg.heading_weight = 0.5         # reward for walking straight
cfg.termination_height = 0.9     # terminate if torso drops below this

# Run 2048 envs in parallel on GPU
cfg.scene.num_envs = 2048

print("=" * 60)
print("ISAAC LAB HUMANOID — Upright Walking")
print(f"Physics: NVIDIA PhysX (Isaac Sim 5.1)")
print(f"Parallel envs: {cfg.scene.num_envs}")
print(f"Uprightness weight: {cfg.up_weight} (20x default)")
print("=" * 60)

env = gym.make("Isaac-Humanoid-Direct-v0", cfg=cfg)

# Wrap for rsl_rl
from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper
env = RslRlVecEnvWrapper(env)

# Load PPO runner config
from isaaclab_tasks.direct.humanoid.agents.rsl_rl_ppo_cfg import HumanoidPPORunnerCfg
from rsl_rl.runners import OnPolicyRunner

runner_cfg = HumanoidPPORunnerCfg()
runner_cfg.max_iterations = 3000

runner = OnPolicyRunner(
    env,
    runner_cfg.to_dict(),
    log_dir="./logs/humanoid_upright",
    device="cuda"
)

print("\nStarting training...")
print("With 2048 parallel envs, this should be MUCH faster than MuJoCo\n")

runner.learn(
    num_learning_iterations=runner_cfg.max_iterations,
    init_at_random_ep_len=True
)

print("\nTraining complete!")
print("Model saved in ./logs/humanoid_upright/")

sim_app.close()