编写高效的 RL 环境#

当数千个环境在单张 GPU 上并行运行时，决定吞吐量的与其说是 env.step() 计算了什么，不如说是它迫使 GPU 停下来等待什么。一个在 n_envs 个状态上运行批量物理内核的 step 是很快的；但只要它把某个值拷回 CPU、分配一个新张量，或者在 Python 里逐个环境循环，同一步就会变慢。本页解释这一性能模型，以及让步进循环始终留在设备上的那些模式。

全文的参考环境是四足运动示例 examples/locomotion/go2_env.py。它构建了 n_envs 份 Go2 机器人副本（批量构建的工作方式见 并行仿真），其 step 在执行过程中没有触发任何一次主机与设备之间的同步。

性能模型#

Genesis World 的物理仿真在 GPU 上运行。你的环境代码则以 Python 在 CPU 上运行，并异步地向 GPU 派发工作：scene.step() 和每一个张量操作都会被排入队列并立即返回，此时 GPU 尚未完成计算。有三类操作会打断这条流水线，而它们损耗的吞吐量是更快的物理仿真也补不回来的。

  • 主机与设备同步： 任何需要在 CPU 上取得张量数值的操作（.item()、.cpu()、.tolist()、bool(t)、print(t)，以及返回动态尺寸张量的 .nonzero()）都会阻塞 CPU，直到 GPU 排空整个队列。每步一次这样的调用，乘以数千步，就足以主导实际耗时。

  • 内存分配： 在循环内创建张量（torch.zeros(...)、torch.tensor(...)，以及大多数新建的读取）会走 CUDA 缓存分配器，而它在寻找或释放内存时需要与尚未完成的工作同步。

  • Python 开销： 逐个环境遍历的 for 循环会发起 n_envs 倍的内核启动次数，却只得到单个批量操作就能产生的结果。此时占主导的是每次启动的固定开销，而不是算术运算。

这三类问题指向同一个方向：一次性对整个 ([n_envs,] ...) 张量做操作，让每个张量都留在设备上，并且绝不让任何数值在步进循环内回到主机。

让张量留在设备上#

所有缓冲区都应在 Genesis World 的设备上、用它的 dtype 分配，这样后续就不会发生隐式的搬迁或类型转换。gs.device 是当前后端对应的设备，gs.tc_float、gs.tc_int 和 gs.tc_bool 则是与 gs.init 时所选精度相匹配的 torch dtype。

self.base_lin_vel = torch.empty((self.num_envs, 3), dtype=gs.tc_float, device=gs.device)
self.reset_buf = torch.ones((self.num_envs,), dtype=gs.tc_bool, device=gs.device)
self.episode_length_buf = torch.empty((self.num_envs,), dtype=gs.tc_int, device=gs.device)

读取实体状态时，请使用批量访问器，而不是把它取出来再查看。每个访问器都返回形状为 ([n_envs,] ...) 的设备张量，可以直接送入下一步操作：

self.base_pos = self.robot.get_pos()                          # shape (n_envs, 3)
self.dof_pos = self.robot.get_dofs_position(self.motors_dof_idx)  # shape (n_envs, n_motors)
self.dof_vel = self.robot.get_dofs_velocity(self.motors_dof_idx)

热路径上常用的刚体实体读取包括 get_pos / get_quat（基连杆的世界位姿）、get_vel / get_ang（基连杆的线速度与角速度），以及 get_dofs_position / get_dofs_velocity（逐自由度状态）。对于传感器输出，当你需要同时观测多个传感器时，优先使用批量接口 scene.read_sensors() 或 entity.read_sensors()，而不是逐个调用 sensor.read()：每个传感器类别只返回一个批量张量，摊销效果比 N 次单独调用更好。

预分配缓冲区，原地写入#

step 和 reset 会写入的每个张量都应在 __init__ 中按最终形状和 dtype 一次性分配好。对于每步都会被完整覆盖的缓冲区，使用 torch.empty(...)，因为没有理由为清零付出代价；只有初始值有意义的地方（例如累加器）才使用 torch.zeros(...)。

此后，请写入已有的存储，而不是替换它。原地操作可以避免每步分配，并让张量的身份保持稳定——当日志器、记录器或观测构造器持有它的视图时，这一点很重要。

  • buf.zero_() / buf.fill_(x)： 原地重置缓冲区。

  • dst.copy_(src)： 在不分配新内存的情况下覆盖 dst。

  • buf.masked_fill_(mask, x)： 在布尔掩码为真的位置写入 x。

  • torch.where(mask, a, b, out=dst)： 选择结果写入已有张量，而不是返回新张量。

Go2 的奖励累加展示了这一模式：rew_buf 先被原地清零，随后每个奖励项累加进去。这个循环遍历的是奖励项（数量很少），而不是环境；并且每个奖励项本身就是作用于全部 n_envs 的批量操作。

self.rew_buf.zero_()
for name, reward_func in self.reward_functions.items():
    rew = reward_func() * self.reward_scales[name]
    self.rew_buf += rew
    self.episode_sums[name] += rew

