Warp 实验性环境

Warp 实验性环境#

备注

Warp 环境基础设施位于 isaaclab_experimental 和 isaaclab_tasks_experimental ，属于实验性功能。

实验性扩展引入了 以 Warp 为核心 、支持 CUDA 图捕获的环境基础设施。环境侧计算（观测、奖励、重置、动作）全部以纯 Warp 内核运行，消除 Python 开销，并通过 CUDA 图捕获获得更高吞吐量。

本页中的 稳定版 指 isaaclab 和 isaaclab_tasks 包中基于 Torch 的标准实现，即未传入 --frontend warp 时使用的默认运行时。对应的 Warp 实现位于 isaaclab_experimental 和 isaaclab_tasks_experimental 包。

工作流程#

支持两种环境工作流程：

直接式工作流程 — 使用 DirectRLEnvWarp 基类，在环境类中通过 Warp 内核直接实现步进循环、观测、奖励和重置。

基于管理器的工作流程 — 使用 ManagerBasedRLEnvWarp 基类，将 MDP 项定义为独立的 Warp 内核函数，再通过配置组合。

可用环境#

直接式 Warp 环境#

直接式任务共享稳定版任务配置。 --frontend warp 将稳定版 <task>_direct_env:<Name>Env 入口解析为对应的 <task>_warp_env:<Name>WarpEnv 实现，只替换环境类。实现无法遵循此约定时，注册可以使用可选的 warp_entry_point 覆盖。以下稳定版任务具有 Warp 实现：

  • Isaac-Cartpole-Direct — Cartpole 平衡

  • Isaac-Ant-Direct — Ant 运动

  • Isaac-Humanoid-Direct — 人形机器人运动

  • Isaac-Reorient-Cube-Allegro-Direct — Allegro 手立方体重定向

基于管理器的 Warp 执行（ --frontend warp ）#

基于管理器的任务无需另行注册 Warp 版本：共享强化学习命令行提供 --frontend {torch,warp} ， --frontend warp 在构建时将 稳定版 任务配置适配到 Warp 运行时，将各 MDP 项替换为对应的 Warp 实现。请通过 presets=newton_mjwarp 显式选择 Newton 求解器。以下稳定版任务的全部项都有对应实现：

  • Isaac-Cartpole

  • Isaac-Ant

  • Isaac-Humanoid

  • Isaac-Reach-Franka

  • Isaac-Reach-UR10

  • Isaac-Velocity-Flat-AnymalD

  • Isaac-Velocity-Flat-Cassie

  • Isaac-Velocity-Flat-G1

  • Isaac-Velocity-Flat-H1

  • Isaac-Velocity-Flat-UnitreeGo2

以下贡献任务也已完整覆盖：

  • IsaacContrib-Velocity-Flat-AnymalB

  • IsaacContrib-Velocity-Flat-AnymalC

  • IsaacContrib-Velocity-Flat-UnitreeA1

  • IsaacContrib-Velocity-Flat-UnitreeGo1

缺少对应项会直接报错，并列出受影响的项，因此仅部分覆盖的任务在构建时就会失败，而不会静默改变行为。目前无法适配使用 height_scan 的崎岖地形速度配置，因为该观测项没有 Warp 对应实现。应检查任务实际使用的项，不要假定所有地形配置都有同样限制。

快速开始#

# Direct workflow: stable task, warp env implementation
uv run isaaclab train --rl_library rsl_rl \
    --task Isaac-Cartpole-Direct --frontend warp presets=newton_mjwarp --num_envs 4096
uv run isaaclab play --rl_library rsl_rl \
    --task Isaac-Cartpole-Direct --frontend warp presets=newton_mjwarp \
    --num_envs 32 --checkpoint latest --visualizer newton_gl

# Manager-based workflow: stable task on the warp runtime
uv run isaaclab train --rl_library rsl_rl \
    --task Isaac-Velocity-Flat-AnymalD --frontend warp presets=newton_mjwarp --num_envs 4096
uv run isaaclab play --rl_library rsl_rl \
    --task Isaac-Velocity-Flat-AnymalD --frontend warp presets=newton_mjwarp \
    --num_envs 32 --checkpoint latest --visualizer newton_gl

支持所有具有 Warp 兼容包装器的强化学习库：RSL-RL、RL Games、SKRL 和 Stable-Baselines3。

备注

Warp 路径的 train 和 play 均不接受 --video ，视频录制需要标准 Torch 前端。要录制轨迹，请通过 --frontend torch 回放同一检查点；见 录制视频 。

性能对比#

下表是稳定版（Torch/管理器）与 Warp（CUDA 图捕获）变体的历史步进时间对比，两者都运行于 Newton 物理后端。测量使用 4096 个环境、300 次迭代，覆盖所测任务，而非当前所有支持的任务。这些数值作为原始基准记录保留，不代表最新 develop 版本的测量结果。更新结果时，应记录硬件及 Isaac Lab、Newton、Warp 的版本，并在相同版本上对比两个前端。共享基础库的改进可能同时惠及两者，不能全部归因于 Warp 前端。

备注

向 Warp 迁移仍在进行中。部分组件（例如场景写入、执行器模型）尚未迁移到 Warp 内核，仍通过 Torch 运行。随着这些组件迁移，预计性能会进一步提高。

环境

类型

稳定版步进时间（μs）

Warp 步进时间（μs）

变化

Isaac-Cartpole-Direct

直接

5,274

4,331

-17.88%

Isaac-Ant-Direct

直接

6,368

3,128

-50.88%

Isaac-Humanoid-Direct

直接

13,937

10,783

-22.63%

Isaac-Reorient-Cube-Allegro-Direct

直接

82,950

74,570

-10.10%

Isaac-Cartpole

管理器

7,971

3,642

-54.31%

Isaac-Ant

管理器

9,781

4,672

-52.23%

Isaac-Humanoid

管理器

17,653

12,505

-29.16%

Isaac-Reach-Franka

管理器

11,458

7,813

-31.83%

Isaac-Velocity-Flat-AnymalD

管理器

32,294

23,977

-25.75%

Isaac-Velocity-Flat-Cassie

管理器

17,320

10,706

-38.19%

Isaac-Velocity-Flat-G1

管理器

34,487

27,300

-20.84%

Isaac-Velocity-Flat-H1

管理器

22,202

15,864

-28.55%

Isaac-Velocity-Flat-UnitreeGo2

管理器

15,221

9,966

-34.52%

哪些工作流程收益最大#

性能收益来自消除环境步进循环中的 Python/Torch 开销，因此原先逐内核 CPU 开销在步进时间中占比越高，环境收益越大。上表可解读为：

  • 基于管理器的经典强化学习 （Cartpole、Ant）— 收益最大，步进时间减少 52%–54%。其奖励/观测项多而小，每项计算量低，逐次启动的 CPU 开销主导了稳定版基线。

  • 基于管理器的运动任务 （AnymalD、G1、H1、Cassie、UnitreeGo2）— 稳定减少 21%–38%。MDP 项更多，但底层物理步也更重，因此节省 Python 开销的相对比例减小。

  • 直接式工作流程 — 收益随环境步进主体中的 Python 占比变化（Ant 减少 51%、Cartpole 减少 18%、Allegro 手减少 10%）。多数工作已写为 GPU 内核的直接式环境收益较小，包含大量 Python 状态处理的环境收益较大。

  • 计算密集 / 场景写入密集的环境 （Allegro 手、大型人形机器人）— 相对收益较小，因为 Warp 侧节省的时间在较重的步进中占比较低。仍通过 Torch 运行的组件（场景写入、执行器模型）目前限制了最低步进时间；其余组件迁移到 Warp 后预计会改善。

如果环境步进时间主要消耗在物理或场景 I/O 上，预期收益较小；如果包含大量小型 MDP 项或大量 Python 步进逻辑，预期收益较大。决定迁移前，请使用下述基准测试流程测量自己的任务。

限制#

Warp 环境路径处于实验阶段，具有以下已知约束。这些约束仅针对 Warp 环境；Newton 物理的成熟度和专题指南见 Newton 。

物理后端

  • 仅支持 Newton。 Warp 环境路径不支持 PhysX。资产和传感器的 class_type 字段解析为 isaaclab_physx.* 类，它们依赖 omni.physics.tensors （Warp 运行时不初始化的 Kit 模块）；多个 Warp API（环境掩码重置、CUDA 图捕获）还要求使用 Newton 关节体。请为配置设置 Newton 物理部分，或使用带类型选择器 physics=newton_mjwarp ；任务没有 Newton 物理预设时，后者会明确报错。

