newton/solvers/mujoco/solver_mujoco.py (2)

1622-1631: Remove or document unused equality constraint variables

Several variables are extracted but immediately commented out. This creates unnecessary overhead and makes the code harder to maintain. Either remove these lines entirely or add a TODO comment explaining their future purpose.

 eq_constraint_body1 = model.equality_constraint_body1.numpy()
 eq_constraint_body2 = model.equality_constraint_body2.numpy()
 eq_constraint_anchor1 = model.equality_constraint_anchor1.numpy()
-# eq_constraint_anchor2 = model.equality_constraint_anchor2.numpy()
-# eq_constraint_distance = model.equality_constraint_distance.numpy()
 eq_constraint_joint1 = model.equality_constraint_joint1.numpy()
 eq_constraint_joint2 = model.equality_constraint_joint2.numpy()
 eq_constraint_polycoef = model.equality_constraint_polycoef.numpy()
-# eq_constraint_enabled = model.equality_constraint_enabled.numpy()
+# TODO: Add support for anchor2, distance, and enabled properties when implementing
+# weld, distance, and other constraint types

---

2020-2039: Consider environment filtering and enabled flag for constraints

The equality constraints should respect the environment filtering logic used elsewhere in the code when separate_envs_to_worlds is True. Also, the equality_constraint_enabled property exists but is commented out - consider using it to allow disabling constraints.

Would you like me to implement the environment filtering logic for equality constraints to ensure they're only added for the selected environment?

newton/sim/builder.py (1)

3530-3556: Consider polycoef array structure and approve the rest.

The equality constraint finalization is well-implemented and follows the established patterns. However, there's a potential issue with the polycoef attribute:

m.equality_constraint_polycoef = wp.array(self.equality_constraint_polycoef, dtype=wp.float32)

Since polycoef is a list of lists (each constraint can have multiple polynomial coefficients), flattening it into a single Warp array might lose the structure. Consider either:

1. Using a 2D array if all constraints have the same number of coefficients
2. Storing it as a list of Warp arrays
3. Adding metadata to track coefficient counts per constraint

The rest of the implementation correctly converts the equality constraint data to appropriate Warp array types and follows the established patterns for model finalization.

Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro

Learnt from: Kenny-Vilella
PR: newton-physics/newton#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.

Learnt from: Kenny-Vilella
PR: newton-physics/newton#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.

newton/tests/test_equality_constraints.py (2)

25-36: Consider improving test organization and method naming.

The test setup is functionally correct, but could be improved:

1. Move timing parameters to setUp(): The simulation timing parameters should be initialized in a setUp() method rather than within the test method for better test organization.

2. Test method name could be more descriptive: The method name test_connect_constraint suggests it's testing connect constraints specifically, but the test actually verifies general equality constraint parsing.

Consider this refactor:

class TestEqualityConstraints(unittest.TestCase):
+    def setUp(self):
+        self.sim_time = 0.0
+        self.frame_dt = 1 / 60
+        self.sim_dt = self.frame_dt / 10
+
-    def test_connect_constraint(self):
+    def test_equality_constraints_parsing(self):
-        self.sim_time = 0.0
-        self.frame_dt = 1 / 60
-        self.sim_dt = self.frame_dt / 10
-

---

62-69: Consider enhancing test coverage with additional assertions.

The current assertion only verifies that equality constraints exist but doesn't validate their correctness. For a more comprehensive test, consider adding assertions to verify:

• Specific constraint types (connect, joint, etc.)
• Constraint properties and parameters
• Expected number of constraints for the 4_bar.xml model

Example additional assertions:

        self.assertGreater(
            self.solver.mj_model.eq_type.shape[0], 0
        )  # check if number of equality constraints in mjModel > 0
+
+        # Verify specific constraint properties
+        eq_types = self.solver.mj_model.eq_type
+        self.assertIn(mujoco.mjtEq.mjEQ_CONNECT, eq_types)  # Verify connect constraint exists
+        self.assertEqual(len([t for t in eq_types if t == mujoco.mjtEq.mjEQ_CONNECT]), 1)  # Expected count

The main execution block with kernel cache clearing is appropriate.

Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro

Learnt from: Kenny-Vilella
PR: newton-physics/newton#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.

Learnt from: Kenny-Vilella
PR: newton-physics/newton#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.

newton/tests/test_equality_constraints.py (1)

42-51: Verify that use_mujoco=True is now supported.

According to previous learnings, the use_mujoco option was unsupported and caused crashes. Please verify that this has been properly implemented and no longer causes issues.

#!/bin/bash
# Description: Check if use_mujoco is still marked as unsupported in the codebase

# Search for any warnings or unsupported messages related to use_mujoco
rg -A 5 "use_mujoco.*unsupported|unsupported.*use_mujoco" 

# Check the MuJoCoSolver implementation for use_mujoco handling
ast-grep --pattern 'class MuJoCoSolver {
  $$$
}'

# Look for any crash-related comments or warnings about use_mujoco
rg -A 5 "use_mujoco.*crash|crash.*use_mujoco"

newton/utils/import_mjcf.py (2)

768-773: Apply consistent body name sanitization.

The connect constraint replaces hyphens with underscores in body names, but other constraint types should do the same for consistency.

---

817-817: Add verbose check for debug print statement.

The weld constraint debug message should be wrapped with verbose check for consistency.

newton/sim/model.py (1)

277-296: Initialize Warp arrays with None.

According to past review comments, Model attributes that are Warp arrays should be initialized with None.

-        self.equality_constraint_body1 = None
-        self.equality_constraint_body2 = None
-        self.equality_constraint_anchor = None
-        self.equality_constraint_torquescale = None
-        self.equality_constraint_relpose = None
-        self.equality_constraint_joint1 = None
-        self.equality_constraint_joint2 = None
-        self.equality_constraint_polycoef = None
-        self.equality_constraint_enabled = None
+        self.equality_constraint_body1 = None
+        self.equality_constraint_body2 = None
+        self.equality_constraint_anchor = None
+        self.equality_constraint_torquescale = None
+        self.equality_constraint_relpose = None
+        self.equality_constraint_joint1 = None
+        self.equality_constraint_joint2 = None
+        self.equality_constraint_polycoef = None
+        self.equality_constraint_enabled = None

Note: The attributes are already initialized to None, so the code is correct as-is.

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

2054-2080: Add error handling and validation for equality constraints

The implementation still lacks proper error handling for invalid indices and unsupported constraint types as previously identified.

newton/sim/model.py (1)

275-297: Consider adding equality constraints to attribute_frequency mapping.

The new equality constraint attributes should be classified in the attribute_frequency dictionary to maintain consistency with other model attributes.

Add the following after line 385:

# attributes per equality constraint
self.attribute_frequency["equality_constraint_type"] = "equality_constraint"
self.attribute_frequency["equality_constraint_body1"] = "equality_constraint"
self.attribute_frequency["equality_constraint_body2"] = "equality_constraint"
self.attribute_frequency["equality_constraint_anchor"] = "equality_constraint"
self.attribute_frequency["equality_constraint_torquescale"] = "equality_constraint"
self.attribute_frequency["equality_constraint_relpose"] = "equality_constraint"
self.attribute_frequency["equality_constraint_joint1"] = "equality_constraint"
self.attribute_frequency["equality_constraint_joint2"] = "equality_constraint"
self.attribute_frequency["equality_constraint_polycoef"] = "equality_constraint"
self.attribute_frequency["equality_constraint_enabled"] = "equality_constraint"

Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro

Learnt from: Kenny-Vilella
PR: newton-physics/newton#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.

newton/utils/import_mjcf.py (4)

814-815: Add verbose check for consistency.

The debug print statement should be wrapped with a verbose check for consistency with the rest of the codebase.

-            if verbose:
-                print(f"Weld constraint: {body1_name} to {body2_name}")
+            if verbose:
+                print(f"Weld constraint: {body1_name} to {body2_name}")

Wait, this is already correctly wrapped. Let me check if there are other missing verbose checks in the function.

---

775-802: Fix undefined variable and improve consistency.

Multiple issues in the connect constraint parsing:

1. Line 787 uses undefined variable anchor instead of anchor_vec
2. Body name sanitization should be consistent with other constraint types

Apply this fix:

         if body1_name and anchor:
             if verbose:
-                print(f"Connect constraint: {body1_name} to {body2_name} at anchor {anchor}")
+                print(f"Connect constraint: {body1_name} to {body2_name} at anchor {anchor_vec}")

Also, consider moving the anchor_vec calculation before the verbose print for better code organization.

---

803-832: Fix incorrect attribute reference and apply consistent naming.

Critical issue: Lines 805-806 incorrectly reference connect.attrib instead of weld.attrib when parsing weld constraints.

Apply this fix:

-        body1_name = connect.attrib.get("body1", "").replace("-", "_") if connect.attrib.get("body1") else None
+        body1_name = weld.attrib.get("body1", "").replace("-", "_") if weld.attrib.get("body1") else None
         body2_name = (
-            connect.attrib.get("body2", "worldbody").replace("-", "_") if connect.attrib.get("body2") else None
+            weld.attrib.get("body2", "worldbody").replace("-", "_") if weld.attrib.get("body2") else None
         )

---

833-851: Fix undefined variable in joint constraint parsing.

The variable common is used but not defined in the joint constraint parsing section.

Add the missing line:

     for joint in equality.findall("joint"):
+        common = parse_common_attributes(joint)
         joint1_name = joint.attrib.get("joint1")

newton/utils/import_mjcf.py (1)

612-613: Remove leftover debug code.

This debug print statement appears to be leftover development code and should be removed.

