MuJoCo Warp (MJWarp)#

MuJoCo Warp (MJWarp) 是 MuJoCo 的一种实现，使用 Warp 编写，并针对 NVIDIA 硬件和平行仿真进行了优化。MJWarp 托管在 google-deepmind/mujoco_warp GitHub 仓库中。

MJWarp 是由 NVIDIA 和 Google DeepMind 共同开发和维护的。

教程笔记本#

MJWarp 的基础知识涵盖在教程 [notebook] [在 colab 中打开] 中。

何时使用 MJWarp？#

高吞吐量#

MuJoCo 生态系统为批量仿真提供了多种选择。

mujoco.rollout：用于在 CPU 上多线程调用 mj_step 的 Python API。使用具有快速核心和大线程数的硬件可以实现高吞吐量，但对于需要频繁进行主机与设备（host<>device）间数据传输的应用（例如，在 CPU 上进行仿真而学习在 GPU 上进行的强化学习），整体性能可能会受到传输开销的限制。
mjx.step：jax.vmap 和 jax.pmap 支持在 CPU、GPU 或 TPU 上使用 JAX 进行多线程和多设备仿真。
mujoco_warp.step：用于通过 NVIDIA GPU 上的 Warp 进行 CUDA 多线程和多设备仿真的 Python API。与 MJX JAX 实现相比，在接触密集的场景中具有更好的扩展性。

低延迟#

MJWarp 针对吞吐量进行了优化：即单位时间内的总仿真步数；而 MuJoCo 针对延迟进行了优化：即单次仿真步所需的时间。通常情况下，在相同的仿真任务下，MJWarp 的单步性能会低于 MuJoCo。

因此，MJWarp 非常适合需要大量样本的应用（如强化学习），而 MuJoCo 可能更适用于实时应用（如在线控制，例如模型预测控制）或交互式图形界面（如基于仿真的遥操作）。

复杂场景#

在包含许多几何体或自由度 (DoF) 的场景中，MJWarp 的扩展性比 MJX 好，但不如 MuJoCo。对于超过 60 个自由度的场景，MJWarp 的性能可能会显著下降。支持这些大型场景是优先级很高的工作，进展可在 GitHub issue 中跟踪：稀疏雅可比矩阵 #88，分块 Cholesky 分解与求解 #320，约束岛 #886，以及休眠岛 #887。

可微性#

MJX 中的动力学 API 通过 JAX 支持自动微分。我们正在考虑是否通过 Warp 在 MJWarp 中支持此功能——如果此功能对您很重要，请在此 issue 中留言。

安装#

从 PyPI 安装

pip install mujoco-warp

从源码安装

git clone https://github.com/google-deepmind/mujoco_warp.git
cd mujoco_warp
uv sync --all-extras

确保一切正常工作

uv run pytest -n 8

基本用法#

安装完成后，可以通过 import mujoco_warp as mjw 导入该包。结构体、函数和枚举均可直接从顶级 mjw 模块获取。

结构体#

在 NVIDIA GPU 上运行 MJWarp 函数之前，必须通过 mjw.put_model 和 mjw.make_data 或 mjw.put_data 函数将结构体复制到设备上。将 mjModel 置于设备上会产生 mjw.Model。将 mjData 置于设备上会产生 mjw.Data。

mjm = mujoco.MjModel.from_xml_string("...")
mjd = mujoco.MjData(mjm)
m = mjw.put_model(mjm)
d = mjw.put_data(mjm, mjd)

这些 MJWarp 变体镜像了它们的 MuJoCo 对应物，但有几个关键区别

mjw.Model 和 mjw.Data 包含已复制到设备上的 Warp 数组。
mjw.Model 和 mjw.Data 缺少一些用于不支持特性的字段。

批量大小#

MJWarp 针对并行仿真进行了优化。可以通过三个参数指定批量仿真

nworld：要仿真的世界数量。
nconmax：每个世界的预期接触数。所有世界的最大接触数为 nconmax * nworld。
naconmax：nconmax 的替代方案，所有世界的最大接触总数。如果同时设置了 nconmax 和 naconmax，则忽略 nconmax。
njmax：每个世界的最大约束数。

nconmax 和 njmax 的语义差异。

如果所有世界的接触总数不超过 nworld x nconmax，则每个世界的接触数可能超过 nconmax。然而，每个世界的约束数严格受限于 njmax。

XML 解析

nconmax 和 njmax 的值不会从 size/nconmax 和 size/njmax 解析（这些参数已被弃用）。这些参数的值必须提供给 mjw.make_data 或 mjw.put_data。

函数#

MuJoCo 函数作为同名的 MJWarp 函数暴露出来，但遵循 PEP 8 命名规范。大部分主仿真函数和一些正向仿真的子组件可从顶级 mjw 模块获取。

最小示例#

# Throw a ball at 100 different velocities.

import mujoco
import mujoco_warp as mjw
import warp as wp

_MJCF=r"""
<mujoco>
  <worldbody>
    <body>
      <freejoint/>
      <geom size=".15" mass="1" type="sphere"/>
    </body>
  </worldbody>
</mujoco>
"""

mjm = mujoco.MjModel.from_xml_string(_MJCF)
m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=100)

# initialize velocities
wp.copy(d.qvel, wp.array([[float(i) / 100, 0, 0, 0, 0, 0] for i in range(100)], dtype=float))

# simulate physics
mjw.step(m, d)

print(f'qpos:\n{d.qpos.numpy()}')

命令行脚本#

使用 testspeed 基准测试环境

mjwarp-testspeed benchmark/humanoid/humanoid.xml

使用 MJWarp 进行交互式环境仿真

mjwarp-viewer benchmark/humanoid/humanoid.xml

功能对等性#

MJWarp 支持 MuJoCo 的大部分主要仿真功能，仅有少数例外。如果要求将字段值引用不支持功能的 mjModel 复制到设备，MJWarp 将抛出异常。有关最新的功能可用性，请参阅 MuJoCo API 兼容性。

性能调优#

以下是优化 MJWarp 性能的注意事项。

图捕获 (Graph capture)#

MJWarp 函数（例如 mjw.step）通常包含一系列内核启动。如果直接调用函数，Warp 将单独启动这些内核。为了提高性能，尤其是当函数将被多次调用时，建议将构成该函数的运算捕获为 CUDA 图

with wp.ScopedCapture() as capture:
  mjw.step(m, d)

然后可以启动或重新启动该图

wp.capture_launch(capture.graph)

与直接调用函数相比，这通常会显著更快。有关详细信息，请参阅 Warp 图 API 参考。

批量大小#

最大接触数和约束数（分别为 nconmax / naconmax 和 njmax）是在使用 mjw.make_data 或 mjw.put_data 创建 mjw.Data 时指定的。内存和计算量随这些参数的值而变化。为获得最佳性能，应尽可能将这些参数设置得小，同时确保仿真不超过这些限制。

预期这些限制的最佳值将取决于具体环境。在实践中，选择合适的值通常涉及反复试验。在调用 mjwarp-testspeed 时使用 --measure_alloc 标志来打印每次仿真步的接触数和约束数，并配合 mjwarp-viewer 检查溢出错误，这些都是迭代测试这些参数值的有效技巧。

求解器迭代#

MuJoCo 关于求解器迭代和线性搜索迭代最大数的默认设置应能提供合理的性能。降低 MJWarp 的 Option.iterations 和/或 Option.ls_iterations 限制可能会提高性能，但这应作为调优 nconmax / naconmax 和 njmax 之后的次要考量。

如果将这些限制降低得过多，可能会阻止约束求解器收敛，并导致仿真不准确或不稳定。

对性能的影响：MJX (JAX) 与 MJWarp

在 MJX 中，这些求解器参数是控制仿真性能的关键。相比之下，对于 MJWarp，一旦所有世界都已收敛，求解器可以提前退出并避免不必要的计算。因此，这些设置的值对性能的影响相对较小。

接触传感器匹配#

包含接触传感器的场景有一个指定每个传感器最大匹配接触数的参数 Option.contact_sensor_max_match。为获得最佳性能，该参数的值应尽可能小，同时确保仿真不超过限制。超过此限制的匹配接触将被忽略。

此参数的值可以直接设置，例如 model.opt.contact_sensor_maxmatch = 16，或者通过 XML 自定义数值字段设置。

<custom>
  <numeric name="contact_sensor_maxmatch" data="16"/>
</custom>

与最大接触数和约束数类似，该设置的理想值预计因环境而异。mjwarp-testspeed 和 mjwarp-viewer 可能有助于调优此参数的值。

并行线性搜索#

除了约束求解器的迭代线性搜索外，MJWarp 还提供了一种并行线性搜索例程，可以并行评估一组步长并选择最佳步长。步长从 Model.opt.ls_parallel_min_step 到 1 对数分布，要评估的步长数量通过 Model.opt.ls_iterations 设置。

在某些情况下，与约束求解器的默认迭代线性搜索相比，并行例程可能会提供更好的性能。

要启用此例程，请设置 Model.opt.ls_parallel=True 或向 XML 添加一个自定义数值字段。

<custom>
  <numeric name="ls_parallel" data="1"/>
</custom>

实验性功能

并行线性搜索目前是一个实验性功能。

内存#

仿真吞吐量通常受到大量世界所需的内存限制。优化内存利用率的注意事项包括：

CCD（连续碰撞检测）碰撞体比基元碰撞体需要更多内存，请参阅 MuJoCo 的成对碰撞体表以了解有关碰撞体的信息。
multiccd 比 CCD 需要更多内存。
CCD 内存需求随 Option.ccd_iterations 线性增长。
包含至少一个网格几何体并使用 multiccd 的场景，其内存需求将随每个面的最大顶点数和每个顶点的最大边数线性增长（在所有网格上计算得出）。

testspeed 提供了 --memory 标志，用于报告仿真的总内存利用率，以及有关需要大量内存的 mjw.Model 和 mjw.Data 字段的信息。内联分配的内存（包括用于 CCD 和约束求解器的内存）也可能很大，报告为 Other memory。

每个碰撞体的最大接触数

某些 MJWarp 碰撞体与 MuJoCo 相比具有不同的最大接触数

PLANE<>MESH：4 对 3
HFieldCCD：4 对 mjMAXCONPAIR

稀疏性

稀疏雅可比矩阵可以显著节省内存。此功能的更新在 GitHub issue #88 中跟踪。

