本博客是基于mjlab rsl_rl的G1/H2 AMP 运动控制同时支持走/跑、跌倒恢复 训练开源代码https://github.com/ccrpRepo/AMP_mjlab部署集成位置部署集成代码位于 https://github.com/ccrpRepo/wbc_fsm 项目中的MJAmp State。其中AMP的算法思路、代码分析参考我这篇博客《VLA 系列》AMP 算法与代码分析技术示例效果跌到、恢复H2的效果项目核心特点⭐统一策略使用同一个policy同时学习locomotion走/跑与recovery跌倒恢复AMP 风格约束通过AMP discriminator约束动作风格与运动先验H2 新增适配原开源版本只支持G1可以通过修改任务配置、机器人资产、动作数据和训练入口等代码新增适配H2训练部署一致训练与导出链路保持一致支持直接导出ONNX policy参考项目unitreerobotics/unitree_rl_mjlab 项目的开源工作与启发。Open-X-Humanoid/TienKung-Lab本博客在rsl_rl的AMP部分参考了该实现。 目录 1. 核心思路 2. 环境要求 3. 下载代码 4. 创建并激活 conda 环境 5. 安装项目依赖⚡ 6. 安装指定版本的 PyTorch 7. 安装 Warp 和 mujoco-warp 兼容版本 8. 应用 mjlab 观测历史补丁️ 9. 修复 mjlab CUDA Graph 兼容问题 10. 查看可用任务 11. G1 训练 12. H2 训练 13. 评估与可视化 14. 运动数据准备️ 15. 核心代码说明⭐ 16. 总结1. 核心思路传统做法常把“走跑策略”和“恢复策略”分开训练并做切换本博客将两类能力放入同一个策略中统一学习。实现要点运动数据分组Walk/Run数据目录src/assets/motions/g1/amp/WalkandRunRecovery数据目录src/assets/motions/g1/amp/RecoveryH2数据目录src/assets/motions/h2/amp/WalkandRun与src/assets/motions/h2/amp/Recovery⏱️延迟重置机制Delayed Termination一部分环境在触发终止后不立即 reset而是给定恢复窗口该子集环境优先从Recovery片段采样 reset 状态统一 AMP 训练单一actor-critic 单一AMP discriminator在同一训练过程中学习速度跟踪、抗扰动与恢复能力✅核心收益这样可以减少策略切换带来的状态不连续问题得到更一致的行为。2. 环境要求系统LinuxPythonPython 3.11建议仿真 / GPU已可用的MuJoCo/ GPU 驱动环境CUDA推荐CUDA 12.8/CUDA 12.9相关环境3. 下载代码gitclone https://github.com/ccrpRepo/AMP_mjlab.gitcdAMP_mjlab4. 创建并激活 conda 环境conda create-nmjlabpython3.11conda activate mjlab5. 安装项目依赖在项目根目录执行python-mpipinstall-e.✅说明该命令会安装本仓库并根据setup.py安装mjlab1.2.0。6. 安装指定版本的 PyTorchpipinstalltorch2.11.0torchvision0.26.0torchaudio2.11.0 --index-url https://download.pytorch.org/whl/cu1287. 安装 Warp 和 mujoco-warp 兼容版本pipinstallwarp-lang1.14.0 mujoco-warp3.9.08. 应用 mjlab 观测历史补丁⚠️注意如果不打这个补丁则需要在代码中去掉history_ordering配置。补丁作用说明增加历史观测的展开方式选项可选择按时间维time或按观测项term展开mjlab 默认仅支持按term展开补丁文件mjlab_patch/mjlab/managers/observation_manager.py示例覆盖命令cpmjlab_patch/mjlab/managers/observation_manager.py\/home/liguopu/miniconda3/envs/mjlab/lib/python3.11/site-packages/mjlab/managers/observation_manager.py路径提示如果环境路径不同请把/home/liguopu/miniconda3/envs/mjlab替换为实际的conda环境路径。也可以用下面命令查看site-packages位置python -PY import site print(site.getsitepackages()) PY9. 修复 mjlab CUDA Graph 兼容问题在torch2.11.0cu128环境中torch.cuda.get_driver_version()不存在而mjlab1.2.0默认会调用它判断CUDA Graph是否可用因此训练会报错AttributeError: module torch.cuda has no attribute get_driver_version推荐修改mjlab安装包里的文件/home/liguopu/miniconda3/envs/mjlab/lib/python3.11/site-packages/mjlab/sim/sim.py找到_should_use_cuda_graph()把原来的 PyTorch driver version 获取逻辑importtorch driver_ver_inttorch.cuda.get_driver_version()driver_ver(driver_ver_int//1000,(driver_ver_int%1000)//10)改为使用 Warpdriver_verwp.get_cuda_driver_version()修改后的关键代码应类似def_should_use_cuda_graph(self)-bool:Determine if CUDA graphs can be used based on device and driver version.ifnotself.wp_device.is_cuda:returnFalsedriver_verwp.get_cuda_driver_version()has_mempoolwp.is_mempool_enabled(self.wp_device)ifdriver_verisNone:print([WARNING] CUDA Graphs disabled: driver version unavailable)returnFalseifhas_mempoolanddriver_ver_GRAPH_CAPTURE_MIN_DRIVER:returnTruereasons[]ifnothas_mempool:reasons.append(mempool disabled)ifdriver_ver_GRAPH_CAPTURE_MIN_DRIVER:reasons.append(fdriver{driver_ver[0]}.{driver_ver[1]} 12.4)print(f[WARNING] CUDA Graphs disabled:{, .join(reasons)})returnFalse✅原因当前warp-lang1.14.0的wp.get_cuda_driver_version()返回类似(12, 9)的tuple正好能和mjlab后续的_GRAPH_CAPTURE_MIN_DRIVER (12, 4)比较。可用下面命令验证python -PY import warp as wp wp.init() print(wp.get_cuda_driver_version()) PY10. 查看可用任务python scripts/list_envs.py--keywordAMP主要任务Unitree-G1-AMP-RoughUnitree-G1-AMP-FlatUnitree-H2-AMP-RoughUnitree-H2-AMP-Flat⭐备注原开源版本只支持G1H2相关任务是本博客通过修改代码后新增适配的任务。临时运行方式如果不想安装项目也可以临时把项目根目录加入Python模块搜索路径仅当前终端生效exportPYTHONPATH/home/liguopu/lgp_dev/project/AMP_mjlab:$PYTHONPATHpython scripts/list_envs.py--keywordAMP命令写法注意python-mscripts.list_envs--keywordAMP⚠️不要写错python -m scripts.list_envs.py是错误写法-m后面应跟模块名不能带.py后缀。11. G1 训练python scripts/train.py Unitree-G1-AMP-Flat --env.scene.num-envs4096日志默认保存到logs/rsl_rl/g1_amp_locomotion/time_stamp_run/ 训练曲线说明对应训练信息在约2w轮约20k iterations附近策略通常会突然学会“跌倒后恢复”行为。✅正常现象对应地logs中多个指标可能出现明显突变阶跃式变化这是正常现象不一定是训练异常。训练过程中终端会持续打印统计信息。下面这一段是G1 Flat训练日志示例用于说明各类指标的含义和正常输出格式################################################################################ Learning iteration : 100000/100001 Computation : 53904 steps/s (collection: 1.071s, learning 0.752s) Mean action noise std : 0.28 Mean value_function loss : 0.0065 Mean surrogate loss : -0.0100 Mean entropy loss : 3.3100 Mean amp loss : 0.3794 Mean amp_grad_pen loss : 0.2058 Mean amp_policy_pred loss : -0.4078 Mean amp_expert_pred loss : 0.4068 Mean skipped_non_finite_batches loss : 0.0000 Mean reward : 41.46 Mean episode length : 995.78 Episode_Reward/track_anchor_linear_velocity : 0.7154 Episode_Reward/track_anchor_angular_velocity: 0.9129 Episode_Reward/track_root_height : 0.0038 Episode_Reward/body_ang_vel_xy_l2 : 0.4666 Episode_Reward/is_terminated : -0.0014 Episode_Reward/joint_acc_l2 : -0.0761 Episode_Reward/joint_pos_limits : -0.0039 Episode_Reward/action_rate_l2 : -0.0616 Episode_Reward/foot_slip : -0.0174 Episode_Reward/self_collisions : -0.0039 Episode_Metrics/mean_action_acc : 0.5422 Episode_Metrics/mean_delay_steps : 0.5537 Metrics/twist/error_vel_xy : 1.5915 Metrics/twist/error_vel_yaw : 1.7172 Episode_Termination/time_out : 5.1667 Episode_Termination/bad_orientation : 0.0000 Episode_Termination/bad_base_height : 0.0417 Metrics/slip_velocity_mean : 0.1289 -------------------------------------------------------------------------------- Total timesteps : 9830498304 Iteration time : 1.82s Time elapsed : 02:57:11 ETA : 00:00:01怎么看这类输出可以快速观察训练速度、reward、AMP loss、终止原因、速度误差和剩余时间。比如上面的G1示例中Mean reward: 41.46Computation: 53904 steps/sMean skipped_non_finite_batches loss: 0.0000说明该阶段训练正在正常推进且没有跳过非有限数值batch。H2 Flat的实测训练日志单独记录在第 12 节。12. H2 训练⚠️训练前置条件H2训练前需要先准备H2 AMP NPZ数据。可由现有G1 AMP NPZ重定位生成详见第 14 节。确认src/assets/motions/h2/amp/WalkandRun中已经存在.npz后执行PYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/train.py Unitree-H2-AMP-Flat\--env.scene.num-envs4096日志默认保存到logs/rsl_rl/h2_amp_locomotion/time_stamp_run/如果要训练粗糙地形版本任务名改为PYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/train.py Unitree-H2-AMP-Rough\--env.scene.num-envs4096 H2 最小连通性测试命令 这里的“冒烟测试”指用最小环境数和最少训练轮数快速跑一遍不追求训练效果只验证H2资产、H2 NPZ数据、任务注册和训练入口是否能正常连通。PYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/train.py Unitree-H2-AMP-Flat\--env.scene.num-envs1\--agent.max-iterations1✅通过标准如果能进入runner并保存model_0.pt说明H2资产、H2 NPZ和训练入口都已打通。 H2 Flat 实测训练记录本次H2 Flat训练命令PYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/train.py Unitree-H2-AMP-Flat\--env.scene.num-envs4096训练末尾输出################################################################################ Learning iteration : 100000/100001 Computation : 32648 steps/s (collection: 1.685s, learning 1.326s) Mean action noise std : 0.26 Mean value_function loss : 0.0067 Mean surrogate loss : -0.0107 Mean entropy loss : -1.8587 Mean amp loss : 0.2792 Mean amp_grad_pen loss : 0.2033 Mean amp_policy_pred loss : -0.5074 Mean amp_expert_pred loss : 0.5010 Mean skipped_non_finite_batches loss : 0.0000 Mean reward : 39.14 Mean episode length : 1000.00 Episode_Reward/track_anchor_linear_velocity : 0.6854 Episode_Reward/track_anchor_angular_velocity: 0.9154 Episode_Reward/track_root_height : 0.0019 Episode_Reward/body_ang_vel_xy_l2 : 0.4594 Episode_Reward/is_terminated : 0.0000 Episode_Reward/joint_acc_l2 : -0.0881 Episode_Reward/joint_pos_limits : -0.0056 Episode_Reward/action_rate_l2 : -0.0622 Episode_Reward/foot_slip : -0.0239 Episode_Reward/self_collisions : -0.0019 Episode_Metrics/mean_action_acc : 0.5028 Episode_Metrics/mean_delay_steps : 0.8181 Metrics/twist/error_vel_xy : 1.6271 Metrics/twist/error_vel_yaw : 2.0770 Episode_Termination/time_out : 4.3333 Episode_Termination/bad_orientation : 0.0000 Episode_Termination/bad_base_height : 0.0000 Metrics/slip_velocity_mean : 0.1238 -------------------------------------------------------------------------------- Total timesteps : 9830498304 Iteration time : 3.01s Time elapsed : 01:21:54 ETA : 00:00:0113. 评估与可视化说明训练与回放阶段都支持ONNX导出默认开启。13.1 通用回放命令使用已训练权重回放时需要指定任务名和checkpoint文件G1 回放命令python scripts/play.py Unitree-G1-AMP-Rough\--checkpoint-file logs/rsl_rl/g1_amp_locomotion/run_dir/model_iter.ptH2 回放命令PYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/play.py Unitree-H2-AMP-Flat\--checkpoint-file logs/rsl_rl/h2_amp_locomotion/run_dir/model_iter.pt\--num-envs1\--viewernative如果评估H2 粗糙地形版本任务名改为Unitree-H2-AMP-Rough⚠️注意G1 checkpoint不能直接用于H2。H2当前策略只控制29 个非头部关节头部两个关节保持固定但H2机器人模型和观测空间仍与G1不同。H2必须使用logs/rsl_rl/h2_amp_locomotion/...下训练得到的checkpoint。13.2 G1 实践示例以下示例使用2026-07-12_11-43-37这次训练中的model_23900.pt进行回放与可视化。先确认权重文件存在findlogs/rsl_rl/g1_amp_locomotion-namemodel_23900.pt示例输出logs/rsl_rl/g1_amp_locomotion/2026-07-12_11-43-37/model_23900.pt13.3 方式一Viser 浏览器可视化启动Viser可视化回放source~/miniconda3/etc/profile.d/conda.shconda activate mjlabPYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/play.py Unitree-G1-AMP-Flat\--checkpoint-file logs/rsl_rl/g1_amp_locomotion/2026-07-12_11-43-37/model_23900.pt\--num-envs1\--viewerviser启动成功后终端会显示类似[INFO]: Loading checkpoint: model_23900.pt [INFO]: Exported ONNX policy to: /home/liguopu/lgp_dev/project/AMP_mjlab/logs/rsl_rl/g1_amp_locomotion/2026-07-12_11-43-37/export/Unitree-G1-AMP-Flat_model_23900.onnx ╭────── viser (listening *:8080) ───────╮ │ ╷ │ │ HTTP │ http://localhost:8080 │ │ Websocket │ ws://localhost:8080 │ │ ╵ │ ╰───────────────────────────────────────╯浏览器打开http://localhost:8080端口提示如果已有其他Viser进程占用了8080终端可能显示8081或其他端口以实际输出为准。13.4 方式二native viewer 本地窗口如果浏览器提示WebGL unavailable或者Viser页面能连接但画面空白可以改用native viewerPYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/play.py Unitree-G1-AMP-Flat\--checkpoint-file logs/rsl_rl/g1_amp_locomotion/2026-07-12_11-43-37/model_36600.pt\--num-envs1\--viewernativeH2native viewer示例PYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/play.py Unitree-H2-AMP-Flat\--checkpoint-file logs/rsl_rl/h2_amp_locomotion/run_dir/model_iter.pt\--num-envs1\--viewernative示例效果native viewer常用按键按键作用s切换模式提高帧率加速播放-减速播放R开关 debug visualizationP开关 reward plotsSpace暂停 / 继续Enter重置环境,切换到上一个环境.切换到下一个环境→暂停时单步前进操作建议可以先按R关闭debug visualization再按P确认reward plots关闭需要更快播放时按加速。13.5 H2 实践示例本小节使用已经训练完成的H2 Flatcheckpoint 进行回放主要用于确认H2 策略是否能正常驱动机器人走、跑以及跌倒恢复。✅示例权重这里使用的是logs/rsl_rl/h2_amp_locomotion/2026-07-13_18-59-23/model_100000.pt。如果使用自己的训练结果请把run_dir和model_iter.pt替换为实际路径。H2 Flat native viewer 回放命令PYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/play.py Unitree-H2-AMP-Flat\--checkpoint-file logs/rsl_rl/h2_amp_locomotion/2026-07-13_18-59-23/model_100000.pt\--num-envs1\--viewernative参数说明Unitree-H2-AMP-Flat使用H2 Flat评估任务--checkpoint-file指定训练完成后的H2权重文件--num-envs 1只启动 1 个环境便于观察单个机器人的动作效果--viewer native使用 MuJoCo 原生窗口查看回放效果⚠️注意这里必须使用h2_amp_locomotion目录下的H2 checkpoint。不要把G1的checkpoint直接拿来回放H2两者机器人模型和观测空间不同。走、跑效果跌倒、恢复效果观察重点回放时可以重点看H2是否能稳定跟踪速度命令、走跑动作是否连续、跌倒后是否能进入恢复动作并重新站起。13.6 常见问题rsl_rl 导入路径如果直接运行python scripts/play.py ...时出现下面错误ModuleNotFoundError: No module named rsl_rl.runners.amp_on_policy_runner✅原因与处理通常是因为没有优先导入本仓库自带的rsl_rl。在命令前加上下面这一段即可PYTHONPATH$PWD:${PYTHONPATH:-}完整命令可以写成PYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/play.py Unitree-G1-AMP-Flat\--checkpoint-file logs/rsl_rl/g1_amp_locomotion/2026-07-12_11-43-37/model_23900.pt\--num-envs1\--viewerviser13.7 ONNX 导出结果这次实践中导出的ONNX文件为logs/rsl_rl/g1_amp_locomotion/2026-07-12_11-43-37/export/Unitree-G1-AMP-Flat_model_23900.onnx14. 运动数据准备仓库提供CSV到NPZ的转换脚本python scripts/csv_to_npz.py--help推荐目录组织原始 CSVmotion_data_csv/amp转换后 NPZsrc/assets/motions/g1/amp/WalkandRun与src/assets/motions/g1/amp/RecoveryH2 重定位 NPZsrc/assets/motions/h2/amp/WalkandRun与src/assets/motions/h2/amp/Recovery✅自动加载只要上述目录中存在可用NPZ训练配置会自动加载。14.1 G1 AMP NPZ 可视化如果想直接查看G1 NPZ动作本身不需要训练checkpoint可以用MuJoCo原生窗口播放PYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/visualize_g1_npz.py\--npz-file src/assets/motions/g1/amp/Recovery/fallAndGetUp1_subject1.npz\--devicecuda:0播放提示脚本默认会循环播放关闭MuJoCo窗口即可退出。常用参数# 无窗口快速校验前三帧适合远程或无显示环境PYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/visualize_g1_npz.py\--npz-file src/assets/motions/g1/amp/Recovery/fallAndGetUp1_subject1.npz\--no-viewer\--no-loop\--max-frames3# 慢速播放PYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/visualize_g1_npz.py\--npz-file src/assets/motions/g1/amp/Recovery/fallAndGetUp1_subject1.npz\--speed0.5该脚本会读取G1 NPZ中的root状态、29 维关节位置/速度并逐帧写入G1 MuJoCo/mjlab模型。14.2 由 G1 AMP NPZ 生成 H2 AMP NPZH2可以直接使用脚本从现有G1 AMP NPZ重定位生成PYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/retarget_g1_amp_to_h2_npz.py\--input-root src/assets/motions/g1/amp\--output-root src/assets/motions/h2/amp\--devicecuda:0如果使用 CPUPYTHONPATH$PWD:${PYTHONPATH:-}\python scripts/retarget_g1_amp_to_h2_npz.py\--input-root src/assets/motions/g1/amp\--output-root src/assets/motions/h2/amp\--devicecpu当前重定位策略关节映射G1的29 个同名关节映射到H2的29 个非头部关节头部关节H2的head_pitch_joint与head_yaw_joint保持默认0位姿和0速度数据格式输出仍然保留H2 31 DoF数据格式其中头部两维为默认0训练策略只控制29 个非头部关节重新计算 body 状态脚本会用H2 MuJoCo/mjlab模型forward重新生成H2的body_pos_w、body_quat_w、body_lin_vel_w、body_ang_vel_w生成后可验证findsrc/assets/motions/h2/amp-name*.npz|sort|wc-l检查 shapePYTHONPATH$PWD:${PYTHONPATH:-}\python -PY import pathlib import numpy as np files sorted(pathlib.Path(src/assets/motions/h2/amp).glob(*/*.npz)) print(files, len(files)) joint_dims set() body_counts set() head_max 0.0 for path in files: data np.load(path) joint_dims.add(data[joint_pos].shape[-1]) body_counts.add(data[body_pos_w].shape[1]) head_max max(head_max, float(abs(data[joint_pos][:, 15:17]).max())) print(joint_dims, joint_dims) print(body_counts, body_counts) print(head_abs_max, head_max) PY期望输出类似files 18 joint_dims {31} body_counts {32} head_abs_max 0.0更详细说明见H2_NPZ生成指导.md15. 核心代码——目录说明src/tasks/amp_locoAMP locomotion/recovery任务实现src/tasks/amp_loco/config/g1G1任务注册、环境与RL配置src/tasks/amp_loco/config/h2H2任务注册、环境与RL配置src/tasks/amp_loco/mdp奖励、观测、事件、终止逻辑scripts/train.py训练入口scripts/play.py回放入口scripts/csv_to_npz.py动作数据转换工具scripts/visualize_g1_npz.pyG1 AMP NPZ的MuJoCo原生窗口可视化工具scripts/visualize_h2_npz.pyH2 AMP NPZ的MuJoCo原生窗口可视化工具scripts/retarget_g1_amp_to_h2_npz.pyG1 AMP NPZ到H2 AMP NPZ的重定位生成工具H2_NPZ生成指导.mdH2 AMP NPZ生成说明mjlab_patch依赖的mjlab本地补丁16. 总结⭐单一策略统一覆盖走跑与跌倒恢复AMP 速度任务联合优化兼顾风格与任务性能️延迟重置与 recovery 采样机制显式强化恢复能力训练到部署链路完整支持ONNX导出✅ 分享完成