🤖 roboto_origin_03 Wiki
首页 / RSL-RL / 向量化环境抽象接口

在 rsl_rl 框架中,所有强化学习算法与外部仿真环境之间的交互都通过单一的抽象边界完成,即 VecEnv。该接口的设计目标是将策略优化逻辑与具体物理引擎或游戏引擎的实现细节彻底解耦,使同一套 PPO、AMP 或蒸馏算法能够在 Isaac Gym、Isaac Lab 乃至自定义环境中无修改运行。本文将从接口契约、观测数据模型、环境交互协议三个维度,系统解析这一抽象层的设计原理与实现约束。

Sources: vec_env.py

接口架构与核心契约

VecEnv 被定义为一个 Python 抽象基类(ABC),仅规定行为契约而不包含任何具体实现。这种设计强制所有外部环境实现者遵循统一的类型签名,从而在编译前(或启动前)即可发现接口不匹配问题。

接口的核心由 六个属性两个抽象方法 构成:

属性/方法 类型/签名 语义与约束
num_envs int 向量化环境中并行实例的数量,决定张量第一维大小
num_actions int 单个环境实例的动作维度,用于构建策略网络输出层
max_episode_length int | torch.Tensor 每轮最大步数;支持标量(全局统一)或张量(每个环境独立),用于动态回合长度场景
episode_length_buf torch.Tensor 当前各环境实例已执行的步数,Runner 在训练初始化阶段会对此进行随机化以增强探索
device torch.device | str 环境内部张量所在设备,Runner 据此进行跨设备数据传输
cfg dict | object 环境配置对象,Runner 将其传递给日志系统以记录实验元信息
get_observations() → TensorDict 返回当前时刻所有并行环境的观测;不可产生副作用
step(actions) → (TensorDict, Tensor, Tensor, dict) 接收 (num_envs, num_actions) 动作张量,推进所有环境一步

Sources: vec_env.py

classDiagram
    class VecEnv {
        <<abstract>>
        +int num_envs
        +int num_actions
        +int|Tensor max_episode_length
        +Tensor episode_length_buf
        +device|str device
        +dict|object cfg
        +get_observations() TensorDict
        +step(actions) tuple
    }
    class OnPolicyRunner {
        +learn(iterations, init_at_random_ep_len)
    }
    class AMPRunner {
        +learn(iterations, init_at_random_ep_len)
    }
    class DistillationRunner {
        +learn(iterations, init_at_random_ep_len)
    }
    class CustomEnv {
        +get_observations()
        +step(actions)
    }
    VecEnv <|-- CustomEnv
    VecEnv <-- OnPolicyRunner : 依赖
    VecEnv <-- AMPRunner : 依赖
    VecEnv <-- DistillationRunner : 依赖
    AMPRunner --|> OnPolicyRunner
    DistillationRunner --|> OnPolicyRunner

观测数据模型:TensorDict 与观测分组

rsl_rl 采用 tensordict.TensorDict 作为 get_observations()step() 返回的观测载体,而非原生字典或裸张量。这一选择的架构意义在于:TensorDict 不仅支持类字典的键值访问,还完整兼容 PyTorch 的张量操作语义(如 .to(device).cat()、广播),使得 Runner 可以在不感知具体观测内容的情况下,批量处理跨设备迁移与存储拼接。

更关键的设计是观测分组(Observation Groups)机制。环境返回的 TensorDict 中,每个键代表一个逻辑观测组(如 base_lin_velprojected_gravitycommands 等)。Runner 的配置项 obs_groups 将这些原始组映射到不同的观测集(Observation Sets),以支持多网络架构:

观测集键名 消费者 典型用途
policy Actor / Student 策略网络输入,直接决定动作输出
critic Critic 价值网络输入,通常包含特权信息(privileged observations)
teacher Teacher 蒸馏框架中教师策略的输入
rnd_state RND Network 随机网络蒸馏模块用于计算内在奖励的表征

obs_groups 配置中缺少某个观测集时,resolve_obs_groups() 会执行降级策略:首先检查是否存在与观测集同名的观测组,若存在则直接复用;否则回退到 policy 的观测配置。这种设计在保证灵活性的同时,降低了简单场景下的配置负担。

Sources: vec_env.py, utils.py

环境交互协议:Step 语义与 Extras 字典

step() 方法是环境状态转移的唯一通道,其返回的四元组具有严格的形状与语义约定:

observations, rewards, dones, extras = env.step(actions)

其中 actions 的形状必须为 (num_envs, num_actions),返回的 rewardsdones 形状必须为 (num_envs,)observationsTensorDict,其批处理维度同样为 num_envs

extras 字典承担了标准接口之外的扩展职责,当前框架约定两个关键键:

Sources: vec_env.py

与训练运行器的协作模式

VecEnv 并非孤立存在,而是深度嵌入 Runner 的生命周期。理解其协作模式有助于把握接口设计的上下文。

初始化阶段:环境探查

OnPolicyRunner.__init__ 中,算法与策略网络的构建完全依赖于环境自描述。Runner 首先调用 env.get_observations() 获取一个样本 TensorDict,随后通过 resolve_obs_groups() 校验配置,并最终将观测结构传入 ActorCritic 构造函数以自动推断各层输入维度。这意味着环境实现必须在首次调用 get_observations() 时即返回完整且形状正确的张量,即使内部模拟尚未开始。

sequenceDiagram
    participant Runner as OnPolicyRunner
    participant VecEnv as VecEnv (外部实现)
    participant Utils as resolve_obs_groups
    participant Policy as ActorCritic

    Runner->>VecEnv: get_observations()
    VecEnv-->>Runner: sample_obs (TensorDict)
    Runner->>Utils: resolve_obs_groups(sample_obs, cfg, default_sets)
    Utils-->>Runner: 校验后的 obs_groups
    Runner->>Policy: __init__(obs, obs_groups, num_actions, ...)
    Policy-->>Runner: policy instance

训练阶段:Rollout 循环

learn() 的主循环中,Runner 与环境的交互呈现经典的"采样-存储-更新"节奏。step() 的返回值被立即分发到三个下游组件:alg.process_env_step() 将 transition 写入 RolloutStorage;logger.process_env_step() 聚合 episode 级统计;dones 信号则可能触发环境内部的自动重置逻辑(由具体实现决定,而非接口强制)。

sequenceDiagram
    participant Runner as OnPolicyRunner
    participant Alg as PPO
    participant Env as VecEnv
    participant Logger as Logger

    loop num_steps_per_env
        Runner->>Alg: act(obs)
        Alg-->>Runner: actions
        Runner->>Env: step(actions)
        Env-->>Runner: obs, rewards, dones, extras
        Runner->>Alg: process_env_step(...)
        Runner->>Logger: process_env_step(...)
    end
    Runner->>Alg: compute_returns(obs)
    Runner->>Alg: update()

此外,episode_length_buf 在训练初始化时支持随机化(init_at_random_ep_len=True),避免所有并行环境在同一时刻同步重置,从而平滑批处理的样本分布。这要求环境实现正确维护并暴露该缓冲区。

Sources: on_policy_runner.py, on_policy_runner.py

设计原则与实现约束

基于以上分析,实现 VecEnv 接口时需遵守以下约束:

  1. 张量形状严格性step() 输入输出必须满足 (num_envs, ...) 的批处理维度。任何形状不匹配都会在 Runner 的存储分配阶段引发运行时错误。
  2. 设备一致性:环境应保证自身张量位于 self.device 上;Runner 会显式调用 .to(self.env.device).to(self.device) 进行双向传输,但跨设备拷贝的隐性开销需实现者留意。
  3. 观测不可变性get_observations() 应仅读取当前状态,不应触发物理步进或改变 episode_length_buf,以保证 Runner 在初始化阶段的多次调用是安全的。
  4. 终止信号双通道:若环境支持固定回合长度,应在 extras["time_outs"] 中准确标记时间截断,避免值估计偏差。
  5. 配置可序列化self.cfg 将被直接传递给日志系统,建议避免包含不可序列化的对象(如函数句柄或张量)。

下一步阅读

掌握 VecEnv 的抽象接口后,建议按以下路径深入: