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

263-271: Consider enhanced handling for global shapes.

The current implementation handles shape indices of -1 but may need additional logic for global shapes that are shared across environments, as mentioned in past review comments.

Also applies to: 294-304

---

1144-1146: Add safeguards for CPU version of MuJoCo.

While the warning for use_mujoco=True with use_mujoco_contacts=False is good, additional safeguards should be added to prevent runtime errors when using CPU MuJoCo with Newton contacts, as mentioned in past review comments.

         if use_mujoco and not use_mujoco_contacts:
             print("Setting use_mujoco_contacts to False has no effect when use_mujoco is True")
+            use_mujoco_contacts = True  # Force to True to prevent issues
 
         disableflags = 0

Also applies to: 1176-1178

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

2120-2142: Good implementation of multi-environment shape mapping.

The code correctly expands the shape-to-geom mapping to handle all environments when they are separated into different MuJoCo worlds. The warning about non-divisible shape counts is helpful for debugging.

Consider adding a debug log to report the number of shapes mapped:

         else:
             full_shape_mapping = shape_mapping
+        
+        if len(full_shape_mapping) > 0:
+            print(f"Mapped {len(full_shape_mapping)} Newton shapes to {len(shape_mapping)} MuJoCo geoms")

Also applies to: 2188-2194

newton/tests/test_mujoco_solver.py (1)

1116-1146: Improve test robustness by adding velocity checks and clarifying tolerance values.

The test logic is sound and properly validates the core Newton contact integration functionality. However, consider these improvements:

1. Magic numbers: The tolerance multipliers (0.9 and 1.2) should be better documented or defined as named constants
2. Settlement verification: The test only checks final position but doesn't verify the sphere has actually settled (velocity near zero)

Consider this enhanced version:

    def test_sphere_on_plane_with_newton_contacts(self):
        """Test that a sphere correctly collides with a plane using Newton contacts."""
        try:
            solver = MuJoCoSolver(self.model, use_mujoco_contacts=False)
        except ImportError as e:
            self.skipTest(f"MuJoCo or deps not installed. Skipping test: {e}")
            return

        sim_dt = 1.0 / 240.0
        num_steps = 120  # Simulate for 0.5 seconds to ensure it settles
+       
+       # Define tolerance constants for better readability
+       SPHERE_HEIGHT_TOLERANCE_LOWER = 0.9  # Allow 10% compression
+       SPHERE_HEIGHT_TOLERANCE_UPPER = 1.2  # Allow 20% above expected height
+       VELOCITY_SETTLING_THRESHOLD = 0.01   # Consider settled if velocity below this

        for _ in range(num_steps):
            self.contacts = self.model.collide(self.state_in)
            solver.step(self.state_in, self.state_out, self.control, self.contacts, sim_dt)
            self.state_in, self.state_out = self.state_out, self.state_in

        final_pos = self.state_in.body_q.numpy()[self.sphere_body_idx, :3]
        final_height = final_pos[2]  # Z-up in MuJoCo
+       final_velocity = self.state_in.body_qd.numpy()[self.sphere_body_idx, :3]
+       final_speed = np.linalg.norm(final_velocity)

        # The sphere should settle on the plane, with its center at its radius's height
        self.assertGreater(
            final_height,
-           self.sphere_radius * 0.9,
+           self.sphere_radius * SPHERE_HEIGHT_TOLERANCE_LOWER,
            f"Sphere fell through the plane. Final height: {final_height}",
        )
        self.assertLess(
            final_height,
-           self.sphere_radius * 1.2,
+           self.sphere_radius * SPHERE_HEIGHT_TOLERANCE_UPPER,
            f"Sphere is floating above the plane. Final height: {final_height}",
        )
+       
+       # Verify the sphere has settled (low velocity)
+       self.assertLess(
+           final_speed,
+           VELOCITY_SETTLING_THRESHOLD,
+           f"Sphere should have settled but still has velocity: {final_speed}",
+       )

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

Learnt from: shi-eric
PR: newton-physics/newton#386
File: .github/workflows/aws-gpu-tests.yml:65-82
Timestamp: 2025-07-12T19:46:47.510Z
Learning: In the newton-physics/newton repository AWS GPU workflow, the workflow only triggers on `pull_request_target` and `push` events, not on `pull_request` events. The `pull_request_target` event is used for testing PRs from forks safely on AWS infrastructure.

Learnt from: Kenny-Vilella
PR: newton-physics/newton#398
File: newton/examples/example_mujoco.py:352-352
Timestamp: 2025-07-14T03:57:29.659Z
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.659Z
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 (3)

89-198: Consider optimizing the contact parameter computation.

