Skip to content

SolverMuJoCo: custom attribute for eq_solref#1245

Merged
eric-heiden merged 2 commits into
newton-physics:mainfrom
adenzler-nvidia:dev/adenzler/equality-solref
Dec 12, 2025
Merged

SolverMuJoCo: custom attribute for eq_solref#1245
eric-heiden merged 2 commits into
newton-physics:mainfrom
adenzler-nvidia:dev/adenzler/equality-solref

Conversation

@adenzler-nvidia

@adenzler-nvidia adenzler-nvidia commented Dec 11, 2025

Copy link
Copy Markdown
Member

This one does not include USD parsing because we don't know about equality constraints in USD yet. We can always add later if 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

Release Notes

  • New Features

    • Added custom attribute support for equality constraints (connect, joint, and weld variants), enabling per-constraint customization.
    • Introduced dynamic updates for MuJoCo equality constraint solver parameters at runtime.
    • Enhanced MJCF parser to properly handle equality constraint attributes.
  • Tests

    • Added tests for equality constraint attribute parsing and runtime solver updates.

✏️ 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>
@coderabbitai

coderabbitai Bot commented Dec 11, 2025

Copy link
Copy Markdown
Contributor
📝 Walkthrough

Walkthrough

This PR introduces comprehensive support for custom per-entity attributes on equality constraints, specifically the eq_solref solver reference parameter. Changes span the builder, model, solver, kernels, and import systems to enable per-constraint customization in MuJoCo through frequency-based attribute tracking and runtime property synchronization.

Changes

Cohort / File(s) Summary
Equality constraint infrastructure
newton/_src/sim/builder.py
Adds custom_attributes parameter to add_equality_constraint, add_equality_constraint_connect, add_equality_constraint_joint, and add_equality_constraint_weld methods. Implements tracking and remapping of equality constraints during builder replication, with proper offset handling and world scoping. Extends finalization logic for custom attribute processing.
Model attribute frequency
newton/_src/sim/model.py
Introduces EQUALITY_CONSTRAINT = 7 enum member to ModelAttributeFrequency to track attribute counts based on equality constraint cardinality.
Solver notification flags
newton/_src/solvers/flags.py
Adds EQUALITY_CONSTRAINT_PROPERTIES = 1 << 6 flag to SolverNotifyFlags enum to enable property update notifications for equality constraints.
MuJoCo kernel
newton/_src/solvers/mujoco/kernels.py
Implements update_eq_properties_kernel Warp kernel to map Newton equality constraints to MuJoCo constraints and propagate solver reference (solref) properties.
MuJoCo solver integration
newton/_src/solvers/mujoco/solver_mujoco.py
Registers eq_solref custom attribute in MuJoCo namespace. Extends model synchronization to handle equality constraint properties, adds mjc_eq_to_newton_eq mapping, implements update_eq_properties method, and expands model fields for multi-world support.
MJCF import
newton/_src/utils/import_mjcf.py
Adds equality constraint custom attribute parsing via frequency-based lookup. Propagates parsed attributes to add_equality_constraint_connect, add_equality_constraint_weld, and add_equality_constraint_joint calls.
Tests
newton/tests/test_import_mjcf.py
Adds test_eq_solref_parsing to verify custom attribute parsing and value assignment for equality constraints during MJCF import.
MuJoCo solver tests
newton/tests/test_mujoco_solver.py
Introduces TestMuJoCoSolverEqualityConstraintProperties test class with test_eq_solref_conversion_and_update to validate mapping, per-world propagation, and runtime updates of eq_solref across multi-world models. (Note: test class definition appears duplicated in diff)

Sequence Diagram(s)

