Upload extract_trajectory.py with huggingface_hub

extract_trajectory.py

@@ -0,0 +1,43 @@
+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()