Skip to content

SolverMuJoCo: make root of fixed-based articulations kinematic to allow root position updates#1104

Merged
eric-heiden merged 10 commits into
newton-physics:mainfrom
adenzler-nvidia:dev/adenzler/mocap-bodies
Nov 25, 2025
Merged

SolverMuJoCo: make root of fixed-based articulations kinematic to allow root position updates#1104
eric-heiden merged 10 commits into
newton-physics:mainfrom
adenzler-nvidia:dev/adenzler/mocap-bodies

Conversation

@adenzler-nvidia

@adenzler-nvidia adenzler-nvidia commented Nov 20, 2025

Copy link
Copy Markdown
Member

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"

  • Necessary tests have been added and new examples are tested (see newton/tests/test_examples.py)
  • Documentation is up-to-date
  • Code passes formatting and linting checks with pre-commit run -a

Summary by CodeRabbit

  • New Features

    • Added mocap (fixed-base) body support with per-body mocap position/orientation outputs and a mapping exposed on the solver; simulation updates now propagate mocap transforms.
  • Bug Fixes

    • Collision geometry and transform updates correctly reflect moved fixed/mocap bodies; avoids duplicate writes and ensures correct per-body storage.
  • Tests

    • New integration tests validating mocap transform updates, collision responses, and dynamics (includes a duplicated test block).
  • Examples

    • Example robot validation split to separately test hand and cube body bounds.

✏️ Tip: You can customize this high-level summary in your review settings.

Signed-off-by: Alain Denzler <adenzler@nvidia.com>
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
@coderabbitai

coderabbitai Bot commented Nov 20, 2025

Copy link
Copy Markdown
Contributor
📝 Walkthrough

Walkthrough

Adds 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

Cohort / File(s) Summary
Kernel interface & logic
newton/_src/solvers/mujoco/kernels.py
update_joint_transforms_kernel signature expanded with input newton_body_to_mocap_index and outputs mocap_pos, mocap_quat. Transform computation now conditionally writes to mocap outputs when mocap_index >= 0, otherwise to body outputs; removed unconditional body writes.
Solver wiring & model mapping
newton/_src/solvers/mujoco/solver_mujoco.py
Added public attribute newton_body_to_mocap_index and logic to detect fixed-base (parent == -1 & FIXED) bodies as mocap, assign sequential mocap indices, build WP array mapping, pass it into kernel calls, and reorder model expansion/notify to populate mocap mapping before finalizing model fields.
Integration tests
newton/tests/test_mujoco_solver.py
Added TestMuJoCoMocapBodies.test_mocap_body_transform_updates_collision_geoms to validate mocap mapping, mocap pose outputs, collision-geometry updates, and dynamic response after applying mocap transforms (test appears duplicated).
Example adjustment
newton/examples/robot/example_robot_allegro_hand.py
Added cube_body_offset per world and split bounding-box assertions into separate checks for hand bodies (tight bounds) and the cube body (looser XY bounds and Z > 0).
Benchmark config
asv.conf.json
Updated mujoco_warp Git commit hash used in installation line.

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
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

  • Areas needing extra attention:
    • Kernel signature change: verify all call sites and argument ordering/types.
    • Device placement and lifetime of newton_body_to_mocap_index WP array and mocap outputs.
    • Reordered model expansion/notify_model_changed—ensure no race/initialization regressions.
    • New tests: remove duplicate test block and confirm collision expectations.
    • Example index arithmetic around per-world offsets.

Possibly related PRs