mjw.make_data 或 mjw.put_data 参数 nccdmax / naccdmax 可以设置为小于 nconmax / naconmax 的值，以减少 CCD 的内存需求。该参数的值应分别为每个世界或所有世界中 CCD 碰撞体产生的最大接触数。例如，一个 10 个世界的批量仿真产生了 80 个总接触，其中每个碰撞体的接触为：mesh-mesh: 30 (CCD), ellipsoid-ellipsoid: 10 (CCD), sphere-sphere: 40 (基元)，则应将 nconmax / naconmax 设置为至少 8 / 80（宽阶段可能需要更多），并将 nccdmax / naccdmax 设置为 3 / 30。

批量 `Model` 字段#

为了支持具有不同模型参数值的批量仿真，许多 mjw.Model 字段都有一个前导批处理维度。默认情况下，前导维度为 1（即 field.shape[0] == 1），相同的值将应用于所有世界。可以使用具有大于 1 的前导维度的 wp.array 覆盖这些字段。该字段将通过世界 ID 和批处理维度的取模运算进行索引：field[worldid % field.shape[0]]。

图捕获

字段数组应在图捕获（即 wp.ScopedCapture）之前覆盖，因为更新不会应用于现有的图。

# override shape and values
m.dof_damping = wp.array([[0.1], [0.2]], dtype=float)

with wp.ScopedCapture() as capture:
  mjw.step(m, d)

可以在图捕获后覆盖字段形状并设置字段值

# override shape
m.dof_damping = wp.empty((2, 1), dtype=float)

with wp.ScopedCapture() as capture:
  mjw.step(m, d)

# set batched values
dof_damping = wp.array([[0.1], [0.2]], dtype=float)
wp.copy(m.dof_damping, dof_damping)  # m.dof = dof_damping will not update the captured graph

修改字段#

修改 mjModel 字段的推荐工作流程是先修改对应的 mjSpec，然后编译以创建带有更新字段的新 mjModel。然而，编译目前需要主机调用：每个新字段实例 1 次调用，即对于 nworld 个实例需要 nworld 次主机调用。

某些字段可以直接修改而无需编译，从而实现设备端更新。请参阅 mjModel 更改以获取有关特定字段的详细信息。此外，GitHub issue 893 正在跟踪为部分字段添加设备端更新的功能。

每世界网格 (Per-world meshes)#

每世界网格支持异构世界，其中不同的世界可以模拟不同的网格。工作流程如下：

创建一个包含所有网格资源和变体间所需的最大几何体槽位的 mjSpec。
通过改变 spec 并调用 spec.compile() 来编译每个变体。
编译一个基础模型并从中创建 mjw.Model。
用从已编译变体构建的每世界数组覆盖相关的 mjw.Model 字段。

示例 1 — 几何体级别随机化（1 个物体，1 个几何体，2 个网格资源）

基础场景包含所有网格资源。几何体引用一个网格（mesh_a）；第二个网格（mesh_b）可用于每世界替换。

<mujoco>
  <asset>
    <mesh name="mesh_a" vertex="0 0 0 1 0 0 0 1 0 0 0 1"/>
    <mesh name="mesh_b" vertex="0 0 0 2 0 0 0 2 0 0 0 2"/>
  </asset>
  <worldbody>
    <body pos="0 0 1">
      <freejoint/>
      <geom name="obj" type="mesh" mesh="mesh_a"/>
    </body>
  </worldbody>
</mujoco>

nworld = 4

# base spec: 1 body with 1 mesh geom, all mesh assets
spec = mujoco.MjSpec()
mesh_a = spec.add_mesh()
mesh_a.name = "mesh_a"
mesh_a.uservert = [0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1]

mesh_b = spec.add_mesh()
mesh_b.name = "mesh_b"
mesh_b.uservert = [0, 0, 0, 2, 0, 0, 0, 2, 0, 0, 0, 2]

body = spec.worldbody.add_body()
body.pos = [0, 0, 1]
body.add_freejoint()
geom = body.add_geom()
geom.name = "obj"
geom.type = mujoco.mjtGeom.mjGEOM_MESH
geom.meshname = "mesh_a"

# compile each variant
geom.meshname = "mesh_a"
mjm_a = spec.compile()
geom.meshname = "mesh_b"
mjm_b = spec.compile()

# restore and compile base
geom.meshname = "mesh_a"
mjm = spec.compile()

m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=nworld)

# build per-world arrays: worlds 0-1 use mesh_a, worlds 2-3 use mesh_b
geom_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_GEOM, "obj")
variants = [mjm_a, mjm_b]
assignment = [0, 0, 1, 1]  # variant index per world

# build per-world arrays
dataid = np.tile(mjm.geom_dataid, (nworld, 1))
geom_size = np.zeros((nworld, mjm.ngeom, 3))
geom_aabb = np.zeros((nworld, mjm.ngeom, 2, 3))
geom_rbound = np.zeros((nworld, mjm.ngeom))
geom_pos = np.zeros((nworld, mjm.ngeom, 3))
body_mass = np.zeros((nworld, mjm.nbody))
body_subtreemass = np.zeros((nworld, mjm.nbody))
body_inertia = np.zeros((nworld, mjm.nbody, 3))
body_invweight0 = np.zeros((nworld, mjm.nbody, 2))
body_ipos = np.zeros((nworld, mjm.nbody, 3))
body_iquat = np.zeros((nworld, mjm.nbody, 4))

