Python: per-environment MJCF insertion + collision-capacity override by haixuanTao · Pull Request #16 · dimforge/nexus

haixuanTao · 2026-07-15T13:39:59Z

insert_mjcf gains an env= kwarg (default 0 — fully backward compatible), making one-robot-per-environment reachable from Python: the robot, its auto-floor and the camera-fit AABB target rbd_world_mut(env), while viewer registration (visual meshes, camera, light) stays env-0-only since the viewer draws environment 0 — batch environments are physics-only. add_environment() was already exposed.

Why

Packing N robots into env 0 is quadratic: every contact-constraint slot stores a dense M⁻¹ column sized to the whole environment's DOF count (contact_cons_col_cap = num_mbs × 192 × dofs_cap). At 2,048 robots that requests ~58 GiB twice (columns + Jacobians), and single-env packing caps at 96 robots — 121.5 MiB, just under wgpu's 128 MiB binding default. One robot per env keeps dofs_cap at the robot's own DOFs and scales linearly with batches, which is the layout the batched solver is designed around.

Second scaling fix: new NexusState::set_rbd_collisions_capacity(). The 4,096-pairs/env default is sized for one busy scene, not thousands of small envs — pair-keyed workspaces bind capacity × envs × sizeof(manifold) ≈ 9.4 GiB at 2,048 envs, for a scene generating ~7 pairs/env. Lowering it to 64 also halved wall-clock at 256 envs (17.6k → 36.1k env-steps/s): the sweep kernels were paying for the slack.

Measured (RTX 5090, headless WebGPU, 12-DOF LeRobot biped, passive, 200-step window, one sync)

envs	env-steps/s
256	36,056
2,048	138,640

Linear scaling; cross-checks against the zealot training stack's physics-only rate at 2,048 envs (~127k on the same engine family).

Notes

Batched finalize requires identical collider counts in every environment — worth a friendlier error eventually.
Callers batching many envs should strip visual geoms from the inserted MJCF: each insert re-parses every referenced mesh, and env>0 discards the visuals anyway (0.5 s → 24 ms per insert on the LeRobot asset).

🤖 Generated with Claude Code

…rride insert_mjcf gains an env= kwarg (default 0, fully backward compatible): the robot, its auto-floor and the camera-fit AABB all target rbd_world_mut(env), and viewer registration (visual meshes, camera, light) only runs for env 0 since the viewer draws environment 0 — batch environments are physics-only. add_environment() was already exposed, so one robot per env is now reachable from Python. Why: loading N robots into env 0 is quadratic — every contact-constraint slot stores a dense M^-1 column sized to the WHOLE env's DOF count (contact_cons_col_cap = num_mbs * 192 * dofs_cap), which at 2048 robots requests ~58 GiB twice and caps single-env packing at 96 robots (121 MiB, just under wgpu's 128 MiB binding default). One robot per env keeps dofs_cap at the robot's own DOFs and scales linearly with batches. Also adds NexusState::set_rbd_collisions_capacity(): the 4096-pairs/env default is sized for one busy scene, not thousands of small envs — pair-keyed workspaces bind capacity x envs x sizeof(manifold) ~ 9.4 GiB at 2048 envs for a scene that generates ~7 pairs per env. Lowering it to 64 also halved wall-clock at 256 envs (17.6k -> 36.1k env-steps/s). Measured (RTX 5090, headless WebGPU, 12-DOF LeRobot biped, passive, 200-step window): 2,048 envs = 138,640 env-steps/s; linear from 256 (36k). Note batched finalize requires identical collider counts per env. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

sebcrozet

That looks good, thank you. Just a few comments that got strangely merged.

sebcrozet · 2026-07-16T13:04:05Z

    /// Sets the number of rigid-body solver steps advanced per
    /// [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call (default 1). Acts as a simulation-speed control.
+    /// Overrides the per-environment collision-pair capacity used when the
+    /// GPU rigid-body state is (re)allocated at `finalize`. The default (4096)
+    /// is sized for one busy scene, not thousands of small batched envs —
+    /// pair-keyed workspaces scale as `capacity x num_envs x sizeof(manifold)`,
+    /// which at 2048 envs binds ~9 GiB unless this is lowered.


The comments got mixed up (the first two sentences should apply to set_rbd_steps_per_frame).

Good catch — fixed in 1aa4958: the set_rbd_steps_per_frame doc is back on its own function.

sebcrozet · 2026-07-16T13:04:34Z

    /// Loads a MuJoCo MJCF scene into environment 0 as multibodies, registering
    /// its render shapes (and a sized floor) with `viewer`. Returns scene info
    /// (suggested camera + whether the scene is Z-up). Call `finalize` after.
-    #[pyo3(signature = (viewer, scene_path, render_colliders=false))]
+    /// Per-environment collision-pair capacity (default 4096). Lower this
+    /// before `finalize` when batching many small environments — pair-keyed
+    /// GPU workspaces scale with `capacity x num_envs`.
+    fn set_rbd_collisions_capacity(&mut self, capacity: u32) {
+        self.0.set_rbd_collisions_capacity(capacity);
+    }
+
+    #[pyo3(signature = (viewer, scene_path, render_colliders=false, env=0))]


The comments got mixed up here too.

Fixed in 1aa4958 — moved the insert_mjcf doc back onto insert_mjcf (and updated its wording for the new env parameter).

Review feedback: the set_rbd_steps_per_frame doc was merged onto set_rbd_collisions_capacity, and insert_mjcf's doc onto the Python set_rbd_collisions_capacity binding. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

sebcrozet · 2026-08-28T12:00:06Z

This will be merged as part of #36

