Move all SolverMuJoCo Warp kernels to kernels.py file #1027 by vreutskyy · Pull Request #1040 · newton-physics/newton

vreutskyy · 2025-11-06T11:02:35Z

Description

This PR refactors the SolverMuJoCo implementation by extracting all Warp kernels and functions into a separate kernels.py file. This improves code organization and maintainability by separating the GPU kernel implementations from the main solver logic.

Changes:

Created newton/_src/solvers/mujoco/kernels.py containing:
- 18 Warp kernels (decorated with @wp.kernel)
- 5 Warp functions (decorated with @wp.func)
- Custom type definitions (vec5, vec10f)
- Required constant (MJ_MINVAL)
Updated solver_mujoco.py to import kernels from the new module
Removed ~1,070 lines of kernel definitions from solver_mujoco.py
Cleaned up imports to only include kernels directly used in solver_mujoco.py

This is a pure refactoring change with no functional modifications. All existing functionality remains intact.

Closes #1027

Newton Migration Guide

Please ensure the migration guide for warp.sim users is up-to-date with the changes made in this PR.

The migration guide in docs/migration.rst is up-to date
- Note: This is an internal refactoring with no API changes, so no migration guide updates are needed

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

Refactor
- Internal consolidation of MuJoCo/Warp physics kernels into a new modular kernel module (including a comprehensive suite of conversion, propagation, inertia, and geometry kernels).
- Maintains existing public interfaces and visible behavior; improvements are internal for maintainability and performance.

…1027 Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

coderabbitai · 2025-11-06T11:02:44Z

📝 Walkthrough

Walkthrough

A new kernels.py module centralizes Warp-based MuJoCo kernel implementations formerly defined inline in solver_mujoco.py. solver_mujoco.py now imports these kernels and retains the solver public interface unchanged.

Changes

Cohort / File(s)	Change Summary
New Kernels Module `newton/_src/solvers/mujoco/kernels.py`	Added a new module containing extracted Warp-based MuJoCo kernels and helpers: contact handling (write_contact, contact_params, convert_newton_contacts_to_mjwarp_kernel), coordinate conversions (convert_mj_coords_to_warp_kernel, convert_warp_coords_to_mj_kernel, convert_mjw_contact_to_warp_kernel), forward kinematics (eval_single_articulation_fk, eval_articulation_fk), body/joint transforms (convert_body_xforms_to_warp_kernel, update_joint_transforms_kernel), mass/inertia updates (update_body_mass_ipos_kernel, update_body_inertia_kernel), model/geom synchronization (update_model_properties_kernel, update_geom_properties_kernel, update_incoming_shape_xform_kernel, update_shape_mappings_kernel), utilities (repeat_array_kernel) and many supporting kernels.
Solver Refactor / Imports `newton/_src/solvers/mujoco/solver_mujoco.py`	Removed numerous inline kernel and helper definitions (vec helpers, orthogonals/make_frame, write_contact, contact_params, convert_/apply_/eval_* kernels, repeat_array_kernel, update_* kernels). Replaced them with imports from the new `kernels.py` (e.g., convert_newton_contacts_to_mjwarp_kernel, convert_mj_coords_to_warp_kernel, convert_warp_coords_to_mj_kernel, convert_mjw_contact_to_warp_kernel, eval_articulation_fk, convert_body_xforms_to_warp_kernel, update_body_mass_ipos_kernel, update_body_inertia_kernel, repeat_array_kernel, update_axis_properties_kernel, update_dof_properties_kernel, update_joint_transforms_kernel, update_incoming_shape_xform_kernel, update_model_properties_kernel, update_geom_properties_kernel, update_shape_mappings_kernel). Public SolverMuJoCo interface unchanged.

Sequence Diagram(s)

sequenceDiagram
    autonumber
    participant Solver as SolverMuJoCo
    participant Kernels as kernels.py
    participant Warp as Warp Arrays

    Note over Solver,Kernels: Kernel implementations moved from solver to kernels.py
    Solver->>Kernels: convert_body_xforms_to_warp_kernel(inputs)
    Kernels-->>Warp: write body transforms, quaternions, DOF indices
    Solver->>Kernels: eval_articulation_fk(articulation_state)
    Kernels-->>Warp: propagate joint transforms & velocities
    Solver->>Kernels: convert_newton_contacts_to_mjwarp_kernel(contacts)
    Kernels-->>Warp: write contact entries (distance, normals, frames, params)
    Note over Warp: Warp arrays consumed by subsequent physics kernels

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Focus review areas:
- Verify signatures and call sites for all moved kernels match prior inline definitions.
- Confirm quaternion ordering, inertia eigen-decomposition, and contact param derivation logic preserved.
- Ensure imports in solver_mujoco.py cover every previously-inline kernel and that tests/compilation pass.

Possibly related PRs

Add possibility to forward newton generated contacts into mujoco warp #370 — Directly related refactor that extracts MuJoCo/Warp kernels (convert_newton_contacts_to_mjwarp_kernel, write_contact, contact_params, orthogonals/make_frame) into a separate module.

Suggested reviewers

adenzler-nvidia
eric-heiden
mmacklin

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 13.04% which is insufficient. The required threshold is 80.00%.	You can run `@coderabbitai generate docstrings` to improve docstring coverage.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately and specifically describes the main change: moving Warp kernels from solver_mujoco.py to a new kernels.py file, with reference to issue #1027.
Linked Issues check	✅ Passed	The PR fully satisfies issue #1027 by creating kernels.py with all Warp kernels, updating solver_mujoco.py imports, and preserving existing functionality as required.
Out of Scope Changes check	✅ Passed	All changes are directly related to the stated objective of refactoring kernels into a separate module; no out-of-scope modifications detected.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

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.}

codecov · 2025-11-06T11:26:10Z

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!

Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

…ernelspy-file

Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

coderabbitai

Actionable comments posted: 0

🧹 Nitpick comments (1)

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

196-586: Consider removing unused parameters in a future cleanup.

Several kernel functions have unused parameters that were carried over from the original implementation:

bodies_per_world (line 196) in convert_newton_contacts_to_mjwarp_kernel

nworld_in (line 212) in convert_newton_contacts_to_mjwarp_kernel

up_axis in multiple kernels (lines 317, 393, 524)

body_q (line 525) in apply_mjc_body_f_kernel

body_com (line 546) and joint_q_start (line 548) in apply_mjc_qfrc_kernel

These parameters can be removed to simplify the function signatures, though this would require updating all call sites.

📜 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 2f7222b and ce58f85.

📒 Files selected for processing (2)

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

🧰 Additional context used

🧠 Learnings (12)

📓 Common learnings

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.

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-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/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, despite the general warp documentation suggesting copies are made by default.

Applied to files:

newton/_src/solvers/mujoco/kernels.py
newton/_src/solvers/mujoco/solver_mujoco.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/_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-08-12T17:51:37.474Z

Learnt from: nvlukasz
Repo: newton-physics/newton PR: 519
File: newton/geometry.py:16-22
Timestamp: 2025-08-12T17:51:37.474Z
Learning: In the Newton public API refactor, common geometry symbols like ParticleFlags and ShapeFlags are now exposed at the top level in newton/__init__.py rather than through newton.geometry. The newton/geometry.py module is intentionally focused only on broad-phase collision detection classes (BroadPhaseAllPairs, BroadPhaseExplicit, BroadPhaseSAP).

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/_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/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: 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

📚 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-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-07-14T03:57:29.670Z

Learnt from: Kenny-Vilella
Repo: newton-physics/newton PR: 398
File: newton/examples/example_mujoco.py:352-352
Timestamp: 2025-07-14T03:57:29.670Z
Learning: The use_mujoco option in newton/examples/example_mujoco.py is currently unsupported and causes crashes. The code automatically disables this option with a warning message when users attempt to enable it. This is intentionally kept as a placeholder for future implementation.

Applied to files:

newton/_src/solvers/mujoco/solver_mujoco.py

🧬 Code graph analysis (2)

newton/_src/solvers/mujoco/kernels.py (2)

newton/_src/sim/joints.py (2)

JointMode (81-95)

JointType (20-44)

newton/_src/sim/collide_unified.py (1)

write_contact (53-149)

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

newton/_src/solvers/mujoco/kernels.py (18)

apply_mjc_body_f_kernel (523-538)

apply_mjc_control_kernel (505-519)

apply_mjc_qfrc_kernel (542-585)

convert_body_xforms_to_warp_kernel (762-780)

convert_mj_coords_to_warp_kernel (313-385)

convert_mjw_contact_to_warp_kernel (460-501)

convert_newton_contacts_to_mjwarp_kernel (175-309)

convert_warp_coords_to_mj_kernel (389-456)

eval_articulation_fk (717-758)

repeat_array_kernel (862-869)

update_axis_properties_kernel (873-916)

update_body_inertia_kernel (812-858)

update_body_mass_ipos_kernel (784-808)

update_dof_properties_kernel (920-937)

update_geom_properties_kernel (1060-1133)

update_joint_transforms_kernel (941-1008)

update_model_properties_kernel (1049-1056)

update_shape_mappings_kernel (1012-1045)

🪛 Ruff (0.14.3)

newton/_src/solvers/mujoco/kernels.py

196-196: Unused function argument: bodies_per_world

(ARG001)

212-212: Unused function argument: nworld_in

(ARG001)

317-317: Unused function argument: up_axis

(ARG001)

393-393: Unused function argument: up_axis

(ARG001)

524-524: Unused function argument: up_axis

(ARG001)

525-525: Unused function argument: body_q

(ARG001)

546-546: Unused function argument: body_com

(ARG001)

548-548: Unused function argument: joint_q_start

(ARG001)

828-828: Ambiguous variable name: I

(E741)

⏰ 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 (Pull Request)
GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)

🔇 Additional comments (2)

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

1-38: LGTM!

The module header, imports, and type definitions are well-structured. The custom vector types (vec5, vec10f) and constant (MJ_MINVAL) are appropriate for MuJoCo integration.

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

46-65: LGTM! Refactoring successfully completed.

The kernel imports from .kernels are correct and complete. All 18 kernels that were moved to kernels.py are now properly imported and used throughout the solver implementation. The refactoring achieves its goal of improving code organization while preserving all functionality.

adenzler-nvidia

Thank you, looks straightforward!

Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

adenzler-nvidia

approving again, thanks!

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (3)

newton/_src/solvers/mujoco/kernels.py (3)
770-775: Remove commented debug code.

Lines 772-774 contain commented-out quaternion conversion attempts that should be removed to keep the codebase clean.

Apply this diff:
     # convert from wxyz to xyzw
     quat = wp.quat(quat[1], quat[2], quat[3], quat[0])
-    # quat = wp.quat(quat[3], quat[0], quat[1], quat[2])
-    # quat = wp.quat_identity()
-    # quat = wp.quat_inverse(quat)
     body_q[wbi] = wp.transform(pos, quat)
823-823: Consider using a more descriptive variable name.

While I is standard physics notation for inertia tensor, Python style guides discourage single-letter uppercase variable names to avoid confusion with built-ins. Consider using inertia_tensor or inertia for better clarity.

Apply this diff:
-    I = body_inertia[tid]
+    inertia_tensor = body_inertia[tid]
 
     # Calculate eigenvalues and eigenvectors
-    eigenvectors, eigenvalues = wp.eig3(I)
+    eigenvectors, eigenvalues = wp.eig3(inertia_tensor)
312-312: Several kernels contain unused parameters.

Multiple kernels have unused parameters flagged by static analysis:

up_axis in convert_mj_coords_to_warp_kernel (line 312)

up_axis in convert_warp_coords_to_mj_kernel (line 388)

up_axis and body_q in apply_mjc_body_f_kernel (lines 519-520)

body_com and joint_q_start in apply_mjc_qfrc_kernel (lines 541, 543)

These parameters should either be used in the implementation, removed, or prefixed with underscore if kept for API consistency with other kernels.

Also applies to: 388-388, 519-520, 541-543