MDP 覆盖范围

  • 目前仅实现了 可用的 Warp MDP 项 中列出的项。稳定版环境如果依赖尚未迁移的项，需要先移植这些项，才能使用 Warp 路径运行。

  • 部分场景侧操作（资产写入、执行器模型、某些传感器类型）仍通过 Torch 运行。它们参与步进，但尚未捕获进图，限制了观测到的最低步进时间。

  • 依赖 Kit RTX 渲染器的传感器（基于相机的观测）不能与 Warp 环境路径组合，因为它们需要 Kit，而 Warp 运行时不初始化 Kit。

与稳定版的 API 差异

  • 重置事件使用布尔 env_mask （ wp.array(dtype=wp.bool) ），而非 env_ids 列表。这是图捕获安全性的要求：可变长度索引会改变图拓扑并破坏回放。

  • 所有缓冲区必须在 __init__ 中预分配。捕获的步进循环中不允许动态分配，因此观测、奖励和终止输出的维度必须在环境初始化时确定。

  • 项函数将结果写入预分配的 out 缓冲区，而不是返回张量。内核与启动函数的编写模式见 迁移已有环境 。

  • 捕获的步进循环内，代码必须遵循图捕获安全规则：不能使用 wp.to_torch、Torch 运算、惰性求值属性，也不能根据 GPU 数据进行 Python 分支判断。完整规则见 图捕获安全性 。

对环境进行基准测试#

上方性能表由 isaaclab benchmark training 生成，它运行固定迭代次数并报告步进时间统计。决定迁移前，可用同一命令估计自己的任务收益。

单任务 A/B 对比

对同一稳定版任务 ID 运行两次，仅改变 --frontend 。两次 运行都传入 presets=newton_mjwarp ，确保物理后端一致，使测得的差异只反映前端（管理器流程与 CUDA 图捕获）：

# Torch frontend (stable managers)
uv run isaaclab benchmark training \
    --rl_library rsl_rl \
    --task <Task-Name> \
    --num_envs 4096 \
    --max_iterations 500 \
    --benchmark_formatter summary \
    --output_path benchmarks/torch \
    presets=newton_mjwarp

# Warp frontend (warp managers, CUDA-graph captured) — same task id
uv run isaaclab benchmark training \
    --rl_library rsl_rl \
    --task <Task-Name> \
    --frontend warp \
    --num_envs 4096 \
    --max_iterations 500 \
    --benchmark_formatter summary \
    --output_path benchmarks/warp \
    presets=newton_mjwarp

summary 格式化器输出步进时间（最小/平均/最大）和总吞吐量。比较两次运行的“step time”，估计每个环境步的收益。

遍历所有可用任务

对稳定版任务集合中的各任务（Cartpole、Ant、人形机器人、运动、操作）运行 isaaclab benchmark training ，再加入 --frontend warp 重新运行，然后对比两个输出目录。

输出中应关注的内容

  • 步进时间（最小/平均/最大）：主要指标，表示每个环境步的开销。

  • 迭代时间：包含策略更新，可用于评估端到端训练吞吐量。

  • 捕获开销：Warp 运行的前几次迭代包含 CUDA 图捕获成本；比较稳态数值时应排除它们。

迁移前估计收益

如果还不能运行 Warp 变体（例如任务尚未移植），可测量稳定版步进时间，并查看时间消耗的位置：

  • num_envs * step_time 主要消耗在物理计算上 → Warp 收益预计较小。

  • step_time 主要消耗在 manager.compute_* 调用上 → 收益预计较大，因为 Warp 管理器正是将这些调用替换为捕获的内核启动。

需要区分环境开销与策略计算时，可在 runtime.py 上使用 --num_steps ，进行无策略的步进时间微基准测试。

迁移已有环境#

本节介绍 Warp 优先环境基础设施使用的主要约定和编写模式，供迁移已有 Torch 环境或直接创建 Warp 环境时参考，包括项目布局，观测、奖励、事件、终止和动作共用的内核与启动函数模式，以及图捕获安全规则和行为一致性测试。

设计依据#

Warp 环境路径围绕 CUDA 图捕获 构建。CUDA 图在捕获阶段记录一系列 GPU 操作（内核启动、内存拷贝），随后通过一次启动回放整个序列。这消除了逐内核的 CPU 开销：通常每次操作需要 20–200 μs 的参数验证、内核选择和缓冲区设置，在图实例化时只执行一次，之后每次回放复用（总计约 10 μs）。捕获期间执行的全部 CPU 端代码（Python 逻辑、Torch 分发）在回放时都会被跳过。Warp 图捕获 API（ wp.ScopedCapture ）见 Warp 并发文档 。