for w in range(nworld):
  ref = variants[assignment[w]]
  dataid[w, geom_id] = ref.geom_dataid[geom_id]
  geom_size[w] = ref.geom_size
  geom_aabb[w] = ref.geom_aabb.reshape(mjm.ngeom, 2, 3)
  geom_rbound[w] = ref.geom_rbound
  geom_pos[w] = ref.geom_pos
  body_mass[w] = ref.body_mass
  body_subtreemass[w] = ref.body_subtreemass
  body_inertia[w] = ref.body_inertia
  body_invweight0[w] = ref.body_invweight0
  body_ipos[w] = ref.body_ipos
  body_iquat[w] = ref.body_iquat

m.geom_dataid = wp.array(dataid, dtype=int)
m.geom_size = wp.array(geom_size, dtype=wp.vec3)
m.geom_aabb = wp.array(geom_aabb, dtype=wp.vec3)
m.geom_rbound = wp.array(geom_rbound, dtype=float)
m.geom_pos = wp.array(geom_pos, dtype=wp.vec3)
m.body_mass = wp.array(body_mass, dtype=float)
m.body_subtreemass = wp.array(body_subtreemass, dtype=float)
m.body_inertia = wp.array(body_inertia, dtype=wp.vec3)
m.body_invweight0 = wp.array(body_invweight0, dtype=wp.vec2)
m.body_ipos = wp.array(body_ipos, dtype=wp.vec3)
m.body_iquat = wp.array(body_iquat, dtype=wp.quat)

示例 2 — 物体级别随机化（1 个物体，1 或 2 个几何体，3 个网格资源）

最大几何体数量

对于物体级别的随机化，提供给 mjw.put_model 的基础 mjModel 应指定所有变体中所需的最大几何体数量。在特定变体中未使用的几何体槽位可以禁用（例如 contype=0，conaffinity=0，dataid=-1），但它们仍应作为基础模型中物体的一部分存在。

<mujoco>
  <asset>
    <mesh name="mA" vertex="0 0 0 1 0 0 0 1 0 0 0 1"/>
    <mesh name="mB" vertex="0 0 0 2 0 0 0 2 0 0 0 2"/>
    <mesh name="mC" vertex="0 0 0 3 0 0 0 3 0 0 0 3"/>
  </asset>
  <worldbody>
    <body name="obj" pos="0 0 1">
      <freejoint/>
      <geom name="obj_0" type="mesh" mesh="mA"/>
      <geom name="obj_1" size=".001" contype="0" conaffinity="0" mass="0"/>
    </body>
  </worldbody>
</mujoco>

nworld = 6

# base spec: body with 2 geom slots (max across variants), all mesh assets
spec = mujoco.MjSpec()
for name, scale in [("mA", 1), ("mB", 2), ("mC", 3)]:
  mesh = spec.add_mesh()
  mesh.name = name
  mesh.uservert = [0, 0, 0, scale, 0, 0, 0, scale, 0, 0, 0, scale]

body = spec.worldbody.add_body()
body.name = "obj"
body.pos = [0, 0, 1]
body.add_freejoint()

g0 = body.add_geom()
g0.name = "obj_0"
g0.type = mujoco.mjtGeom.mjGEOM_MESH
g0.meshname = "mA"

# null geom slot: disabled collision, no mesh
g1 = body.add_geom()
g1.name = "obj_1"
g1.size = [0.001, 0, 0]
g1.contype = 0
g1.conaffinity = 0
g1.mass = 0

# variant A: 1 geom (mesh mA), g1 stays null
mjm_a = spec.compile()

# variant B: 2 geoms (mesh mB + mC)
g0.meshname = "mB"
g1.type = mujoco.mjtGeom.mjGEOM_MESH
g1.meshname = "mC"
g1.contype = 1
g1.conaffinity = 1
mjm_b = spec.compile()

# restore base and compile
g0.meshname = "mA"
g1.type = mujoco.mjtGeom.mjGEOM_SPHERE
g1.contype = 0
g1.conaffinity = 0
mjm = spec.compile()

m = mjw.put_model(mjm)
d = mjw.make_data(mjm, nworld=nworld)

# worlds 0-2: variant A (1 active geom), worlds 3-5: variant B (2 active geoms)
variants = [mjm_a, mjm_b]
assignment = [0, 0, 0, 1, 1, 1]

geom0_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_GEOM, "obj_0")
geom1_id = mujoco.mj_name2id(mjm, mujoco.mjtObj.mjOBJ_GEOM, "obj_1")
body_id = mjm.geom_bodyid[geom0_id]

# build per-world arrays
dataid = np.tile(mjm.geom_dataid, (nworld, 1))
geom_size = np.zeros((nworld, mjm.ngeom, 3))
geom_rbound = np.zeros((nworld, mjm.ngeom))
geom_aabb = np.zeros((nworld, mjm.ngeom, 2, 3))
geom_pos = np.zeros((nworld, mjm.ngeom, 3))
body_mass = np.zeros((nworld, mjm.nbody))
body_subtreemass = np.zeros((nworld, mjm.nbody))
body_inertia = np.zeros((nworld, mjm.nbody, 3))
body_invweight0 = np.zeros((nworld, mjm.nbody, 2))
body_ipos = np.zeros((nworld, mjm.nbody, 3))
body_iquat = np.zeros((nworld, mjm.nbody, 4))

