rsl_rl 的训练循环并非针对单环境编写，而是面向批量同步执行的向量化环境（Vectorized Environment）设计的。本页面向具备 PyTorch 基础的中级开发者，系统讲解如何让自己的仿真环境满足 VecEnv 抽象接口，从而无缝接入 PPO、AMP、RND 等算法的训练管线。核心思路非常直接：你的环境只需要返回符合约定的 TensorDict 观测、接收 (num_envs, num_actions) 动作张量，并暴露若干关键属性，运行器（Runner）会自动完成策略网络构建、Rollout 采集与参数更新。

核心交互架构

在 rsl_rl 中，环境不是被算法直接调用的工具对象，而是与 Runner 平级、通过严格接口契约交互的协作方。训练启动时，OnPolicyRunner（或其子类如 AMPRunner）首先向环境查询一次观测以推断网络输入维度，随后进入固定的 Rollout 循环：策略采样动作 → 环境批量步进 → 算法处理 Transition。这一交互模式要求环境内部维持 num_envs 个并行实例的状态，并在每次 step() 中同步推进。

graph LR
    subgraph 训练管线
        R[OnPolicyRunner]
        A[Algorithm<br/>PPO / PPOAMP]
        S[RolloutStorage]
    end
    subgraph 环境侧
        E[VecEnv<br/>你的环境实现]
    end
    R -->|get_observations| E
    R -->|构造算法| A
    E -->|TensorDict obs| R
    R -->|act| A
    A -->|actions| R
    R -->|step(actions)| E
    E -->|obs, rewards, dones, extras| R
    R -->|process_env_step| A
    A -->|存储| S

Sources: vec_env.py, on_policy_runner.py

必须实现的接口与属性

VecEnv 是一个抽象基类，定义在 rsl_rl.env.vec_env 中。你的环境类必须继承它，并实现两个抽象方法，同时确保以下属性在实例化后可用。

Sources: vec_env.py

观测数据组织与 obs_groups

rsl_rl 采用 tensordict.TensorDict 作为观测容器，以支持多模态、多分组的复杂观测结构。你的 get_observations() 和 step() 返回的观测必须是一个 TensorDict，其中每个键代表一个观测组（observation group），值为形状 (num_envs, ...) 的张量。

训练配置中的 obs_groups 负责把“观测组”映射到“观测集（observation set）”。算法会依据观测集决定哪些数据喂给策略网络、哪些喂给价值网络。系统预留的观测集名称包括：

OnPolicyRunner 初始化时会调用 resolve_obs_groups() 校验配置：若 obs_groups 中缺少 policy 键会直接报错；若缺少 critic 等默认集，则会尝试自动回退（优先使用同名的观测组，否则复用 policy 的组）。建议在环境侧暴露语义清晰的观测组名称，并在训练配置中显式声明映射关系，以避免隐式回退带来的调试困扰。

Sources: vec_env.py, utils.py

环境步进协议详解

step(actions) 是训练循环中最高频的调用点，其签名与返回值必须严格遵守以下约定：

def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]:

• actions: 形状 (num_envs, num_actions)，由策略网络输出。注意 runner 会在调用前将动作张量转移到 self.env.device，因此环境内部应确保该设备与自身计算设备一致。
• observations: TensorDict，与 get_observations() 的结构一致。
• rewards: 形状 (num_envs,) 的 torch.Tensor。
• dones: 形状 (num_envs,) 的 torch.Tensor，布尔或浮点类型均可。
• extras: dict，用于传递辅助信息。系统识别以下两个特殊键：
  • "time_outs" (torch.Tensor): 标记因达到 max_episode_length 而终止的环境（区别于任务失败导致的终止）。这对 bootstrap 价值估计至关重要。
  • "log" (dict[str, float | torch.Tensor]): 供日志系统记录的自定义指标，键名建议以 / 开头做命名空间隔离。

Sources: vec_env.py, on_policy_runner.py

设备管理与张量流向

rsl_rl 支持 runner 与环境位于不同设备的场景（例如环境在 cuda:0 进行物理仿真，而策略优化在 cuda:1）。OnPolicyRunner 的 learn() 循环遵循固定的设备转换协议：

1. 策略网络在 runner.device 上采样动作；
2. 动作被显式转移到 env.device 后传入 env.step()；
3. 环境返回的观测、奖励、dones 再被转移回 runner.device 供算法处理。

因此，你的环境必须确保 self.device 准确反映内部张量所在的设备，并且 step() 返回的张量已经在该设备上。若环境内部使用 GPU 加速（如 Isaac Lab、Isaac Gym），通常将 device 设为对应 CUDA 设备即可；若环境在 CPU 上运行，runner 会自动完成 cpu → cuda 的搬运，但会带来额外的同步开销。

Sources: on_policy_runner.py

最小接入实现步骤

