import sys
sys.argv = ['play', '--headless']
from isaaclab.app import AppLauncher
launcher = AppLauncher({"headless": True})
sim_app = launcher.app

import torch
import json
import gymnasium as gym
import isaaclab_tasks

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

from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper
env = RslRlVecEnvWrapper(env)

from rsl_rl.runners import OnPolicyRunner
from isaaclab_tasks.direct.humanoid.agents.rsl_rl_ppo_cfg import HumanoidPPORunnerCfg

runner_cfg = HumanoidPPORunnerCfg()
runner = OnPolicyRunner(env, runner_cfg.to_dict(), log_dir=None, device="cuda")
runner.load("/workspace/IsaacLab/logs/rsl_rl/humanoid_direct/2026-04-13_00-52-00/model_2999.pt")

policy = runner.get_inference_policy(device="cuda")
obs, _ = env.get_observations()

trajectory = []
for step in range(500):
    with torch.no_grad():
        actions = policy(obs)
    obs, _, dones, infos = env.step(actions)
    qpos = env.unwrapped.scene["robot"].data.joint_pos[0].cpu().numpy().tolist()
    qvel = env.unwrapped.scene["robot"].data.joint_vel[0].cpu().numpy().tolist()
    root_pos = env.unwrapped.scene["robot"].data.root_pos_w[0].cpu().numpy().tolist()
    root_quat = env.unwrapped.scene["robot"].data.root_quat_w[0].cpu().numpy().tolist()
    trajectory.append({"step": step, "joint_pos": qpos, "joint_vel": qvel, "root_pos": root_pos, "root_quat": root_quat, "actions": actions[0].cpu().numpy().tolist()})
    if step % 100 == 0:
        print(f"Step {step}/500 - root_pos: {root_pos[:3]}")

with open("/workspace/isaac_trajectory.json", "w") as f:
    json.dump(trajectory, f)
print(f"Saved {len(trajectory)} steps to /workspace/isaac_trajectory.json")
sim_app.close()