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

2244-2253: Use existing shape_map to expand across environments instead of modulo arithmetic

The current logic derives a local_shape_idx via template_shape_idx % shapes_per_env and expands to all envs. This presumes uniform shape layouts. Since you already compute self.shape_map earlier with precise (worldid, mj_geom) tuples per Newton shape (including -1 for global shapes), you can build full_shape_mapping by iterating that structure directly, eliminating layout assumptions and edge cases.

Here’s a minimal adjustment using self.shape_map to build full_shape_mapping:

-        full_shape_mapping = {}
-        # `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():
-            # The local index is consistent for a given part of the model across all environments.
-            local_shape_idx = template_shape_idx % shapes_per_env
-            for env_idx in range(model.num_envs):
-                # Calculate the global Newton shape index for the current environment.
-                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] = (env_idx, geom_idx)
+        full_shape_mapping = {}
+        # self.shape_map: {newton_shape_idx: (worldid, mj_geom_idx_or_None)}
+        for shape_idx, (worldid, mj_geom_idx) in self.shape_map.items():
+            if mj_geom_idx is None:
+                continue
+            wid = worldid if worldid != -1 else 0  # replicate globals across envs below
+            if worldid == -1:
+                for env_idx in range(model.num_envs):
+                    full_shape_mapping[shape_idx] = (env_idx, mj_geom_idx)
+            else:
+                full_shape_mapping[shape_idx] = (wid, mj_geom_idx)

This keeps semantics consistent with create_newton_contact_geom_mapping and handles global (-1) shapes uniformly.

---

2297-2317: Fix to_mjc_geom_index construction: remove transient dict assignment & size array by model.shape_count

Short: Confirmed both issues — the code assigns model.to_mjc_geom_index to a dict then overwrites it, and to_mjc_geom_array is sized by len(full_shape_mapping) instead of model.shape_count which can drop high-index shapes.

Files/locations:

• newton/solvers/mujoco/solver_mujoco.py — around lines 2298–2316 (see the two consecutive assignments to model.to_mjc_geom_index and the to_mjc_geom_array construction).

Recommended diff (apply to the block shown in the review):

-            # build the geom index mappings now that we have the actual indices
-            model.to_mjc_geom_index = shape_mapping  # pyright: ignore[reportAttributeAccessIssue]
-            num_shapes = len(list(full_shape_mapping))
+            # build the geom index mappings now that we have the actual indices
+            # Do not assign a transient dict to model.to_mjc_geom_index (keep it as a dense array below).
+            num_shapes = model.shape_count
@@
-            # create mapping from Newton shape index to MuJoCo geom index (for all envs)
-            to_mjc_geom_array = np.full(num_shapes, -1, dtype=np.int32)
-            if len(full_shape_mapping) > 0:
-                for shape_idx, mjc_indices in full_shape_mapping.items():
-                    if shape_idx < len(to_mjc_geom_array):
-                        to_mjc_geom_array[shape_idx] = mjc_indices[1]
-            model.to_mjc_geom_index = wp.array(to_mjc_geom_array, dtype=wp.int32)  # pyright: ignore[reportAttributeAccessIssue]
+            # create mapping from Newton shape index to MuJoCo geom index (for all envs)
+            to_mjc_geom_array = np.full(model.shape_count, -1, dtype=np.int32)
+            if len(full_shape_mapping) > 0:
+                for shape_idx, mjc_indices in full_shape_mapping.items():
+                    to_mjc_geom_array[shape_idx] = mjc_indices[1]
+            model.to_mjc_geom_index = wp.array(to_mjc_geom_array, dtype=wp.int32)  # pyright: ignore[reportAttributeAccessIssue]

Extra suggestion: add an assertion after construction to catch regressions:

• assert model.to_mjc_geom_index.shape[0] == model.shape_count

Reasoning (concise): assigning a dict temporarily can break callers expecting an array; sizing by len(full_shape_mapping) loses slots for shapes that have no MuJoCo geom entry in the mapping and leads to skipped/out-of-range indices.

---