for w in range(nworld):
  ref = variants[assignment[w]]
  dataid[w] = ref.geom_dataid
  # disable unused geom slot for variant A
  if assignment[w] == 0:
    dataid[w, geom1_id] = -1
  geom_size[w] = ref.geom_size
  geom_rbound[w] = ref.geom_rbound
  geom_aabb[w] = ref.geom_aabb.reshape(mjm.ngeom, 2, 3)
  geom_pos[w] = ref.geom_pos
  body_mass[w] = ref.body_mass
  body_subtreemass[w] = ref.body_subtreemass
  body_inertia[w] = ref.body_inertia
  body_invweight0[w] = ref.body_invweight0
  body_ipos[w] = ref.body_ipos
  body_iquat[w] = ref.body_iquat

m.geom_dataid = wp.array(dataid, dtype=int)
m.geom_size = wp.array(geom_size, dtype=wp.vec3)
m.geom_rbound = wp.array(geom_rbound, dtype=float)
m.geom_aabb = wp.array(geom_aabb, dtype=wp.vec3)
m.geom_pos = wp.array(geom_pos, dtype=wp.vec3)
m.body_mass = wp.array(body_mass, dtype=float)
m.body_subtreemass = wp.array(body_subtreemass, dtype=float)
m.body_inertia = wp.array(body_inertia, dtype=wp.vec3)
m.body_invweight0 = wp.array(body_invweight0, dtype=wp.vec2)
m.body_ipos = wp.array(body_ipos, dtype=wp.vec3)
m.body_iquat = wp.array(body_iquat, dtype=wp.quat)

批量字段 — 对于每世界网格必须覆盖的字段

字段	dtype	形状
`geom_dataid`	`int`	`(nworld, ngeom)`
`geom_size`	`wp.vec3`	`(nworld, ngeom)`
`geom_aabb`	`wp.vec3`	`(nworld, ngeom, 2)`
`geom_rbound`	`float`	`(nworld, ngeom)`
`geom_pos`	`wp.vec3`	`(nworld, ngeom)`
`body_mass`	`float`	`(nworld, nbody)`
`body_subtreemass`	`float`	`(nworld, nbody)`
`body_inertia`	`wp.vec3`	`(nworld, nbody)`
`body_invweight0`	`wp.vec2`	`(nworld, nbody)`
`body_ipos`	`wp.vec3`	`(nworld, nbody)`
`body_iquat`	`wp.quat`	`(nworld, nbody)`

批量渲染#

MJWarp 提供了一个基于 Warp 加速 BVH 构建的批量渲染器，用于高吞吐量光线追踪，支持并行渲染具有多个摄像机的世界。

主要功能

带纹理的网格渲染：具有完全纹理支持的 BVH 加速网格渲染。
高度场渲染：针对高度场的优化渲染。
Flex 渲染：渲染 flex 对象。
照明和阴影：具有可配置阴影的动态照明；支持域随机化：light_active, light_type, light_castshadow, light_xpos, light_xdir。
异构多摄像机：每个世界支持多个摄像机，每个摄像机可以具有不同的分辨率 (cam_resolution)、视野 (cam_fovy, cam_sensorsize, cam_intrinsic) 和输出模式 (cam_output)。
域随机化：每世界 mjw.Model 字段（见上文批量 Model 字段）：geom_matid, geom_size, geom_rgba, mat_texid, mat_texrepeat, mat_rgba。
BVH 加速射线/多射线 API：光线投射：通过 Warp 的 BVH 进行加速的 mjw.ray、mjw.rays 和测距传感器。

基本用法#

渲染或光线投射需要一个 mjw.RenderContext，其中包含 BVH 结构、特定于渲染的字段和输出缓冲区。

rc = mjw.create_render_context(
    mjm,
    nworld=1,
    cam_res=(256, 256),           # Override camera resolution (or per-camera list)
    render_rgb=True,              # Enable RGB output (or per-camera list)
    render_depth=True,            # Enable depth output (or per-camera list)
    use_textures=True,            # Apply material textures
    use_shadows=False,            # Enable shadow casting (slower)
    enabled_geom_groups=[0, 1],   # Only render geoms in groups 0 and 1
    cam_active=[True, False],     # Selectively enable/disable cameras
    flex_render_smooth=True,      # Smooth shading for soft bodies
)

每个 mjw.RenderContext 参数都可以全局应用或按摄像机应用。此外，mjw.RenderContext 参数的值可以从 XML 解析

<camera name="front_camera" pos="3 0 2" xyaxes="0 1 0 -0.6 0 0.8" resolution="64 64" output="rgb depth"/>

或通过 mjSpec 设置以实现摄像机自定义。

要进行渲染，首先调用 mjw.refit_bvh 来更新 BVH 树，然后调用 mjw.render 写入输出缓冲区。

mjw.refit_bvh(m, d, rc)
mjw.render(m, d, rc)

输出缓冲区包含所有摄像机的堆叠像素，形状为 (nworld, npixel)，RGB 数据打包成一个 uint32 变量。RenderContext.rgb_adr 和 RenderContext.depth_adr 提供按摄像机的索引。为方便起见，mjw.get_rgb 和 mjw.get_depth 返回针对给定摄像机处理和重塑后的 RGB 和深度数据，并对所有世界进行了批处理。

