🤖 roboto_origin_03 Wiki
首页 / 训练 / 自定义任务与环境扩展指南

本文面向需要在 RoboLab 框架中新增足式机器人训练任务、修改现有环境行为或接入全新机器人资产的进阶开发者。RoboLab 的任务层采用 Isaac Lab 的双轨架构:Direct RLManager-based 并行共存,二者分别适用于需要高度自定义步进逻辑的场景与偏好声明式配置的场景。掌握本文介绍的扩展模式后,你可以在不改动算法库(rsl_rl)的前提下,完成从场景定义、观测设计、奖励函数编写到 gym 注册与训练启动的全链路扩展。

Sources: base_env.py, vec_env.py, train.py

环境扩展架构总览

RoboLab 的任务代码位于 robolab/tasks/ 目录,其顶层按范式划分为 direct/manager_based/ 两大分支。Direct RL 环境通过 Python 类继承显式重写 step()compute_current_observations() 等方法实现行为定制;Manager-based 环境则通过配置类组合 ObservationManager、RewardManager、TerminationManager 等模块,将逻辑委托给 Isaac Lab 的通用调度器。无论哪种范式,最终都需向 gymnasium 注册,并通过 RslRlVecEnvWrapper 封装为 RSL-RL 算法库可识别的 VecEnv 接口。

下图为两类扩展路径的继承与交互关系:

graph TB
    subgraph DirectRL["Direct RL 分支"]
        DirectRLEnv["isaaclab.envs.DirectRLEnv"]
        BaseEnv["BaseEnv<br>robolab/tasks/direct/base"]
        AttnEncEnv["AttnEncEnv<br>感知增强"]
        InterruptEnv["InterruptEnv<br>关节干预"]
        BaseEnv --> DirectRLEnv
        AttnEncEnv --> BaseEnv
        InterruptEnv --> BaseEnv
    end

    subgraph ManagerBased["Manager-based 分支"]
        ManagerBasedRLEnv["isaaclab.envs.ManagerBasedRLEnv"]
        AnimationEnv["AnimationEnv<br>动画驱动"]
        AmpEnv["AmpEnv<br>AMP 对抗先验"]
        BeyondMimic["BeyondMimic<br>配置驱动"]
        AnimationEnv --> ManagerBasedRLEnv
        AmpEnv --> AnimationEnv
        BeyondMimic --> ManagerBasedRLEnv
    end

    subgraph RSLRL["RSL-RL 算法端"]
        VecEnv["rsl_rl.env.VecEnv<br>抽象接口"]
        RslRlVecEnvWrapper["RslRlVecEnvWrapper<br>适配封装"]
        OnPolicyRunner["OnPolicyRunner"]
    end

    BaseEnv -->|封装为| RslRlVecEnvWrapper
    AmpEnv -->|封装为| RslRlVecEnvWrapper
    BeyondMimic -->|封装为| RslRlVecEnvWrapper
    RslRlVecEnvWrapper --> VecEnv
    OnPolicyRunner --> VecEnv

Sources: base_env.py, animation_env.py, amp_env.py, train.py

Direct RL 环境扩展

Direct RL 模式适合需要在单步 step() 中插入自定义物理前处理、动作干预或复杂观测拼接逻辑的场景。BaseEnv 已经封装了标准的足式机器人环境初始化、命令生成器(UniformVelocityCommand)、接触传感器解析与基础奖励计算,扩展时通常只需继承 BaseEnv 并选择性重写关键方法。

1. 继承 BaseEnv 与配置类

新建任务目录时,首先创建环境类与对应的配置类。配置类应继承 BaseEnvCfg,并在 __post_init__ 中绑定机器人资产、场景、观测/动作维度与地形。以 ATOM01FlatEnvCfg 为例,它在 __post_init__ 中设定 action_space = 23observation_space = 78state_space = 139,并指定 ATOM01_CFG 作为机器人模型,同时实例化 SceneCfg 以注入地形、光照与射线传感器。

Sources: base_config.py, atom01_env_cfg.py

