import sys
import os
# 确保当前目录 (demo) 在 sys.path 中
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
import argparse
from isaaclab.app import AppLauncher

# 启动参数，默认关闭 headless 模式以便观察
parser = argparse.ArgumentParser(description="Play trained T1 policy.")
parser.add_argument("--checkpoint", type=str, help="Path to the trained model checkpoint.")
AppLauncher.add_app_launcher_args(parser)
args_cli = parser.parse_args()

app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

import torch
import gymnasium as gym
from isaaclab_rl.rl_games import RlGamesVecEnvWrapper
from rl_games.torch_runner import Runner
import yaml
import os

from rl_game.demo.config.t1_env_cfg import T1EnvCfg

gym.register(
    id="Isaac-T1-Walking-v0",
    entry_point="isaaclab.envs:ManagerBasedRLEnv", # Isaac Lab 统一的强化学习环境入口
    kwargs={
        "cfg": T1EnvCfg(),
    },
)


def main():
    # 1. 创建环境 (只给 1 个机器人看演示)
    env = gym.make("Isaac-T1-Walking-v0", num_envs=1, render_mode="rgb_array")
    wrapped_env = RlGamesVecEnvWrapper(env)

    # 2. 加载配置并推理
    config_path = os.path.join(os.path.dirname(__file__), "config", "ppo_cfg.yaml")
    with open(config_path, "r") as f:
        rl_config = yaml.safe_load(f)

    # 注入 checkpoint 路径
    rl_config['params']['config']['checkpoint'] = args_cli.checkpoint

    runner = Runner()
    runner.load(rl_config)
    # 运行推理模式
    runner.run({"vecenv": wrapped_env, "train": False})

    simulation_app.close()


if __name__ == "__main__":
    main()