下面给出将自定义环境接入 rsl_rl 的标准流程。你可以把它当作一份 checklist 逐条落实。

flowchart TD
    A[继承 VecEnv 抽象类] --> B[定义 num_envs / num_actions / device / cfg]
    B --> C[初始化 episode_length_buf 与 max_episode_length]
    C --> D[实现 get_observations 返回 TensorDict]
    D --> E[实现 step 接收 actions 并返回四元组]
    E --> F[在 extras 中提供 time_outs 与 log]
    F --> G[配置 obs_groups 映射观测组到 policy/critic 等集合]
    G --> H[实例化 OnPolicyRunner 并传入环境]
    H --> I[调用 runner.learn 开始训练]

骨架代码示例

from rsl_rl.env import VecEnv
from tensordict import TensorDict
import torch

class MyVecEnv(VecEnv):
    def __init__(self, cfg, num_envs=4096, device="cuda:0"):
        self.cfg = cfg
        self.num_envs = num_envs
        self.num_actions = 12
        self.device = device
        self.max_episode_length = 1000
        # 记录每个环境当前回合已走步数
        self.episode_length_buf = torch.zeros(num_envs, device=device, dtype=torch.long)
        # ... 初始化物理仿真、状态缓冲区等 ...

    def get_observations(self) -> TensorDict:
        # 假设环境内部已维护当前状态
        return TensorDict({
            "base_lin_vel": self.base_lin_vel,      # [num_envs, 3]
            "base_ang_vel": self.base_ang_vel,      # [num_envs, 3]
            "projected_gravity": self.projected_gravity,  # [num_envs, 3]
            "joint_pos": self.joint_pos,            # [num_envs, 12]
            "joint_vel": self.joint_vel,            # [num_envs, 12]
            "commands": self.commands,              # [num_envs, 3]
        }, batch_size=[self.num_envs])

    def step(self, actions: torch.Tensor):
        # actions: [num_envs, 12]，已在 self.device 上
        # ... 执行物理步进、计算奖励、判定终止 ...
        obs = self.get_observations()
        rewards = self.compute_rewards()            # [num_envs]
        dones = self.check_terminations()           # [num_envs]
        time_outs = self.check_timeouts()           # [num_envs]
        extras = {
            "time_outs": time_outs,
            "log": {"/episode/reward": rewards.mean()},
        }
        # 更新回合长度计数（runner 会负责重置）
        self.episode_length_buf += 1
        return obs, rewards, dones, extras

实际集成时，观测组的命名不必与示例完全一致，但必须保证训练配置中的 obs_groups 引用的键名与环境返回的 TensorDict 键名一一对应。

Sources: vec_env.py

配置属性与生命周期属性

cfg 和 episode_length_buf 虽然看似是“辅助信息”，但在训练流程中扮演明确角色：

• cfg: OnPolicyRunner 在构造 Logger 时会将 env.cfg 传入，用于在日志目录中序列化环境配置，方便实验复现。它可以是任意对象（如 omegaconf.DictConfig），不要求特定结构。
• episode_length_buf: 若调用 learn(..., init_at_random_ep_len=True)，runner 会直接原地修改该张量，将各环境的初始回合长度随机化，以避免所有环境在同一时刻重置造成的数据相关性峰值。你的环境 step() 实现应当依赖或兼容该缓冲区。
• max_episode_length: 当为标量时，所有环境共享同一上限；当为 torch.Tensor 时，支持每个环境实例拥有独立的动态上限。这在课程学习或地形难度渐进场景中非常有用。

Sources: on_policy_runner.py

算法特定扩展要求

对于基础 PPO 训练，上述最小实现已足够。若计划使用 AMP、RND 或对称性增强等高级特性，runner 与配置解析函数会尝试读取环境的更多内部属性：

这些访问点并未写入 VecEnv 基类，属于“约定俗成”的实现细节。若你的环境使用标准 wrapper 模式（如 Isaac Lab 的 ManagerBasedRLEnv），通常通过 unwrapped 链或 env.env 链即可暴露底层属性；若自行从零实现，建议在类中提供 .unwrapped 属性指向自身，以兼容上述解析逻辑。

Sources: rnd.py, amp.py, symmetry.py, amp_runner.py

常见问题排查

Sources: utils.py, vec_env.py

下一步阅读建议

完成环境接入后，你可以继续深入以下主题：

• 若需要理解 Runner 如何驱动环境完成 Rollout 并更新策略，请阅读 PPO 算法实现与训练流程。
• 若对 VecEnv 抽象接口的设计哲学与多 GPU 场景下的环境复用感兴趣，请阅读 向量化环境抽象接口。
• 若已接入 AMP 并需要理解判别器观测与演示数据的要求，请阅读 AMP 对抗动作先验算法。
• 若需配置观测归一化、网络激活函数等策略细节，请阅读 观测归一化与网络工具函数。