📜 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 ce58f85 and 1c40493.

📒 Files selected for processing (1)

newton/_src/solvers/mujoco/kernels.py (1 hunks)

🧰 Additional context used

🧠 Learnings (7)

📓 Common learnings

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.

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-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/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, despite the general warp documentation suggesting copies are made by default.

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/_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

📚 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

📚 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

🧬 Code graph analysis (1)

newton/_src/solvers/mujoco/kernels.py (2)

newton/_src/sim/joints.py (2)

JointMode (81-95)

JointType (20-44)

newton/_src/sim/collide_unified.py (1)

write_contact (53-149)

🪛 Ruff (0.14.3)

newton/_src/solvers/mujoco/kernels.py

191-191: Unused function argument: bodies_per_world

(ARG001)

207-207: Unused function argument: nworld_in

(ARG001)

312-312: Unused function argument: up_axis

(ARG001)

388-388: Unused function argument: up_axis

(ARG001)

519-519: Unused function argument: up_axis

(ARG001)

520-520: Unused function argument: body_q

(ARG001)

541-541: Unused function argument: body_com

(ARG001)

543-543: Unused function argument: joint_q_start

(ARG001)

823-823: Ambiguous variable name: I

(E741)

⏰ 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 (Pull Request)
GitHub Check: Run GPU Unit Tests on AWS EC2 (Pull Request)

🔇 Additional comments (7)

newton/_src/solvers/mujoco/kernels.py (7)

27-28: Custom vector type definitions look good.

Both vec5 and vec10 use consistent definition patterns with wp.float32. While a past review mentioned alignment, the current definitions are already consistent in structure and dtype.

335-344: Quaternion conversions are correctly implemented.

The coordinate conversion kernels properly handle quaternion format differences between MuJoCo (wxyz) and Warp (xyzw). The conversions are consistent across FREE and BALL joint types in both directions.

Also applies to: 359-369, 410-420, 436-440

353-357: Angular velocity transformations are correct.

For FREE joints, the angular velocities are properly transformed between coordinate frames using wp.quat_rotate (MuJoCo→Warp) and wp.quat_rotate_inv (Warp→MuJoCo), maintaining consistency with the quaternion rotations.

Also applies to: 429-433

583-708: Forward kinematics implementation is comprehensive and correct.

The eval_single_articulation_fk function properly handles all joint types (PRISMATIC, REVOLUTE, BALL, FREE, DISTANCE, D6) with correct transform composition and velocity propagation through the kinematic tree. The D6 joint implementation correctly separates linear and angular DOF processing.

169-304: Contact conversion kernel correctly transforms Newton contacts to MuJoCo format.

The kernel properly handles coordinate transformations, contact distance computation with thickness, contact frame construction, and parameter derivation. The use of the midpoint for contact position (line 249) is consistent with XPBD solver conventions.

831-842: Eigenvalue sorting logic is correct.

The bubble sort implementation correctly sorts 3 eigenvalues and their corresponding eigenvectors in descending order. For a fixed size of 3 elements, this approach is simple and efficient.

1-24: File structure and organization are well-designed.

The file has a clean layout with proper license header, logical grouping of custom types, constants, utility functions, and kernels. The separation of @wp.func helpers from @wp.kernel implementations is clear and maintainable.

