SolverMuJoCo: make root of fixed-based articulations kinematic to allow root position updates#1104
Conversation
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
📝 WalkthroughWalkthroughAdds mocap (fixed-base) body support to the MuJoCo solver: kernel signature and logic now accept a per-body mocap index and write mocap pose outputs; SolverMuJoCo builds and exposes a Newton→MuJoCo mocap mapping and wires it into kernel calls; tests and an example are updated to exercise mocap behavior. Changes
Sequence Diagram(s)sequenceDiagram
participant Solver as SolverMuJoCo
participant Kernel as update_joint_transforms_kernel
participant MJC as mjw_data
rect rgb(248,249,255)
note over Solver: Model construction
Solver->>Solver: Detect fixed-base (parent==-1 & FIXED)
Solver->>Solver: Assign sequential mocap indices -> newton_body_to_mocap_index
Solver->>MJC: Reserve/prepare mocap slots
end
rect rgb(240,255,240)
note over Solver,Kernel: Per-frame update flow
Solver->>Kernel: Call with newton_body_to_mocap_index
Kernel->>Kernel: Compute transform tf for joint/body
alt mocap_index >= 0
Kernel->>MJC: Write mocap_pos / mocap_quat
else not mocap
Kernel->>MJC: Write body_pos / body_quat
end
MJC-->>Solver: Updated poses available for observers/collision
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Possibly related PRs
Suggested reviewers
Pre-merge checks and finishing touches❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: Path: .coderabbit.yml Review profile: CHILL Plan: Pro 📒 Files selected for processing (3)
🧰 Additional context used🧠 Learnings (15)📓 Common learnings📚 Learning: 2025-11-24T08:05:21.380ZApplied to files:
📚 Learning: 2025-08-19T12:27:30.629ZApplied to files:
📚 Learning: 2025-08-18T15:56:26.587ZApplied to files:
📚 Learning: 2025-08-19T12:27:30.629ZApplied to files:
📚 Learning: 2025-09-22T21:08:31.901ZApplied to files:
📚 Learning: 2025-08-25T11:07:47.818ZApplied to files:
📚 Learning: 2025-09-22T21:03:39.624ZApplied to files:
📚 Learning: 2025-09-22T21:08:31.901ZApplied to files:
📚 Learning: 2025-09-25T16:14:22.033ZApplied to files:
📚 Learning: 2025-09-25T16:14:22.033ZApplied to files:
📚 Learning: 2025-09-22T21:03:18.367ZApplied to files:
📚 Learning: 2025-08-12T18:04:06.577ZApplied to files:
📚 Learning: 2025-08-18T15:56:26.587ZApplied to files:
📚 Learning: 2025-09-24T00:26:54.486ZApplied to files:
🧬 Code graph analysis (1)newton/tests/test_mujoco_solver.py (1)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
🔇 Additional comments (2)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
adding an unresolved comment to remind me that this needs a dependency update as well. MjWarp PR is google-deepmind/mujoco_warp#858 |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
newton/examples/robot/example_robot_allegro_hand.py (1)
129-131: Assumption that cube is the last body is a bit brittleUsing
allegro_hand.body_count - 1assumes the cube remains the last body in the USD template. If the asset changes (extra bodies added, reordering),cube_body_offsetandhand_body_indicesmay no longer match the cube.If you want this to be more future‑proof, consider deriving the cube index from
allegro_hand.body_key(e.g., by matching the DexCube body name) instead of relying on ordering.newton/_src/solvers/mujoco/solver_mujoco.py (3)
239-240: Clarifynewton_body_to_mocap_indexmapping semantics in the docstringThe field is documented as:
self.newton_body_to_mocap_index: wp.array(dtype=int) | None = None """Mapping from Newton body index to MuJoCo mocap body index. -1 if body is not mocap. Shape [bodies_per_world], dtype int32."""This is accurate conceptually, but given the way it’s constructed later (from
body_mapping.keys()for the first world only), it would help future readers to explicitly mention that:
- It is defined per world, using the same local body index convention as
to_mjc_body_index/bodies_per_world.- Only bodies that appear in the MuJoCo model (i.e., in
body_mapping) participate; others are implicitly non‑mocap (-1).No functional changes needed here—just tightening the docstring to match the actual construction.
1502-1508:newton_body_to_mocap_indexarray construction relies on insertion order; consider making the indexing more explicitYou currently build the Warp array via:
newton_body_indices = np.fromiter(body_mapping.keys(), dtype=int)[1:] body_to_mocap = np.array( [newton_body_to_mocap_index.get(idx, -1) for idx in newton_body_indices], dtype=np.int32, ) self.newton_body_to_mocap_index = wp.array(body_to_mocap, dtype=wp.int32)and later index it in the kernel as:
mocap_index = newton_body_to_mocap_index[child] # in update_joint_transforms_kernelThis is correct as long as:
body_mapping’s non‑world keys are exactly the local Newton body indices[0, ..., bodies_per_world-1]for the first world, and- They are inserted into the dict in that order, so that
newton_body_indices[i] == i.That’s true today for the models used in your tests (since bodies are built in order and
body_mappingis filled once per child), but it’s a fairly implicit coupling between:
- how
body_mappingis populated, and- how
newton_body_to_mocap_indexis indexed.To make this more robust and self‑documenting, consider one of:
- Building
body_to_mocapexplicitly by local body index rather than relying on dict key iteration, e.g.:bodies_per_world = self.model.body_count // self.model.num_worlds body_to_mocap = np.full(bodies_per_world, -1, dtype=np.int32) for newton_body, mj_body in body_mapping.items(): if newton_body >= 0 and newton_body < bodies_per_world: body_to_mocap[newton_body] = newton_body_to_mocap_index.get(newton_body, -1) self.newton_body_to_mocap_index = wp.array(body_to_mocap, dtype=wp.int32)
- Or, if you intend
newton_body_to_mocap_indexto be indexed in the same “local body index” space asto_mjc_body_index, compute both arrays from a shared, ordered list of local body indices (e.g.,selected_bodies).That would avoid surprises if the body construction order or
body_mappingimplementation ever changes, and makes the expected indexing scheme much clearer.
1745-1770: Propagating mocap mapping intoupdate_joint_propertiesis correct; just ensure it’s always initializedThe updated
update_joint_propertieslaunch:wp.launch( update_joint_transforms_kernel, dim=self.model.joint_count, inputs=[ self.model.joint_X_p, self.model.joint_X_c, self.model.joint_qd_start, self.model.joint_dof_dim, self.model.joint_axis, self.model.joint_child, self.model.joint_type, self.model.joint_limit_ke, self.model.joint_limit_kd, self.joint_mjc_dof_start, self.to_mjc_body_index, self.newton_body_to_mocap_index, joints_per_world, ], outputs=[ self.mjw_model.jnt_pos, self.mjw_model.jnt_axis, self.mjw_model.jnt_solref, self.mjw_model.body_pos, self.mjw_model.body_quat, self.mjw_data.mocap_pos, self.mjw_data.mocap_quat, ], )correctly threads the new mapping and mocap outputs into the kernel, and is consistent with how other properties are updated via
notify_model_changed(SolverNotifyFlags.JOINT_PROPERTIES).One thing to double‑check:
self.newton_body_to_mocap_indexis only initialized in_convert_to_mjc(the mujoco_warp path). If someone ever constructs a solver with an externalmjw_model/mjw_datapair or in a CPU‑only configuration and still callsnotify_model_changed(SolverNotifyFlags.JOINT_PROPERTIES), this field will beNoneandwp.launchwill fail.Given that the kernel itself is specific to the mujoco_warp backend, it would be safer to either:
- Guard
update_joint_propertieswithif self.use_mujoco_cpu: return, or- Initialize
self.newton_body_to_mocap_indexto an all‑(-1) array in any code path that may call this method.Functionally fine for the common mjwarp path, but worth tightening for edge configurations.
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
newton/_src/solvers/mujoco/kernels.py(2 hunks)newton/_src/solvers/mujoco/solver_mujoco.py(8 hunks)newton/examples/robot/example_robot_allegro_hand.py(2 hunks)newton/tests/test_mujoco_solver.py(1 hunks)
🧰 Additional context used
🧠 Learnings (14)
📓 Common learnings
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.
📚 Learning: 2025-08-18T15:56:26.587Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.
Applied to files:
newton/_src/solvers/mujoco/kernels.pynewton/examples/robot/example_robot_allegro_hand.pynewton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.
Applied to files:
newton/_src/solvers/mujoco/kernels.pynewton/examples/robot/example_robot_allegro_hand.pynewton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:03:39.624Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.
Applied to files:
newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, Example.use_mujoco is a high-level attribute that controls whether to use MuJoCo solver vs other solvers (like XPBD), while SolverMuJoCo.use_mujoco_cpu controls the backend within MuJoCo (CPU vs Warp). These are separate concepts serving different purposes - the PR rename only applies to the solver parameter, not the Example class attributes.
Applied to files:
newton/_src/solvers/mujoco/kernels.pynewton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.
Applied to files:
newton/_src/solvers/mujoco/kernels.pynewton/examples/robot/example_robot_allegro_hand.pynewton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: In NVIDIA Warp's Newton physics library, wp.transform supports direct numerical indexing (e.g., body_tf[0], body_tf[1], body_tf[2] for position and body_tf[3], body_tf[4], body_tf[5], body_tf[6] for quaternion components) to access its elements. This is the standard API used throughout the codebase.
Applied to files:
newton/_src/solvers/mujoco/kernels.pynewton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: wp.transform in NVIDIA Warp supports indexing using numerical indices (e.g., body_tf[0], body_tf[3], etc.) to access translation and rotation components. The indexing approach used in the Newton codebase is valid and compiles correctly.
Applied to files:
newton/_src/solvers/mujoco/kernels.pynewton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:03:18.367Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:638-648
Timestamp: 2025-09-22T21:03:18.367Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently only supports additive updates, meaning it assumes n_coords == n_dofs. Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work.
Applied to files:
newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-25T20:10:59.536Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 634
File: newton/examples/ik/example_ik_franka.py:0-0
Timestamp: 2025-08-25T20:10:59.536Z
Learning: In Newton's IK examples, when creating solver arrays with `wp.array(source_array, shape=new_shape)`, this creates a view into the same underlying memory rather than a copy. This means updates to the reshaped array are automatically reflected in the original array without needing explicit copy operations like `wp.copy()`.
Applied to files:
newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
Repo: newton-physics/newton PR: 521
File: newton/examples/example_cloth_hanging.py:36-36
Timestamp: 2025-08-12T05:17:34.423Z
Learning: The Newton migration guide (docs/migration.rst) is specifically for documenting how to migrate existing warp.sim functionality to Newton equivalents. New Newton-only features that didn't exist in warp.sim do not need migration documentation.
Applied to files:
newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, there's a distinction between solver parameters and Example class attributes. The Example class can have its own use_mujoco attribute for controlling example-level behavior (like CUDA graphs, rendering logic), while the solver uses use_mujoco_cpu for backend selection. These serve different purposes and should not be conflated during API renames.
Applied to files:
newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
Repo: newton-physics/newton PR: 521
File: newton/examples/example_cloth_hanging.py:36-36
Timestamp: 2025-08-12T05:17:34.423Z
Learning: The STYLE3D solver in Newton is a new solver that did not exist in warp.sim, so it does not require any migration documentation.
Applied to files:
newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-24T00:26:54.486Z
Learnt from: eric-heiden
Repo: newton-physics/newton PR: 710
File: newton/_src/utils/import_usd.py:597-599
Timestamp: 2025-09-24T00:26:54.486Z
Learning: In the Newton ModelBuilder class, `joint_count`, `body_count`, `articulation_count`, and `shape_count` are implemented as property decorated methods, so they should be accessed as attributes (e.g., `builder.joint_count`) rather than called as methods (e.g., `builder.joint_count()`).
Applied to files:
newton/tests/test_mujoco_solver.py
🧬 Code graph analysis (3)
newton/examples/robot/example_robot_allegro_hand.py (3)
newton/_src/sim/builder.py (1)
body_count(875-879)newton/examples/__init__.py (1)
test_body_state(38-113)newton/_src/utils/__init__.py (1)
vec_inside_limits(213-220)
newton/_src/solvers/mujoco/solver_mujoco.py (3)
newton/_src/sim/joints.py (1)
JointType(20-44)newton/_src/solvers/solver.py (1)
notify_model_changed(267-294)newton/_src/solvers/flags.py (1)
SolverNotifyFlags(22-56)
newton/tests/test_mujoco_solver.py (3)
newton/_src/sim/builder.py (6)
add_body(1618-1690)add_joint_fixed(2042-2086)add_shape_box(3099-3138)add_joint_free(2088-2142)add_shape_sphere(3058-3097)finalize(4780-5168)newton/_src/solvers/mujoco/solver_mujoco.py (1)
notify_model_changed(373-383)newton/_src/sim/model.py (2)
state(511-548)control(550-589)
🪛 GitHub Actions: Pull Request
newton/tests/test_mujoco_solver.py
[error] 2119-2119: Mocap body transform test failed. Test 'test_mocap_body_transform_updates_collision_geoms' expected downward velocity but got approximately 0 m/s (np.float32(-0.00024268524)). This indicates ball did not fall as expected during the simulation.
🪛 Ruff (0.14.5)
newton/examples/robot/example_robot_allegro_hand.py
220-220: Unused lambda argument: qd
(ARG005)
233-233: Unused lambda argument: qd
(ARG005)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2
- GitHub Check: Run GPU Tests / Run GPU Unit Tests on AWS EC2
🔇 Additional comments (3)
newton/_src/solvers/mujoco/kernels.py (1)
939-948: Mocap routing inupdate_joint_transforms_kernellooks correct but depends on existing body-index conventionsThe new signature and routing:
- Adds
newton_body_to_mocap_index: wp.array(dtype=wp.int32).- Adds
mocap_pos/mocap_quatoutputs.- For each non‑FREE joint, computes
tf = parent_xform * inv(child_xform)and then:child = joint_child[joint_in_world] # local Newton body id (per-world) body_id = body_mapping[child] # MuJoCo body id mocap_index = newton_body_to_mocap_index[child] rotation = wp.quat(tf.q.w, tf.q.x, tf.q.y, tf.q.z) # xyzw -> wxyz if mocap_index >= 0: mocap_pos[worldid, mocap_index] = tf.p mocap_quat[worldid, mocap_index] = rotation else: body_pos[worldid, body_id] = tf.p body_quat[worldid, body_id] = rotationThis matches the rest of the file’s quat conventions (Warp uses xyzw, MuJoCo fields use wxyz) and preserves previous behavior for non‑mocap bodies while giving mocap roots dedicated storage.
One thing to be aware of: this assumes that:
joint_child[joint_in_world]is a per‑world local body index in[0, bodies_per_world), andnewton_body_to_mocap_indexis indexed with the same local body index convention asbody_mapping/to_mjc_body_index.Given how
to_mjc_body_indexis constructed insolver_mujoco.py, that’s consistent for single‑world models and for multi‑world models using the existing “per‑world slice” layout, but the coupling is fairly implicit.I’d keep the implementation as‑is but consider documenting these indexing assumptions in the kernel docstring (and perhaps asserting valid ranges in debug builds) to avoid future regressions if body indexing ever changes.
Also applies to: 970-986
newton/_src/solvers/mujoco/solver_mujoco.py (2)
988-992: Mocap root detection and bookkeeping look correct for fixed-base articulationsThe new logic in
_convert_to_mjc:
- Detects mocap roots via:
fixed_base = parent == -1 and joint_type[j] == JointType.FIXED ... body = mj_bodies[body_mapping[parent]].add_body(..., mocap=fixed_base) ... if fixed_base: newton_body_to_mocap_index[child] = next_mocap_index next_mocap_index += 1This matches the intended use: treat fixed joints from world (
parent == -1) as mocap bodies so subsequent joint frame updates can drive them kinematically.The bookkeeping (
next_mocap_indexand per‑body dict) is straightforward and integrates cleanly with the later array construction.I don’t see functional issues here; this is a reasonable way to identify mocap roots for fixed‑base trees.
Also applies to: 1219-1225
1552-1557: Delayedexpand_model_fieldsandnotify_model_changedsequencing makes senseMoving the expansion and initial
notify_model_changedto the end of_convert_to_mjc:self.mjw_data = mujoco_warp.put_data(...) # expand model fields that can be expanded: self.expand_model_fields(self.mjw_model, nworld) # now complete the data from the Newton model self.notify_model_changed(SolverNotifyFlags.ALL)is a good change:
- All derived mappings (
to_mjc_body_index,to_newton_shape_index,newton_body_to_mocap_index, etc.) are fully initialized before data expansion and updates.- The
SolverNotifyFlags.ALLcall ensures that model-, joint-, DOF-, body-, and geom‑related MuJoCo/Warp arrays are brought in sync with the Newton model right after creation, which is especially important now that mocap fields are part of the joint/property update path.No issues here; this is a sensible ordering.
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
newton/_src/solvers/mujoco/solver_mujoco.py (1)
1189-1225: Ensure mocap mapping is initialized for external mjw_model/mjw_data pathWhen
mjw_modelandmjw_dataare provided externally,self.newton_body_to_mocap_indexremainsNone. However,update_joint_properties()unconditionally passes it to the kernel at line 1759, which expects awp.array(dtype=wp.int32), notNone. This will cause a runtime crash if any code path triggers the JOINT_PROPERTIES notify flag after initialization.Initialize
newton_body_to_mocap_indexto a safe default (array of -1 values) in the external model path to prevent the crash:if mjw_model is not None and mjw_data is not None: self.mjw_model = mjw_model self.mjw_data = mjw_data self.use_mujoco_cpu = False # No Newton→mocap mapping available in this path; treat all bodies as non‑mocap bodies_per_world = model.body_count // max(model.num_worlds, 1) self.newton_body_to_mocap_index = wp.full(bodies_per_world, -1, dtype=wp.int32, device=model.device)
🧹 Nitpick comments (1)
newton/examples/robot/example_robot_allegro_hand.py (1)
129-131: Cube body offset assumption is brittle to asset changesUsing
self.cube_body_offset = allegro_hand.body_count - 1assumes the cube remains the last body in the template. If the USD changes (extra bodies added, ordering altered), this silently points at the wrong body. Consider resolving the cube by key (e.g., matchingbody_keyagainst aDexCubepattern) and storing its index, then using that for per-world offsets.
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
newton/_src/solvers/mujoco/kernels.py(2 hunks)newton/_src/solvers/mujoco/solver_mujoco.py(8 hunks)newton/examples/robot/example_robot_allegro_hand.py(2 hunks)newton/tests/test_mujoco_solver.py(1 hunks)
🧰 Additional context used
🧠 Learnings (13)
📓 Common learnings
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.
Learnt from: gdaviet
Repo: newton-physics/newton PR: 631
File: newton/examples/mpm/example_anymal_c_walk_on_sand.py:164-169
Timestamp: 2025-08-25T11:07:47.818Z
Learning: In Newton's MuJoCo solver implementation, setting TARGET_POSITION mode on floating base DOFs (first 6 DOFs) does not overconstrain the base - the solver handles this appropriately and maintains expected behavior for floating robots.
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, Example.use_mujoco is a high-level attribute that controls whether to use MuJoCo solver vs other solvers (like XPBD), while SolverMuJoCo.use_mujoco_cpu controls the backend within MuJoCo (CPU vs Warp). These are separate concepts serving different purposes - the PR rename only applies to the solver parameter, not the Example class attributes.
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:638-648
Timestamp: 2025-09-22T21:03:18.367Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently only supports additive updates, meaning it assumes n_coords == n_dofs. Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work.
📚 Learning: 2025-08-18T15:56:26.587Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.
Applied to files:
newton/_src/solvers/mujoco/kernels.pynewton/examples/robot/example_robot_allegro_hand.pynewton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.
Applied to files:
newton/_src/solvers/mujoco/kernels.pynewton/examples/robot/example_robot_allegro_hand.pynewton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:03:39.624Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.
Applied to files:
newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.
Applied to files:
newton/_src/solvers/mujoco/kernels.pynewton/examples/robot/example_robot_allegro_hand.pynewton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: In NVIDIA Warp's Newton physics library, wp.transform supports direct numerical indexing (e.g., body_tf[0], body_tf[1], body_tf[2] for position and body_tf[3], body_tf[4], body_tf[5], body_tf[6] for quaternion components) to access its elements. This is the standard API used throughout the codebase.
Applied to files:
newton/_src/solvers/mujoco/kernels.pynewton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-25T16:14:22.033Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_objectives.py:0-0
Timestamp: 2025-09-25T16:14:22.033Z
Learning: wp.transform in NVIDIA Warp supports indexing using numerical indices (e.g., body_tf[0], body_tf[3], etc.) to access translation and rotation components. The indexing approach used in the Newton codebase is valid and compiles correctly.
Applied to files:
newton/_src/solvers/mujoco/kernels.pynewton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-09-22T21:03:18.367Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:638-648
Timestamp: 2025-09-22T21:03:18.367Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently only supports additive updates, meaning it assumes n_coords == n_dofs. Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work.
Applied to files:
newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-25T20:10:59.536Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 634
File: newton/examples/ik/example_ik_franka.py:0-0
Timestamp: 2025-08-25T20:10:59.536Z
Learning: In Newton's IK examples, when creating solver arrays with `wp.array(source_array, shape=new_shape)`, this creates a view into the same underlying memory rather than a copy. This means updates to the reshaped array are automatically reflected in the original array without needing explicit copy operations like `wp.copy()`.
Applied to files:
newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-09-24T00:26:54.486Z
Learnt from: eric-heiden
Repo: newton-physics/newton PR: 710
File: newton/_src/utils/import_usd.py:597-599
Timestamp: 2025-09-24T00:26:54.486Z
Learning: In the Newton ModelBuilder class, `joint_count`, `body_count`, `articulation_count`, and `shape_count` are implemented as property decorated methods, so they should be accessed as attributes (e.g., `builder.joint_count`) rather than called as methods (e.g., `builder.joint_count()`).
Applied to files:
newton/tests/test_mujoco_solver.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, Example.use_mujoco is a high-level attribute that controls whether to use MuJoCo solver vs other solvers (like XPBD), while SolverMuJoCo.use_mujoco_cpu controls the backend within MuJoCo (CPU vs Warp). These are separate concepts serving different purposes - the PR rename only applies to the solver parameter, not the Example class attributes.
Applied to files:
newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-19T12:27:30.629Z
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, there's a distinction between solver parameters and Example class attributes. The Example class can have its own use_mujoco attribute for controlling example-level behavior (like CUDA graphs, rendering logic), while the solver uses use_mujoco_cpu for backend selection. These serve different purposes and should not be conflated during API renames.
Applied to files:
newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-08-12T05:17:34.423Z
Learnt from: shi-eric
Repo: newton-physics/newton PR: 521
File: newton/examples/example_cloth_hanging.py:36-36
Timestamp: 2025-08-12T05:17:34.423Z
Learning: The Newton migration guide (docs/migration.rst) is specifically for documenting how to migrate existing warp.sim functionality to Newton equivalents. New Newton-only features that didn't exist in warp.sim do not need migration documentation.
Applied to files:
newton/_src/solvers/mujoco/solver_mujoco.py
🧬 Code graph analysis (3)
newton/tests/test_mujoco_solver.py (2)
newton/_src/sim/builder.py (7)
ModelBuilder(70-5272)add_body(1618-1690)add_joint_fixed(2042-2086)add_shape_box(3099-3138)add_joint_free(2088-2142)add_shape_sphere(3058-3097)copy(228-229)newton/_src/solvers/mujoco/solver_mujoco.py (1)
notify_model_changed(373-383)
newton/examples/robot/example_robot_allegro_hand.py (3)
newton/_src/sim/builder.py (1)
body_count(875-879)newton/examples/__init__.py (1)
test_body_state(38-113)newton/_src/utils/__init__.py (1)
vec_inside_limits(213-220)
newton/_src/solvers/mujoco/solver_mujoco.py (3)
newton/_src/sim/joints.py (1)
JointType(20-44)newton/_src/solvers/solver.py (1)
notify_model_changed(267-294)newton/_src/solvers/flags.py (1)
SolverNotifyFlags(22-56)
🪛 GitHub Actions: Pull Request
newton/tests/test_mujoco_solver.py
[error] 2119-2119: Test mocap_body_transform_updates_collision_geoms failed. Ball should be falling with downward velocity, but velocity was 0 (np.float32(-0.00024268524) not less than -0.2).
🪛 GitHub Actions: Pull Request - AWS GPU
newton/tests/test_mujoco_solver.py
[error] 2119-2119: MuJoCo mocap body transform test failed: ball velocity not downward (expected negative velocity).
newton/examples/robot/example_robot_allegro_hand.py
[error] 3-3: IndentationError: unexpected indent encountered while parsing generated warp function for test body state.
🪛 Ruff (0.14.5)
newton/examples/robot/example_robot_allegro_hand.py
220-220: Unused lambda argument: qd
(ARG005)
233-233: Unused lambda argument: qd
(ARG005)
🔇 Additional comments (5)
newton/_src/solvers/mujoco/kernels.py (1)
927-986: Mocap routing inupdate_joint_transforms_kernellooks consistent with existing conventionsThe new
newton_body_to_mocap_indexparameter and mocap branch correctly:
- reuse the existing
child = joint_child[joint_in_world]+body_mapping[child]pattern so indices stay per‑world,- compute
tfonce and share it between mocap/body paths, and- preserve the established wxyz storage convention by using
rotation = wp.quat(tf.q.w, tf.q.x, tf.q.y, tf.q.z)before writing tobody_quat/mocap_quat(aligns withtest_joint_frame_updateexpectations).No correctness issues stand out here; the changes are a straightforward extension of the prior body transform logic to mocap outputs.
newton/_src/solvers/mujoco/solver_mujoco.py (4)
233-245: New mocap-body mapping attribute is well scopedThe new
self.newton_body_to_mocap_indexattribute and its docstring make the mocap mapping explicit and keep it parallel toself.to_mjc_body_index. This is a good place to expose the Newton→mocap index mapping to tests and higher-level code.
983-992: Mocap index tracking in_convert_to_mjcmatches fixed-base semanticsThe
newton_body_to_mocap_indexdict is only filled whenparent == -1 and joint_type[j] == JointType.FIXED, which is exactly the fixed‑base articulation root you want to treat as a mocap body. Coupled withmocap=fixed_basein theadd_bodycall, this correctly identifies and indexes mocap bodies at conversion time.
1502-1508: Mocap mapping array construction matches body index conventionsThis block:
newton_body_indices = np.fromiter(body_mapping.keys(), dtype=int)[1:] body_to_mocap = np.array( [newton_body_to_mocap_index.get(idx, -1) for idx in newton_body_indices], dtype=np.int32 ) self.newton_body_to_mocap_index = wp.array(body_to_mocap, dtype=wp.int32)produces a length‑
bodies_per_worldarray aligned withto_mjc_body_index(samebody_mapping.keys()[1:]source). Using.get(idx, -1)correctly marks non‑mocap bodies as-1. This layout matches howupdate_joint_transforms_kernelexpects to indexnewton_body_to_mocap_index[child].
1736-1770: Mocap data wiring intoupdate_joint_transforms_kernelis consistentPassing
self.newton_body_to_mocap_indexand wiringself.mjw_data.mocap_pos/mocap_quatas outputs integrates the kernel’s mocap support cleanly with the rest of the solver:
- For non‑mocap bodies (
mocap_index == -1), behavior is unchanged (body_pos/body_quat are written).- For mocap bodies, the kernel now keeps
body_pos/body_quatuntouched and writes into the mocap buffers, which is what MuJoCo expects for kinematic mocap roots.This matches the new integration test’s expectations and doesn’t disturb existing joint/limit logic.
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
|
Mjwarp Update is in now, should be good to go. |
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
newton/examples/robot/example_robot_allegro_hand.py (1)
233-234: Simplify lambda to use closure and single-line expression.The lambda has two issues:
Non-standard signature with default arguments: The signature
lambda q, _qd, lower=cube_lower, upper=cube_upper: ...creates a 4-parameter function (2 required + 2 with defaults), buttest_body_stateexpects a 2-argument callable. Default arguments may also confuse Warp'screate_warp_function.Multi-line expression: Splitting the lambda body across lines 233-234 can cause
IndentationErrorduring Warp kernel codegen.For consistency with the hand test lambda (line 220), use closure to capture
cube_lowerandcube_upper, collapse to a single line, and add the# noqa: B023comment:Apply this diff:
- lambda q, _qd, lower=cube_lower, upper=cube_upper: newton.utils.vec_inside_limits(q.p, lower, upper) - and q.p[2] > 0.0, + lambda q, _qd: newton.utils.vec_inside_limits(q.p, cube_lower, cube_upper) and q.p[2] > 0.0, # noqa: B023
🧹 Nitpick comments (1)
newton/examples/robot/example_robot_allegro_hand.py (1)
220-220: Rename unused parameterqdto_qdfor consistency.The lambda parameter
qdis unused, triggering a Ruff warning. For consistency with the cube test lambda (line 233) and to indicate intentional non-use, rename it to_qd.Apply this diff:
- lambda q, qd: newton.utils.vec_inside_limits(q.p, hand_lower, hand_upper), # noqa: B023 + lambda q, _qd: newton.utils.vec_inside_limits(q.p, hand_lower, hand_upper), # noqa: B023
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
newton/examples/robot/example_robot_allegro_hand.py(2 hunks)
🧰 Additional context used
🧠 Learnings (5)
📓 Common learnings
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/_src/sim/ik/ik_lbfgs_optimizer.py:739-752
Timestamp: 2025-09-22T21:03:39.624Z
Learning: The L-BFGS optimizer in newton/_src/sim/ik/ik_lbfgs_optimizer.py currently intentionally only supports additive updates (assuming n_coords == n_dofs). Velocity space integration for joints with mismatched coordinate/DOF dimensions (like free/ball joints) is planned for future work and should not be flagged as an issue in current reviews.
Learnt from: gdaviet
Repo: newton-physics/newton PR: 631
File: newton/examples/mpm/example_anymal_c_walk_on_sand.py:164-169
Timestamp: 2025-08-25T11:07:47.818Z
Learning: In Newton's MuJoCo solver implementation, setting TARGET_POSITION mode on floating base DOFs (first 6 DOFs) does not overconstrain the base - the solver handles this appropriately and maintains expected behavior for floating robots.
Learnt from: preist-nvidia
Repo: newton-physics/newton PR: 579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, Example.use_mujoco is a high-level attribute that controls whether to use MuJoCo solver vs other solvers (like XPBD), while SolverMuJoCo.use_mujoco_cpu controls the backend within MuJoCo (CPU vs Warp). These are separate concepts serving different purposes - the PR rename only applies to the solver parameter, not the Example class attributes.
📚 Learning: 2025-09-26T05:58:21.013Z
Learnt from: WenchaoHuang
Repo: newton-physics/newton PR: 835
File: newton/_src/solvers/style3d/collision/collision.py:49-53
Timestamp: 2025-09-26T05:58:21.013Z
Learning: In newton/_src/solvers/style3d/collision/collision.py, the broad_phase buffers (broad_phase_ee, broad_phase_ef, broad_phase_vf) created with wp.array() don't need zero-initialization because they follow a write-first-read-later pattern: BVH query kernels in frame_begin() populate these buffers completely (write-only operations), and collision kernels only read from them afterward in accumulate_contact_force(), ensuring no undefined reads occur.
Applied to files:
newton/examples/robot/example_robot_allegro_hand.py
📚 Learning: 2025-08-18T15:56:26.587Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 552
File: newton/_src/solvers/mujoco/solver_mujoco.py:0-0
Timestamp: 2025-08-18T15:56:26.587Z
Learning: In Newton's MuJoCo solver, when transforming joint axes from Newton's internal frame to MuJoCo's expected frame, use wp.quat_rotate(joint_rot, axis) not wp.quat_rotate_inv(joint_rot, axis). The joint_rot represents rotation from joint-local to body frame, so forward rotation is correct.
Applied to files:
newton/examples/robot/example_robot_allegro_hand.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering.
Applied to files:
newton/examples/robot/example_robot_allegro_hand.py
📚 Learning: 2025-09-22T21:08:31.901Z
Learnt from: dylanturpin
Repo: newton-physics/newton PR: 806
File: newton/examples/ik/example_ik_franka.py:121-123
Timestamp: 2025-09-22T21:08:31.901Z
Learning: In the newton physics framework, when creating warp arrays for IK solver joint variables using wp.array(self.model.joint_q, shape=(1, coord_count)), the resulting array acts as a reference/pointer to the original model's joint coordinates, so updates from the IK solver automatically reflect in the model's joint_q buffer used for rendering, despite the general warp documentation suggesting copies are made by default.
Applied to files:
newton/examples/robot/example_robot_allegro_hand.py
🧬 Code graph analysis (1)
newton/examples/robot/example_robot_allegro_hand.py (2)
newton/examples/__init__.py (1)
test_body_state(38-113)newton/_src/utils/__init__.py (1)
vec_inside_limits(213-220)
🪛 Ruff (0.14.5)
newton/examples/robot/example_robot_allegro_hand.py
220-220: Unused lambda argument: qd
(ARG005)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2
- GitHub Check: Run GPU Tests / Run GPU Unit Tests on AWS EC2
🔇 Additional comments (2)
newton/examples/robot/example_robot_allegro_hand.py (2)
129-130: LGTM!The
cube_body_offsetattribute correctly identifies the cube as the last body in each replicated world, enabling targeted testing.
227-228: Good fix for cube Z bounds!The cube lower Z bound is now correctly set to
0.0, allowing the cube to fall to the ground plane while keeping X/Y bounds tight. This properly addresses the previous review concern and aligns with the PR objective to allow the cube to occasionally fall from the hand.
eric-heiden
left a comment
There was a problem hiding this comment.
Thanks @adenzler-nvidia, this is great to have this support for fixed-based articulations through mocap bodies! This will allow us to support kinematically driven bodies in MuJoCo by setting the transforms of fixed joints. Just one question on general massless bodies in MuJoCo, see comment.
…ow root position updates (newton-physics#1104) Signed-off-by: Alain Denzler <adenzler@nvidia.com>
…ow root position updates (newton-physics#1104) Signed-off-by: Alain Denzler <adenzler@nvidia.com>
Reason for this change: Allegro hand example test was flaky #909
The reason for the flakyness was that we animate the position of the fixed-base root. This does not work well with MuJoCo as-is, because the collision geometries attached to static bodies are not being updated, and thus we lost collisions with the palm of the hand.
This PR makes the fixed-base root a "mocap body", which is MuJoCo-speak for a kinematically driven body. This requires us to change some things in the plumbing, most notably updating mocap_pos/quat instead of body_pos/quat. While working on this, I also found a bug in MuJoCo Warp which is a requirement for this PR to get in: google-deepmind/mujoco_warp#858
Other than that, the change should be straightforward.
I'm not sure if that's all that's needed to close #909, but I also changed the test criteria a bit to allow the cube to fall out of the hand. We don't really have a smart controller so I think it can happen randomly that one of the cubes falls.
Before your PR is "Ready for review"
newton/tests/test_examples.py)pre-commit run -aSummary by CodeRabbit
New Features
Bug Fixes
Tests
Examples
✏️ Tip: You can customize this high-level summary in your review settings.