newton/tests/test_mujoco_solver.py (1)

844-847: Friction pattern change is fine; consider a tiny readability improvement

The new pattern raises the baseline friction; assertions compute expected values from new_mu so behavior remains consistent.

Optional: vectorize for clarity and slightly faster test runs.

-        new_mu = np.zeros(shape_count)
-        for i in range(shape_count):
-            new_mu[i] = 1.0 + (i + 1) * 0.05  # Pattern: 1.05, 1.10, ...
+        # Pattern: 1.05, 1.10, ...
+        new_mu = 1.0 + 0.05 * (np.arange(shape_count) + 1)

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

1838-1840: Fix typo and clarify intent of shape replication comment.

• Typo: “exludes” → “excludes”.
• Consider clarifying that “negative collision group” means global shapes attached to body -1 that are not replicated per env.

Apply this diff:

-        # number of shapes which are replicated per env (exludes singular static shapes from a negative group)
+        # number of shapes replicated per env (excludes global/static shapes in a negative collision group)

---

1913-1986: Incoming geom transform parameter is unused; either apply it or remove it.

add_geoms(newton_body_id, incoming_xform) accepts incoming_xform but never applies it. If the intent is to pre-adjust geom frames, multiply it in when computing geom_params.

Apply this diff:

-            for shape in shapes:
+            for shape in shapes:
                 if skip_visual_only_geoms and not (shape_flags[shape] & ShapeFlags.COLLIDE_SHAPES):
                     continue
                 stype = shape_type[shape]
                 name = f"{geom_type_name[stype]}_{shape}"
                 if stype == GeoType.PLANE and newton_body_id != -1:
                     raise ValueError("Planes can only be attached to static bodies")
                 geom_params = {
                     "type": geom_type_mapping[stype],
                     "name": name,
                 }
                 tf = wp.transform(*shape_transform[shape])
+                if incoming_xform is not None:
+                    tf = incoming_xform * tf

---

2201-2202: Replace hard assert on geometry count with a warning.

This assert can be tripped by benign differences (e.g., if we later add non-colliding geoms conditionally). Prefer a warning to avoid aborting conversion in release configurations.

Apply this diff:

-        assert colliding_shapes_per_env == len(spec.geoms)
+        if colliding_shapes_per_env != len(spec.geoms):
+            warnings.warn(
+                f"Per-env colliding shape count ({colliding_shapes_per_env}) != spec geom count ({len(spec.geoms)}). "
+                "Verify collision flags and skip_visual_only_geoms.",
+                stacklevel=2,
+            )

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

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

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

1852-1854: Fallback for shape_range_len is wrong when all shapes are global (-1).

Setting shape_range_len = model.shape_count incorrectly inflates per-env stride, leading to out-of-bounds indices downstream. There are no per-env replicated shapes in this case, so the stride must be zero.

Apply this diff:

             else:
                 first_collision_group = -1
-                shape_range_len = model.shape_count
+                # No per-env replicated shapes; only global (-1) shapes exist
+                shape_range_len = 0

---

1902-1907: shape_incoming_xform initialization is invalid; use a sequence of wp.transform objects.

np.tile(np.array(wp.transform_identity()), (N, 1)) creates a 2D float array, not a list of wp.transform. This can cause dtype issues when converting to wp.array(dtype=wp.transform).

Apply this diff:

-        shape_incoming_xform = np.tile(np.array(wp.transform_identity()), (model.shape_count, 1))
+        # one wp.transform per Newton shape
+        shape_incoming_xform = [wp.transform_identity() for _ in range(model.shape_count)]

---

2244-2272: Global shape index computation across environments is incorrect and brittle; build maps from self.shape_map.

The arithmetic global_shape_idx = env_idx * shape_range_len + template_shape_idx assumes contiguous env-local shape indices and a stable stride. This is not guaranteed and causes mis-mappings or OOB indices. Use self.shape_map, which already encodes (worldid, geom) per Newton shape; treat worldid == -1 as “replicate across all envs”.

Apply this diff to replace the entire block:

-        if separate_envs_to_worlds and model.num_envs > 1:
-            full_shape_mapping = {}
-            reverse_shape_mapping = {}
-            for env_idx in range(model.num_envs):
-                # `geom_to_shape_idx` provides the reverse mapping for the template env: {mj_geom_idx: newton_shape_idx}
-                for geom_idx, template_shape_idx in geom_to_shape_idx.items():
-                    # Compute the Newton shape index for the given geom index in the current environment.
-                    # If the shape is a template shape, it will be used in all environments.
-                    # If the shape is not a template shape, it will be used only in the current environment.
-                    if shape_collision_group[template_shape_idx] < 0:
-                        global_shape_idx = template_shape_idx
-                    else:
-                        global_shape_idx = env_idx * shape_range_len + template_shape_idx
-                        if global_shape_idx >= model.shape_count:
-                            print()
-                        assert global_shape_idx < model.shape_count, (
-                            f"Global shape index {global_shape_idx} exceeds model.shape_count {model.shape_count}"
-                        )
-                    full_shape_mapping[global_shape_idx] = (env_idx, geom_idx)
-                    reverse_shape_mapping[(env_idx, geom_idx)] = global_shape_idx
-                    # Update incoming shape transforms
-                    # compute the difference between the original shape transform
-                    # and the transform after applying the joint child transform
-                    # and the transform MuJoCo does on mesh geoms
-                    original_tf = wp.transform(*shape_transform[global_shape_idx])
-                    mjc_p = self.mj_model.geom_pos[geom_idx]
-                    mjc_q = self.mj_model.geom_quat[geom_idx]
-                    mjc_tf = wp.transform(mjc_p, quat_from_mjc(mjc_q))
-                    shape_incoming_xform[global_shape_idx] = mjc_tf * wp.transform_inverse(original_tf)
+        if separate_envs_to_worlds and model.num_envs > 1:
+            full_shape_mapping = {}
+            reverse_shape_mapping = {}
+            # Build mappings from the authoritative shape_map
+            for shape_idx, (worldid, geom_idx) in self.shape_map.items():
+                if geom_idx is None:
+                    continue
+                if worldid == -1:
+                    # replicate across all envs in the reverse map
+                    for env_idx in range(model.num_envs):
+                        reverse_shape_mapping[(env_idx, geom_idx)] = shape_idx
+                    # forward map doesn't depend on env id (geom_idx is the same)
+                    full_shape_mapping[shape_idx] = (0, geom_idx)
+                else:
+                    reverse_shape_mapping[(worldid, geom_idx)] = shape_idx
+                    full_shape_mapping[shape_idx] = (worldid, geom_idx)
+            # Compute per-shape incoming transform once using the compiled geom frame
+            for shape_idx, (_, geom_idx) in full_shape_mapping.items():
+                original_tf = wp.transform(*shape_transform[shape_idx])
+                mjc_p = self.mj_model.geom_pos[geom_idx]
+                mjc_q = self.mj_model.geom_quat[geom_idx]
+                mjc_tf = wp.transform(mjc_p, quat_from_mjc(mjc_q))
+                shape_incoming_xform[shape_idx] = mjc_tf * wp.transform_inverse(original_tf)

This removes the fragile stride arithmetic and the stray print() used for debugging.

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

1839-1840: Typo in comment ("exludes" → "excludes").

Minor wording fix to keep comments clean.

Apply this diff:

-        # number of shapes which are replicated per env (exludes singular static shapes from a negative group)
+        # number of shapes that are replicated per env (excludes singular static shapes from a negative group)

---

1914-1916: Remove unused parameter from add_geoms().

incoming_xform is never used; keep the signature minimal.

Apply this diff:

-        def add_geoms(newton_body_id: int, incoming_xform: wp.transform | None = None):
+        def add_geoms(newton_body_id: int):

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

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

newton/tests/test_mujoco_solver.py (4)

320-327: Same mapping-shape assertion as mass test

The same assumption about to_mjc_body_index applies here; reuse the shape assertions added above to catch regressions early.

---

407-415: Same mapping-shape assertion as mass/COM tests

Ditto for inertia: ensure to_mjc_body_index is 1-D and sized to bodies_per_env.

---

605-621: Repeat cached mapping usage for updated checks

Reuse the cached axis_to_actuator and the same shape assertions to keep consistency and avoid repeated .numpy() calls here as well.

---

824-826: Same mapping and incoming-transform shape validations in update test

Mirror the to_newton_shape_index and shape_incoming_xform shape checks here to keep both tests robust to future changes in solver outputs.

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

914-914: Module ID for kernel should be unique to avoid collisions.

Using module="unique" as the module identifier seems unusual. If this is meant to ensure unique module compilation, consider using a more descriptive name or letting Warp handle module uniqueness automatically.

-@wp.kernel(module="unique", enable_backward=False)
+@wp.kernel(enable_backward=False)

If a unique module is specifically required for performance reasons, consider documenting why and using a more descriptive identifier.

---

305-315: Consider validating worldid bounds.

While the code handles negative worldid by using world_geom_b, there's no validation that the positive worldid values are within bounds of the number of environments.

 worldid = world_geom_a[0]
 if worldid < 0:
     worldid = world_geom_b[0]
+else:
+    # Ensure worldid is valid (though this should be guaranteed by the mapping)
+    assert worldid < nworld_in, f"Invalid worldid {worldid} >= {nworld_in}"

newton/tests/test_mujoco_solver.py (4)

270-276: Assert per-env body mapping shape to prevent silent indexing bugs

You treat to_mjc_body_index as a 1-D, per-env-local mapping (length = bodies_per_env). Add a quick shape check to make failures obvious if the mapping semantics change.

Apply this diff near where the mapping is read:

 bodies_per_env = self.model.body_count // self.model.num_envs
-body_mapping = solver.to_mjc_body_index.numpy()
+body_mapping = solver.to_mjc_body_index.numpy()
+self.assertEqual(
+    body_mapping.ndim, 1, "Expected 1-D per-env body mapping (local body -> MuJoCo body)"
+)
+self.assertEqual(
+    body_mapping.shape[0], bodies_per_env,
+    f"Body mapping length should equal bodies_per_env ({bodies_per_env})"
+)

---

523-539: Cache axis→actuator mapping and validate its shape (avoids per-iteration .numpy() and clarifies semantics)

Calling .numpy() inside the tight loop is avoidable and may allocate repeatedly on some backends. Also, we assume a per-env-local 1-D mapping of length dofs_per_env. Cache it once and assert the expected shape.

Apply this minimal change at the usage site:

-                actuator_idx = solver.mjc_axis_to_actuator.numpy()[axis_idx]
+                actuator_idx = axis_to_actuator[axis_idx]

And add this setup once, just before the verification loops (outside the changed line range):

# Cache and validate axis->actuator mapping
axis_to_actuator = solver.mjc_axis_to_actuator.numpy()
self.assertEqual(axis_to_actuator.ndim, 1, "Expected 1-D per-env axis->actuator mapping")
self.assertEqual(
    axis_to_actuator.shape[0], dofs_per_env,
    f"axis->actuator mapping length should equal dofs_per_env ({dofs_per_env})"
)

---

661-665: Strengthen mapping check: assert shape instead of hasattr

Rather than just checking for the attribute, validate its shape matches [num_envs, ngeom] right after converting to NumPy.

Apply this diff:

-        # Verify to_newton_shape_index mapping exists
-        self.assertTrue(hasattr(solver, "to_newton_shape_index"))
-
-        # Get mappings and arrays
-        to_newton_shape_index = solver.to_newton_shape_index.numpy()
+        # Get mappings and arrays
+        to_newton_shape_index = solver.to_newton_shape_index.numpy()
+        self.assertEqual(
+            to_newton_shape_index.shape, (self.model.num_envs, solver.mj_model.ngeom),
+            "to_newton_shape_index must be shaped [num_envs, ngeom]"
+        )

---

675-679: Validate shape_incoming_xform length against model.shape_count

This guards against accidental per-env duplication or AoS/SoA mismatches.

Apply this diff:

-        shape_incoming_xform = solver.shape_incoming_xform.numpy()
+        shape_incoming_xform = solver.shape_incoming_xform.numpy()
+        self.assertEqual(
+            shape_incoming_xform.shape[0], self.model.shape_count,
+            "shape_incoming_xform should be per-shape (length = model.shape_count)"
+        )

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

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: eric-heiden
PR: newton-physics/newton#587
File: newton/_src/sim/builder.py:656-661
Timestamp: 2025-08-20T18:02:36.042Z
Learning: In Newton physics, collisions between shapes with body index -1 (world-attached/global shapes) are automatically skipped by the collision system, so no manual collision filter pairs need to be added between global shapes from different builders.

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

656-661: Good fix for preserving global shapes during builder merging.

The logic correctly handles global shapes (body == -1) by:

1. Initializing body_shapes[-1] if not present
2. Extending with offset shape indices to preserve shape references

This aligns with the PR objective of fixing shape index mapping for environments.

---

870-874: Good restoration of global shapes after DFS reconstruction.

The logic correctly preserves static/global shapes during the body_shapes reconstruction by:

1. Retrieving static shapes before clearing
2. Restoring them after clearing

This ensures global shapes persist through topological reorganization.

---

3727-3727: Good change to always export environment group arrays.

Removing conditional emission of None and always creating group arrays aligns with the PR objective to improve multi-environment handling and simplifies downstream processing.

Also applies to: 3770-3770, 3925-3925, 3938-3938, 3977-3977

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

993-1028: Well-structured kernel for computing environment-aware shape transforms.

The kernel correctly:

1. Maps template shapes to global indices based on environment
2. Handles static shapes (group < 0) by preserving their index
3. Computes and stores the incoming transform differences

This is a key improvement for the multi-environment shape mapping.

---

1218-1235: Good addition of public fields for environment-aware mappings.

The new fields provide clear access to the internal mappings needed for multi-environment simulations. The docstrings are helpful for understanding each field's purpose.

---

1247-1265: Good use of ScopedTimer for performance monitoring.

Wrapping the expensive model conversion in a timer (even if inactive by default) helps with profiling and debugging performance issues.

---

1756-1765: Efficient array filling from dictionary with tuple support.

The helper function efficiently handles both single integers and tuples as keys/values, which is useful for the multi-dimensional mappings in this PR.

---

1906-1927: Correct selection of shapes/bodies/joints for first environment.

The logic properly identifies entities from the first non-negative group or all negative groups for building the template environment in MuJoCo.

---

1953-1967: Good addition of shape filtering in add_geoms.

The function now correctly:

1. Skips shapes not in the selected environment
2. Validates that planes are only attached to static bodies

This ensures only the appropriate shapes are added to each MuJoCo environment.

---

2254-2257: Good validation of geom count consistency.

The assertion ensures that the number of MuJoCo geoms matches the expected number of colliding shapes, catching potential mapping issues early.

---

2294-2328: Well-implemented shape mapping initialization using kernel.

The code efficiently:

1. Creates the geom-to-shape index mapping
2. Initializes the bidirectional mappings with proper dimensions
3. Launches a kernel to compute incoming transforms in parallel

This is a significant performance improvement over the previous CPU-based approach mentioned in the PR description.

newton/tests/test_mujoco_solver.py (2)

261-262: Docstring wording LGTM

Minor grammar tweak reads fine and matches the intent of the test.

---

846-847: Friction pattern change LGTM

Starting at 1.05 avoids zero/near-zero friction edge cases and keeps monotonicity. The downstream expectations reflect this; good.

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

1913-1919: Wrong fallback inflates per-env shape stride when all shapes are global; use 0

Using shape_range_len = model.shape_count when there are no non-negative groups overestimates per-env stride and breaks global index arithmetic. If all shapes are in negative groups, there are no per-env replicated shapes, so the stride must be 0.

Apply this diff:

-            else:
-                first_group = -1
-                shape_range_len = model.shape_count
+            else:
+                first_group = -1
+                # No per-env replicated shapes; only global (-1) shapes exist
+                shape_range_len = 0

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

1223-1224: Fix docstring: mapping direction is inverted

The docstring says “Mapping from MuJoCo body index to Newton body index”, but the array is used as Newton body index → MuJoCo body index (env-local). This confuses readers and led to a broken construction below.

Apply this diff:

-        """Mapping from MuJoCo body index to Newton body index (skip world body index -1). Shape [bodies_per_env], dtype int32."""
+        """Mapping from Newton body index (env-local) to MuJoCo body index (skip world body index -1). Shape [bodies_per_env], dtype int32."""

---

1902-1905: Track body names for post-compile mapping

Introduce a body_names dictionary to capture the unique names per Newton body id.

Apply this diff:

-        # ensure unique names
+        # ensure unique names
         body_name_counts = {}
         joint_names = {}
+        body_names = {}

---

1018-1021: Misleading comment about “joint child transform”

The corrective transform only accounts for MuJoCo’s mesh compiler frame adjustment (pos/quat), not joint child transforms—which are already in shape_transform.

Apply this diff:

-    # compute the difference between the original shape transform
-    # and the transform after applying the joint child transform
-    # and the transform MuJoCo does on mesh geoms
+    # Compute corrective transform between Newton's shape frame and MuJoCo's compiled geom frame
+    # (accounts for MuJoCo mesh compiler pos/quat adjustment only; joint child transforms are already in shape_transform)

---

1953-1961: Remove unused incoming_xform parameter in add_geoms

incoming_xform is never used; keep the signature lean.

Proposed change (outside diff range): drop the parameter from the closure signature and call sites.

---

224-255: Kernel zeroing may overrun when rigid_contact_count[0] < dim; confirm bounds

You zero collision_hftri_index_out[tid] for all tid, but guard tid >= rigid_contact_count[0] later. Ensure the zeroing arrays are at least dim length, or move zeroing before the early return behind a tid < collision_hftri_index_out.shape[0] check. Given mujoco_warp shapes this is likely safe, but it’s a sharp edge.

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

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

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

2571-2585: Confirm that num_worlds equals self.mjw_model.nworld when available

Hard-coding num_worlds = self.model.num_envs works for the “separate envs to worlds” path but will drift if a future code path sets nworld differently. Prefer reading from self.mjw_model.nworld when the mjwarp model is active.

Would you like me to patch this to: num_worlds = getattr(self.mjw_model, "nworld", self.model.num_envs)?