2. 观测空间重写模式

观测空间是 Direct RL 扩展中最常见的修改点。BaseEnv 将观测分为 actor_obscritic_obs,并通过 compute_current_observations() 计算单帧观测,再于 _get_observations() 中拼接历史缓冲。若需要新增感知模态(如注意力编码器的高程扫描)或改变拼接顺序,应重写这两个方法。

AttnEncEnv 的扩展示范了三种典型改写:第一,在 compute_current_observations() 中根据 cfg.attn_enc.vel_in_obs 条件拼接线速度;第二,在 critic 观测中加入足端相对机体坐标系的位置 feet_pos;第三,在 _get_observations() 中当启用注意力编码器时,将高程扫描以独立键 perception_aperception_c 返回,而非直接拼接到主观测向量。这种分离使得网络侧可在 ActorCritic 中为感知模态配置独立的编码器分支。

Sources: attn_enc_env.py, attn_enc_env.py

3. 步进与动作干预

若任务需要在物理步进前后插入外部扰动或动作覆写,可重写 _pre_physics_step()step()InterruptEnv 提供了关节干预(interrupt)的典型实现:它在 init_buffers() 中创建干预掩码 interrupt_mask 与课程难度标量 interrupt_rad_curriculum;在 _pre_physics_step() 中按概率切换掩码,并将干预后的动作夹紧在课程半径内;在 step() 中每 interrupt_update_step 步重新采样干预目标,同时在 GUI 模式下更新可视化标记。该设计将干预逻辑与基础环境完全解耦,保留了 BaseEnv 的奖励与终止判定不变。

Sources: interrupt_env.py, interrupt_env.py, interrupt_env.py

4. 配置扩展与噪声向量对齐

当观测维度发生变化时,必须同步修改 init_obs_buffer() 中的噪声向量 noise_scale_vec,否则 add_noise 阶段会发生越界。AttnEncEnv 中根据 vel_in_obs 的布尔条件分别计算噪声向量的切片位置:若线速度参与 actor 观测,则 noise_vec 的索引从 0:3(角速度)推移到 3:6(线速度),后续命令、关节位置、关节速度、动作依次后移。任何新增观测项都必须在此显式声明对应的噪声尺度,否则噪声注入将错位。

Sources: attn_enc_env.py, base_env.py

Manager-based 环境扩展

Manager-based 模式将环境逻辑拆解为可插拔的 Manager 集合,适合以配置声明方式定义任务,且便于利用 Isaac Lab 的 ObservationTermCfgRewardTermCfg 等基础设施进行快速实验。

1. 配置驱动架构

Manager-based 环境的核心是一个继承自 ManagerBasedRLEnvCfg 的配置类,其内部聚合了 sceneobservationsactionsrewardsterminationseventscurriculum 等子配置。运行时,Isaac Lab 根据这些配置自动实例化对应的 Manager。以 BeyondMimicEnvCfg 为例,它通过 MySceneCfg 定义地形与机器人资产,通过 ObservationsCfg 分别声明 policy 与 critic 的观测项,通过 ActionsCfg 将动作映射为关节位置偏移,所有逻辑均以 ObsTermRewTermDoneTerm 的实例化参数表达。

Sources: beyondmimic_env_cfg.py

2. 场景配置

场景配置继承 InteractiveSceneCfg,负责声明地形导入器、机器人资产、光照、接触传感器与射线传感器。BeyondMimicEnvCfg 中的 MySceneCfg 展示了最小可用场景:地形采用 plane 类型,接触传感器通过正则 "{ENV_REGEX_NS}/Robot/.*" 批量挂载,所有环境在初始化时由 InteractiveSceneCfg 自动克隆。若需使用生成式地形,可将 terrain_type 改为 "generator" 并传入 TerrainGeneratorCfg 实例。

Sources: beyondmimic_env_cfg.py, scene_cfg.py

3. 观测、动作、奖励与终止条件

