1025 lines
41 KiB
Python
1025 lines
41 KiB
Python
# 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()
|