sequenceDiagram
    participant MJCF as MJCF Parser
    participant Builder as ModelBuilder
    participant Model as Model
    participant Solver as SolverMuJoCo
    participant Kernel as update_eq_properties_kernel
    
    MJCF->>MJCF: Parse equality constraints<br/>with custom attrs
    MJCF->>Builder: add_equality_constraint_*<br/>(custom_attributes=eq_solref)
    Builder->>Builder: Validate & process<br/>custom attributes
    Builder->>Model: Store attributes<br/>by EQUALITY_CONSTRAINT freq
    
    Note over Builder,Model: Model construction & finalization
    
    Solver->>Model: Load model with<br/>eq_solref attributes
    Solver->>Solver: Create mjc_eq_to_newton_eq<br/>index mapping
    Solver->>Solver: Set initial MuJoCo<br/>eq_solref values
    
    Note over Solver: Runtime updates
    
    Solver->>Solver: Notify:<br/>EQUALITY_CONSTRAINT_PROPERTIES
    Solver->>Solver: Call update_eq_properties()
    Solver->>Kernel: update_eq_properties_kernel<br/>(eq_solref, mapping)
    Kernel->>Kernel: Map Newton idx →<br/>MuJoCo idx
    Kernel->>Kernel: Copy solref to<br/>MuJoCo state
    Kernel-->>Solver: Updated solver state
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

  • Builder changes (builder.py): Custom attribute parameter threading and tracking/remapping logic for equality constraints during replication requires careful verification of offset handling and world scoping consistency
  • MuJoCo solver integration (solver_mujoco.py): Index mapping (mjc_eq_to_newton_eq), property synchronization flow, and multi-world expansion logic need validation for correctness
  • Test duplication (test_mujoco_solver.py): The noted duplicate test class insertion should be flagged and consolidated
  • Cross-file attribute flow: Verify custom attributes flow correctly from MJCF import → builder → model → solver → kernel execution

Possibly related PRs

Suggested reviewers

  • eric-heiden
  • vreutskyy

Pre-merge checks and finishing touches

✅ 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 summarizes the main change: adding a custom attribute (eq_solref) for equality constraints in the MuJoCo solver. It is concise, clear, and directly reflects the primary objective of the pull request.
Docstring Coverage ✅ Passed Docstring coverage is 80.77% which is sufficient. The required threshold is 80.00%.
✨ 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 56105c0 and 77d0184.

📒 Files selected for processing (8)
  • newton/_src/sim/builder.py (15 hunks)
  • newton/_src/sim/model.py (1 hunks)
  • newton/_src/solvers/flags.py (1 hunks)
  • newton/_src/solvers/mujoco/kernels.py (1 hunks)
  • newton/_src/solvers/mujoco/solver_mujoco.py (10 hunks)
  • newton/_src/utils/import_mjcf.py (5 hunks)
  • newton/tests/test_import_mjcf.py (1 hunks)
  • newton/tests/test_mujoco_solver.py (1 hunks)
🧰 Additional context used
🧠 Learnings (13)
📓 Common learnings
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: Milad-Rakhsha-NV
Repo: newton-physics/newton PR: 710
File: newton/_src/utils/schema_resolver.py:164-168
Timestamp: 2025-10-27T18:45:18.446Z
Learning: In newton/_src/utils/schema_resolver.py, when using USD's GetAuthoredPropertiesInNamespace() to collect solver-specific attributes, prop.GetName() already returns the full namespaced attribute name (e.g., "physxJoint:armature", not just "armature"), so no additional namespace prefix is needed when storing the keys.
📚 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/flags.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/_src/sim/builder.py
  • newton/_src/sim/model.py
  • newton/_src/utils/import_mjcf.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/tests/test_mujoco_solver.py
  • 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, 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/tests/test_mujoco_solver.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-11-24T08:05:21.390Z
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.390Z
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/tests/test_mujoco_solver.py
  • newton/_src/solvers/mujoco/kernels.py
  • newton/_src/utils/import_mjcf.py
  • newton/_src/solvers/mujoco/solver_mujoco.py
📚 Learning: 2025-10-10T10:47:41.082Z
Learnt from: nvtw
Repo: newton-physics/newton PR: 899
File: newton/tests/test_rigid_contact.py:214-268
Timestamp: 2025-10-10T10:47:41.082Z
Learning: In `newton/tests/test_rigid_contact.py`, using `add_shape_plane(width=0, length=0)` for an infinite plane works correctly with the collision system and does not cause degenerate AABB issues in the test scenario `test_shape_collisions_gjk_mpr_multicontact`.