各项 MDP 术语通过配置对象中的函数引用绑定到 Python 实现。观测项支持噪声腐蚀(AdditiveUniformNoiseCfg)与历史缓冲;奖励项支持权重与超参数;终止项区分 time_out 与硬性终止。下表对比了 Direct RL 与 Manager-based 中各类 MDP 术语的声明位置与扩展方式:

术语类型 Direct RL 声明位置 Manager-based 声明位置 扩展方式
观测 compute_current_observations() 返回值 ObservationsCfg 中的 ObsTerm 重写方法 / 新增 ObsTerm 并注册函数
动作 _pre_physics_step() 中的张量变换 ActionsCfg 中的 JointPositionActionCfg 重写方法 / 替换 action term
奖励 RewardCfg 子类中的 RewTerm(由 RewardManager 计算) RewardsCfg 中的 RewTerm 在配置类中新增字段并编写函数
终止 _get_dones() 中的布尔张量 TerminationsCfg 中的 DoneTerm 重写方法 / 新增 DoneTerm 并注册函数
事件 EventCfg 中的 EventTerm EventCfg 中的 EventTerm 在配置类中新增字段并引用函数
课程 update_terrain_levels() 等手动实现 CurriculumCfg 中的 CurriculumTerm 在配置类中声明

Sources: base_env.py, beyondmimic_env_cfg.py

4. 自定义 Manager 与步进覆写

当 Manager-based 的通用步进逻辑无法满足需求时,可通过继承 ManagerBasedRLEnv 并覆写 step() 实现定制。AmpEnv 的父类 AnimationEnv 已经展示了如何注入 AnimationManagerMotionDataManagerAmpEnv 进一步覆写 step(),在环境重置前保留 AMP 观测,避免重置导致的观测丢失。自定义 Manager 的加载应在 load_managers() 中完成,并确保在调用 super().load_managers() 之前实例化,以便后续 Manager 依赖其状态。

Sources: animation_env.py, amp_env.py

MDP 术语函数扩展

无论 Direct RL 还是 Manager-based,新增奖励、观测或终止函数时都遵循同一套函数签名约定。函数应接收 env(当前环境实例)与配置参数,返回形状为 (num_envs,)torch.Tensor

1. 奖励函数编写规范

track_lin_vel_xy_yaw_frame_exp 为例,其签名接收 env: BaseEnvstd: float,内部通过 env.scene[asset_cfg.name] 获取 Articulation 资产,并利用 math_utils.quat_apply_inverse 将世界系速度转换到偏航对齐的机体坐标系,最终与 env.command_generator.command 比较并返回指数型奖励。所有自定义奖励函数应放在任务目录的 mdp/rewards.py 中,并通过 mdp/__init__.pyfrom .rewards import * 暴露,使配置类能够以 mdp.xxx 形式引用。

Sources: rewards.py, mdp/init.py

2. 场景实体解析

MDP 函数中访问机器人躯体或传感器时,通常使用 SceneEntityCfg 进行延迟解析。例如 undesired_contacts 通过 env.scene.sensors[sensor_cfg.name] 获取接触传感器,再利用 sensor_cfg.body_ids 筛选目标躯体。配置阶段只需在 RewTerm.params 中传入 SceneEntityCfg("contact_sensor", body_names="(?!.*ankle_roll.*).*"),运行时 Isaac Lab 会自动将正则匹配解析为具体的躯体索引数组。这种机制使得奖励函数无需硬编码躯体序号,具备跨机器人复用能力。

Sources: rewards.py, atom01_env_cfg.py

地形扩展

地形配置通过 TerrainGeneratorCfg 实例定义,包含子地形比例、尺寸、水平/垂直精度等参数。terrain_generator_cfg.py 中预定义了 GRAVEL_TERRAINS_CFGROUGH_TERRAINS_CFGROUGH_HARD_TERRAINS_CFG 三种配置,分别对应平坦碎石地、课程化 rough 地型与高难度混合地型。新增地形时,只需实例化新的 TerrainGeneratorCfg 并在环境配置的 __post_init__ 中赋值给 scene_context.terrain_generator,同时确保 terrain_type = "generator"。若启用课程学习(curriculum=True),BaseEnv.update_terrain_levels() 会根据机器人行走距离自动提升或降低环境的地形等级。