* fix(rbd): apply the per-batch stride to collider_parent reads in the narrow phase Replaces #21 Co-Authored-By: Haixuan Xavier Tao <tao.xavier@outlook.com> * feat: per-environment collision-pair capacity override Replaces #24 Co-Authored-By: Haixuan Xavier Tao <tao.xavier@outlook.com> * feat(rbd): make the narrow-phase contact prediction distance configurable Replaces #28 Co-Authored-By: Haixuan Xavier Tao <tao.xavier@outlook.com> * fix(rbd): thread the configurable prediction distance through the brute-force broad phase Completes #28 * feat(python): per-environment MJCF insertion Replaces #16 Co-Authored-By: Haixuan Xavier Tao <tao.xavier@outlook.com> * fix(python): drop the duplicated collisions-capacity setter and pass the RbdCoupling to insert_rigid_body_in Completes #16 * feat(python): per-step MJCF actuator control + multibody state readback Replaces #12 Co-Authored-By: Haixuan Xavier Tao <tao.xavier@1ms.ai> * fix(python): gate multibody control/readback on dim3, add the missing PyArray2 import Replaces #12 Co-Authored-By: Haixuan Xavier Tao <tao.xavier@outlook.com> * fix(rbd): decode the SoA link workspace for the multibody readback and drop the stale set_gravity copy Completes #12 * perf(rbd): dedupe shared TriMesh uploads in from_rapier Replaces #19 Co-Authored-By: Haixuan Xavier Tao <tao.xavier@outlook.com> * perf(rbd): optional GPU contact reduction, merging per-pair manifolds to <=4 points Replaces #17 Co-Authored-By: Haixuan Xavier Tao <tao.xavier@outlook.com> * fix(rbd): pass the prediction distance to manifold_reduction in the contact-reduction kernel Completes #17 * perf(rbd): flat 1-D narrow-phase dispatch, packing warps across batches Replaces #21 Co-Authored-By: Haixuan Xavier Tao <tao.xavier@outlook.com> * fix(rbd): restore the contacts capacity binding and import atomic_load_u32 for the flat dispatch Completes #21 * fix(rbd): drop the stale 2mm PREDICTION constant reintroduced by the flat-dispatch port Completes #21 * fix mpm feature-gating * feat(rbd): expose dof_state_mut, links_static, joint_constraints and link_of_body * feat(rbd): env-reset primitives, GPU motor scatter, contact sensors, actuator delay, encoded step, substep-refresh cadence and per-DoF armature/frictionloss * fix(rbd): guard against implicit-coriolis drifting from the batch_indices uniform * feat(rbd): cluster contact manifolds by normal, matching rapier, with a tunable threshold * feat(rbd): model multibody joint frictionloss as a constraint instead of a force * feat(rbd): seed per-DoF joint friction from rapier's Multibody::frictions * refactor(rbd): read the contact prediction distance from RbdSimParams instead of a dedicated uniform * chore: cargo fmt * refactor(rbd): move the contact merge cosine into RbdSimParams * refactor: move read_multibody_links onto NexusState and drive every env from control_multibody_motors * test(rbd): add a headless many-small-environments step-timing harness * revert(rbd): drop the flat 1-D narrow-phase dispatch Measured 7-21% slower on Metal. * chore: cleanup comments * fix: gate control_multibody_motors on dim3 so the 2D build still compiles * fix instability in joint-ball3 demo * chore: remove debug test files * chore: cleanups * fix(rbd): build the bench harness without the metal feature and only in 3D * fix(rbd): silence the clippy needless-borrow and unnecessary-mut lints * fix(rbd): split the joint-constraint back-solve into its own dispatch to fit 8 storage buffers * test(rbd): keep the bench harness under wgpu's default buffer-size limit * fix(rbd): split the batched env reset into pose and DoF passes to fit 8 storage buffers * chore: switch to the published rapier version * fix: make all envs share the same RbdSimParams * chore: clippy fixes --------- Co-authored-by: Haixuan Xavier Tao <tao.xavier@outlook.com> Co-authored-by: Haixuan Xavier Tao <tao.xavier@1ms.ai>

haixuanTao force-pushed the feat/per-env-mjcf-upstream branch from e839c74 to 0447bf0 Compare July 16, 2026 12:17

sebcrozet requested changes Jul 16, 2026

View reviewed changes

docs: reattach doc comments to the right functions

1aa4958

Review feedback: the set_rbd_steps_per_frame doc was merged onto set_rbd_collisions_capacity, and insert_mjcf's doc onto the Python set_rbd_collisions_capacity binding. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

sebcrozet mentioned this pull request Aug 28, 2026

feat: add features required by RL #36

Merged

sebcrozet closed this in #36 Aug 29, 2026

Navigation Menu

Uh oh!

Python: per-environment MJCF insertion + collision-capacity override - #16

Python: per-environment MJCF insertion + collision-capacity override#16
haixuanTao wants to merge 2 commits into
dimforge:mainfrom
haixuanTao:feat/per-env-mjcf-upstream

haixuanTao commented Jul 15, 2026

Uh oh!

sebcrozet left a comment

Uh oh!

sebcrozet Jul 16, 2026

Uh oh!

haixuanTao Jul 16, 2026

Uh oh!

sebcrozet Jul 16, 2026

Uh oh!

haixuanTao Jul 16, 2026

Uh oh!

sebcrozet commented Aug 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

haixuanTao commented Jul 15, 2026

Why

Measured (RTX 5090, headless WebGPU, 12-DOF LeRobot biped, passive, 200-step window, one sync)

Notes

Uh oh!

sebcrozet left a comment

Choose a reason for hiding this comment

Uh oh!

sebcrozet Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

haixuanTao Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

sebcrozet Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

haixuanTao Jul 16, 2026

Choose a reason for hiding this comment

Uh oh!

sebcrozet commented Aug 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants