Python: per-environment MJCF insertion + collision-capacity override#16
Open
haixuanTao wants to merge 2 commits into
Open
Python: per-environment MJCF insertion + collision-capacity override#16haixuanTao wants to merge 2 commits into
haixuanTao wants to merge 2 commits into
Conversation
…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>
haixuanTao
force-pushed
the
feat/per-env-mjcf-upstream
branch
from
July 16, 2026 12:17
e839c74 to
0447bf0
Compare
sebcrozet
requested changes
Jul 16, 2026
sebcrozet
left a comment
sebcrozet
left a comment
Member
There was a problem hiding this comment.
That looks good, thank you. Just a few comments that got strangely merged.
Comment on lines
+157
to
+163
| /// 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. |
Member
There was a problem hiding this comment.
The comments got mixed up (the first two sentences should apply to set_rbd_steps_per_frame).
Contributor
Author
There was a problem hiding this comment.
Good catch — fixed in 1aa4958: the set_rbd_steps_per_frame doc is back on its own function.
Comment on lines
+266
to
+276
| /// 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))] |
Member
There was a problem hiding this comment.
The comments got mixed up here too.
Contributor
Author
There was a problem hiding this comment.
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
insert_mjcfgains anenv=kwarg (default 0 — fully backward compatible), making one-robot-per-environment reachable from Python: the robot, its auto-floor and the camera-fit AABB targetrbd_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 keepsdofs_capat the robot's own DOFs and scales linearly withbatches, 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 bindcapacity × 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)
Linear scaling; cross-checks against the zealot training stack's physics-only rate at 2,048 envs (~127k on the same engine family).
Notes
finalizerequires identical collider counts in every environment — worth a friendlier error eventually.🤖 Generated with Claude Code