envs_idx 使用布尔掩码，而不是索引张量#

终止判定是同步最常偷偷混入的经典位置。用 (condition).nonzero() 计算哪些环境需要重置会强制一次同步，因为主机必须知道产生了多少个索引才能确定结果尺寸。请让选择结果全程保持为形状 (n_envs,) 的布尔掩码：用批量比较构造它，用 |= 合并条件，然后直接传给各个 setter。

self.reset_buf = self.episode_length_buf > self.max_episode_length
self.reset_buf |= torch.abs(self.base_euler[:, 1]) > self.env_cfg["termination_if_pitch_greater_than"]
self.reset_buf |= torch.abs(self.base_euler[:, 0]) > self.env_cfg["termination_if_roll_greater_than"]
self.reset_buf |= self.scene.rigid_solver.get_error_envs_mask()

最后一行把数值发散也纳入其中。当求解器在某个环境中遇到 NaN 或约束求解失败时，scene.rigid_solver.get_error_envs_mask() 会返回一个标记受影响环境的 (n_envs,) 布尔掩码。用与正常回合结束相同的机制重置这些环境，这样单个发散的环境会自行终止，而不会让整批环境崩溃。

Genesis World 的各种 setter 和求解器都直接接受布尔 envs_idx 掩码，因此你无需为了使用掩码而先构造出索引。

无同步地重置一批环境#

重置会把初始状态写入选中的环境。无同步的组合是：布尔 envs_idx 掩码、预分配的源张量，以及 setter 上的两个关键字开关：

self.robot.set_qpos(self.init_qpos, envs_idx=envs_idx, zero_velocity=True, skip_forward=True)

  • zero_velocity=True： 在同一次写入中顺带清零速度，而不必再调用一次 setter。

  • skip_forward=True： 推迟正向运动学，使其在下一次 scene.step() 时统一执行一次，而不是在重置期间的每次 setter 调用中执行。

只需根据是全量重置还是部分重置分支一次。重置所有环境时，传入 envs_idx=None；实现会走更快的全量覆盖路径，你也可以对自己的缓冲区使用更廉价的 zero_() / fill_()，而不必做掩码写入。

if envs_idx is None:
    self.actions.zero_()
    self.episode_length_buf.zero_()
    self.reset_buf.fill_(True)
else:
    self.actions.masked_fill_(envs_idx[:, None], 0.0)
    self.episode_length_buf.masked_fill_(envs_idx, 0)
    self.reset_buf.masked_fill_(envs_idx, True)

如果希望用一次粗粒度重置恢复所有求解器状态，scene.rigid_solver.set_state(f, state, envs_idx=mask, partial=True) 是等价的批量接口。请保留 partial=True：partial=False 会重建整个场景的辅助状态，速度明显更慢。

高效地下发命令#

控制下发是写入侧的热路径。Go2 策略为每个被驱动的关节输出一个动作，将其缩放，再叠加默认姿态作为偏置，得到位置目标：

target_dof_pos = exec_actions * self.env_cfg["action_scale"] + self.default_dof_pos
self.robot.control_dofs_position(target_dof_pos[:, self.actions_dof_idx], slice(6, 18))

有两个细节能让这一过程保持廉价。

  • 用 slice 而非索引张量来指定 DOF。 control_dofs_position 的第三个参数选择要驱动哪些 DOF。slice(6, 18) 是一个零成本的视图；索引张量则会强制一次 gather。Go2 让被驱动的 DOF 保持连续，并在初始化时一次性预计算 actions_dof_idx = torch.argsort(self.motors_dof_idx)，这样按关节名排序的策略输出可以在调用前被重排成这种连续布局。

  • 尽可能复用目标缓冲区。 每次 a * scale + b 都会分配一个新张量。在缓冲区的身份必须保持稳定的地方，请改用 out= 写入，把目标构造进预分配的张量。（会改变形状的操作，例如拼接观测向量时的 torch.concatenate，必然要分配内存；这是可以接受的例外，而不是需要去对抗的规则。）

刚体实体上的零拷贝命令写入接口包括：用于 PD 目标和直接力的 control_dofs_position、control_dofs_velocity 和 control_dofs_force，以及用于直接写入状态的 set_dofs_position / set_dofs_velocity / set_qpos。

训练时开启性能模式#

环境定稿之后，gs.init(performance_mode=True) 会把此时已固定的张量形状编译进内核，仿真大约快 30%。代价是任何对场景的改动都会触发重新编译，可能耗时数分钟。在研究、调试和交互式开发时请保持关闭；在长时间训练和生产运行中开启它——此时场景已经固定，一次重新编译换来的更快步进会惠及整个运行的剩余部分。其他 gs.init 选项见 你好，Genesis World。

用性能分析器验证#

不要猜同步发生在哪里。在若干步周围挂上 PyTorch 性能分析器，观察内核启动之间的空隙，这些空隙就是 CPU 在等待 GPU。性能分析的配置方法以及如何解读 Quadrants 的内核耗时，见 仿真性能分析。

另见#