2235-2242: Don't rely on integer-division to expand template shapes — derive full_shape_mapping from self.shape_map

The code currently assumes shapes are evenly partitioned (shapes_per_env = model.shape_count // model.num_envs) and only warns if not divisible. self.shape_map is already built from model.body_shapes and encodes per-shape world ids; use it instead to avoid silent mis-mapping.

Files/locations to fix

• newton/solvers/mujoco/solver_mujoco.py
  • self.shape_map construction: ~lines 2202–2219
  • problematic expansion using shapes_per_env and geom_to_shape_idx: ~lines 2234–2256

Suggested replacement (conceptual diff)

• Replace the shapes_per_env / modulo expansion with a direct iteration over self.shape_map and lookups into the computed template→geom index map (shape_to_geom_idx):

-    shapes_per_env = model.shape_count // model.num_envs
-    ...  # warning + arithmetic expansion using template_shape_idx % shapes_per_env
+    # Build full mapping from the authoritative self.shape_map to avoid assuming even partitions.
+    full_shape_mapping = {}
+    for global_shape_idx, (worldid, template_shape) in self.shape_map.items():
+        # template_shape should correspond to the template/newton shape key used earlier;
+        # use shape_to_geom_idx (created above) to get the actual MuJoCo geom index.
+        mjc_geom_idx = shape_to_geom_idx.get(template_shape, -1)
+        full_shape_mapping[global_shape_idx] = (worldid if worldid >= 0 else 0, mjc_geom_idx)
+        if mjc_geom_idx >= 0:
+            original_tf = wp.transform(*shape_transform[global_shape_idx])
+            mjc_p = self.mj_model.geom_pos[mjc_geom_idx]
+            mjc_q = self.mj_model.geom_quat[mjc_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)

Why this is required

• The current arithmetic approach breaks when global/static shapes exist or when environments have unequal shape counts; a warning is insufficient because downstream code will silently map shapes incorrectly.

Action

• Please replace the arithmetic expansion with the self.shape_map-based approach above (or raise an error instead of warning if uneven partitioning is unacceptable). I can draft the concrete patch if you want.

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

2244-2246: Typo in comment

“environemnt” → “environment”.

-            # geom_to_shape_idx is a mapping from MuJoCo template geom index to Newton shape indices of the last environemnt
+            # geom_to_shape_idx is a mapping from MuJoCo template geom index to Newton shape indices of the last environment

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 settings in your CodeRabbit configuration.

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

2242-2271: Modulo-based expansion can mis-map shapes across envs (static shapes, non-divisible counts).

The expansion via local_shape_idx = template_shape_idx % shapes_per_env assumes:

• shape indices are perfectly chunked per environment; and
• there are no globally-scoped/static shapes breaking divisibility.

If shape_count % num_envs != 0 (e.g., due to worldbody/static shapes), this code can index an unrelated shape (global_shape_idx) and compute a wrong shape_incoming_xform for it. You warn, but still proceed and may write out-of-bounds or incorrect data.

Minimum hardening within this block:

• Guard against out-of-range global_shape_idx.
• Skip expansion when divisibility doesn’t hold (or compute mapping from an explicit per-shape world mapping rather than modulo).

Suggested guard to prevent misindexing:

-                for env_idx in range(model.num_envs):
-                    # Calculate the global Newton shape index for the current environment.
-                    global_shape_idx = env_idx * shapes_per_env + local_shape_idx
+                for env_idx in range(model.num_envs):
+                    # Calculate the global Newton shape index for the current environment.
+                    global_shape_idx = env_idx * shapes_per_env + local_shape_idx
+                    if global_shape_idx < 0 or global_shape_idx >= model.shape_count:
+                        # Skip invalid indices (can happen if shape_count is not divisible by num_envs)
+                        continue

Stronger (recommended) refactor: derive expansion from self.shape_map (already built) instead of modulo. It already carries (worldid, geom_idx) for every Newton shape (including static/worldbody with worldid == -1). That avoids assumptions about shape indexing and divisibility. Example approach (not a drop-in diff; illustrating intent):

full_shape_mapping = {}
for shape_idx, (worldid, geom_idx) in self.shape_map.items():
    if geom_idx is None or geom_idx < 0:
        continue
    # compute incoming transform once per real shape
    original_tf = wp.transform(*shape_transform[shape_idx])
    mjc_tf = wp.transform(self.mj_model.geom_pos[geom_idx], quat_from_mjc(self.mj_model.geom_quat[geom_idx]))
    shape_incoming_xform[shape_idx] = mjc_tf * wp.transform_inverse(original_tf)

    # expand mapping across worlds if static (-1), otherwise just the owning world
    worlds = range(model.num_envs) if worldid == -1 else (worldid,)
    for w in worlds:
        # Note: maintaining downstream expectations may require also adjusting how
        # reverse mapping is constructed (to handle static shapes across all worlds).
        full_shape_mapping[shape_idx] = (w, geom_idx)

If you prefer, I can rewrite this block and the subsequent reverse-mapping construction to use self.shape_map end-to-end.

---

2252-2253: Typo and wording in comment.

• Typo: “environemnt” → “environment”.
• The “last environment” phrasing is misleading; these are “template” geoms.

-            # geom_to_shape_idx is a mapping from MuJoCo template geom index to Newton shape indices of the last environemnt
-            # `geom_to_shape_idx` provides the reverse mapping for the template env: {mj_geom_idx: newton_shape_idx}
+            # geom_to_shape_idx maps MuJoCo template geom index -> Newton shape index (for the template environment)
+            # `geom_to_shape_idx` provides the reverse mapping for the template env: {mj_geom_idx: newton_shape_idx}

---

2242-2261: Please add a regression test that exercises multi-env incoming xform population.

A simple coverage target:

• Build a minimal model with num_envs > 1 including at least one MESH geom.
• Construct SolverMuJoCo with separate_envs_to_worlds=True and False.
• Assert model.shape_incoming_xform is populated for all shapes (not identity for the mesh) in both branches.
• Assert per-env shapes that correspond to the same template geom yield consistent incoming xforms.

I can draft this test in newton/tests (e.g., test_mujoco_incoming_xform.py) if helpful.

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)

2326-2331: to_mjc_geom_index array is sized incorrectly; risk of OOB access in kernels

to_mjc_geom_array is allocated with length num_shapes = len(full_shape_mapping), then indexed by absolute Newton shape_idx. This can cause out-of-bounds or stale -1 entries for valid shape_idx values when used by kernels like convert_newton_contacts_to_mjwarp_kernel, which index by the global Newton shape id.

Allocate it to model.shape_count and drop the guard.

-            # create mapping from Newton shape index to MuJoCo geom index (for all envs)
-            to_mjc_geom_array = np.full(num_shapes, -1, dtype=np.int32)
-            if len(full_shape_mapping) > 0:
-                for shape_idx, mjc_indices in full_shape_mapping.items():
-                    if shape_idx < len(to_mjc_geom_array):
-                        to_mjc_geom_array[shape_idx] = mjc_indices[1]
+            # create mapping from Newton shape index to MuJoCo geom index (for all envs)
+            to_mjc_geom_array = np.full(model.shape_count, -1, dtype=np.int32)
+            if full_shape_mapping:
+                for shape_idx, mjc_indices in full_shape_mapping.items():
+                    to_mjc_geom_array[shape_idx] = mjc_indices[1]

This also aligns with how to_mjc_geom_index[shape_id] is consumed elsewhere.

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

2273-2280: Single-env path now populates incoming transforms — resolves prior regression

This loop mirrors the multi-env computation and fixes the earlier issue where shape_incoming_xform stayed identity for single-env builds. Good catch.

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

2243-2271: Multi-env incoming xforms: logic is correct and addresses the original gap

The expansion via full_shape_mapping and computing shape_incoming_xform[global_shape_idx] for every env is sound. This ensures non-identity incoming transforms for mesh geoms across all environments rather than just the template env.

Minor perf nit: avoid recomputing mjc_tf inside the inner env loop.

