在 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_vel、projected_gravity、commands 等)。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),返回的 rewards 与 dones 形状必须为 (num_envs,)。observations 为 TensorDict,其批处理维度同样为 num_envs。
extras 字典承担了标准接口之外的扩展职责,当前框架约定两个关键键:
"time_outs"(torch.Tensor, shape(num_envs,)): 标识因达到max_episode_length而被强制截断的终止,而非环境本身进入失效状态。这一区分对值函数 bootstrap 至关重要——时间截断不应触发值函数的终止假设。"log"(dict[str, float | torch.Tensor]): 供环境向日志系统透传诊断信息。键名应以/开头以实现命名空间隔离,值为张量时 Runner 会自动取均值记录。
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 接口时需遵守以下约束:
- 张量形状严格性:
step()输入输出必须满足(num_envs, ...)的批处理维度。任何形状不匹配都会在 Runner 的存储分配阶段引发运行时错误。 - 设备一致性:环境应保证自身张量位于
self.device上;Runner 会显式调用.to(self.env.device)和.to(self.device)进行双向传输,但跨设备拷贝的隐性开销需实现者留意。 - 观测不可变性:
get_observations()应仅读取当前状态,不应触发物理步进或改变episode_length_buf,以保证 Runner 在初始化阶段的多次调用是安全的。 - 终止信号双通道:若环境支持固定回合长度,应在
extras["time_outs"]中准确标记时间截断,避免值估计偏差。 - 配置可序列化:
self.cfg将被直接传递给日志系统,建议避免包含不可序列化的对象(如函数句柄或张量)。
下一步阅读
掌握 VecEnv 的抽象接口后,建议按以下路径深入:
- 若需实现自定义环境接入,请参阅 向量化环境接入指南,了解具体的继承与调试方法。
- 若关注算法如何消费这些观测与奖励,可阅读 PPO 算法实现与训练流程。
- 若希望理解 Runner 如何编排完整训练生命周期,请参阅 训练运行器生命周期管理。
- Rollout 数据的内部组织方式记载于 Rollout 数据存储与 Transition。