-    if body_name == "thigh_right":
-        print(geoms)

Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro

Learnt from: Kenny-Vilella
PR: newton-physics/newton#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.

newton/solvers/mujoco/solver_mujoco.py (2)

1656-1665: Add validation for equality constraint attributes before accessing them.

The code directly accesses multiple equality constraint attributes without checking if they exist on the model object. This could cause runtime exceptions if any attributes are missing or not properly initialized.

+        # Validate equality constraint attributes exist
+        constraint_attrs = [
+            'equality_constraint_type', 'equality_constraint_body1', 'equality_constraint_body2',
+            'equality_constraint_anchor', 'equality_constraint_torquescale',
+            'equality_constraint_relpose', 'equality_constraint_joint1',
+            'equality_constraint_joint2', 'equality_constraint_polycoef',
+            'equality_constraint_enabled'
+        ]
+        for attr in constraint_attrs:
+            if not hasattr(model, attr):
+                raise AttributeError(f"Model missing required equality constraint attribute: {attr}")
+
         eq_constraint_type = model.equality_constraint_type.numpy()
         eq_constraint_body1 = model.equality_constraint_body1.numpy()
         eq_constraint_body2 = model.equality_constraint_body2.numpy()

---

2055-2081: Add comprehensive error handling and validation for equality constraints.

The constraint conversion logic lacks proper error handling and validation:

1. Direct dictionary access could raise KeyError if indices are invalid
2. No bounds checking for body/joint indices before accessing model arrays
3. No explicit handling for unsupported constraint types
4. No validation of constraint data before use

 for i, typ in enumerate(eq_constraint_type):
     if typ == newton.EQ_CONNECT:
+        # Validate body indices
+        body1_idx = eq_constraint_body1[i]
+        body2_idx = eq_constraint_body2[i]
+        if body1_idx >= len(model.body_key) or body2_idx >= len(model.body_key):
+            print(f"Warning: Invalid body indices in connect constraint {i}")
+            continue
+            
         eq = spec.add_equality(objtype=mujoco.mjtObj.mjOBJ_BODY)
         eq.type = mujoco.mjtEq.mjEQ_CONNECT
         eq.active = eq_constraint_enabled[i]
         eq.name1 = model.body_key[eq_constraint_body1[i]]
         eq.name2 = model.body_key[eq_constraint_body2[i]]
         eq.data[0:3] = eq_constraint_anchor[i]

     elif typ == newton.EQ_JOINT:
+        # Validate joint indices
+        joint1_idx = eq_constraint_joint1[i]
+        joint2_idx = eq_constraint_joint2[i]
+        if joint1_idx >= len(model.joint_key) or joint2_idx >= len(model.joint_key):
+            print(f"Warning: Invalid joint indices in joint constraint {i}")
+            continue
+            
         eq = spec.add_equality(objtype=mujoco.mjtObj.mjOBJ_JOINT)
         eq.type = mujoco.mjtEq.mjEQ_JOINT
         eq.active = eq_constraint_enabled[i]
         eq.name1 = model.joint_key[eq_constraint_joint1[i]]
         eq.name2 = model.joint_key[eq_constraint_joint2[i]]
         eq.data[0:5] = eq_constraint_polycoef[i]

     elif typ == newton.EQ_WELD:
+        # Validate body indices
+        body1_idx = eq_constraint_body1[i]
+        body2_idx = eq_constraint_body2[i]
+        if body1_idx >= len(model.body_key) or body2_idx >= len(model.body_key):
+            print(f"Warning: Invalid body indices in weld constraint {i}")
+            continue
+            
         eq = spec.add_equality(objtype=mujoco.mjtObj.mjOBJ_BODY)
         eq.type = mujoco.mjtEq.mjEQ_WELD
         eq.active = eq_constraint_enabled[i]
         eq.name1 = model.body_key[eq_constraint_body1[i]]
         eq.name2 = model.body_key[eq_constraint_body2[i]]
         eq.data[0:3] = eq_constraint_anchor[i]
         eq.data[3:10] = eq_constraint_relpose[i]
         eq.data[10] = eq_constraint_torquescale[i]
+    
+    else:
+        print(f"Warning: Unknown or unsupported equality constraint type '{typ}' at index {i}")

newton/utils/import_mjcf.py (1)

612-613: Remove leftover debug code.

This debug print statement appears to be leftover development code that should be removed before merging.

-        if body_name == "thigh_right":
-            print(geoms)

Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro

Learnt from: Kenny-Vilella
PR: newton-physics/newton#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.

newton/utils/import_mjcf.py (1)

612-613: Remove leftover debug code.

This debug print statement appears to be leftover development code that should be removed before merging.

-        if body_name == "thigh_right":
-            print(geoms)

Configuration used: .coderabbit.yml
Review profile: CHILL
Plan: Pro

Learnt from: Kenny-Vilella
PR: newton-physics/newton#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.