# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""Script to play a checkpoint if an RL agent from RSL-RL."""

"""Launch Isaac Sim Simulator first."""

import argparse
import sys
import os
import torch
import torchvision
import gymnasium as gym
import time
from datetime import datetime

from isaaclab.app import AppLauncher

# local imports
import cli_args  # isort: skip

# add argparse arguments
parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.")
parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
parser.add_argument("--video_length", type=int, default=2000, help="Length of the recorded video (in steps).")
parser.add_argument(
    "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
)
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
parser.add_argument(
    "--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point."
)
parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
parser.add_argument(
    "--use_pretrained_checkpoint",
    action="store_true",
    help="Use the pre-trained checkpoint from Nucleus.",
)
parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.")
# append RSL-RL cli arguments
cli_args.add_rsl_rl_args(parser)
# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
# parse the arguments
args_cli, hydra_args = parser.parse_known_args()
# always enable cameras to record video
if args_cli.video:
    args_cli.enable_cameras = True

# clear out sys.argv for Hydra
sys.argv = [sys.argv[0]] + hydra_args

# launch omniverse app
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

"""Rest everything follows."""

from rsl_rl.runners import DistillationRunner, OnPolicyRunner

from isaaclab.envs import (
    DirectMARLEnv,
    DirectMARLEnvCfg,
    DirectRLEnvCfg,
    ManagerBasedRLEnvCfg,
    multi_agent_to_single_agent,
)
from isaaclab.utils.assets import retrieve_file_path
from isaaclab.utils.dict import print_dict

from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg, RslRlVecEnvWrapper, export_policy_as_jit, export_policy_as_onnx
from isaaclab_rl.utils.pretrained_checkpoint import get_published_pretrained_checkpoint

import isaaclab_tasks  # noqa: F401
from isaaclab_tasks.utils import get_checkpoint_path
from isaaclab_tasks.utils.hydra import hydra_task_config

import mindbot.tasks  # noqa: F401

# ==============================================================================
#  [Update] 多相机、多环境录制类
# ==============================================================================
# 1. 确保文件顶部有以下导入
import imageio
import numpy as np

"""use gpu compute to record video"""

# class MultiCameraRecorder:
#     def __init__(self, env, camera_names: list[str], env_indices: list[int], output_dir: str, fps: int = 30):
#         self.env = env
#         self.camera_names = camera_names
#         self.env_indices = env_indices
#         self.output_dir = output_dir
#         self.fps = fps
#         self.frames = {cam_name: {env_idx: [] for env_idx in env_indices} for cam_name in camera_names}
#         os.makedirs(self.output_dir, exist_ok=True)
        
#         self.cameras = {}
#         for name in camera_names:
#             if name in self.env.unwrapped.scene.keys():
#                 self.cameras[name] = self.env.unwrapped.scene[name]
#                 print(f"[INFO] Camera {name} linked.")

#     def record_step(self):
#         """保持在 GPU 上克隆数据"""
#         for cam_name, camera_obj in self.cameras.items():
#             # 获取数据前强制同步一次（防止后端丢失）
#             rgb_data = camera_obj.data.output["rgb"] 
            
#             for env_idx in self.env_indices:
#                 # 使用 .clone() 保持在 GPU，但要注意显存占用
#                 self.frames[cam_name][env_idx].append(rgb_data[env_idx].clone())

#     def save_videos(self, filename_suffix=""):
#         print(f"[INFO] Saving videos from GPU to Disk...")
#         for cam_name, env_dict in self.frames.items():
#             for env_idx, frame_list in env_dict.items():
#                 if not frame_list: continue
                
#                 # 转换为 torchvision 格式 (T, C, H, W)
#                 video_tensor = torch.stack(frame_list)
#                 if video_tensor.shape[-1] == 4: # RGBA -> RGB
#                     video_tensor = video_tensor[..., :3]
                
#                 # 移动到 CPU 并保存
#                 video_cpu = video_tensor.cpu()
#                 output_path = os.path.join(self.output_dir, f"{cam_name}_env{env_idx}_{filename_suffix}.mp4")
                
#                 # 使用 torchvision 写入 (T, H, W, C)
#                 torchvision.io.write_video(output_path, video_cpu, fps=self.fps)
                
#                 # 【关键】保存后立即释放显存
#                 del video_tensor
#                 del video_cpu
#                 frame_list.clear() 
        
#         torch.cuda.empty_cache()

"""use cpu compute to record video"""
# # 2. 修改 MultiCameraRecorder 类
class MultiCameraRecorder:
    def __init__(self, env, camera_names: list[str], env_indices: list[int], output_dir: str, fps: int = 30):
        self.env = env
        self.camera_names = camera_names
        self.env_indices = env_indices
        self.output_dir = output_dir
        self.fps = fps
        self.frames = {cam_name: {env_idx: [] for env_idx in env_indices} for cam_name in camera_names}
        os.makedirs(self.output_dir, exist_ok=True)
        self.cameras = {}
        for name in camera_names:
            try:
                self.cameras[name] = self.env.unwrapped.scene[name]
                print(f"[INFO][MultiCameraRecorder] Found camera: {name}")
            except KeyError:
                print(f"[WARN][MultiCameraRecorder] Camera '{name}' not found!")

    def record_step(self):
        """在每个仿真步调用"""
        for cam_name, camera_obj in self.cameras.items():
            # [关键修改] 获取数据前先确保数据已同步
            # 这可以防止读取到正在渲染中的内存导致 access violation
            rgb_data = camera_obj.data.output["rgb"]
            
            for env_idx in self.env_indices:
                if env_idx >= rgb_data.shape[0]: continue
                
                # 转换为 CPU 上的 numpy，这种方式通常比 torchvision 的 tensor 堆叠更稳
                frame = rgb_data[env_idx].clone().detach().cpu().numpy()
                self.frames[cam_name][env_idx].append(frame)

    def save_videos(self, filename_suffix=""):
        """循环结束后调用"""
        print(f"[INFO][MultiCameraRecorder] Saving videos...")
        for cam_name, env_dict in self.frames.items():
            for env_idx, frame_list in env_dict.items():
                if not frame_list: continue
                
                print(f"  -> Saving {cam_name} (Env {env_idx})...")
                
                # 处理格式并使用 imageio 保存
                processed_frames = []
                for img in frame_list:
                    # [0, 1] -> [0, 255]
                    if img.dtype != np.uint8:
                        if img.max() <= 1.01: img = (img * 255).astype(np.uint8)
                        else: img = img.astype(np.uint8)
                    # 去掉 Alpha 通道
                    if img.shape[-1] == 4: img = img[:, :, :3]
                    processed_frames.append(img)

                fname = f"{cam_name}_env{env_idx}_{filename_suffix}.mp4"
                output_path = os.path.join(self.output_dir, fname)
                
                try:
                    # 使用 imageio 写入视频
                    imageio.mimsave(output_path, processed_frames, fps=self.fps)
                    print(f"     Saved: {output_path}")
                except Exception as e:
                    print(f"     [ERROR] Failed to save {fname}: {e}")



# class MultiCameraRecorder:
#     """
#     用于从 Isaac Lab 环境中录制多个相机的图像数据，并支持多个环境实例。
#     """
#     def __init__(self, env, camera_names: list[str], env_indices: list[int], output_dir: str, fps: int = 30):
#         self.env = env
#         self.camera_names = camera_names
#         self.env_indices = env_indices # 要录制的环境索引列表，例如 [0, 1]
#         self.output_dir = output_dir
#         self.fps = fps
        
#         # 数据结构: self.frames[camera_name][env_idx] = [list of tensors]
#         self.frames = {
#             cam_name: {env_idx: [] for env_idx in env_indices} 
#             for cam_name in camera_names
#         }

#         os.makedirs(self.output_dir, exist_ok=True)

#         # 获取场景中的相机对象
#         self.cameras = {}
#         for name in camera_names:
#             try:
#                 # 尝试获取相机句柄
#                 self.cameras[name] = self.env.unwrapped.scene[name]
#                 print(f"[INFO][MultiCameraRecorder] Found camera: {name}")
#             except KeyError:
#                 print(f"[WARN][MultiCameraRecorder] Camera '{name}' not found in scene! Skipping.")

#     def record_step(self):
#         """在每个仿真步调用"""
#         for cam_name, camera_obj in self.cameras.items():
#             # 获取 RGB 数据 (Shape: [num_envs, H, W, C])
#             # 注意：如果数据类型是 RGBA，后面处理
#             rgb_data = camera_obj.data.output["rgb"]

#             for env_idx in self.env_indices:
#                 # 边界检查
#                 if env_idx >= rgb_data.shape[0]:
#                     continue

#                 # [关键] 立即 clone 并转到 CPU。
#                 # 如果不转 CPU，4个相机 * 长时间录制会迅速耗尽 GPU 显存。
#                 frame = rgb_data[env_idx].clone().detach().cpu()
                