Sources: terrain_generator_cfg.py, base_env.py

智能体配置扩展

智能体配置决定网络架构、PPO 超参数与训练管线。Direct RL 与 Manager-based 环境共用同一套 RslRlOnPolicyRunnerCfg 体系。扩展时,新建配置类继承 BaseAgentCfg(或直接使用 RslRlOnPolicyRunnerCfg),在 __post_init__ 中覆写实验名、学习率、网络隐藏层维度、对称增强(symmetry_cfg)与 RND 配置(rnd_cfg)等字段。ATOM01FlatAgentCfg 中设置了 actor_hidden_dims=[512, 256, 128]learning_rate=1.0e-3gamma=0.99,而 ATOM01RoughAgentCfg 通过继承仅修改了 experiment_namewandb_project,展示了最小差异扩展策略。

Sources: base_config.py, atom01_agent_cfg.py

完整注册与训练启动流程

完成环境与配置类的代码编写后,必须向 gymnasium 注册,才能在 train.py 中通过 --task 参数调用。注册在任务包根目录的 __init__.py 中完成,需指定唯一 ID、入口类、环境配置入口点与智能体配置入口点。

flowchart LR
    A[创建 Env 类<br>继承 BaseEnv/ManagerBasedRLEnv] --> B[创建 Cfg 类<br>继承 BaseEnvCfg/ManagerBasedRLEnvCfg]
    B --> C[编写 AgentCfg 类<br>继承 BaseAgentCfg]
    C --> D[在 __init__.py 中<br>gym.register]
    D --> E[import robolab.tasks<br>触发注册]
    E --> F[运行 train.py<br>--task Atom01-XXX]
    F --> G[hydra_task_config<br>自动解析配置]
    G --> H[gym.make<br>实例化环境]
    H --> I[RslRlVecEnvWrapper<br>封装为 VecEnv]
    I --> J[OnPolicyRunner/AMPRunner<br>开始训练]

AttnEncEnv 的注册为例,gym.registerid 设为 "Atom01-AttnEnc"entry_point 指向 attn_enc_env:AttnEncEnvenv_cfg_entry_pointrsl_rl_cfg_entry_point 分别指向对应配置类的完整模块路径。训练脚本通过 import robolab.tasks 触发所有子包的注册逻辑,随后 hydra_task_config(args_cli.task, args_cli.agent) 根据任务名自动加载并实例化配置。

Sources: init.py, train.py, list_envs.py

最佳实践与调试要点

问题场景 建议做法
观测维度不匹配导致 RuntimeError 同步修改 __post_init__ 中的 observation_space/state_spaceinit_obs_buffer() 中的噪声向量长度
新增奖励项未生效 检查配置类是否正确继承 RewardCfg 并在 __post_init__ 中赋值给 env_cfg.reward;确认 mdp/__init__.py 已导出函数
地形生成器缓存异常 TerrainGeneratorCfg.use_cache 设为 False,避免修改子地形后读取旧缓存
Manager-based 观测缺失 检查 ObservationsCfg 中各 ObsGroupenable_corruptionconcatenate_terms 设置;确认 observation_manager.compute() 调用时机
gym 注册后找不到任务 确认 import robolab.tasks 已被执行;使用 python -m scripts.tools.list_envs 查看已注册环境列表
动作干预导致关节限位错误 _pre_physics_step() 的最终动作输出前增加 torch.clip,并在干预逻辑中考虑 default_joint_pos_limits

Sources: base_env.py, train.py, list_envs.py


完成自定义环境扩展后,建议先阅读 Direct RL 环境架构Manager-based 环境架构 深入理解底层机制,再参考 奖励函数与 MDP 设计 优化奖励 shaping。若需将训练好的策略部署到 MuJoCo 或真机,可继续阅读 MuJoCo Sim2Sim 部署与真机迁移