-        for geom_idx, template_shape_idx in geom_to_shape_idx.items():
-            # The local index is consistent for a given part of the model across all environments.
-            local_shape_idx = template_shape_idx % shapes_per_env
-            for env_idx in range(model.num_envs):
+        for geom_idx, template_shape_idx in geom_to_shape_idx.items():
+            # The local index is consistent for a given part of the model across all environments.
+            local_shape_idx = template_shape_idx % shapes_per_env
+            # Precompute MuJoCo transform once per geom
+            if geom_idx >= 0:
+                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))
+            for env_idx in range(model.num_envs):
                 # Calculate the global Newton shape index for the current environment.
                 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] = (env_idx, geom_idx)
-                if geom_idx >= 0:
-                    # 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 geom_idx >= 0:
+                    # compute corrective transform from original Newton shape frame to MuJoCo geom frame
+                    original_tf = wp.transform(*shape_transform[global_shape_idx])
+                    shape_incoming_xform[global_shape_idx] = mjc_tf * wp.transform_inverse(original_tf)

---

2252-2253: Fix comment typo and wording

• Typo: “environemnt” -> “environment”.
• “last environment” is misleading; this mapping is built from the template env (typically env0).

-            # geom_to_shape_idx is a mapping from MuJoCo template geom index to Newton shape indices of the last environemnt
-            # `geom_to_shape_idx` provides the reverse mapping for the template env: {mj_geom_idx: newton_shape_idx}
+            # geom_to_shape_idx maps from a MuJoCo template geom index to a Newton shape index in the template environment
+            # (usually env0), i.e., {mj_geom_idx: newton_shape_idx}.

---

2245-2249: Non-divisible shape counts: warn may hide a hard-to-debug mapping bug

If model.shape_count is not divisible by model.num_envs, local_shape_idx = template_shape_idx % shapes_per_env becomes ambiguous and can generate incorrect cross-env mappings. Consider failing fast.

Do you expect non-divisible shape counts in practice? If not, please raise instead of warn:

-            if model.shape_count % model.num_envs != 0:
-                warnings.warn(
-                    f"Total shape count {model.shape_count} is not divisible by number of environments {model.num_envs}. "
-                    "Shape mapping to MuJoCo geoms may be incorrect.",
-                    stacklevel=2,
-                )
+            if model.shape_count % model.num_envs != 0:
+                raise ValueError(
+                    f"Shape count ({model.shape_count}) must be divisible by num_envs ({model.num_envs}) "
+                    "to expand template geom mappings across environments."
+                )

If non-divisible layouts are legitimate, we should derive local_shape_idx from the model’s per-env indexing scheme (e.g., via self.shape_map or body/shape tables) rather than modulo arithmetic. I can help patch that path if needed.

---

2239-2242: Minor: clarify “template” vs “expanded” mappings

shape_mapping is reassigned to the resolved template indices. The subsequent code builds full_shape_mapping for all envs. Consider a short comment clarifying these two layers to avoid confusion during maintenance.

-        shape_mapping = shape_to_geom_idx  # Replace with actual indices
+        # Template mapping: Newton shape (template env) -> MuJoCo geom index
+        shape_mapping = shape_to_geom_idx

---

2243-2280: Please add a regression test covering both multi-env and single-env paths

The reviewer requested a test. Suggested coverage:

• Multi-env (separate_envs_to_worlds=True, num_envs > 1): verify that
  • model.to_newton_shape_index.shape == (num_envs, mj_model.ngeom), and for each template geom, to_newton_shape_index[env, geom] points to the expected per-env shape id.
  • For a mesh geom with a non-identity Newton shape_transform, the corresponding model.shape_incoming_xform[shape_idx] is non-identity and satisfies incoming * shape_tf ≈ mj_geom_tf.
• Single-env (separate_envs_to_worlds=False or num_envs == 1): same assertions for the one env; ensure shape_incoming_xform is populated.

If helpful, I can scaffold a minimal model builder and a test file under newton/tests/ that exercises these paths. Let me know.

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.