Suggested reviewers

  • eric-heiden
  • vreutskyy
  • Kenny-Vilella

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)
Check name Status Explanation Resolution
Out of Scope Changes check ⚠️ Warning One change is out of scope: the asv.conf.json commit hash update is unrelated to fixing the mocap body handling or the Allegro hand test flakiness. Remove or justify the asv.conf.json commit hash change, as it appears to be a separate dependency update unrelated to the mocap body fix.
Docstring Coverage ⚠️ Warning Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. You can run @coderabbitai generate docstrings to improve docstring coverage.
✅ Passed checks (3 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately describes the main change: making the root of fixed-base articulations kinematic (as mocap bodies) to enable root position updates, which directly addresses the core problem.
Linked Issues check ✅ Passed The PR implements the fix for issue #909 by making fixed-base articulation roots mocap bodies to enable position updates and prevent collision geometry issues.
✨ Finishing touches
  • 📝 Generate docstrings
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

📜 Recent review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between ce6b1c2 and 6768d13.

📒 Files selected for processing (3)
  • newton/_src/solvers/mujoco/kernels.py (2 hunks)
  • newton/_src/solvers/mujoco/solver_mujoco.py (7 hunks)
  • newton/tests/test_mujoco_solver.py (1 hunks)
🧰 Additional context used
🧠 Learnings (15)
📓 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: adenzler-nvidia
Repo: newton-physics/newton PR: 1107
File: newton/_src/solvers/mujoco/kernels.py:973-974
Timestamp: 2025-11-24T08:05:21.380Z
Learning: In Newton's MuJoCo solver integration (newton/_src/solvers/mujoco/), the mapping between Newton joints and MuJoCo joints is not 1-to-1. Instead, each Newton DOF maps to a distinct MuJoCo joint. This means that for multi-DOF joints (like D6 with 6 DOFs), there will be 6 corresponding MuJoCo joints, each with its own properties (margin, solimp, solref, etc.). The mapping is done via dof_to_mjc_joint array, ensuring each DOF's properties are written to its own MuJoCo joint without overwriting.
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-11-24T08:05:21.380Z
Learnt from: adenzler-nvidia
Repo: newton-physics/newton PR: 1107
File: newton/_src/solvers/mujoco/kernels.py:973-974
Timestamp: 2025-11-24T08:05:21.380Z
Learning: In Newton's MuJoCo solver integration (newton/_src/solvers/mujoco/), the mapping between Newton joints and MuJoCo joints is not 1-to-1. Instead, each Newton DOF maps to a distinct MuJoCo joint. This means that for multi-DOF joints (like D6 with 6 DOFs), there will be 6 corresponding MuJoCo joints, each with its own properties (margin, solimp, solref, etc.). The mapping is done via dof_to_mjc_joint array, ensuring each DOF's properties are written to its own MuJoCo joint without overwriting.

Applied to files:

  • newton/_src/solvers/mujoco/solver_mujoco.py
  • 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/solver_mujoco.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/_src/solvers/mujoco/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
  • 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, 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-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/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
📚 Learning: 2025-08-25T11:07:47.818Z
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.

Applied to files:

  • newton/_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/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.py
  • newton/tests/test_mujoco_solver.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/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.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/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.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/solver_mujoco.py
  • newton/_src/solvers/mujoco/kernels.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-12T18:04:06.577Z
Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/_src/solvers/featherstone/kernels.py:75-75
Timestamp: 2025-08-12T18:04:06.577Z
Learning: The Newton physics framework requires nightly Warp builds, which means compatibility concerns with older stable Warp versions (like missing functions such as wp.spatial_adjoint) are not relevant for this project.

Applied to files:

  • newton/tests/test_mujoco_solver.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: When a user reports that a suggested code change breaks their simulation, acknowledge their empirical evidence and avoid being insistent about theoretical analysis without practical verification. Ask for clarification about the intended behavior rather than doubling down on the original suggestion.

Applied to files:

  • newton/tests/test_mujoco_solver.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 (1)
newton/tests/test_mujoco_solver.py (1)
newton/_src/solvers/mujoco/solver_mujoco.py (3)
  • SolverMuJoCo (77-1953)
  • step (311-334)
  • notify_model_changed (383-393)
⏰ 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 Tests / Run GPU Unit Tests on AWS EC2
  • GitHub Check: Run GPU Benchmarks / Run GPU Benchmarks on AWS EC2
🔇 Additional comments (2)
newton/_src/solvers/mujoco/kernels.py (1)

1004-1052: Mocap branch in joint transform kernel looks consistent with existing mapping conventions

Routing transforms through newton_body_to_mocap_index[child] and selectively writing either mocap_pos/mocap_quat (for fixed-base mocap bodies) or body_pos/body_quat (for regular bodies) fits cleanly with the existing dof_to_mjc_joint / to_mjc_body_index pattern. The per-world indexing via worldid and joint_in_world stays intact, and non-mocap behavior is unchanged apart from the small refactor.

To be safe for multi-world setups, please verify that a model with multiple fixed-base articulations across worlds updates each world’s mocap_pos/mocap_quat correctly (e.g., by extending the new test to num_worlds > 1 or adding a small dedicated test).

newton/_src/solvers/mujoco/solver_mujoco.py (1)

249-250: Mocap body plumbing in solver is coherent and matches kernel expectations

The end-to-end wiring for fixed-base mocap bodies looks solid:

  • Detecting mocap candidates as parent == -1 and joint_type[j] == JointType.FIXED and passing mocap=fixed_base into add_body matches the intended “fixed-base articulation root” semantics without touching other fixed joints.
  • Tracking mocap indices in newton_body_to_mocap_index during body creation and then densifying to a [bodies_per_world] Warp array aligns with the kernel’s template-world indexing (joint_in_world / child) and keeps runtime lookups cheap.
  • Feeding self.newton_body_to_mocap_index plus self.mjw_data.mocap_pos/mocap_quat into update_joint_transforms_kernel gives a clean separation: mocap bodies go through mocap arrays, while non-mocap bodies continue to use body_pos/body_quat.
  • The new expand_model_fields(..., nworld) followed by notify_model_changed(SolverNotifyFlags.ALL) ensures multi-world mjw_model fields are properly tiled before any update kernels run, which should prevent subtle shape/stride issues when scaling to multiple worlds.

Overall, this design cleanly introduces mocap support without perturbing existing non-mocap flows.

As a follow-up sanity check, I’d recommend:

  • Verifying that a model with both floating-base and fixed-base articulations still behaves as before for the floating bases.
  • Running the existing mass/geom/joint property tests together with the new mocap test on a multi-world model to confirm there are no regressions from the added expand_model_fields + notify_model_changed(SolverNotifyFlags.ALL) calls.

Also applies to: 999-1003, 1208-1245, 1543-1549, 1593-1599, 1796-1827


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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@adenzler-nvidia

Copy link
Copy Markdown
Member Author

adding an unresolved comment to remind me that this needs a dependency update as well. MjWarp PR is google-deepmind/mujoco_warp#858

@adenzler-nvidia adenzler-nvidia self-assigned this Nov 20, 2025
@adenzler-nvidia adenzler-nvidia linked an issue Nov 20, 2025 that may be closed by this pull request

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 brittle

Using allegro_hand.body_count - 1 assumes the cube remains the last body in the USD template. If the asset changes (extra bodies added, reordering), cube_body_offset and hand_body_indices may 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: Clarify newton_body_to_mocap_index mapping semantics in the docstring

The 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_index array construction relies on insertion order; consider making the indexing more explicit

You 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_kernel

This 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_mapping is filled once per child), but it’s a fairly implicit coupling between:

  • how body_mapping is populated, and
  • how newton_body_to_mocap_index is indexed.

To make this more robust and self‑documenting, consider one of:

  • Building body_to_mocap explicitly 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_index to be indexed in the same “local body index” space as to_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_mapping implementation ever changes, and makes the expected indexing scheme much clearer.


1745-1770: Propagating mocap mapping into update_joint_properties is correct; just ensure it’s always initialized

The updated update_joint_properties launch:

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_index is only initialized in _convert_to_mjc (the mujoco_warp path). If someone ever constructs a solver with an external mjw_model/mjw_data pair or in a CPU‑only configuration and still calls notify_model_changed(SolverNotifyFlags.JOINT_PROPERTIES), this field will be None and wp.launch will fail.

Given that the kernel itself is specific to the mujoco_warp backend, it would be safer to either:

  • Guard update_joint_properties with if self.use_mujoco_cpu: return, or
  • Initialize self.newton_body_to_mocap_index to 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

📥 Commits

Reviewing files that changed from the base of the PR and between a546871 and 1583392.

📒 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.py
  • newton/examples/robot/example_robot_allegro_hand.py
  • newton/_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.py
  • newton/examples/robot/example_robot_allegro_hand.py
  • newton/_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.py
  • newton/_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.py
  • newton/examples/robot/example_robot_allegro_hand.py
  • newton/_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.py
  • newton/_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.py
  • newton/_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 in update_joint_transforms_kernel looks correct but depends on existing body-index conventions

The new signature and routing:

  • Adds newton_body_to_mocap_index: wp.array(dtype=wp.int32).
  • Adds mocap_pos / mocap_quat outputs.
  • 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] = rotation

This 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), and
  • newton_body_to_mocap_index is indexed with the same local body index convention as body_mapping / to_mjc_body_index.

Given how to_mjc_body_index is constructed in solver_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 articulations

The 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 += 1

This 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_index and 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: Delayed expand_model_fields and notify_model_changed sequencing makes sense