Warp 基础设施的所有设计都遵循这一约束：步进循环中的每个操作都必须是使用稳定内存指针的 GPU 内核启动，使捕获的图无需修改即可回放。

由此带来的主要要求：

  • 所有缓冲区都 预先分配 — 步进循环内不动态分配

  • 数据通过 持久的 wp.array 指针 传递 — 只覆盖内容，不替换数组

  • MDP 项是 纯 @wp.kernel 函数 — 不在 Python 中根据 GPU 数据分支

  • 重置使用 布尔掩码 （ env_mask ）代替索引列表（ env_ids ），避免可变长度索引改变图拓扑

项目结构#

与基于 Torch 的管理器和环境类不同的 Warp 专属实现，位于 _experimental 包中：

  • isaaclab_experimental — Warp 管理器、环境基类和 Warp MDP 项

  • isaaclab_tasks_experimental — Warp 任务配置及任务特有的 MDP 项

凡是与基于 Torch 的管理器或环境类不同的新 Warp 实现，都应放入这些包。Warp 路径仅支持 Newton，因此 Warp 任务配置直接引用 Newton 物理配置，不使用 PresetCfg 。

编写 Warp MDP 项#

导入#

Warp 任务配置从实验性包导入：

# Warp
from isaaclab_experimental.managers import ObservationTermCfg, RewardTermCfg, SceneEntityCfg
import isaaclab_experimental.envs.mdp as mdp

项配置类具有相同接口，只改变导入路径。

通用模式#

所有 Warp MDP 项（观测、奖励、终止、事件、动作）均遵循 内核 + 启动 模式。Torch 项使用 Torch 张量并返回结果；Warp 项通过 @wp.kernel 函数写入预分配的 wp.array 输出缓冲区：

# Torch — returns a tensor
def lin_vel_z_l2(env, asset_cfg) -> torch.Tensor:
    return torch.square(asset.data.root_lin_vel_b[:, 2])

# Warp — writes into pre-allocated output
@wp.kernel
def _lin_vel_z_l2_kernel(vel: wp.array(...), out: wp.array(dtype=wp.float32)):
    i = wp.tid()
    out[i] = vel[i][2] * vel[i][2]

def lin_vel_z_l2(env, out, asset_cfg) -> None:
    wp.launch(_lin_vel_z_l2_kernel, dim=env.num_envs, inputs=[..., out])

输出缓冲区形状因项类型而异：

  • 观测 ： (num_envs, D) ，其中 D 为观测维度

  • 奖励 ： (num_envs,)

  • 终止 ： (num_envs,) ，数据类型为 bool

  • 事件 ： (num_envs,) 掩码 — 事件不生成输出，而是修改仿真状态

观测项#

Warp 项写入预分配缓冲区，因此观测管理器必须在初始化时知道各项的输出维度，才能分配正确的 (num_envs, D) 输出数组。维度通过逐级回退机制确定（见 isaaclab_experimental/managers/observation_manager.py 中的 ObservationManager._infer_term_dim_scalar ）：

警告

IO 描述符装饰器已在 Isaac Lab 3.0 中弃用，将在 Isaac Lab 3.2 中移除。在替代的运行时配置开发完成前，其输出维度元数据暂时保留，供以 Warp 为核心的环境使用。

  1. 装饰器中显式设置 out_dim （推荐）：

    @generic_io_descriptor_warp(out_dim=3, observation_type="RootState")
def base_lin_vel(env, out, asset_cfg) -> None: ...

    out_dim 可以是整数，也可以是在初始化时解析的字符串：

    • "joint" — asset_cfg 中所选关节的数量

    • "body:N" — asset_cfg 中每个所选刚体的 N 个分量

    • "command" — 从指令管理器获取维度

    • "action" — 从动作管理器获取维度

  2. axes 元数据 ：维度等于列出的轴数量：

    @generic_io_descriptor_warp(axes=["X", "Y", "Z"], observation_type="RootState")
def projected_gravity(env, out, asset_cfg) -> None: ...
# → dimension = 3

  3. 旧式参数 ： term_cfg.params 中的 term_dim 、 out_dim 或 obs_dim 键。

  4. 资产配置回退 ：对于关节级项，使用 asset_cfg.joint_ids （或 joint_ids_wp ）的数量。