#                 self.frames[cam_name][env_idx].append(frame)

#     def save_videos(self, filename_suffix=""):
#         """循环结束后调用，保存所有视频"""
#         print(f"[INFO][MultiCameraRecorder] Saving videos for {len(self.cameras)} cameras...")
        
#         for cam_name, env_dict in self.frames.items():
#             for env_idx, frame_list in env_dict.items():
#                 if not frame_list:
#                     continue
                
#                 print(f"  -> Processing {cam_name} (Env {env_idx}) with {len(frame_list)} frames...")
                
#                 # 堆叠 (T, H, W, C)
#                 video_tensor = torch.stack(frame_list)

#                 # 去除 Alpha 通道
#                 if video_tensor.shape[-1] == 4:
#                     video_tensor = video_tensor[..., :3]

#                 # (T, H, W, C) -> (T, C, H, W)
#                 video_tensor = video_tensor.permute(0, 3, 1, 2)

#                 # 确保是 uint8
#                 if video_tensor.dtype != torch.uint8:
#                     if video_tensor.max() <= 1.0:
#                         video_tensor = (video_tensor * 255).to(torch.uint8)
#                     else:
#                         video_tensor = video_tensor.to(torch.uint8)

#                 # 文件名: left_hand_camera_env0_2024-xx-xx.mp4
#                 fname = f"{cam_name}_env{env_idx}_{filename_suffix}.mp4"
#                 output_path = os.path.join(self.output_dir, fname)
                
#                 torchvision.io.write_video(output_path, video_tensor, fps=self.fps)
#                 print(f"     Saved: {output_path}")

# ==============================================================================


