Python: per-step MJCF actuator control + multibody state readback - #12
Closed
haixuanTao wants to merge 4 commits into
Closed
Python: per-step MJCF actuator control + multibody state readback#12haixuanTao wants to merge 4 commits into
haixuanTao wants to merge 4 commits into
Conversation
- GpuMultibodySet::sync_joint_data_from_rapier: refresh every link's joint data (motor targets/gains, limits) from the rapier multibody set in one buffer write, mirroring from_rapier's traversal order. - NexusState::control_multibody_motors: runtime actuation entry point that mutates the rapier joints (e.g. rapier3d-mjcf's apply_controls_multibody) and pushes the refreshed joint data to the GPU, without marking the world dirty (no rebuild). - Viewer::read_multibody_links + links_workspace COPY_SRC: one-readback joint coordinates, link world poses and world-space velocities per env. - Python: NexusState keeps the MjcfRobotHandles from insert_mjcf and exposes actuator_names() / apply_actuator_controls(viewer, ctrl, env) with full MJCF actuator semantics; NexusViewer.read_multibody_links(state, env) returns (coords, positions, quats, linvels, angvels) numpy arrays. Together these make MJCF robots drivable per control step from Python (position-servo PD runs inside the solver) with full state observation -- the two pieces sim-to-sim eval loops need. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
sebcrozet
force-pushed
the
py-mjcf-actuation
branch
from
July 16, 2026 12:14
43c5c98 to
a7064e3
Compare
…port The rebase onto upstream main left control_multibody_motors and read_multibody_links ungated while RbdState::multibodies(_mut) is dim3-only, breaking the 2D crates. Also fixes the missing numpy PyArray2 import and an unused_assignments warning in the MJCF loader. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Stock naga 29's MSL writer re-evaluates a loop's break_if condition after the continuing block has advanced the loop phis, dropping the final body iteration on Metal. In the multibody solve kernels the per-lane J.v loops run exactly one iteration per lane, so they executed zero times: zero impulses from the contact and joint/PD sweeps while gravity kept integrating -> robots free-fell through the floor on macOS. Pin the fixed fork (see zealot/docs/metal-contact-bug-proposal.md). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Member
|
This will be merged as part of #36 |
sebcrozet
added a commit
that referenced
this pull request
Aug 29, 2026
* 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>
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.
Makes MJCF robots drivable per control step from Python, with full state observation — the two pieces an RL / sim-to-sim evaluation loop needs. Follow-up to #7/#8, standalone on main.
What
Control —
NexusState.apply_actuator_controls(viewer, ctrl, env=0)applies one MJCF control vector (one entry per<actuator>, in actuator order —NexusState.actuator_names()reports the order) with full MJCF actuator semantics viarapier3d-mjcf'sapply_controls_multibody(<position>servos with kp/kv,<motor>force/gear, force limits), then pushes the refreshed joint-motor state to the GPU in one buffer write. The PD servo itself runs inside the GPU solver, so there is no per-step solver round-trip beyond that single upload.Observation —
NexusViewer.read_multibody_links(state, env=0)reads every link's state back in one readback:(coords, positions, quats, linvels, angvels)numpy arrays (one row per link, GPU traversal order).coordsare the generalized joint coordinates (a revolute joint's angle iscoords[3], the first angular slot); poses/velocities are world-space.How
GpuMultibodySet::sync_joint_data_from_rapierrefreshes each link'sGenericJointdata (motors, limits) from the rapier multibody set, mirroringfrom_rapier's traversal order, then uploadslinks_staticonce. Only joint data is touched — coordinates/velocities/mass properties are untouched, so this can't teleport links.NexusState::control_multibody_motorswraps "mutate rapier joints, then sync" without marking the world dirty (motor updates are control, not topology — no GPU rebuild).links_workspacegainsCOPY_SRC+ public accessors for the readback.NexusStatekeeps theMjcfRobotHandlesfrominsert_mjcf.Validation
On the LeRobot bipedal platform MJCF (12 position actuators, 13 links): a passive run collapses to torso z ≈ −0.34 within 1 s; driving the servos each step holds it at z ≈ −0.02, and
read_multibody_linkstracks joint angles/limits throughout.cargo checkclean onnexus3d,nexus_viewer3d,nexus_python3d.🤖 Generated with Claude Code