…1027 (newton-physics#1040) Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

Move all SolverMuJoCo Warp kernels to kernels.py file newton-physics#…

2f7222b

…1027 Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

vreutskyy added 3 commits November 6, 2025 12:27

Update branch

7104dd9

Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

Merge branch 'main' into 1027-move-all-solvermujoco-warp-kernels-to-k…

ebe08b2

…ernelspy-file

Update impelmentation

ce58f85

Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

coderabbitai Bot reviewed Nov 6, 2025

View reviewed changes

adenzler-nvidia previously approved these changes Nov 6, 2025

View reviewed changes

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

Small clean up

1c40493

Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

vreutskyy dismissed adenzler-nvidia’s stale review via 1c40493 November 7, 2025 09:07

adenzler-nvidia approved these changes Nov 7, 2025

View reviewed changes

adenzler-nvidia enabled auto-merge (squash) November 7, 2025 09:25

coderabbitai Bot reviewed Nov 7, 2025

View reviewed changes

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

adenzler-nvidia merged commit a39fc71 into newton-physics:main Nov 7, 2025
18 checks passed

coderabbitai Bot mentioned this pull request Nov 11, 2025

Solver data fields #1022

Closed

3 tasks

This was referenced Nov 20, 2025

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

Merged

SolverMuJoCo: allow runtime updates for joint limits #1110

Merged

coderabbitai Bot mentioned this pull request Dec 14, 2025

SDF Hydroelastic Collisions #1248

Merged

4 tasks

coderabbitai Bot mentioned this pull request Jan 14, 2026

Fix SolverMuJoCo to follow Newton convention around spatial velocities #1350

Merged

coderabbitai Bot mentioned this pull request Jan 21, 2026

Implement randomization and runtime update of equality constraint data + active #1192 #1430

Merged

4 tasks

eric-heiden pushed a commit to eric-heiden/newton that referenced this pull request Jan 28, 2026

Move all SolverMuJoCo Warp kernels to kernels.py file newton-physics#…

cc52c7e

…1027 (newton-physics#1040) Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

This was referenced Jan 30, 2026

Add runtime updates for MuJoCo contact pair custom attributes #1501

Merged

Introduce Extended Contact Attributes & Refactor SensorContact #1512

Merged

coderabbitai Bot mentioned this pull request Feb 16, 2026

Fix free joint body_pos and add ref/qpos0 support for MuJoCo solver #1645

Merged

3 tasks

This was referenced Feb 17, 2026

SolverMuJoCo: Support kinematic joints via high-armature DOFs #1667

Closed

Expose qfrc_actuator from mujoco #1698

Merged

Fix xyzw-to-wxyz quaternion conversion in body inertia kernel #1694

Merged

coderabbitai Bot mentioned this pull request Mar 13, 2026

Fix loop joint coordinate mapping and constraint types #2075

Merged

3 tasks

This was referenced Mar 24, 2026

Fix MuJoCo-to-Newton contact force conversion #2208

Merged

Defer import of mujoco_warp for use in kernels #2263

Merged

Attempt to fix issue 2071 #2073

Merged

This was referenced Apr 1, 2026

Fix and test for connect constraints with non-zero joint reference positions #2288

Merged

Migrate wp.array(dtype=X) annotations to bracket syntax #2282

Merged

mmacklin pushed a commit to mmacklin/newton that referenced this pull request Apr 7, 2026

Move all SolverMuJoCo Warp kernels to kernels.py file newton-physics#…

cd1a6c6

…1027 (newton-physics#1040) Signed-off-by: Viktor Reutskyy <vreutskyy@nvidia.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Move all SolverMuJoCo Warp kernels to kernels.py file #1027#1040

Move all SolverMuJoCo Warp kernels to kernels.py file #1027#1040
adenzler-nvidia merged 5 commits into
newton-physics:mainfrom
vreutskyy:1027-move-all-solvermujoco-warp-kernels-to-kernelspy-file

vreutskyy commented Nov 6, 2025 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Nov 6, 2025 •

edited

Loading

Walkthrough

Changes

Sequence Diagram(s)

Estimated code review effort

Possibly related PRs

Suggested reviewers

Uh oh!

codecov Bot commented Nov 6, 2025 •

edited

Loading

Uh oh!

coderabbitai Bot left a comment

Uh oh!

adenzler-nvidia left a comment

Uh oh!

Uh oh!

adenzler-nvidia left a comment

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

vreutskyy commented Nov 6, 2025 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Changes:

Newton Migration Guide

Before your PR is "Ready for review"

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Nov 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Estimated code review effort

Possibly related PRs

Suggested reviewers

Pre-merge checks and finishing touches

Uh oh!

codecov Bot commented Nov 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

adenzler-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

adenzler-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

vreutskyy commented Nov 6, 2025 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Nov 6, 2025 •

edited

Loading

codecov Bot commented Nov 6, 2025 •

edited

Loading