nworld = 1
cam_index = 0
resolution = rc.cam_res.numpy()[cam_index]
rgb_data = wp.zeros((nworld, resolution[1], resolution[0]), dtype=wp.vec3)
mjw.get_rgb(rc, rgb_data=rgb_data, cam_id=cam_index)

完整的示例可以在 MJWarp 教程 [notebook] [在 colab 中打开] 中找到。

基准测试#

可以使用 testspeed 对渲染进行基准测试

mjwarp-testspeed benchmarks/primitives.xml --function=render

有关各种场景的基准测试结果，请参阅发布的基准测试。

说明#

网格：渲染计算随网格复杂度（即顶点数和面数）线性增长。与同等尺寸的网格或高度场相比，基元预期具有更好的性能（即更高的吞吐量）。
扩展性：渲染随分辨率（总像素数）和摄像机数量线性增长。

常见问题解答#

学习框架#

MJWarp 能与 JAX 配合使用吗？

是的。MJWarp 可与 JAX 互操作。详情请参阅 Warp 互操作性文档。

此外，MJX 为 MJWarp 的部分 API 提供了 JAX API。该实现通过 impl='warp' 指定。

MJWarp 能与 PyTorch 配合使用吗？

是的。MJWarp 可与 PyTorch 互操作。详情请参阅 Warp 互操作性文档。

如何使用 MJWarp 物理引擎训练策略？

有关使用 MJWarp 物理引擎训练策略的示例，请参阅

Isaac Lab：通过 Newton API 进行训练。
mjlab：使用 PyTorch 直接通过 MJWarp 训练。
MuJoCo Playground：通过 MJX API 训练。

功能#

MJWarp 可微分吗？

否。MJWarp 目前无法通过 Warp 的自动微分功能进行微分。团队关于在 MJWarp 中启用自动微分的相关更新在此 GitHub issue 中跟踪。

MJWarp 支持多 GPU 吗？

是的。Warp 的 wp.ScopedDevice 支持多 GPU 计算

# create a graph for each device
graph = {}
for device in wp.get_cuda_devices():
  with wp.ScopedDevice(device):
    m = mjw.put_model(mjm)
    d = mjw.make_data(mjm)
    with wp.ScopedCapture(device) as capture:
      mjw.step(m, d)
    graph[device] = capture.graph

# launch a graph on each device
for device in wp.get_cuda_devices():
  wp.capture_launch(graph[device])

有关详细信息，请参阅 Warp 文档，有关强化学习示例，请参阅 mjlab 分布式训练。

MJWarp 在 GPU 上是确定性的吗？

否。在同一代码的不同执行计算结果之间，可能存在排序或微小的数值差异。这是 GPU 上非确定性原子操作的特性。若要获得确定性结果，请使用 wp.set_device("cpu") 将设备设置为 CPU。

有关在 GPU 上获得确定性结果的进展在此 GitHub issue 中跟踪。

方向是如何表示的？

方向表示为单位四元数，并遵循 MuJoCo 的约定：w, x, y, z 或 scalar, vector。

wp.quaternion

MJWarp 使用 Warp 的内置类型 wp.quaternion。但重要的是，MJWarp 不使用 Warp 的 x, y, z, w 四元数约定或运算，而是实现了遵循 MuJoCo 约定的四元数例程。请参阅 math.py 以获取实现。

MJWarp 有命名访问 API/绑定吗？

否。此功能的更新在此 GitHub issue 中跟踪。

为什么没有碰撞时会报告接触？

对于对碰撞传感器有贡献的每一个唯一几何体对，即使该几何体对没有发生碰撞，也会报告 1 个接触。与 MuJoCo 或 MJX 不同（在 MuJoCo 或 MJX 中，碰撞传感器在计算传感器数据的同时对碰撞例程进行单独调用），MJWarp 在运行其主碰撞流水线时在接触中计算并存储这些传感器的数据。

接触传感器将报告影响物理的接触的正确信息。

为什么雅可比矩阵总是稠密的？

稀疏雅可比矩阵尚未实现，Data 字段：ten_J, actuator_moment, flexedge_J 和 efc.J 始终表示为稠密矩阵。对稀疏雅可比矩阵的支持在 GitHub issue #88 中跟踪。

为什么某些数组的形状与 mjModel 或 mjData 相比不同？

对于批量仿真，默认情况下，许多 mjw.Data 字段具有大小为 Data.nworld 的前导批处理维度。一些 mjw.Model 字段具有大小为 1 的前导批处理维度，表明此字段可以用批量参数数组覆盖以进行域随机化。

此外，包括 Model.qM, Data.efc.J 和 Data.efc.D 在内的特定字段已填充，以实现 GPU 上的快速加载。

为什么 MJWarp 和 MuJoCo 的数值结果不同？

MJWarp 使用 float，而 MuJoCo 默认使用 double 表示 mjtNum。求解器设置（包括迭代、碰撞检测和较小的摩擦值）可能对浮点表示的差异敏感。

如果您遇到意料之外的结果（包括 NaN），请提交 GitHub issue。

为什么惯性矩阵 qM 的稀疏性与 MuJoCo / MJX 不一致？

mjtJacobian 语义

MuJoCo 的惯性矩阵始终是稀疏的，而 mjtJacobian 会影响约束雅可比矩阵及相关量
MJWarp（和 MJX）的约束雅可比矩阵始终是稠密的，而 mjtJacobian 被重新调整用途以影响可以表示为稠密或稀疏的惯性矩阵