事件项#

事件使用 env_mask （布尔 wp.array ）代替 env_ids ，各内核检查掩码以跳过未选中的环境：

def reset_joints_by_offset(env, env_mask, ...):
    wp.launch(_kernel, dim=env.num_envs, inputs=[env_mask, ...])

@wp.kernel
def _kernel(env_mask: wp.array(dtype=wp.bool), ...):
    i = wp.tid()
    if not env_mask[i]:
        return
    # ... modify state for selected envs only

  • 随机数生成使用逐环境 env.rng_state_wp （ wp.uint32 ），而非 torch.rand

  • 启动/预启动 事件使用 Torch 约定 (env, env_ids, **params)

  • 重置/间隔 事件使用 Warp 约定 (env, env_mask, **params)

动作项#

动作采用 两阶段执行 ： process_actions 每个环境步调用一次，对原始动作缩放和限幅； apply_actions 每个仿真步调用一次，将目标写入资产。两个阶段都使用 Warp 内核，以及预分配的 _raw_actions 和 _processed_actions 缓冲区。

图捕获安全性#

编写在捕获的步进循环内运行的项时，应注意：

  • 不使用 wp.to_torch 或 Torch 运算 — 全程使用 Warp

  • 不访问惰性求值属性 — 直接使用与仿真绑定的第一层（Tier 1）数据；需要派生量时，在内核内计算

  • 不动态分配 — 所有缓冲区必须在 __init__ 中预分配

一致性测试#

通过两个层次的一致性测试验证 Warp 项：

1. 实现一致性（Torch 与 Warp） — 验证 Warp 内核与 Torch 实现产生相同结果。没有 Torch 对应实现的项（例如直接用 Warp 编写的新项）可以不进行此测试。

import isaaclab.envs.mdp.observations as torch_obs
import isaaclab_experimental.envs.mdp.observations as warp_obs

# Torch baseline
expected = torch_obs.joint_pos(torch_env, asset_cfg=cfg)

# Warp (uncaptured)
out = wp.zeros((num_envs, num_joints), dtype=wp.float32, device=device)
warp_obs.joint_pos(warp_env, out, asset_cfg=cfg)
actual = wp.to_torch(out)

torch.testing.assert_close(actual, expected)

2. 捕获一致性（Warp 与图捕获后的 Warp） — 验证 CUDA 图回放和直接启动产生完全相同的结果。不一致说明代码不适合捕获，例如使用了过时指针、动态分配，或回放时不会执行的惰性属性访问。即使没有 Torch 对应实现，也应始终运行此测试。

# Warp uncaptured
out_uncaptured = wp.zeros((num_envs, num_joints), dtype=wp.float32, device=device)
warp_obs.joint_pos(warp_env, out_uncaptured, asset_cfg=cfg)

# Warp captured (graph replay)
out_captured = wp.zeros((num_envs, num_joints), dtype=wp.float32, device=device)
with wp.ScopedCapture() as cap:
    warp_obs.joint_pos(warp_env, out_captured, asset_cfg=cfg)
wp.capture_launch(cap.graph)

torch.testing.assert_close(wp.to_torch(out_captured), wp.to_torch(out_uncaptured))

完整的一致性测试示例见 source/isaaclab_experimental/test/envs/mdp/ 。

可用的 Warp MDP 项#

类别

可用项

观测（11）
base_pos_z
base_lin_vel
base_ang_vel
projected_gravity
joint_pos
joint_pos_rel
joint_pos_limit_normalized
joint_vel
joint_vel_rel
last_action
generated_commands

奖励（16）
is_alive
is_terminated
lin_vel_z_l2
ang_vel_xy_l2
flat_orientation_l2
joint_torques_l2
joint_vel_l1
joint_vel_l2
joint_acc_l2
joint_deviation_l1
joint_pos_limits
action_rate_l2
action_l2
undesired_contacts
track_lin_vel_xy_exp
track_ang_vel_z_exp

事件（6）
reset_joints_by_offset
reset_joints_by_scale
reset_root_state_uniform
push_by_setting_velocity
apply_external_force_torque
randomize_rigid_body_com

终止（4）
time_out
root_height_below_minimum
joint_pos_out_of_manual_limit
illegal_contact

动作（2）
JointPositionAction
JointEffortAction

未列出的项仍仅有 Torch 实现。环境需要未列出的项时，必须先用 Warp 实现这些项。

目录