robot_lab中rsl_rl的replay_amp_data.py简洁解析

引言

这个replay_amp_data.py脚本是用于在Isaac Sim中可视化训练好的RL策略或预存运动数据的播放工具。让我们逐步解析其主要功能：

核心功能

策略可视化：加载训练好的RL策略检查点，在仿真环境中直观展示机器人运动运动数据回放：通过AMP模块回放预存参考运动轨迹摄像头跟随：自动调整视角跟踪机器人运动性能优化：支持无头模式(Headless)和实时模式(real-time)

1. 脚本架构概览

1.1 模块依赖关系

• Isaac Sim核心模块: AppLauncher, parse_env_cfg
• 强化学习框架: gymnasium, rsl_rl
• 数据处理: torch, numpy
• 机器人控制: robot_lab.tasks
• 工具函数: quat_rotate, reorder_from_isaacgym_to_isaacsim_tool

1.2 执行流程

graph TDA[启动Isaac Sim] --> B[解析配置]B --> C[创建仿真环境]C --> D[初始化AMP数据加载器]D --> E{仿真运行?}E --> |Yes| F[加载运动数据帧]F --> G[更新机器人状态]G --> H[环境步进]H --> I[更新摄像头]I --> EE --> |No| J[关闭环境]

2. 核心功能模块解析

2.1 Isaac Sim环境初始化

（1）启动Omniverse应用

app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

功能：初始化Isaac Sim的Omniverse内核关键技术：通过AppLauncher处理图形渲染后端支持Headless模式（无GUI）和实时模式自动加载USD场景文件

2.2 环境配置解析

env_cfg = parse_env_cfg(args_cli.task, device=args_cli.device,num_envs=args_cli.num_envs,use_fabric=not args_cli.disable_fabric
)

参数调整策略：强制设置num_envs=1以降低计算负载禁用地形生成器的课程学习(curriculum=False)关闭观测随机化(enable_corruption=False)

2.3 AMP运动数据加载

frames = env.unwrapped.amp_loader.get_full_frame_at_time_batch(np.array([traj_idx]), np.array([t])

数据结构：traj_idx: 轨迹索引（不同运动片段）t: 时间戳（基于仿真步长dt）关键方法：get_root_pos_batch(): 根关节位置（x,y,z）get_root_rot_batch(): 根关节旋转（四元数xyzw格式）get_joint_pose_batch(): 各关节位置

2.4 关节顺序转换

joint_pos = reorder_from_isaacgym_to_isaacsim_tool(joint_pos)

背景：Isaac Gym使用深度优先关节排序（父关节→子关节）Isaac Sim使用广度优先排序（同层级关节连续排列）

转换示例：

# 原始顺序: [hip0, thigh0, calf0, hip1, thigh1, calf1,...]
# 转换后顺序: [hip0, hip1, hip2, hip3, thigh0, thigh1,...]

2.5 物理约束处理

joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1])
joint_vel.clamp_(-joint_vel_limits, joint_vel_limits)

约束来源：joint_pos_limits: 来自URDF定义的机械限位joint_vel_limits: 根据电机性能设定的最大速度工程意义：防止仿真中的超限导致学习到不可行的动作确保策略在sim-to-real时的安全性

3. 摄像头控制子系统

3.1 视角定位算法

lookat = [root_pos_x, root_pos_y, root_pos_z]
eye = [lookat[0]+2, lookat[1]+2, lookat[2]+2]
viewport_camera_controller.update_view_location(eye, lookat)

参数解析：lookat: 机器人根关节的世界坐标eye: 摄像机位置（相对于lookat的偏移量）运动平滑：每帧更新实现动态跟踪偏移量(2,2,2)提供45度俯视角

3.2 摄像机控制类方法

class ViewportCameraController:def set_view_env_index(self, env_index): ...def update_view_location(self, eye, lookat): ...

功能扩展：支持多环境视角切换（本脚本固定env_index=0）可扩展第三人称/第一人称视角切换

4. 关键参数详解

4.1 运动播放参数

4.2 性能调优参数

env_cfg.sim.dt = 0.0167  # 仿真步长(对应60Hz)
env_cfg.sim.render_interval = 1  # 渲染间隔(每物理步都渲染)

实时性保证：当--real-time启用时，仿真时钟与真实时间同步需要GPU渲染速度 > 60FPS

5. 调试与验证技巧

5.1 运动数据可视化

（1）添加调试代码查看足端位置

foot_pos_amp = amp_loader.get_tar_toe_pos_local_batch(frames)
print(foot_pos_amp[0])  # 输出第一个环境的足端局部坐标

5.2 状态同步验证

（2）对比仿真关节位置与AMP数据

sim_joint_pos = env.unwrapped.robot.data.joint_pos[0]
amp_joint_pos = joint_pos[0]
print(f"Sim: {sim_joint_pos}\nAMP: {amp_joint_pos}")

5.3 性能监控

运行前设置Torch性能分析

export TORCH_CPP_LOG_LEVEL=INFO
export TORCH_CUDA_APPTAINER=1

6. 工程实践建议

6.1 扩展功能

多视角支持：通过添加--camera_angle参数切换不同视角数据导出：添加--export_path参数保存运动状态数据交互控制：集成键盘事件监听实现手动控制

6.2 常见问题排查

6.3 性能优化

在循环外预分配内存

preallocated_frames = torch.empty((1, amp_loader.frame_size), device='cuda')
while ...:amp_loader.get_full_frame_at_time_batch(..., out=preallocated_frames)

