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

2530-2536: update_geom_properties: use mjw_model’s world dimension, not model.num_envs

When envs aren’t separated, mjw_model holds nworld=1 fields. Launching kernels with dim=(model.num_envs, ...) will cause out-of-bounds writes to mjw_model arrays. Use the world dimension from mjw_model’s field shapes.

Apply this diff:

-        num_geoms = self.mj_model.ngeom
-        num_worlds = self.model.num_envs  # why is there no 'self.mjw_model.nworld'?
+        num_geoms = self.mj_model.ngeom
+        # Use the actual world dimension of the mjw model arrays.
+        # This is 1 when environments are not separated.
+        num_worlds = int(self.mjw_model.geom_pos.shape[0])

newton/tests/test_mujoco_solver.py (2)

688-696: Add an invariant check to catch the original cross-env mapping bug.

Right now, “expected” values are derived from the same mapping under test, which could mask regressions (e.g., all worlds mapping to env 0 shapes). Assert that, for each world, geom→shape indices differ from world 0 by a constant offset across geoms.

Apply this diff inside the loop after computing shape_idx to ensure the per-env offset is consistent (skip unmapped geoms):

 for world_idx in range(self.model.num_envs):
     for geom_idx in range(num_geoms):
-        shape_idx = to_newton_shape_index[world_idx, geom_idx]
+        shape_idx = to_newton_shape_index[world_idx, geom_idx]
+        # Invariant: for replicated envs, mapping should shift by a constant per-env offset
+        if world_idx > 0:
+            base_shape_idx = to_newton_shape_index[0, geom_idx]
+            if base_shape_idx >= 0 and shape_idx >= 0:
+                if env_offsets[world_idx] is None:
+                    env_offsets[world_idx] = int(shape_idx - base_shape_idx)
+                else:
+                    self.assertEqual(
+                        int(shape_idx - base_shape_idx),
+                        env_offsets[world_idx],
+                        f"Geom→shape mapping offset must be constant within world {world_idx}",
+                    )
         if shape_idx < 0:  # No mapping for this geom
             continue

Additionally, define env_offsets before the loop:

# Place right after `tested_count = 0`
env_offsets = [None] * self.model.num_envs

---

893-901: Repeat the invariant check in the update test to guard against regressions.

Ensure we don’t silently reintroduce the “all envs map to env 0” bug when updating properties.

Apply this diff inside the loop after computing shape_idx:

 for world_idx in range(self.model.num_envs):
     for geom_idx in range(num_geoms):
-        shape_idx = to_newton_shape_index[world_idx, geom_idx]
+        shape_idx = to_newton_shape_index[world_idx, geom_idx]
+        # Invariant: ensure a constant per-env offset relative to world 0
+        if world_idx > 0:
+            base_shape_idx = to_newton_shape_index[0, geom_idx]
+            if base_shape_idx >= 0 and shape_idx >= 0:
+                if env_offsets[world_idx] is None:
+                    env_offsets[world_idx] = int(shape_idx - base_shape_idx)
+                else:
+                    self.assertEqual(
+                        int(shape_idx - base_shape_idx),
+                        env_offsets[world_idx],
+                        f"Geom→shape mapping offset must be constant within world {world_idx}",
+                    )
         if shape_idx < 0:  # No mapping
             continue

And define env_offsets before the loop:

# Place right before the `for world_idx` loop
env_offsets = [None] * self.model.num_envs

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

992-1009: Migration note: public field shape change

to_newton_shape_index is now 2D. Please update the migration guide and any downstream docs/usages accordingly.

I can draft a short migration snippet (old vs new indexing) for warp.sim users if helpful.

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

Learnt from: adenzler-nvidia
PR: newton-physics/newton#479
File: newton/solvers/mujoco/solver_mujoco.py:2147-2149
Timestamp: 2025-07-28T15:04:54.360Z
Learning: The ls_parallel option is specific to MuJoCo Warp and is not available in the pure MuJoCo mjOption structure. It should only be set on mjw_model.opt, not mj_model.opt.

Learnt from: dylanturpin
PR: newton-physics/newton#394
File: newton/tests/test_ik.py:119-119
Timestamp: 2025-07-15T21:00:03.709Z
Learning: In the Newton physics engine codebase, it's acceptable to test private API functions (those prefixed with underscore) in unit and component tests, as the maintainers consider this appropriate for testing internal functionality.