Applied to files:

  • 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, 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-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-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-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-10-27T18:45:18.446Z
Learnt from: Milad-Rakhsha-NV
Repo: newton-physics/newton PR: 710
File: newton/_src/utils/schema_resolver.py:164-168
Timestamp: 2025-10-27T18:45:18.446Z
Learning: In newton/_src/utils/schema_resolver.py, when using USD's GetAuthoredPropertiesInNamespace() to collect solver-specific attributes, prop.GetName() already returns the full namespaced attribute name (e.g., "physxJoint:armature", not just "armature"), so no additional namespace prefix is needed when storing the keys.

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
🧬 Code graph analysis (3)
newton/_src/sim/builder.py (1)
newton/_src/sim/model.py (1)
  • ModelAttributeFrequency (49-72)
newton/tests/test_mujoco_solver.py (3)
newton/_src/solvers/mujoco/solver_mujoco.py (3)
  • SolverMuJoCo (81-2554)
  • register_custom_attributes (155-324)
  • notify_model_changed (544-556)
newton/_src/solvers/solver.py (2)
  • register_custom_attributes (307-314)
  • notify_model_changed (267-294)
newton/_src/solvers/flags.py (1)
  • SolverNotifyFlags (22-60)
newton/_src/solvers/mujoco/solver_mujoco.py (3)
newton/_src/solvers/mujoco/kernels.py (1)
  • update_eq_properties_kernel (1292-1309)
newton/_src/sim/builder.py (3)
  • add_custom_attribute (640-685)
  • ModelBuilder (70-5960)
  • CustomAttribute (312-402)
newton/_src/sim/model.py (2)
  • ModelAttributeFrequency (49-72)
  • ModelAttributeAssignment (31-46)
⏰ 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 (18)
newton/_src/solvers/flags.py (1)

48-59: LGTM!

The new EQUALITY_CONSTRAINT_PROPERTIES flag correctly follows the bit-shifting pattern (1 << 6) and is properly included in the ALL bitmask. The docstring accurately describes the eq_solref property.

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

1291-1309: LGTM!

The new kernel follows the established pattern used by similar kernels like update_dof_properties_kernel and update_jnt_properties_kernel. The mapping lookup, early return for unmapped constraints, and conditional null check on eq_solref are all correct.

newton/tests/test_import_mjcf.py (1)

1592-1641: LGTM!

The test is well-structured and follows the established pattern of other custom attribute parsing tests (e.g., test_geom_solimp_parsing). Good use of assertAlmostEqual for float comparisons and clear documentation of the expected constraint ordering (connect constraints parsed before weld constraints).

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

71-72: LGTM!

The new EQUALITY_CONSTRAINT = 7 enum member correctly follows the sequential numbering pattern and the docstring appropriately references the existing equality_constraint_count attribute.

newton/tests/test_mujoco_solver.py (1)

2598-2708: eq_solref multi‑world test is consistent and thorough

The new TestMuJoCoSolverEqualityConstraintProperties.test_eq_solref_conversion_and_update mirrors the existing solimp/solref tests: it builds a replicated multi‑world model, assigns per‑constraint eq_solref values via model.mujoco.eq_solref, validates them through mjc_eq_to_newton_eq, and re‑checks after notify_model_changed(SolverNotifyFlags.EQUALITY_CONSTRAINT_PROPERTIES). Indexing and mapping logic look sound, and the final sanity check that values changed is a good guard. No issues from my side.

newton/_src/utils/import_mjcf.py (1)

138-140: Equality‑constraint custom attributes are wired correctly into MJCF import

Using builder.get_custom_attributes_by_frequency([ModelAttributeFrequency.EQUALITY_CONSTRAINT]) to build builder_custom_attr_eq and then threading custom_attrs = parse_custom_attributes(..., builder_custom_attr_eq, parsing_mode="mjcf") into all add_equality_constraint_* calls cleanly extends the existing custom‑attribute machinery to equality constraints. This should allow attributes like solref to populate eq_solref without affecting non‑equality paths, and the change respects skip_equality_constraints and existing name/active handling. Looks good as‑is.

