本页文档聚焦于 ATOM01-Train 中两项进阶训练能力:BeyondMimic 多技能动作模仿与基于注意力机制的地形感知策略。前者在 IsaacLab 的 ManagerBasedRLEnv 框架下,通过自适应采样从长时运动剪辑中学习多个技能段;后者在 DirectRLEnv 框架下,利用高度扫描仪获取地形高程图,并通过 CNN+多头注意力编码器将地形特征融入策略与值函数网络。两项能力均基于 RSL-RL 的 PPO 训练流程,并共享同一套 Sim2Sim 与部署管线。在阅读本节之前,建议先完成 IsaacLab环境搭建与训练流程 与 AMP模仿学习与动作数据集 的基础配置。
Sources: README_CN.md
架构总览
整体训练架构可分为两条支线:基于运动指令的 ManagerBased 模仿环境与基于地形感知编码器的 Direct 强化学习环境。两者最终都通过 RSL-RL 的 Runner 进行策略迭代。
graph TD
subgraph BeyondMimic多技能模仿
A[运动剪辑 NPZ<br>joint_pos/body_pos_w] --> B(MotionLoader)
B --> C[MotionCommand<br>自适应采样]
C --> D[ManagerBasedRLEnv<br>BeyondMimic/GetupMimic]
end
subgraph 注意力地形感知
E[高度扫描仪 RayCaster] --> F[AttnEncEnv<br>DirectRLEnv]
G[本体观测历史] --> F
end
D --> H[RslRlVecEnvWrapper]
F --> H
H --> I[OnPolicyRunner / DistillationRunner / AMPRunner]
I --> J[ActorCriticAttnEnc<br>或标准 ActorCritic]
J --> K[策略网络 + 值网络]
K -->|动作| D
K -->|动作| F
Sources: init.py, init.py, train.py
1. BeyondMimic 多技能模仿环境
1.1 MotionCommand 与 MotionLoader
BeyondMimic 的核心指令来源是预录制的运动剪辑(.npz),其中包含每一帧的关节位置、关节速度以及各刚体的位姿与速度。MotionLoader 负责在初始化时将文件载入 GPU 张量,并根据配置中的 body_names 进行索引过滤;MotionCommand 则在每个环境中维护一个 time_steps 计数器,通过逐帧推进来提供当前时刻的参考运动状态。策略观测中的 command 项即由当前帧的关节位置与速度拼接而成,作为模仿学习的参考信号。
Sources: commands.py
1.2 观测空间设计
环境采用策略-评论家分离的观测分组。策略网络(Policy)仅接收低维本体信息:运动指令、base_ang_vel、projected_gravity、相对默认位置的关节位置与速度、以及上一时刻动作,且开启观测加噪与拼接。评论家网络(Critic)则享有更完整的状态信息,包括运动锚点在机体坐标系下的位置与朝向、所有被跟踪刚体的相对位姿、基座线速度等,以便更准确地估计值函数。
| 观测组 | 包含项 | 维度特性 |
|---|---|---|
| Policy | command, base_ang_vel, projected_gravity, joint_pos_rel, joint_vel_rel, last_action | 低维本体 + 指令 |
| Critic | Policy 项 + motion_anchor_pos_b, motion_anchor_ori_b, body_pos, body_ori, base_lin_vel, base_ang_vel | 全状态 + 运动参考 |
Sources: beyondmimic_env_cfg.py, observations.py
1.3 奖励函数设计
奖励由跟踪奖励与正则惩罚两部分构成。跟踪奖励对运动锚点的全局位置/朝向、各刚体的相对位置/朝向以及全局线速度/角速度施加指数型误差奖励,权重分别为 0.5~1.0;正则项则对关节加速度、关节力矩、动作变化率及超限关节位置进行惩罚,以避免抖动与能量浪费。
| 奖励项 | 权重 | 参数 |
|---|---|---|
| motion_global_anchor_pos | 0.5 | std=0.3 |
| motion_global_anchor_ori | 0.5 | std=0.4 |
| motion_body_pos | 1.0 | std=0.3 |
| motion_body_ori | 1.0 | std=0.4 |
| motion_body_lin_vel | 1.0 | std=1.0 |
| motion_body_ang_vel | 1.0 | std=3.14 |
| joint_acc_l2 | -2.5e-7 | — |
| joint_torques_l2 | -1e-5 | — |
| action_rate_l2 | -1e-1 | — |
| joint_pos_limits | -10.0 | — |
Sources: beyondmimic_env_cfg.py, rewards.py
1.4 自适应采样机制(多技能核心)
BeyondMimic 的“多技能”能力并非通过显式切换多个独立文件实现,而是依托自适应采样在同一长时剪辑内部对不同困难段落进行聚焦学习。具体而言,MotionCommand 将整条剪辑按时间轴划分为若干等宽 bin(数量为 time_step_total // (decimation * dt))。每当环境触发终止时,系统记录失败所处的 bin 索引,并通过指数移动平均 adaptive_alpha 更新 bin_failed_count。采样概率由 bin_failed_count + adaptive_uniform_ratio / bin_count 构成,再与一个非因果平滑核(核权重为 adaptive_lambda^i)做一维卷积,最终归一化后通过多项式分布采样得到起始 bin。该机制确保智能体在容易失败的复杂动作段上获得更高的初始化概率,从而在长时剪辑中隐式习得多种技能模式。同时系统会记录采样熵 sampling_entropy 与最高概率 bin 指标,便于监控训练分布。
Sources: commands.py, commands.py
1.5 终止条件与域随机化
为了保证策略的鲁棒性,环境设置了基于运动锚点的终止条件:当锚点高度偏差超过 0.25 m 或朝向误差超过 0.8 rad 时触发终止,避免机器人在严重偏离参考姿态后继续无意义训练。域随机化覆盖摩擦系数、restitution、关节默认位置偏移、基座质量(±3 kg)、质心偏移、连杆质量比例(0.8~1.2)、执行器刚度/阻尼比例、关节摩擦与 armature 比例,以及周期性的外部推力(1~3 s 间隔),从而缩小 Sim2Real 差距。
Sources: beyondmimic_env_cfg.py, beyondmimic_env_cfg.py
1.6 环境配置与注册
BeyondMimicEnvCfg 定义了通用的场景与 MDP 参数:4096 个并行环境、decimation=4、仿真步长 0.005 s、单幕最大时长 20 s。针对 ATOM01 机器人的具体实现 Atom01BeyondMimicEnvCfg 指定了运动文件 yundong1.npz、锚点刚体 torso_link 以及需要跟踪的 23 个刚体列表。此外,仓库还提供了 Atom01GetupMimicEnvCfg,使用俯卧起身动作剪辑 getup_supin2prone.npz,并设置 reset_on_motion_end=False,在运动结束后保持末帧姿态并引入 stand_still_after_motion 奖励,展示同一框架下对不同技能的适配能力。两个环境分别注册为 Atom01-BeyondMimic 与 Atom01-Getup-Mimic。
Sources: atom01_beyondmimic_env_cfg.py, atom01_getup_mimic_env_cfg.py, init.py
2. 注意力策略(Attention Policy)
2.1 AttnEncEnv 直接环境
AttnEncEnv 继承自 BaseEnv(DirectRLEnv),在基础 locomotion 观测之上增加了地形高程扫描。环境通过 RayCaster 高度扫描仪获取以机器人为中心、分辨率 0.1 m、尺寸 1.6 m × 1.0 m 的地形高程图(17×11 网格)。在 _get_observations 中,actor 与 critic 的本体观测(角速度、投影重力、指令、关节位置/速度、历史动作等)被分别压入 CircularBuffer 以支持时序堆叠;当启用 use_attn_enc 时,高程图被拆分为 perception_a(供策略网络)与 perception_c(供值网络)两个感知键,随观测字典一并送入网络。
Sources: attn_enc_env.py, atom01_attn_enc_env_cfg.py
2.2 注意力编码器 AttentionEncoder
注意力编码器负责将高程图与本体观测融合为固定长度的地形嵌入。其处理流程如下:首先将高程图 map_scans 通过小型 CNN 提取特征,输出通道数为 embedding_dim - 3;同时构造二维空间位置编码(以地图中心为原点,分辨率 map_resolution),并与原始高度值、CNN 特征在通道维度拼接,得到 embedding_dim 维的局部特征图。该特征图被展平为 L×W 个点级特征,作为多头注意力(MHA)的 Key 与 Value。本体观测经线性层映射为同样维度的 Query,经过层归一化后执行 nn.MultiheadAttention,最终输出地形嵌入 map_encoding 以及空间注意力权重 attn_weights。该设计使策略网络能够显式地“注视”对当前步态重要的地形区域。
Sources: attn_encoder.py
2.3 ActorCriticAttnEnc 网络架构
ActorCriticAttnEnc 将上述编码器嵌入标准的 Actor-Critic 框架。初始化时,它根据观测字典中的 policy 与 critic 维度以及历史长度 actor_history_length / critic_history_length 计算单步观测维度。网络支持两项可选扩展:**观测编码器(obs_encoder)**通过 MLP 将完整历史观测压缩为 latent_dim 维隐向量;**评论家估计器(critic_estimation)**通过 MLP 预测指定观测切片(如 estimation_slice=[78,79,80]),并在辅助损失中约束预测值与真实值的 MSE。在前向传播时,act 方法提取最新单步 actor 观测(及隐向量/估计值)作为 Query,传入 AttentionEncoder 与 perception_a 交互,得到的地形嵌入与原始观测拼接后送入 Actor MLP;evaluate 方法对 critic 执行类似流程,但使用 perception_c。动作分布采用独立高斯分布,支持标量或对数标准差。
Sources: actor_critic_attn_enc.py, actor_critic_attn_enc.py, actor_critic_attn_enc.py, actor_critic_attn_enc.py
2.4 对称数据增强
ATOM01 具备左右对称的肢体结构,因此注意力策略配置中实现了完整的手动镜像规则。data_augmentation_func 对 policy 观测、critic 观测、动作以及高程图 perception_a/perception_c 分别定义了镜像索引与符号翻转表,并在训练时将原批次与镜像批次拼接,实现数据量翻倍。该增强与 RSL-RL 的 symmetry_cfg 配合使用,镜像损失系数设为 0.1,在不增加仿真开销的前提下提升策略的左右泛化能力。
Sources: atom01_attn_enc_agent_cfg.py
2.5 Agent 与训练配置
ATOM01AttnEncAgentCfg 将上述网络与训练超参数集中管理。策略类指定为 ActorCriticAttnEnc,关键网络参数包括 embedding_dim=32、head_num=4、map_size=(17,11)、actor_history_length=5、critic_history_length=5。训练开启 enable_critic_estimation 与 enable_obs_encoder,隐维度 latent_dim=32。算法层面使用 PPO,自适应学习率调度,enable_aux_loss=True 且辅助损失系数为 0.05;对称增强同时开启数据增强与镜像损失。该配置通过 Hydra 入口 rsl_rl_cfg_entry_point 与 ATOM01AttnEncEnvCfg 绑定,注册为 Atom01-AttnEnc 任务。
Sources: atom01_attn_enc_agent_cfg.py, init.py
3. 训练与推理
3.1 启动训练
BeyondMimic 与 AttnEnc 的训练均复用同一套 train.py 脚本。通过 --task 指定注册的环境名称即可切换任务。训练脚本内部使用 hydra_task_config 解析环境配置与智能体配置,自动将环境包装为 RslRlVecEnvWrapper,并根据 agent_cfg.class_name 实例化 OnPolicyRunner、DistillationRunner 或 AMPRunner。典型启动命令如下:
# BeyondMimic 多技能模仿
python robolab/scripts/rsl_rl/train.py --task=Atom01-BeyondMimic --headless --num_envs=8192
# 注意力地形感知策略
python robolab/scripts/rsl_rl/train.py --task=Atom01-AttnEnc --headless --num_envs=4096
Sources: train.py
3.2 策略回放与导出
训练完成后,使用 play_bm.py(或通用的 play.py)加载最近检查点进行可视化回放。脚本会自动导出 JIT 格式的 policy.pt 与 ONNX 格式的 policy.onnx,存放于检查点同级目录的 exported/ 文件夹中,供后续 Sim2Sim迁移与MuJoCo验证 与真机部署使用。回放时可通过 --real-time 以真实时间步长运行,或开启 --video 录制验证视频。
Sources: play_bm.py
3.3 关键超参数对比
| 参数 | BeyondMimic (Atom01-BeyondMimic) | AttnEnc (Atom01-AttnEnc) |
|---|---|---|
| 网络类 | 标准 ActorCritic | ActorCriticAttnEnc |
| Actor 隐层 | [512, 256, 128] | [512, 256, 128] |
| Critic 隐层 | [512, 256, 128] | [512, 256, 128] |
| 学习率 | 1.0e-3 (adaptive) | 1.0e-3 (adaptive) |
| 每环境步数 | 24 | 24 |
| Mini-batch | 4 | 4 |
| PPO epoch | 5 | 5 |
| 历史长度 (actor/critic) | 1 / 1 (默认) | 5 / 5 |
| 注意力嵌入维度 | — | 32 |
| 注意力头数 | — | 4 |
| 地图尺寸 | — | (17, 11) |
| 辅助损失 | — | enable_aux_loss=True, coef=0.05 |
| 对称增强 | — | use_data_augmentation + mirror_loss |
Sources: atom01_beyondmimic_agent_cfg.py, atom01_attn_enc_agent_cfg.py
4. 阅读建议与下一步
BeyondMimic 与注意力策略均建立在 IsaacLab 环境与 RSL-RL 训练流程之上。若尚未完成基础环境搭建,请先回顾 IsaacLab环境搭建与训练流程;若对动作数据集制作与 AMP 模仿学习不熟悉,建议先阅读 AMP模仿学习与动作数据集。掌握本节内容后,下一步可将训练好的策略导出并迁移至 MuJoCo 进行验证,具体流程请参阅 Sim2Sim迁移与MuJoCo验证。