该脚本作为AMP训练流程的终端验证环节，将算法输出与物理仿真深度整合，为研究人员提供了直观的策略评估手段。掌握其实现细节对调试运动控制策略、优化sim-to-real迁移效果具有重要意义。

总体代码

# Copyright (c) 2024-2025 Ziqi Fan
# SPDX-License-Identifier: Apache-2.0"""Script to play a checkpoint if an RL agent from RSL-RL.""""""Launch Isaac Sim Simulator first."""import argparse
import os
import sysfrom isaaclab.app import AppLauncher# local imports
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import cli_args# 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("--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)
args_cli = parser.parse_args()
# always enable cameras to record video
if args_cli.video:args_cli.enable_cameras = True# launch omniverse app
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app"""Rest everything follows."""import gymnasium as gym
import numpy as np
import torchfrom isaaclab.utils.math import quat_rotate
from isaaclab_rl.rsl_rl import RslRlVecEnvWrapper
from isaaclab_tasks.utils import parse_env_cfgimport robot_lab.tasks  # noqa: F401def main():"""Play with RSL-RL agent."""# parse configurationenv_cfg = parse_env_cfg(args_cli.task, device=args_cli.device, num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric)# make a smaller scene for playenv_cfg.scene.num_envs = 1# spawn the robot randomly in the grid (instead of their terrain levels)env_cfg.scene.terrain.max_init_terrain_level = None# reduce the number of terrains to save memoryif env_cfg.scene.terrain.terrain_generator is not None:env_cfg.scene.terrain.terrain_generator.num_rows = 1env_cfg.scene.terrain.terrain_generator.num_cols = 1env_cfg.scene.terrain.terrain_generator.curriculum = False# disable randomization for playenv_cfg.observations.policy.enable_corruption = False# remove random pushingenv_cfg.events.randomize_apply_external_force_torque = Noneenv_cfg.events.push_robot = Noneenv_cfg.amp_num_preload_transitions = 1env_cfg.amp_replay_buffer_size = 2# create isaac environmentenv = gym.make(args_cli.task, cfg=env_cfg)# wrap around environment for rsl-rlenv = RslRlVecEnvWrapper(env)# reset environmentobs, _ = env.get_observations()# simulate environmentwhile simulation_app.is_running():# run everything in inference modewith torch.inference_mode():env_ids = torch.tensor([0], device=env.unwrapped.device)t = 0.0traj_idx = 0while traj_idx < len(env.unwrapped.amp_loader.trajectory_lens) - 1:actions = torch.zeros((env_cfg.scene.num_envs, env.unwrapped.num_actions), device=env.unwrapped.device)if (t + env.unwrapped.amp_loader.time_between_frames + env_cfg.sim.dt * env_cfg.sim.render_interval) >= env.unwrapped.amp_loader.trajectory_lens[traj_idx]:traj_idx += 1t = 0else:t += env_cfg.sim.dt * env_cfg.sim.render_intervalframes = env.unwrapped.amp_loader.get_full_frame_at_time_batch(np.array([traj_idx]), np.array([t]))positions = env.unwrapped.amp_loader.get_root_pos_batch(frames)orientations = env.unwrapped.amp_loader.get_root_rot_batch(frames)  # xyzw# Func quat_rotate() and Isaacsim/IsaacLab all need wxyzorientations = torch.cat((orientations[:, -1].unsqueeze(1), orientations[:, :-1]), dim=1)lin_vel = quat_rotate(orientations, env.unwrapped.amp_loader.get_linear_vel_batch(frames))ang_vel = quat_rotate(orientations, env.unwrapped.amp_loader.get_angular_vel_batch(frames))velocities = torch.cat([lin_vel, ang_vel], dim=-1)env.unwrapped.robot.write_root_pose_to_sim(torch.cat([positions, orientations], dim=-1), env_ids=env_ids)env.unwrapped.robot.write_root_velocity_to_sim(velocities, env_ids=env_ids)joint_pos = env.unwrapped.amp_loader.get_joint_pose_batch(frames)joint_vel = env.unwrapped.amp_loader.get_joint_vel_batch(frames)# Isaac Sim uses breadth-first joint ordering, while Isaac Gym uses depth-first joint orderingjoint_pos = env.unwrapped.amp_loader.reorder_from_isaacgym_to_isaacsim_tool(joint_pos)joint_vel = env.unwrapped.amp_loader.reorder_from_isaacgym_to_isaacsim_tool(joint_vel)joint_pos_limits = env.unwrapped.robot.data.soft_joint_pos_limits[env_ids]joint_pos = joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1])joint_vel_limits = env.unwrapped.robot.data.soft_joint_vel_limits[env_ids]joint_vel = joint_vel.clamp_(-joint_vel_limits, joint_vel_limits)env.unwrapped.robot.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids)# print("---")# foot_pos_amp = env.unwrapped.amp_loader.get_tar_toe_pos_local_batch(frames)# print(env.unwrapped.get_amp_observations()[env_ids.item(), 12:24])# print(foot_pos_amp[0])# env steppingobs, _, _, _ = env.step(actions)# camera followenv.unwrapped.viewport_camera_controller.set_view_env_index(env_index=0)lookat = [env.unwrapped.robot.data.root_pos_w[env_ids.item(), i].cpu().item() for i in range(3)]eye_offset = [2, 2, 2]pairs = zip(lookat, eye_offset)eye = [x + y for x, y in pairs]env.unwrapped.viewport_camera_controller.update_view_location(eye, lookat)# close the simulatorenv.close()if __name__ == "__main__":# run the main functionmain()# close sim appsimulation_app.close()