Fix loop joint coordinate mapping and constraint types#2075
Merged
eric-heiden merged 19 commits intoMar 13, 2026
Merged
Conversation
Loop joints (articulation=-1) become equality constraints in MuJoCo with no qpos/qvel slots, but their DOFs in Newton's joint_q_start offset all subsequent joints' MuJoCo indices. Build mj_q_start and mj_qd_start arrays at solver init that correctly map Newton template joint indices to MuJoCo qpos/qvel start addresses (-1 for loop joints), and pass them to the 5 coordinate conversion kernels. Also fix equality constraint types: use mjEQ_WELD for fixed loop joints (single constraint, all 6 DOFs) and single mjEQ_CONNECT for revolute/ball loop joints (pins anchor point, rotation freedom from kinematic chain). Previously all loop joints used 2x mjEQ_CONNECT which over-constrained non-fixed joints.
Parametrize test_fourbar_linkage with use_loop_joint flag. When true, the loop closure uses a revolute joint with articulation=-1 instead of an explicit equality constraint. This validates that the coordinate mapping fix produces correct dynamics, not just round-trip conversion.
Correct the weld constraint relpose computation for fixed loop joints to use parent_xform * inv(child_xform) instead of the inverse, matching the body-transform convention used elsewhere in the solver. Add the missing child anchor (data[3:6]) on connect constraints for revolute/ball loop joints. Add quat_to_mjc() conversion on weld relpose quaternion. Also fix _total_loop_joint_coords to handle separate_worlds correctly, remove dead _loop_joint_dof_count, fix "qvel" → "DOF" comment in qfrc kernel, and use solver._mujoco instead of direct import in new tests.
Make weld test use asymmetric parent/child xforms to exercise the relpose computation (parent * inv(child)) and quat_to_mjc conversion. Assert eq_data values: anchor, relpose translation, and wxyz quaternion. Remove eq.data[3:6] assignment on connect constraints — MuJoCo's compiler overwrites it by converting body1's anchor to world space then back to body2-local coords. Add parent anchor assertion to connect test instead.
Offset b1 from b0 so the loop joint uses asymmetric local anchors that coincide in world space: (0,0,-0.5) on b1 at (0,0,1) and (0,0,0.5) on b0 at origin both map to world (0,0,0.5). Assert both data[0:3] (parent anchor) and data[3:6] (child anchor auto-computed by MuJoCo compiler) at compile time. Also verify the runtime update kernel preserves loop joint constraint data in mjw_model.eq_data after a round-trip.
Add assertion after mj_q_start/mj_qd_start construction to verify every non-loop joint got a valid MuJoCo start index. This catches future bugs where a non-loop joint accidentally retains the -1 sentinel and silently vanishes from coordinate conversion. Replace hardcoded line number reference in relpose comment with a description of the code location. Improve relpose comment to say "relative transform from child anchor to parent anchor frame" instead of the imprecise "child body pose in parent body frame".
Classify loop joints by their decomposed DOF counts (lin_count, ang_count) and create appropriate MuJoCo equality constraints: - FIXED (0,0) -> WELD with auto-computed relpose - REVOLUTE (0,1) -> 2x CONNECT offset along hinge axis to constrain 5 DOFs (3 translational + 2 rotational) - BALL (0,2+) -> 1x CONNECT for 3 translational DOFs Add dynamics tests for each variant with setups that would fail if the wrong constraint type were used: revolute test uses a four-bar with ball joint in tree (out-of-plane buckling under Z-gravity without 2nd CONNECT), ball test verifies 3D pendulum motion (suppressed with extra CONNECT), fixed test checks rigid L-shape integrity (broken without WELD).
Move the ball joint from coupler->rocker (adjacent to loop joint) to world->crank (diagonal). This lets the crank tilt out of plane when the second CONNECT is missing, causing the entire mechanism to buckle dramatically under Z-gravity instead of just the rocker tilting slightly.
Tilt the pendulum sideways and kick it tangentially to create nonzero angular momentum about the gravity axis. This produces a conical orbit that spans both X and Z, clearly distinguishing a ball joint (1xCONNECT, 3D orbit) from a revolute (2xCONNECT, planar motion). Only check displacement in the second half of the trajectory so the initial X tilt doesn't mask a constraint that immediately collapses the motion to a plane.
Correct WELD relpose description to reference body2's local frame, clarify universal joint underconstrain severity, use "DOF entries" in force-conversion kernel comment, and make plane reference axis-agnostic in ball loop test.
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughReplaces joint-based start indices with MuJoCo-specific Changes
Sequence Diagram(s)sequenceDiagram
participant Solver as SolverMuJoCo
participant Kernels as Mujoco Kernels
participant MuJoCo as MuJoCo Model
Solver->>Solver: build model, compute mj_q_start / mj_qd_start (per-template)
Solver->>Kernels: call convert_warp_coords_to_mj (pass mj_q_start / mj_qd_start)
Kernels->>Kernels: for each joint/DOF: if mj_q_start[i] < 0 then skip DOF
Kernels->>MuJoCo: read/write qpos / qvel / qfrc using mj indices
MuJoCo-->>Kernels: updated qpos / qvel / qfrc
Kernels-->>Solver: converted coords & actuator forces
Solver->>Solver: apply results, update Newton state and constraints
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Use assert np.allclose to match TestMuJoCoArticulationConversion style. Remove unnecessary bound_inertia/bound_mass where bodies already have explicit mass and inertia. Use non-default position and orientation for the free body so the roundtrip cannot pass by accident.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
newton/tests/test_mujoco_solver.py (2)
6117-6142: The new offset tests still miss themj_qd_start/nontrivial quaternion cases.Both round-trips keep velocities at zero, and the free body uses the identity rotation. That gives good coverage for the
qpostranslation offset, but it can still miss a badmj_qd_startmapping or a quaternion-slot offset bug. Please seed a non-zerojoint_qd/free-body velocity and a non-identity free-body quaternion before the round-trip so these tests cover the full regression surface this PR is fixing.Also applies to: 6196-6218
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@newton/tests/test_mujoco_solver.py` around lines 6117 - 6142, The test currently only checks qpos round-trips with zero velocities and identity quaternion; modify the setup before calling solver._update_mjc_data / solver._mujoco_warp.kinematics / solver._update_newton_state to also seed a non-zero joint velocity (set model.joint_qd to a small nonzero vector) and set the free-body quaternion in state.body_q (for b_free) to a non-identity rotation (e.g., a small known quaternion not equal to [1,0,0,0]) so the round-trip validates mj_qd_start and quaternion-slot mapping; ensure you update body_q_before after seeding and adjust the assertions to compare both translational, quaternion (sign-invariant) and velocity components to catch mj_qd_start mapping bugs.
5988-5995: Avoid hard-coding the compiled weldrelpose.The solver now intentionally leaves fixed-loop
data[3:10]forspec.compile()to derive from the compiled body poses. Pinning exact numericrelposevalues here makes the test brittle to legitimate pose/layout changes and re-asserts the hand-derived behavior this PR is removing. I’d keep the type/anchor assertions here and move anyrelposecheck to a case that computes the expected body-frame value independently.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@newton/tests/test_mujoco_solver.py` around lines 5988 - 5995, The test is hard-coding the compiled weld relpose (eq_data[3:10]) which is now left for spec.compile() to compute; remove the two np.testing.assert_allclose checks that assert eq_data[3:6] and eq_data[6:10] and keep only the type/anchor assertion (np.testing.assert_allclose(eq_data[0:3], ...)). If you still need to validate relpose, add a separate test that computes the expected body-frame relpose from the bodies' poses and compares that against solver.mj_model.eq_data[0][3:10] (use the same solver.mj_model.eq_data reference and eq_data variable to locate the code).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@newton/_src/solvers/mujoco/solver_mujoco.py`:
- Around line 3252-3256: The new single-world template sizing introduced in
_update_mjc_data (using effective_coord_count and single_world_template) is not
reflected in the CPU paths, causing _apply_mjc_control and _update_newton_state
to index past the smaller mj_q_start / mj_qd_start buffers when
use_mujoco_cpu=True and separate_worlds=True; update those functions to compute
effective_coord_count and single_world_template the same way and use
expected_qpos (or the template-sized offsets) when launching CPU kernels so
their loops and buffer indexing match the template sizing, or alternatively add
an upfront validation rejecting the combination use_mujoco_cpu=True with
separate_worlds=True; ensure references to mj_q_start and mj_qd_start, and any
loops over model.joint_count, use the template-sized indices to avoid
out-of-bounds access.
- Around line 4761-4779: Change the loop-joint handling so you only map to a
CONNECT when there are exactly 3 angular DOFs: replace the condition `lin_count
== 0 and ang_count >= 2` with `lin_count == 0 and ang_count == 3` in the branch
that creates `eq` (the CONNECT equality), and adjust the surrounding comment to
state it is exact for 3 angular DOFs (ball). Leave the 2-DOF case (universal) to
fall through to the `else` that issues the warning so the code does not silently
underconstrain the mechanism; references: the condition around `eq =
spec.add_equality(...)`, `eq.type = mujoco.mjtEq.mjEQ_CONNECT`,
`mjc_eq_to_newton_jnt[eq.id] = j`, and the `JointType(j_type).name` warning
block.
In `@newton/tests/test_mujoco_solver.py`:
- Around line 6144-6154: The test never exercises the runtime constraint-update
path because it only calls _update_mjc_data, kinematics, and
_update_newton_state and then reads mjw_model.eq_data (which is still the
compiled value); to fix, before asserting eq_data mutate a constraint-relevant
field on the model/joint used by the loop joint (for example change the loop
joint's parent/child anchor or another constraint property) and then call
notify_model_changed(SolverNotifyFlags.CONSTRAINT_PROPERTIES) on the Solver so
the runtime update path runs, then read mjw_model.eq_data and run the existing
assertions.
In `@newton/tests/test_physics_validation.py`:
- Around line 1338-1341: Change the eq_data access to use the mujoco-warp model:
replace solver.mj_model.eq_data[0, 3:6] with solver.mjw_model.eq_data[0, 3:6] so
child_anchor_local is read from the mjw_model; update the reference in the
computation that uses child_anchor_local (the pivot_end calculation that uses
body_q, link, pos and rot derived from wp.quat_to_matrix) to ensure consistency
with other usages in this test.
---
Nitpick comments:
In `@newton/tests/test_mujoco_solver.py`:
- Around line 6117-6142: The test currently only checks qpos round-trips with
zero velocities and identity quaternion; modify the setup before calling
solver._update_mjc_data / solver._mujoco_warp.kinematics /
solver._update_newton_state to also seed a non-zero joint velocity (set
model.joint_qd to a small nonzero vector) and set the free-body quaternion in
state.body_q (for b_free) to a non-identity rotation (e.g., a small known
quaternion not equal to [1,0,0,0]) so the round-trip validates mj_qd_start and
quaternion-slot mapping; ensure you update body_q_before after seeding and
adjust the assertions to compare both translational, quaternion (sign-invariant)
and velocity components to catch mj_qd_start mapping bugs.
- Around line 5988-5995: The test is hard-coding the compiled weld relpose
(eq_data[3:10]) which is now left for spec.compile() to compute; remove the two
np.testing.assert_allclose checks that assert eq_data[3:6] and eq_data[6:10] and
keep only the type/anchor assertion (np.testing.assert_allclose(eq_data[0:3],
...)). If you still need to validate relpose, add a separate test that computes
the expected body-frame relpose from the bodies' poses and compares that against
solver.mj_model.eq_data[0][3:10] (use the same solver.mj_model.eq_data reference
and eq_data variable to locate the code).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: 3550e649-571f-4c94-9e62-a57e0c2db822
📒 Files selected for processing (4)
newton/_src/solvers/mujoco/kernels.pynewton/_src/solvers/mujoco/solver_mujoco.pynewton/tests/test_mujoco_solver.pynewton/tests/test_physics_validation.py
Follow the mujoco_cpu/mujoco_warp solver dict pattern used by the other physics validation tests so loop joint tests run on all devices.
Restrict ball loop joint mapping to ang_count == 3 (was >= 2) so universal joints fall through to the warning path instead of being silently underconstrainted. Add single_world_template handling to _update_newton_state and _apply_mjc_control to match _update_mjc_data, preventing OOB access on mj_q_start/mj_qd_start when separate_worlds=True with CPU. Remove misleading eq_data assertions from the coordinate conversion roundtrip test (they only checked compiled values, not runtime).
…oint-coord-mapping # Conflicts: # newton/tests/test_mujoco_solver.py
eric-heiden
previously approved these changes
Mar 13, 2026
eric-heiden
left a comment
Member
There was a problem hiding this comment.
Thanks for fixing this! I tested this branch with the digit_v4.usd asset from Isaac Sim, it seems to work well. Just some test issues remain, but it seems to be a test setup problem.
- Use solver._mujoco.mjtEq instead of bare mujoco import in test_loop_joints_only and test_mixed_loop_joints_and_equality_constraints - Add inertia to free bodies in coordinate conversion offset tests so MuJoCo spec.compile() does not reject them - Guard wp.ScopedCapture behind device.is_cuda in test_fourbar_linkage since graph capture requires CUDA - Sync eq_solref to mj_model for CPU solver path so stiff constraint parameters take effect in MuJoCo CPU tests
…oint-coord-mapping
Also clarify AGENTS.md changelog instruction to explicitly say entries go at the end of the category list.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@newton/tests/test_mujoco_solver.py`:
- Line 5985: The current assertion compares eq_data[6:10] to [1.0,0.0,0.0,0.0]
which fails for an equivalent quaternion with opposite sign; change the
assertion to be sign-invariant by checking either equality to the target or its
negation (e.g. replace the single np.allclose(eq_data[6:10], [1.0,0.0,0.0,0.0],
atol=1e-6) with a boolean that passes if np.allclose(eq_data[6:10], target,
atol=1e-6) or np.allclose(eq_data[6:10], -target, atol=1e-6)), referencing the
eq_data slice eq_data[6:10] and the target quaternion values.
In `@newton/tests/test_physics_validation.py`:
- Around line 957-962: The test writes to solver.mjw_model unconditionally and
only then conditionally to solver.mj_model, and it also sets gravity only on
mjw_model; fix by branching on solver.use_mujoco_cpu first and writing to the
correct backend model (use solver.mj_model when solver.use_mujoco_cpu is true,
otherwise solver.mjw_model) for all property updates (eq_solref, gravity, etc.),
and ensure gravity is written to mj_model.opt.gravity for CPU mode and to
mjw_model.opt.gravity for Warp mode so _update_model_properties() sees the right
values.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: 9c54ba4a-9409-4142-8e88-125e4c0a2794
📒 Files selected for processing (3)
newton/_src/solvers/mujoco/solver_mujoco.pynewton/tests/test_mujoco_solver.pynewton/tests/test_physics_validation.py
Write eq_solref and gravity to the correct backend model (mj_model for CPU, mjw_model for Warp) instead of unconditionally writing to mjw_model first. The gravity write in test_revolute_loop_joint was missing the CPU branch entirely, so the CPU solver ran with default gravity. Make the weld relpose quaternion assertion sign-invariant since q and -q represent the same rotation.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
CHANGELOG.md (1)
29-29: Prefer user-outcome wording over internal constraint mechanics.This entry is clear but still exposes internal implementation detail (
WELD,CONNECT,2x CONNECT). Consider rephrasing to emphasize observable behavior/compatibility impact for users.As per coding guidelines, changelog entries for user-facing changes should avoid internal implementation details.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@CHANGELOG.md` at line 29, The changelog line exposes internal implementation names (WELD, CONNECT, 2x CONNECT); reword it to focus on user-visible impact instead: replace the current sentence with a user-outcome phrasing that describes behavior and compatibility (e.g., improved joint constraint generation that matches each joint’s degrees of freedom, reduces redundant constraints, and improves compatibility) and remove or avoid mentioning internal symbols like WELD, CONNECT, and “2x CONNECT”.newton/tests/test_physics_validation.py (1)
850-877: Optional: extract duplicated four-bar solve/initialization math.The Freudenstein-based initialization appears in multiple tests; a shared helper would reduce drift risk when thresholds or geometry assumptions evolve.
♻️ Minimal cleanup sketch
+def _fourbar_initial_angles(a_link, b_link, c_link, d_link): + # solve theta3_0 / theta4 / delta_rocker once + ... + return theta3_0, delta_rockerAlso applies to: 1098-1115
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@newton/tests/test_physics_validation.py` around lines 850 - 877, Multiple tests duplicate the Freudenstein four-bar initialization (the math computing K1,K2,K3, A,B,C, denom, arg and solving for theta4/theta3) as implemented in solve_fourbar; extract that logic into a single helper (e.g., compute_freudenstein or fourbar_solve) that returns (theta3, theta4) and replace duplicated blocks in solve_fourbar and the other occurrences (the other test locations using the same math) to call the helper, update any imports/uses accordingly, and ensure the helper preserves the exact calculations for K1,K2,K3, the arctan2/arccos branch used for the open configuration and the returned theta3, theta4 values.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@newton/tests/test_mujoco_solver.py`:
- Around line 6047-6203: The tests currently only exercise the Warp backend;
instantiate the MuJoCo solver with the CPU path to cover the CPU-specific
template-sizing logic by creating a SolverMuJoCo with use_mujoco_cpu=True (e.g.
replace or add the existing SolverMuJoCo(model) call in both
test_loop_joint_coordinate_conversion_offset and
test_ball_loop_joint_coordinate_conversion_offset with/alongside
SolverMuJoCo(model, use_mujoco_cpu=True) so the calls into
SolverMuJoCo._update_mjc_data and SolverMuJoCo._update_newton_state exercise the
CPU template-sized branches (or run the same round-trip sequence twice: once
default and once with use_mujoco_cpu=True).
---
Nitpick comments:
In `@CHANGELOG.md`:
- Line 29: The changelog line exposes internal implementation names (WELD,
CONNECT, 2x CONNECT); reword it to focus on user-visible impact instead: replace
the current sentence with a user-outcome phrasing that describes behavior and
compatibility (e.g., improved joint constraint generation that matches each
joint’s degrees of freedom, reduces redundant constraints, and improves
compatibility) and remove or avoid mentioning internal symbols like WELD,
CONNECT, and “2x CONNECT”.
In `@newton/tests/test_physics_validation.py`:
- Around line 850-877: Multiple tests duplicate the Freudenstein four-bar
initialization (the math computing K1,K2,K3, A,B,C, denom, arg and solving for
theta4/theta3) as implemented in solve_fourbar; extract that logic into a single
helper (e.g., compute_freudenstein or fourbar_solve) that returns (theta3,
theta4) and replace duplicated blocks in solve_fourbar and the other occurrences
(the other test locations using the same math) to call the helper, update any
imports/uses accordingly, and ensure the helper preserves the exact calculations
for K1,K2,K3, the arctan2/arccos branch used for the open configuration and the
returned theta3, theta4 values.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
Run ID: fbb61988-6bc1-47e2-84f4-75442624e668
📒 Files selected for processing (4)
AGENTS.mdCHANGELOG.mdnewton/tests/test_mujoco_solver.pynewton/tests/test_physics_validation.py
eric-heiden
approved these changes
Mar 13, 2026
mzamoramora-nvidia
pushed a commit
to mzamoramora-nvidia/newton
that referenced
this pull request
Mar 16, 2026
mmacklin
pushed a commit
to mmacklin/newton
that referenced
this pull request
Apr 7, 2026
This was referenced Apr 8, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
are present. Loop joints occupy slots in Newton's joint arrays but have
no MuJoCo qpos/qvel entries, so using Newton indices directly caused
all subsequent joints to read/write at wrong offsets. Introduces
separate
mj_q_start/mj_qd_startarrays with -1 sentinels forloop joints.
emitting 2x CONNECT: FIXED -> WELD (0 DOFs free), REVOLUTE -> 2x
CONNECT along hinge axis (1 rotational DOF free), BALL -> 1x CONNECT
(3 rotational DOFs free).
physics validation tests (revolute, ball, fixed loop joints) that each
discriminate between correct and incorrect constraint types.
Test plan
pytest newton/tests/test_mujoco_solver.py -k loop_joint— unit tests for mappings and round-tripspytest newton/tests/test_physics_validation.py -k loop_joint— dynamics tests for all 3 joint typespytest newton/tests/test_physics_validation.py -k fourbar— four-bar linkage with loop joint variantSummary by CodeRabbit
Bug Fixes
New Features
Tests
Documentation