MJWarp 用于 AUTO 的自动稀疏阈值已针对 GPU 进行了优化，并设置为 nv > 32，这与使用 nv >= 60 的 MuJoCo 和 MJX 不同。稠密 DENSE 和稀疏 SPARSE 设置与 MuJoCo 和 MJX 一致。

此功能将来可能会更改。

如何修复仿真运行时警告？

当仿真期间内存需求超过现有分配时，会提供警告

nconmax / njmax：已超过最大接触数/约束数。通过更新相关参数至 mjw.make_data 或 mjw.put_data 来增加设置值。
mjw.Option.ccd_iterations：凸碰撞检测算法已超过最大迭代次数。在 XML / mjSpec / mjModel 中增加此设置的值。重要的是，必须对提供给 mjw.put_model 和 mjw.make_data / mjw.put_data 的 mjModel 实例进行此更改。
mjw.Option.contact_sensor_maxmatch：已超过接触传感器匹配标准的最大匹配数。增加此 MJWarp 专用设置 m.opt.contact_sensor_maxmatch。或者，重构接触传感器匹配标准，例如如果已知感兴趣的 2 个几何体，请指定 geom1 和 geom2。
height field collision overflow：高度场生成的潜在接触数超过了 mjMAXCONPAIR，部分接触被忽略。为解决此警告，请降低高度场分辨率或减小与高度场交互的几何体尺寸。

编译#

如何缩短编译时间？

限制需要通用凸碰撞流水线的唯一碰撞体数量。这些碰撞体在 collision_convex.py 中列为 _CONVEX_COLLISION_PAIRS。此流水线的编译时间改进在此 GitHub issue 中跟踪。

为什么升级 MJWarp 后物理效果不如预期？

Warp 缓存可能与当前代码不兼容，应作为调试过程的一部分进行清理。可以通过删除目录 ~/.cache/warp 或通过 Python 完成。

import warp as wp
wp.clear_kernel_cache()

是否可以提前编译 MJWarp 而不是在运行时编译？

是的。详情请参阅 Warp 的提前编译工作流程文档。

与 MuJoCo 的区别#

本节说明 MJWarp 和 MuJoCo 之间的差异。

热启动 (Warmstart)#

如果热启动未禁用，MJWarp 求解器的热启动始终使用 qacc_warmstart 初始化加速度。相比之下，MuJoCo 会在 qacc_smooth 和 qacc_warmstart 之间进行比较，以确定使用哪一个进行初始化。

惯性矩阵分解#

当使用稠密计算时，MJWarp 的惯性矩阵 qLD 分解是使用 Warp 的 L'L Cholesky 分解 wp.tile_cholesky 计算的，且结果预期不会与 MuJoCo 的对应字段匹配，因为使用了不同的反向模式 L'DL 例程 mj_factorM。

选项#

mjw.Option 字段对应其 mjOption 对应物，但有以下例外：

impratio 以其逆平方根 impratio_invsqrt 形式存储。
约束求解器设置 tolerance 被钳制为最小值 1e-6。
接触覆盖参数 o_margin, o_solref, o_solimp 和 o_friction 不可用。

disableflags 有以下区别

mjDSBL_MIDPHASE 不可用。
mjDSBL_AUTORESET 不可用。
mjDSBL_NATIVECCD 将默认的盒子-盒子碰撞体从 CCD 更改为基元碰撞体。
mjDSBL_ISLAND 当前不可用。约束岛发现功能在 GitHub issue #886 中跟踪。

enableflags 有以下区别

mjENBL_OVERRIDE 不可用。
mjENBL_FWDINV 不可用。
通过 mjENBL_ISLAND 启用的约束岛休眠功能当前不可用。此功能在 GitHub issue #886 和 #887 中跟踪。

提供额外的 MJWarp 专用选项

ls_parallel：在约束求解器中使用并行线性搜索
ls_parallel_min_step：并行线性搜索的最小步长
broadphase：宽阶段算法类型 (mjw.BroadphaseType)
broadphase_filter：宽阶段使用的过滤类型 (mjw.BroadphaseFilter)
graph_conditional：使用 CUDA 图条件
run_collision_detection：使用碰撞检测例程
contact_sensor_maxmatch：接触传感器匹配标准的最大接触数

流体模型

修改流体模型参数：density, viscosity 或 wind 可能需要更新 Model.has_fluid。

图捕获

修改 mjw.Option 字段后，可能需要进行新的图捕获，以便更新后的设置生效。

SDF 插件#

SDF 碰撞支持插件。plugin/sdf/bowl.xml 的以下示例展示了如何在 bowl.cc 中实现 SDF 插件。

import mujoco_warp as mjw
import warp as wp

# distance function
@wp.func
def bowl(p: wp.vec3, attr: wp.vec3) -> float:
  """Signed distance function for a bowl shape.

  attr[0] = height
  attr[1] = radius
  attr[2] = thickness
  """
  height = attr[0]
  radius = attr[1]
  thick = attr[2]
  width = wp.sqrt(radius * radius - height * height)

  # q = (norm_xy(p), p.z)
  q0 = wp.sqrt(p[0] * p[0] + p[1] * p[1])
  q1 = p[2]

  # qdiff = q - (width, height)
  qdiff0 = q0 - width
  qdiff1 = q1 - height

  if height * q0 < width * q1:
    dist = wp.sqrt(qdiff0 * qdiff0 + qdiff1 * qdiff1)
  else:
    q_norm = wp.sqrt(q0 * q0 + q1 * q1)
    dist = wp.abs(q_norm - radius)

  return dist - thick