@hydra_task_config(args_cli.task, args_cli.agent)
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg):
    """Play with RSL-RL agent."""
    # grab task name for checkpoint path
    task_name = args_cli.task.split(":")[-1]
    train_task_name = task_name.replace("-Play", "")

    # override configurations with non-hydra CLI arguments
    agent_cfg: RslRlBaseRunnerCfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
    env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs

    # set the environment seed
    env_cfg.seed = agent_cfg.seed
    env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device

    # specify directory for logging experiments
    log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
    log_root_path = os.path.abspath(log_root_path)
    print(f"[INFO] Loading experiment from directory: {log_root_path}")
    if args_cli.use_pretrained_checkpoint:
        resume_path = get_published_pretrained_checkpoint("rsl_rl", train_task_name)
        if not resume_path:
            print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.")
            return
    elif args_cli.checkpoint:
        resume_path = retrieve_file_path(args_cli.checkpoint)
    else:
        resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)

    log_dir = os.path.dirname(resume_path)

    # set the log directory for the environment (works for all environment types)
    env_cfg.log_dir = log_dir

    # create isaac environment
    env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)

    # convert to single-agent instance if required by the RL algorithm
    if isinstance(env.unwrapped, DirectMARLEnv):
        env = multi_agent_to_single_agent(env)

    # wrap for video recording (standard Gym recording - Viewport)
    if args_cli.video:
        video_kwargs = {
            "video_folder": os.path.join(log_dir, "videos", "play"),
            "step_trigger": lambda step: step == 0,
            "video_length": args_cli.video_length,
            "disable_logger": True,
        }
        print("[INFO] Recording videos during training.")
        print_dict(video_kwargs, nesting=4)
        env = gym.wrappers.RecordVideo(env, **video_kwargs)

    # wrap around environment for rsl-rl
    env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)

    print(f"[INFO]: Loading model checkpoint from: {resume_path}")
    # load previously trained model
    if agent_cfg.class_name == "OnPolicyRunner":
        runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
    elif agent_cfg.class_name == "DistillationRunner":
        runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
    else:
        raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}")
    runner.load(resume_path)

    # obtain the trained policy for inference
    policy = runner.get_inference_policy(device=env.unwrapped.device)

    # extract the neural network module
    try:
        policy_nn = runner.alg.policy
    except AttributeError:
        policy_nn = runner.alg.actor_critic

    # extract the normalizer
    if hasattr(policy_nn, "actor_obs_normalizer"):
        normalizer = policy_nn.actor_obs_normalizer
    elif hasattr(policy_nn, "student_obs_normalizer"):
        normalizer = policy_nn.student_obs_normalizer
    else:
        normalizer = None

    # export policy to onnx/jit
    export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
    export_policy_as_jit(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.pt")
    export_policy_as_onnx(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.onnx")

    dt = env.unwrapped.step_dt

    # ==========================================================================
    # 初始化多相机录制器
    # ==========================================================================
    recorder = None
    if args_cli.video:
        # 1. 定义要录制的相机列表（需与 mindbot_env_cfg.py 中的命名一致）
        target_cameras = [
            "left_hand_camera",
            "right_hand_camera",
            "head_camera",
            "chest_camera"
        ]
        
        # 2. 定义要录制的环境索引 (env_nums)
        # 默认只录制第 0 个环境，如果要录制前 2 个，改为 [0, 1]
        # 注意：不要录制太多环境，否则写入视频时非常慢且占内存
        envs_to_record = [0,1,2,3]
        # 如果你想录制所有环境（不推荐，除非num_envs很小）：
        # envs_to_record = list(range(env.unwrapped.num_envs))

        save_dir = os.path.join(log_dir, "robot_camera_recordings")
        
        recorder = MultiCameraRecorder(
            env=env, 
            camera_names=target_cameras, 
            env_indices=envs_to_record,
            output_dir=save_dir, 
            fps=int(1/dt)
        )

    # reset environment
    obs = env.get_observations()
    timestep = 0
    
    print("[INFO] Starting simulation loop...")
    
    # simulate environment
    while simulation_app.is_running():
        start_time = time.time()
        # run everything in inference mode
        with torch.inference_mode():
            # agent stepping
            actions = policy(obs)
            # env stepping
            obs, _, dones, _ = env.step(actions)
            
            # ==================================================================
            # 每一帧录制数据
            # ==================================================================
            if recorder:
                recorder.record_step()

            # reset recurrent states for episodes that have terminated
            policy_nn.reset(dones)

        if args_cli.video:
            timestep += 1
            # Exit the play loop after recording one video
            if timestep == args_cli.video_length:
                # ==============================================================
                # 循环结束时保存所有视频
                # ==============================================================
                if recorder:
                    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
                    recorder.save_videos(filename_suffix=timestamp)
                break

        # time delay for real-time evaluation
        sleep_time = dt - (time.time() - start_time)
        if args_cli.real_time and sleep_time > 0:
            time.sleep(sleep_time)

    # close the simulator
    env.close()


if __name__ == "__main__":
    # run the main function
    main()
    # close sim app
    simulation_app.close()

# # Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# # All rights reserved.
# #
# # SPDX-License-Identifier: BSD-3-Clause

# """Script to play a checkpoint if an RL agent from RSL-RL."""

# """Launch Isaac Sim Simulator first."""

# import argparse
# import sys
# import os
# import torch
# import torchvision
# import imageio
# import numpy as np
# import gymnasium as gym
# import time
# from datetime import datetime

# from isaaclab.app import AppLauncher

# # local imports
# import cli_args  # isort: skip

# # add argparse arguments
# parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.")
# parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
# parser.add_argument("--video_length", type=int, default=2000, help="Length of the recorded video (in steps).")
# parser.add_argument(
#     "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
# )
# parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
# parser.add_argument("--task", type=str, default=None, help="Name of the task.")
# parser.add_argument(
#     "--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point."
# )
# parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
# parser.add_argument(
#     "--use_pretrained_checkpoint",
#     action="store_true",
#     help="Use the pre-trained checkpoint from Nucleus.",
# )
# parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.")
# # append RSL-RL cli arguments
# cli_args.add_rsl_rl_args(parser)
# # append AppLauncher cli args
# AppLauncher.add_app_launcher_args(parser)
# # parse the arguments
# args_cli, hydra_args = parser.parse_known_args()
# # always enable cameras to record video
# if args_cli.video:
#     args_cli.enable_cameras = True

# # clear out sys.argv for Hydra
# sys.argv = [sys.argv[0]] + hydra_args

# # launch omniverse app
# app_launcher = AppLauncher(args_cli)
# simulation_app = app_launcher.app

# """Rest everything follows."""

# from rsl_rl.runners import DistillationRunner, OnPolicyRunner

# from isaaclab.envs import (
#     DirectMARLEnv,
#     DirectMARLEnvCfg,
#     DirectRLEnvCfg,
#     ManagerBasedRLEnvCfg,
#     multi_agent_to_single_agent,
# )
# from isaaclab.utils.assets import retrieve_file_path
# from isaaclab.utils.dict import print_dict

# from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg, RslRlVecEnvWrapper, export_policy_as_jit, export_policy_as_onnx
# from isaaclab_rl.utils.pretrained_checkpoint import get_published_pretrained_checkpoint

# import isaaclab_tasks  # noqa: F401
# from isaaclab_tasks.utils import get_checkpoint_path
# from isaaclab_tasks.utils.hydra import hydra_task_config

# import mindbot.tasks  # noqa: F401

# # ==============================================================================
# #  新增：自定义相机录制类
# # ==============================================================================
# class CameraRecorder:
#     """
#     用于从 Isaac Lab 环境中录制指定相机的图像数据并保存为视频。
#     """
#     def __init__(self, env, camera_name: str, output_dir: str, fps: int = 30):
#         self.env = env
#         self.camera_name = camera_name
#         self.output_dir = output_dir
#         self.fps = fps
#         self.frames = [] # 用于存储每一帧的 Tensor

#         # 确保输出目录存在
#         os.makedirs(self.output_dir, exist_ok=True)

#         # 尝试从场景中获取相机对象
#         # env.unwrapped 访问原始的 ManagerBasedRLEnv
#         try:
#             self.camera = self.env.unwrapped.scene[camera_name]
#             print(f"[INFO][CameraRecorder] Successfully found camera: {camera_name}")
#         except KeyError:
#             print(f"[ERROR][CameraRecorder] Camera '{camera_name}' not found in scene entities!")
#             print(f"Available entities: {self.env.unwrapped.scene.keys()}")
#             raise

#     def record_step(self):
#         """在每个仿真步调用，获取图像数据"""
#         # 获取 RGB 数据。注意：Isaac Lab 的输出通常是 (num_envs, H, W, C)
#         # 我们这里只取第一个环境 (env_index=0)
#         # output["rgb"] 可能是 RGBA 或 RGB
#         rgb_tensor = self.camera.data.output["rgb"][0] 

#         # 如果数据在 GPU 上，克隆一份并 detach，防止被后续步骤覆盖
#         # 此时先不转 CPU，以免拖慢仿真循环，最后保存时再转
#         self.frames.append(rgb_tensor.clone().detach())

#     def save_video(self, filename="robot_camera_video.mp4"):
#         """循环结束后调用，将内存中的帧保存为视频"""
#         if not self.frames:
#             print("[WARN][CameraRecorder] No frames recorded.")
#             return

#         print(f"[INFO][CameraRecorder] Saving {len(self.frames)} frames to video...")
        
#         # 将 tensor 转换为 numpy 格式
#         video_frames = []
#         for frame in self.frames:
#             # 转换为 CPU 上的 numpy 数组 (H, W, C)
#             img = frame.cpu().numpy()
#             # 处理归一化数据 [0, 1] -> [0, 255]
#             if img.dtype != np.uint8:
#                 if img.max() <= 1.01:
#                     img = (img * 255).astype(np.uint8)
#                 else:
#                     img = img.astype(np.uint8)
#             # 处理 RGBA -> RGB
#             if img.shape[-1] == 4:
#                 img = img[:, :, :3]
#             video_frames.append(img)

#         output_path = os.path.join(self.output_dir, filename)
        
#         # 使用 imageio 写入视频，这种方式在 Windows 上更稳定
#         try:
#             imageio.mimsave(output_path, video_frames, fps=self.fps)
#             print(f"[INFO][CameraRecorder] Video saved to: {output_path}")
#         except Exception as e:
#             print(f"[ERROR][CameraRecorder] Failed to save video using imageio: {e}")
#             # 如果 imageio 也失败，尝试最后的备选方案
#             print("[INFO] Attempting to save frames as images instead...")
#             # ... (可选：保存为一系列图片)
#     # def save_video(self, filename="robot_camera_video.mp4"):
#     #     """循环结束后调用，将内存中的帧保存为视频"""
#     #     if not self.frames:
#     #         print("[WARN][CameraRecorder] No frames recorded.")
#     #         return

#     #     print(f"[INFO][CameraRecorder] Saving {len(self.frames)} frames to video...")
        
#     #     # 堆叠帧 -> (T, H, W, C)
#     #     video_tensor = torch.stack(self.frames)

#     #     # 处理 Alpha 通道：如果是 RGBA (4通道)，只取前3个 RGB
#     #     if video_tensor.shape[-1] == 4:
#     #         video_tensor = video_tensor[..., :3]

#     #     # 转换为 (T, C, H, W) 以符合 torchvision 的输入要求
#     #     # 原始: (T, H, W, C) -> Permute -> (T, C, H, W)
#     #     video_tensor = video_tensor.permute(0, 3, 1, 2)

#     #     # 确保是 uint8 类型 [0, 255]
#     #     if video_tensor.dtype != torch.uint8:
#     #          # 如果是 float [0,1]，则 * 255 并转 uint8
#     #         if video_tensor.max() <= 1.0:
#     #             video_tensor = (video_tensor * 255).to(torch.uint8)
#     #         else:
#     #             video_tensor = video_tensor.to(torch.uint8)

#     #     # 移动到 CPU
#     #     video_tensor = video_tensor.cpu()

#     #     output_path = os.path.join(self.output_dir, filename)
        
#     #     # 使用 torchvision 保存视频
#     #     torchvision.io.write_video(output_path, video_tensor, fps=self.fps)
#     #     print(f"[INFO][CameraRecorder] Video saved to: {output_path}")

# # ==============================================================================


# @hydra_task_config(args_cli.task, args_cli.agent)
# def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg):
#     """Play with RSL-RL agent."""
#     # grab task name for checkpoint path
#     task_name = args_cli.task.split(":")[-1]
#     train_task_name = task_name.replace("-Play", "")

#     # override configurations with non-hydra CLI arguments
#     agent_cfg: RslRlBaseRunnerCfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
#     env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs

#     # set the environment seed
#     env_cfg.seed = agent_cfg.seed
#     env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device

#     # specify directory for logging experiments
#     log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
#     log_root_path = os.path.abspath(log_root_path)
#     print(f"[INFO] Loading experiment from directory: {log_root_path}")
#     if args_cli.use_pretrained_checkpoint:
#         resume_path = get_published_pretrained_checkpoint("rsl_rl", train_task_name)
#         if not resume_path:
#             print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.")
#             return
#     elif args_cli.checkpoint:
#         resume_path = retrieve_file_path(args_cli.checkpoint)
#     else:
#         resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)

#     log_dir = os.path.dirname(resume_path)

#     # set the log directory for the environment (works for all environment types)
#     env_cfg.log_dir = log_dir

#     # create isaac environment
#     env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)

#     # convert to single-agent instance if required by the RL algorithm
#     if isinstance(env.unwrapped, DirectMARLEnv):
#         env = multi_agent_to_single_agent(env)

#     # wrap for video recording (standard Gym recording)
#     # 注意：这只是录制 Viewport 的画面，不是你机器人相机的画面
#     if args_cli.video:
#         video_kwargs = {
#             "video_folder": os.path.join(log_dir, "videos", "play"),
#             "step_trigger": lambda step: step == 0,
#             "video_length": args_cli.video_length,
#             "disable_logger": True,
#         }
#         print("[INFO] Recording videos during training.")
#         print_dict(video_kwargs, nesting=4)
#         env = gym.wrappers.RecordVideo(env, **video_kwargs)

#     # wrap around environment for rsl-rl
#     env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)

#     print(f"[INFO]: Loading model checkpoint from: {resume_path}")
#     # load previously trained model
#     if agent_cfg.class_name == "OnPolicyRunner":
#         runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
#     elif agent_cfg.class_name == "DistillationRunner":
#         runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
#     else:
#         raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}")
#     runner.load(resume_path)

#     # obtain the trained policy for inference
#     policy = runner.get_inference_policy(device=env.unwrapped.device)

#     # extract the neural network module
#     try:
#         policy_nn = runner.alg.policy
#     except AttributeError:
#         policy_nn = runner.alg.actor_critic

#     # extract the normalizer
#     if hasattr(policy_nn, "actor_obs_normalizer"):
#         normalizer = policy_nn.actor_obs_normalizer
#     elif hasattr(policy_nn, "student_obs_normalizer"):
#         normalizer = policy_nn.student_obs_normalizer
#     else:
#         normalizer = None

#     # export policy to onnx/jit
#     export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
#     export_policy_as_jit(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.pt")
#     export_policy_as_onnx(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.onnx")

#     dt = env.unwrapped.step_dt

#     # ==========================================================================
#     # 初始化自定义相机录制器
#     # ==========================================================================
#     recorder = None
#     if args_cli.video:
#         # 这里的 output_dir 设为日志目录下的 hand_camera_videos
#         save_dir = os.path.join(log_dir, "hand_camera_videos")
#         # 这里的 "left_hand_camera" 必须对应 mindbot_env_cfg.py 中定义的名称
#         recorder = CameraRecorder(env, "left_hand_camera", save_dir, fps=int(1/dt))

#     # reset environment
#     obs = env.get_observations()
#     timestep = 0
    
#     print("[INFO] Starting simulation loop...")
    
#     # simulate environment
#     while simulation_app.is_running():
#         start_time = time.time()
#         # run everything in inference mode
#         with torch.inference_mode():
#             # agent stepping
#             actions = policy(obs)
#             # env stepping
#             obs, _, dones, _ = env.step(actions)
            
#             # ==================================================================
#             # 每一帧录制数据
#             # ==================================================================
#             if recorder:
#                 recorder.record_step()

#             # reset recurrent states for episodes that have terminated
#             policy_nn.reset(dones)

#         if args_cli.video:
#             timestep += 1
#             # Exit the play loop after recording one video
#             if timestep == args_cli.video_length:
#                 # ==============================================================
#                 # 循环结束时保存视频
#                 # ==============================================================
#                 if recorder:
#                     timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
#                     recorder.save_video(filename=f"hand_cam_{timestamp}.mp4")
#                 break

#         # time delay for real-time evaluation
#         sleep_time = dt - (time.time() - start_time)
#         if args_cli.real_time and sleep_time > 0:
#             time.sleep(sleep_time)

#     # close the simulator
#     env.close()


# if __name__ == "__main__":
#     # run the main function
#     main()
#     # close sim app
#     simulation_app.close()


# # # Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# # # All rights reserved.
# # #
# # # SPDX-License-Identifier: BSD-3-Clause

# # """Script to play a checkpoint if an RL agent from RSL-RL."""

# # """Launch Isaac Sim Simulator first."""

# # import argparse
# # import sys

# # from isaaclab.app import AppLauncher

# # # local imports
# # import cli_args  # isort: skip

# # # add argparse arguments
# # parser = argparse.ArgumentParser(description="Train an RL agent with RSL-RL.")
# # parser.add_argument("--video", action="store_true", default=False, help="Record videos during training.")
# # parser.add_argument("--video_length", type=int, default=200, help="Length of the recorded video (in steps).")
# # parser.add_argument(
# #     "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
# # )
# # parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
# # parser.add_argument("--task", type=str, default=None, help="Name of the task.")
# # parser.add_argument(
# #     "--agent", type=str, default="rsl_rl_cfg_entry_point", help="Name of the RL agent configuration entry point."
# # )
# # parser.add_argument("--seed", type=int, default=None, help="Seed used for the environment")
# # parser.add_argument(
# #     "--use_pretrained_checkpoint",
# #     action="store_true",
# #     help="Use the pre-trained checkpoint from Nucleus.",
# # )
# # parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.")
# # # append RSL-RL cli arguments
# # cli_args.add_rsl_rl_args(parser)
# # # append AppLauncher cli args
# # AppLauncher.add_app_launcher_args(parser)
# # # parse the arguments
# # args_cli, hydra_args = parser.parse_known_args()
# # # always enable cameras to record video
# # if args_cli.video:
# #     args_cli.enable_cameras = True

# # # clear out sys.argv for Hydra
# # sys.argv = [sys.argv[0]] + hydra_args

# # # launch omniverse app
# # app_launcher = AppLauncher(args_cli)
# # simulation_app = app_launcher.app

# # """Rest everything follows."""

# # import gymnasium as gym
# # import os
# # import time
# # import torch

# # from rsl_rl.runners import DistillationRunner, OnPolicyRunner

# # from isaaclab.envs import (
# #     DirectMARLEnv,
# #     DirectMARLEnvCfg,
# #     DirectRLEnvCfg,
# #     ManagerBasedRLEnvCfg,
# #     multi_agent_to_single_agent,
# # )
# # from isaaclab.utils.assets import retrieve_file_path
# # from isaaclab.utils.dict import print_dict

# # from isaaclab_rl.rsl_rl import RslRlBaseRunnerCfg, RslRlVecEnvWrapper, export_policy_as_jit, export_policy_as_onnx
# # from isaaclab_rl.utils.pretrained_checkpoint import get_published_pretrained_checkpoint

# # import isaaclab_tasks  # noqa: F401
# # from isaaclab_tasks.utils import get_checkpoint_path
# # from isaaclab_tasks.utils.hydra import hydra_task_config

# # import mindbot.tasks  # noqa: F401


# # @hydra_task_config(args_cli.task, args_cli.agent)
# # def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlBaseRunnerCfg):
# #     """Play with RSL-RL agent."""
# #     # grab task name for checkpoint path
# #     task_name = args_cli.task.split(":")[-1]
# #     train_task_name = task_name.replace("-Play", "")

# #     # override configurations with non-hydra CLI arguments
# #     agent_cfg: RslRlBaseRunnerCfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
# #     env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs

# #     # set the environment seed
# #     # note: certain randomizations occur in the environment initialization so we set the seed here
# #     env_cfg.seed = agent_cfg.seed
# #     env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device

# #     # specify directory for logging experiments
# #     log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
# #     log_root_path = os.path.abspath(log_root_path)
# #     print(f"[INFO] Loading experiment from directory: {log_root_path}")
# #     if args_cli.use_pretrained_checkpoint:
# #         resume_path = get_published_pretrained_checkpoint("rsl_rl", train_task_name)
# #         if not resume_path:
# #             print("[INFO] Unfortunately a pre-trained checkpoint is currently unavailable for this task.")
# #             return
# #     elif args_cli.checkpoint:
# #         resume_path = retrieve_file_path(args_cli.checkpoint)
# #     else:
# #         resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)

# #     log_dir = os.path.dirname(resume_path)

# #     # set the log directory for the environment (works for all environment types)
# #     env_cfg.log_dir = log_dir

# #     # create isaac environment
# #     env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)

# #     # convert to single-agent instance if required by the RL algorithm
# #     if isinstance(env.unwrapped, DirectMARLEnv):
# #         env = multi_agent_to_single_agent(env)

# #     # wrap for video recording
# #     if args_cli.video:
# #         video_kwargs = {
# #             "video_folder": os.path.join(log_dir, "videos", "play"),
# #             "step_trigger": lambda step: step == 0,
# #             "video_length": args_cli.video_length,
# #             "disable_logger": True,
# #         }
# #         print("[INFO] Recording videos during training.")
# #         print_dict(video_kwargs, nesting=4)
# #         env = gym.wrappers.RecordVideo(env, **video_kwargs)

# #     # wrap around environment for rsl-rl
# #     env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)

# #     print(f"[INFO]: Loading model checkpoint from: {resume_path}")
# #     # load previously trained model
# #     if agent_cfg.class_name == "OnPolicyRunner":
# #         runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
# #     elif agent_cfg.class_name == "DistillationRunner":
# #         runner = DistillationRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
# #     else:
# #         raise ValueError(f"Unsupported runner class: {agent_cfg.class_name}")
# #     runner.load(resume_path)

# #     # obtain the trained policy for inference
# #     policy = runner.get_inference_policy(device=env.unwrapped.device)

# #     # extract the neural network module
# #     # we do this in a try-except to maintain backwards compatibility.
# #     try:
# #         # version 2.3 onwards
# #         policy_nn = runner.alg.policy
# #     except AttributeError:
# #         # version 2.2 and below
# #         policy_nn = runner.alg.actor_critic

# #     # extract the normalizer
# #     if hasattr(policy_nn, "actor_obs_normalizer"):
# #         normalizer = policy_nn.actor_obs_normalizer
# #     elif hasattr(policy_nn, "student_obs_normalizer"):
# #         normalizer = policy_nn.student_obs_normalizer
# #     else:
# #         normalizer = None

# #     # export policy to onnx/jit
# #     export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
# #     export_policy_as_jit(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.pt")
# #     export_policy_as_onnx(policy_nn, normalizer=normalizer, path=export_model_dir, filename="policy.onnx")

# #     dt = env.unwrapped.step_dt

# #     # reset environment
# #     obs = env.get_observations()
# #     timestep = 0
# #     # simulate environment
# #     while simulation_app.is_running():
# #         start_time = time.time()
# #         # run everything in inference mode
# #         with torch.inference_mode():
# #             # agent stepping
# #             actions = policy(obs)
# #             # env stepping
# #             obs, _, dones, _ = env.step(actions)
# #             # reset recurrent states for episodes that have terminated
# #             policy_nn.reset(dones)
# #         if args_cli.video:
# #             timestep += 1
# #             # Exit the play loop after recording one video
# #             if timestep == args_cli.video_length:
# #                 break

# #         # time delay for real-time evaluation
# #         sleep_time = dt - (time.time() - start_time)
# #         if args_cli.real_time and sleep_time > 0:
# #             time.sleep(sleep_time)

# #     # close the simulator
# #     env.close()


# # if __name__ == "__main__":
# #     # run the main function
# #     main()
# #     # close sim app
# #     simulation_app.close()