The contact_params function contains complex branching logic that could impact GPU performance. The multiple conditional assignments for mix calculation might benefit from simplification.

Consider consolidating the mix calculation logic:

-    mix = solmix1 / (solmix1 + solmix2)
-    mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 < MJ_MINVAL), 0.5, mix)
-    mix = wp.where((solmix1 < MJ_MINVAL) and (solmix2 >= MJ_MINVAL), 0.0, mix)
-    mix = wp.where((solmix1 >= MJ_MINVAL) and (solmix2 < MJ_MINVAL), 1.0, mix)
-    mix = wp.where(p1 == p2, mix, wp.where(p1 > p2, 1.0, 0.0))
+    sum_mix = solmix1 + solmix2
+    mix = wp.where(sum_mix < MJ_MINVAL, 0.5, solmix1 / sum_mix)
+    mix = wp.where(p1 != p2, wp.where(p1 > p2, 1.0, 0.0), mix)

---

1120-1120: Document the default behavior and interaction with other flags.

The use_mujoco_contacts parameter defaults to True, which maintains backward compatibility. However, the interaction with use_mujoco flag could be better documented.

Consider adding more detailed documentation about the parameter interaction:

-            use_mujoco_contacts (bool): If True, use the MuJoCo contact solver. If False, use the Newton contact solver (newton contacts must be passed in through the step function in that case).
+            use_mujoco_contacts (bool): If True, use the MuJoCo contact solver. If False, use the Newton contact solver and newton contacts must be passed in through the step function. Note: This parameter has no effect when use_mujoco=True, as pure MuJoCo backend always uses MuJoCo contacts.

---

1148-1149: The warning message could be more descriptive.

The current warning is functional but could provide more guidance to users about the expected behavior.

-        if use_mujoco and not use_mujoco_contacts:
-            print("Setting use_mujoco_contacts to False has no effect when use_mujoco is True")
+        if use_mujoco and not use_mujoco_contacts:
+            print("WARNING: use_mujoco_contacts=False has no effect when use_mujoco=True. Pure MuJoCo backend always uses MuJoCo's contact solver.")

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

Learnt from: shi-eric
PR: newton-physics/newton#386
File: .github/workflows/aws-gpu-tests.yml:65-82
Timestamp: 2025-07-12T19:46:47.510Z
Learning: In the newton-physics/newton repository AWS GPU workflow, the workflow only triggers on `pull_request_target` and `push` events, not on `pull_request` events. The `pull_request_target` event is used for testing PRs from forks safely on AWS infrastructure.

Learnt from: Kenny-Vilella
PR: newton-physics/newton#398
File: newton/examples/example_mujoco.py:352-352
Timestamp: 2025-07-14T03:57:29.659Z
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.659Z
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 (1)

1199-1203: The conditional logic in the step method needs error handling.

The code correctly handles the two modes but lacks validation for required Newton contacts when use_mujoco_contacts=False.

This matches the previous review feedback about adding error handling for required Newton contacts.

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

Learnt from: shi-eric
PR: newton-physics/newton#386
File: .github/workflows/aws-gpu-tests.yml:65-82
Timestamp: 2025-07-12T19:46:47.510Z
Learning: In the newton-physics/newton repository AWS GPU workflow, the workflow only triggers on `pull_request_target` and `push` events, not on `pull_request` events. The `pull_request_target` event is used for testing PRs from forks safely on AWS infrastructure.

Learnt from: Kenny-Vilella
PR: newton-physics/newton#398
File: newton/examples/example_mujoco.py:352-352
Timestamp: 2025-07-14T03:57:29.659Z
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.659Z
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 (4)

75-86: Add zero-length vector check to make_frame function.

The function calls wp.normalize(a) without checking if the vector has zero length, which can cause runtime errors.

This was flagged in a previous review. Apply the suggested fix:

@wp.func
def make_frame(a: wp.vec3):
+   if wp.length(a) == 0.0:
+       return wp.mat33(
+           1.0, 0.0, 0.0,
+           0.0, 1.0, 0.0,
+           0.0, 0.0, 1.0
+       )
    a = wp.normalize(a)
    b, c = orthogonals(a)
    
    # fmt: off
    return wp.mat33(
    a.x, a.y, a.z,
    b.x, b.y, b.z,
    c.x, c.y, c.z
  )
    # fmt: on

---

299-303: Fix world ID calculation for invalid body pairs.

The current logic defaults to worldid = 0 when both body_a and body_b are invalid (< 0), which incorrectly assigns contacts to environment 0.

This was flagged in previous reviews. Apply the suggested fix:

-   worldid = 0
-   if body_a >= 0:
-       worldid = body_a // bodies_per_env
-   elif body_b >= 0:
-       worldid = body_b // bodies_per_env
+   if body_a >= 0:
+       worldid = body_a // bodies_per_env
+   elif body_b >= 0:
+       worldid = body_b // bodies_per_env
+   else:
+       return  # both bodies invalid – skip this contact

---

1112-1113: Collision detection flag setting needs validation.

The code sets the collision detection flag but doesn't verify if the MuJoCo Warp model supports this configuration.

This was flagged in a previous review. Add validation:

        if self.mjw_model is not None:
+           if not hasattr(self.mjw_model.opt, 'run_collision_detection'):
+               raise ValueError("MuJoCo Warp version does not support disabling collision detection")
            self.mjw_model.opt.run_collision_detection = use_mujoco_contacts

---

1131-1135: Step function modification needs error handling.

The conditional logic properly handles the two contact modes, but lacks validation for required Newton contacts.

This was flagged in a previous review. Add validation:

                if self.mjw_model.opt.run_collision_detection:
                    self.mujoco_warp.step(self.mjw_model, self.mjw_data)
                else:
+                   if contacts is None or contacts.rigid_contact_count is None:
+                       raise ValueError("Newton contacts are required when use_mujoco_contacts=False")
                    self.convert_contacts_to_mjwarp(self.model, state_in, contacts)
                    self.mujoco_warp.step(self.mjw_model, self.mjw_data)

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

1079-1080: Warning message is helpful but could be more specific.

The warning correctly identifies when the parameter has no effect, but could be more explicit about the behavior.

Consider making the warning more specific:

-       print("Setting use_mujoco_contacts to False has no effect when use_mujoco is True")
+       print("Warning: use_mujoco_contacts=False has no effect when use_mujoco=True. MuJoCo CPU backend will always use its own contact solver.")

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.659Z
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.659Z
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 (3)

299-303: Fix world ID calculation for invalid body pairs.

The current logic defaults worldid = 0 when both body_a and body_b are negative, which incorrectly assigns contacts to environment 0. This was flagged in previous reviews but hasn't been addressed.

Apply this fix to handle invalid body pairs correctly:

-    worldid = 0
-    if body_a >= 0:
-        worldid = body_a // bodies_per_env
-    elif body_b >= 0:
-        worldid = body_b // bodies_per_env
+    if body_a >= 0:
+        worldid = body_a // bodies_per_env
+    elif body_b >= 0:
+        worldid = body_b // bodies_per_env
+    else:
+        return  # both bodies invalid – skip this contact

---

1137-1142: Add error handling for missing Newton contacts.

The conditional logic correctly handles the two contact modes, but there's no validation that Newton contacts are provided when use_mujoco_contacts=False. This was flagged in previous reviews but hasn't been addressed.

Add validation for required Newton contacts:

                if self.mjw_model.opt.run_collision_detection:
                    self.mujoco_warp.step(self.mjw_model, self.mjw_data)
                else:
+                   if contacts is None or contacts.rigid_contact_count is None:
+                       raise ValueError("Newton contacts are required when use_mujoco_contacts=False")
                    self.convert_contacts_to_mjwarp(self.model, state_in, contacts)
                    self.mujoco_warp.step(self.mjw_model, self.mjw_data)

---

75-86: Add zero-length vector handling before normalization.

The function normalizes the input vector a without checking if it has zero length, which could cause undefined behavior. Even though orthogonals handles zero-length vectors, the normalization should be protected.

Apply this fix to handle zero-length vectors:

@wp.func
def make_frame(a: wp.vec3):
+   if wp.length(a) == 0.0:
+       return wp.mat33(
+           1.0, 0.0, 0.0,
+           0.0, 1.0, 0.0,
+           0.0, 0.0, 1.0
+       )
    a = wp.normalize(a)
    b, c = orthogonals(a)
    
    # fmt: off
    return wp.mat33(
    a.x, a.y, a.z,
    b.x, b.y, b.z,
    c.x, c.y, c.z
  )
    # fmt: on

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

2038-2113: Complex but correct multi-environment shape mapping logic.

The logic correctly extends the template environment's shape mapping to all environments by calculating global shape indices. The implementation properly handles the case where multiple environments share the same MuJoCo geom indices.

Consider adding more inline comments to explain the mapping logic:

+           # Calculate the global Newton shape index for the current environment.
+           # Each environment has the same local structure, so we offset by env_idx * shapes_per_env
            global_shape_idx = env_idx * shapes_per_env + local_shape_idx
+           # All corresponding shapes map to the same MuJoCo geom index (since mj_model is single-env).
            full_shape_mapping[global_shape_idx] = geom_idx

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.659Z
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.659Z
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.