Also applies to: 967-1064

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

61-61: LGTM!

Import follows the established pattern for property update kernels.


313-324: LGTM!

The custom attribute registration follows established patterns with appropriate defaults matching MuJoCo's standard solref values.


555-556: LGTM!

The notification handler follows the established pattern for model property updates.


1156-1156: LGTM!

Custom attribute retrieval follows the established pattern.


1695-1720: LGTM!

Solref assignment for all three equality constraint types (CONNECT, JOINT, WELD) is consistent and correctly guarded by the None check.


1923-1934: LGTM!

The mapping creation follows the established pattern for other entity mappings (bodies, joints, geoms) with appropriate handling of the per-world index calculation.


2033-2033: LGTM!

Adding eq_solref to the expandable model fields enables proper multi-world replication.


2332-2359: LGTM!

The update_eq_properties method follows the established pattern for property update methods with appropriate early returns and correct kernel invocation.

newton/_src/sim/builder.py (4)

1586-1586: Equality-constraint replication and custom-attribute remapping look consistent

start_equality_constraint_idx plus the dedicated equality-constraint copy block and the new EQUALITY_CONSTRAINT branch in the custom-attribute merge correctly remap sub‑builder indices into the main builder’s index space. World indices are overridden to current_world in the same way as for bodies/joints, so replicated equality constraints integrate cleanly into multi‑world setups. I don’t see index or off‑by‑one issues here.

Also applies to: 1705-1729, 1805-1825


2611-2670: Custom attributes on equality constraints are wired through correctly

The custom_attributes plumbed into add_equality_constraint and its connect/joint/weld helpers are validated and stored via _process_custom_attributes with ModelAttributeFrequency.EQUALITY_CONSTRAINT, using the freshly appended constraint_idx. This aligns with how bodies/shapes/joints handle custom attributes and should give per‑constraint arrays the expected layout without extra bookkeeping.

Also applies to: 2671-2776


5260-5266: Including equality_constraint_world in world-order validation is appropriate

Adding ("equality_constraint_world", self.equality_constraint_world) to the validated world arrays ensures equality constraints now participate in the same monotonic/contiguous world-index checks as other entities. Given how add_equality_constraint and add_builder assign world indices, this should catch misuses without breaking existing, well-ordered patterns (e.g., global constraints at the start or end).


5833-5855: EQUALITY_CONSTRAINT custom-attribute sizing matches model counts

The new elif frequency == ModelAttributeFrequency.EQUALITY_CONSTRAINT: count = m.equality_constraint_count branch in finalize() is consistent with the ModelAttributeFrequency enum and with how equality constraints are populated earlier in finalize. build_array(count, ...) will correctly produce one entry per equality constraint, with defaults for unspecified ones, mirroring the behavior for other frequencies.


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

codecov Bot commented Dec 11, 2025

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 93.47826% with 3 lines in your changes missing coverage. Please review.
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
newton/_src/solvers/mujoco/solver_mujoco.py 89.65% 3 Missing ⚠️

📢 Thoughts on this report? Let us know!

@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!

@eric-heiden eric-heiden added this pull request to the merge queue Dec 12, 2025
Merged via the queue into newton-physics:main with commit 5d8d1be Dec 12, 2025
20 checks passed
@coderabbitai coderabbitai Bot mentioned this pull request Dec 25, 2025
4 tasks
eric-heiden pushed a commit to eric-heiden/newton that referenced this pull request Jan 28, 2026
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
@coderabbitai coderabbitai Bot mentioned this pull request Mar 11, 2026
3 tasks
mmacklin pushed a commit to mmacklin/newton that referenced this pull request Apr 7, 2026
Signed-off-by: Alain Denzler <adenzler@nvidia.com>
@coderabbitai coderabbitai Bot mentioned this pull request May 27, 2026
3 tasks
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.

2 participants