# gradient of distance function
@wp.func
def bowl_sdf_grad(p: wp.vec3, attr: wp.vec3) -> wp.vec3:
  """Gradient of bowl SDF via finite differences."""
  eps = float(1e-6)
  f0 = bowl(p, attr)

  px = wp.vec3(p[0] + eps, p[1], p[2])
  py = wp.vec3(p[0], p[1] + eps, p[2])
  pz = wp.vec3(p[0], p[1], p[2] + eps)

  grad = wp.vec3(
    (bowl(px, attr) - f0) / eps,
    (bowl(py, attr) - f0) / eps,
    (bowl(pz, attr) - f0) / eps,
  )
  return grad


# register the bowl SDF plugin
@wp.func
def user_sdf(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> float:
  return bowl(p, attr)


@wp.func
def user_sdf_grad(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3:
  return bowl_sdf_grad(p, attr)


# override the module-level hooks
mjw._src.collision_sdf.user_sdf = user_sdf
mjw._src.collision_sdf.user_sdf_grad = user_sdf_grad

物理回调#

MuJoCo 提供全局物理回调，允许用户将自定义逻辑注入仿真流水线。MJWarp 支持类似的机制，但回调是设置在每个 mjw.Model 实例上的 Python 函数（通过 Model.callback 设置），而不是作为全局函数指针。

以下回调可用：

回调	描述
`control`	自定义控制律，写入 `Data.ctrl`
`passive`	自定义被动力，写入 `Data.qfrc_passive`
`act_dyn`	自定义执行器动力学，写入 `Data.act_dot`
`act_gain`	自定义执行器增益，写入 `Data.actuator_force`
`act_bias`	自定义执行器偏差，写入 `Data.actuator_force`
`sensor`	自定义传感器，写入 `Data.sensordata`；接收一个额外的 `stage` 参数
`contactfilter`	自定义接触过滤，写入 `Data.contact`

import mujoco
import mujoco_warp as mjw
import warp as wp

_MJCF = r"""
<mujoco>
  <worldbody>
    <body>
      <geom size=".1"/>
      <joint name="hinge"/>
    </body>
  </worldbody>
  <actuator>
    <motor joint="hinge"/>
  </actuator>
</mujoco>
"""

@wp.kernel
def _ctrl_callback(ctrl_out: wp.array2d(dtype=float)):
  worldid = wp.tid()
  ctrl_out[worldid, 0] = 2.0

def ctrl_callback(m, d):
  wp.launch(_ctrl_callback, dim=(d.nworld,), outputs=[d.ctrl])

mjm = mujoco.MjModel.from_xml_string(_MJCF)
m = mjw.put_model(mjm)
d = mjw.make_data(mjm)

m.callback.control = ctrl_callback
mjw.step(m, d)
assert d.ctrl.numpy()[0, 0] == 2.0

盒子-盒子碰撞#

默认情况下，盒子-盒子碰撞使用通用凸碰撞流水线 (GJK/EPA)。通过设置 NATIVECCD 禁用标志，可以使用基于 engine_collision_box.c 的专用基元碰撞体

m.opt.disableflags |= mjw.DisableBit.NATIVECCD

与凸流水线最多 4 个接触点相比，专用碰撞体最多可生成 8 个接触点，并可能提高涉及盒子堆叠或操纵任务的接触稳定性。

CCD 边距#

非零几何体边距或配对边距不支持某些 CCD 碰撞体，在调用 mjw.put_model 时将抛出 NotImplementedError

几何体对	场景	变通方法
box-box, box-mesh, mesh-mesh	MULTICCD 已启用（默认开启）	将边距设置为 `0` 或禁用 `MULTICCD`
box-box	NATIVECCD 已启用（默认开启）	将边距设置为 `0` 或禁用 `NATIVECCD`

渲染#

MJWarp 中包含的批量渲染器的用途与 MuJoCo 的渲染器不同。MJWarp 批量渲染器是一个单次命中光线投射器，针对高吞吐量和低保真度进行了优化。

它支持

简单的 Lambertian 漫反射阴影
基本的点光源和定向光源
纹理
阴影

它不支持

全局照明等高级照明效果
基于物理的材质属性

MuJoCo Warp (MJWarp)#

教程笔记本#

何时使用 MJWarp？#

高吞吐量#

低延迟#

复杂场景#

可微性#

安装#

基本用法#

结构体#

批量大小#

函数#

最小示例#

命令行脚本#

功能对等性#

性能调优#

图捕获 (Graph capture)#

批量大小#

求解器迭代#

接触传感器匹配#

并行线性搜索#

内存#

批量 Model 字段#

修改字段#

每世界网格 (Per-world meshes)#

批量渲染#

基本用法#

基准测试#

说明#

常见问题解答#

学习框架#

功能#

编译#

与 MuJoCo 的区别#

热启动 (Warmstart)#

惯性矩阵分解#

选项#

SDF 插件#

物理回调#

盒子-盒子碰撞#

CCD 边距#

渲染#

批量 `Model` 字段#