Moving the expansion and initial notify_model_changed to 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.ALL call 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.

Comment thread newton/examples/robot/example_robot_allegro_hand.py
Comment thread newton/tests/test_mujoco_solver.py
Signed-off-by: Alain Denzler <adenzler@nvidia.com>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 path

When mjw_model and mjw_data are provided externally, self.newton_body_to_mocap_index remains None. However, update_joint_properties() unconditionally passes it to the kernel at line 1759, which expects a wp.array(dtype=wp.int32), not None. This will cause a runtime crash if any code path triggers the JOINT_PROPERTIES notify flag after initialization.

Initialize newton_body_to_mocap_index to 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 changes

Using self.cube_body_offset = allegro_hand.body_count - 1 assumes 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., matching body_key against a DexCube pattern) and storing its index, then using that for per-world offsets.

📜 Review details

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a546871 and 1583392.

📒 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.py
  • newton/examples/robot/example_robot_allegro_hand.py
  • newton/_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.py
  • newton/examples/robot/example_robot_allegro_hand.py
  • newton/_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.py
  • newton/examples/robot/example_robot_allegro_hand.py
  • newton/_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.py
  • newton/_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.py
  • newton/_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 in update_joint_transforms_kernel looks consistent with existing conventions

The new newton_body_to_mocap_index parameter and mocap branch correctly:

  • reuse the existing child = joint_child[joint_in_world] + body_mapping[child] pattern so indices stay per‑world,
  • compute tf once 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 to body_quat / mocap_quat (aligns with test_joint_frame_update expectations).

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 scoped

The new self.newton_body_to_mocap_index attribute and its docstring make the mocap mapping explicit and keep it parallel to self.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_mjc matches fixed-base semantics

The newton_body_to_mocap_index dict is only filled when parent == -1 and joint_type[j] == JointType.FIXED, which is exactly the fixed‑base articulation root you want to treat as a mocap body. Coupled with mocap=fixed_base in the add_body call, this correctly identifies and indexes mocap bodies at conversion time.


1502-1508: Mocap mapping array construction matches body index conventions

This 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_world array aligned with to_mjc_body_index (same body_mapping.keys()[1:] source). Using .get(idx, -1) correctly marks non‑mocap bodies as -1. This layout matches how update_joint_transforms_kernel expects to index newton_body_to_mocap_index[child].


1736-1770: Mocap data wiring into update_joint_transforms_kernel is consistent

Passing self.newton_body_to_mocap_index and wiring self.mjw_data.mocap_pos / mocap_quat as 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_quat untouched 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.

Comment thread newton/examples/robot/example_robot_allegro_hand.py
Comment thread newton/tests/test_mujoco_solver.py
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
@codecov

codecov Bot commented Nov 21, 2025

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ All tests successful. No failed tests found.

📢 Thoughts on this report? Let us know!

@adenzler-nvidia adenzler-nvidia changed the title Draft: SolverMuJoCo: make root of fixed-based articulations kinematic to allow root position updates SolverMuJoCo: make root of fixed-based articulations kinematic to allow root position updates Nov 21, 2025
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
@adenzler-nvidia

Copy link
Copy Markdown
Member Author

Mjwarp Update is in now, should be good to go.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. 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), but test_body_state expects a 2-argument callable. Default arguments may also confuse Warp's create_warp_function.

  2. Multi-line expression: Splitting the lambda body across lines 233-234 can cause IndentationError during Warp kernel codegen.

For consistency with the hand test lambda (line 220), use closure to capture cube_lower and cube_upper, collapse to a single line, and add the # noqa: B023 comment:

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 parameter qd to _qd for consistency.

The lambda parameter qd is 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

📥 Commits

Reviewing files that changed from the base of the PR and between e717dca and ce6b1c2.

📒 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_offset attribute 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.

Comment thread newton/_src/solvers/mujoco/solver_mujoco.py

@eric-heiden eric-heiden left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@eric-heiden eric-heiden added this pull request to the merge queue Nov 25, 2025
Merged via the queue into newton-physics:main with commit 02ad18a Nov 25, 2025
20 checks passed
eric-heiden pushed a commit to eric-heiden/newton that referenced this pull request Jan 28, 2026
…ow root position updates (newton-physics#1104)

Signed-off-by: Alain Denzler <adenzler@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Flaky test: test_robot.example_robot_allegro_hand_cuda_0

2 participants