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

511-552: Rename kernel parameter for clarity and consistency.

convert_mjw_contact_to_warp_kernel still names the first parameter contact_geom_mapping, but you now pass to_newton_shape_index. Renaming improves readability and reduces future confusion.

Apply this diff:

 @wp.kernel
 def convert_mjw_contact_to_warp_kernel(
     # inputs
-    contact_geom_mapping: wp.array2d(dtype=wp.int32),
+    to_newton_shape_index: wp.array2d(dtype=wp.int32),
     pyramidal_cone: bool,
@@
-    for i in range(2):
-        pair[i] = contact_geom_mapping[worldid, geoms_mjw[i]]
+    for i in range(2):
+        pair[i] = to_newton_shape_index[worldid, geoms_mjw[i]]

---

2014-2017: Storing geom names and resolving indices post-compile — sensible decoupling.

Writing body.add_geom(**geom_params) and deferring to mj_name2id avoids relying on speculative indices. The subsequent name→id resolution (Lines 2274-2284) is robust given the unique naming scheme f"{geom_type_name[stype]}_{shape}".

Optionally assert uniqueness at creation time to catch accidental collisions early:

-                name = f"{geom_type_name[stype]}_{shape}"
+                name = f"{geom_type_name[stype]}_{shape}"
+                # Defensive: ensure name is unique
+                if mujoco.mj_name2id(spec, mujoco.mjtObj.mjOBJ_GEOM, name) >= 0:
+                    raise ValueError(f"Duplicate geom name detected: {name}")

Configuration used: Path: .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)

1617-1636: Switched to to_newton_shape_index for contact mapping — correct direction.

Passing self.to_newton_shape_index into convert_mjw_contact_to_warp_kernel aligns update_contacts with the new post-compile mapping. Shapes: the kernel expects a 2D int32 map [world, geom] → shape, which matches how to_newton_shape_index is constructed at Lines 2295-2321. No functional concerns here.

---

1006-1028: Assumption: template → global shape indexing via env-stride may be invalid if indices aren’t block-strided by group.

The global_shape_idx = env_idx * shape_range_len + template_shape_idx formula assumes per-env shapes occupy contiguous blocks with identical intra-env ordering. If shape indices are interleaved or not strictly block-strided, reverse mapping becomes incorrect.

Consider asserting this precondition right after computing selected_shapes (Line 1937) in convert_to_mjc. Example check (debug-only):

@@
-        self.selected_shapes = wp.array(selected_shapes, dtype=wp.int32, device=model.device)
+        self.selected_shapes = wp.array(selected_shapes, dtype=wp.int32, device=model.device)
+        # Debug: verify per-env shapes are block-strided and consistent
+        if separate_envs_to_worlds:
+            non_neg = np.where(shape_group >= 0)[0]
+            if len(non_neg):
+                first_group = np.min(shape_group[non_neg])
+                per_env = np.where(shape_group == first_group)[0]
+                stride = len(per_env)
+                for env_idx in range(1, model.num_envs):
+                    g = first_group + env_idx
+                    env_shapes = np.where(shape_group == g)[0]
+                    if len(env_shapes) != stride:
+                        raise AssertionError("Per-env shape counts differ; mapping by fixed stride is invalid")
+                    # Optional: verify relative ordering matches template indices

If this assertion fails in real models, we should pass an explicit “template index → per-env global index” lookup into update_incoming_shape_xform_kernel instead of deriving it arithmetically.

---

2539-2570: Use the actual number of worlds from mjwarp data when launching geom-updates.

update_geom_properties uses num_worlds = self.model.num_envs. If separate_envs_to_worlds=False, put_data(..., nworld=1) creates 2D geom arrays with first dim 1, but self.model.num_envs could be >1, leading to OOB writes.

Consider deriving num_worlds from the target arrays or from self.mjw_data.nworld:

-        num_worlds = self.model.num_envs  # why is there no 'self.mjw_model.nworld'?
+        num_worlds = int(self.mjw_data.nworld) if hasattr(self, "mjw_data") else self.model.num_envs

If mjw_data isn’t available (CPU path), ensure this path isn’t invoked or adjust accordingly.

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

1617-1630: Contact mapping switch is correct; add precondition checks and align kernel arg name.

Good change passing to_newton_shape_index into the contact conversion. Two follow-ups:

• Add a guard to fail fast if the mapping wasn’t initialized (e.g., when constructing the solver with pre-existing mjw_model/mjw_data).
• Minor naming nit: the kernel’s parameter is still called contact_geom_mapping, which now conceptually represents to_newton_shape_index. Renaming improves readability.

Apply the precondition check close to the launch:

         wp.launch(
             convert_mjw_contact_to_warp_kernel,
             dim=mj_data.nconmax,
             inputs=[
-                self.to_newton_shape_index,
+                # Preconditions
+                # Ensure mapping exists and matches current MjModel ngeom.
+                # This is especially important if the solver was constructed
+                # with pre-existing mjw_model/mjw_data and skipped convert_to_mjc().
+                (None if self.to_newton_shape_index is not None else (_ for _ in ()).throw(
+                    RuntimeError("to_newton_shape_index is None; mapping not initialized. Call convert_to_mjc or provide mapping.")
+                )),
+                (None if self.to_newton_shape_index is not None and
+                        self.to_newton_shape_index.shape[1] == self.mj_model.ngeom else (_ for _ in ()).throw(
+                    ValueError(f"to_newton_shape_index shape mismatch; expected second dim {self.mj_model.ngeom}")
+                )),
+                self.to_newton_shape_index,
                 self.mjw_model.opt.cone == int(self.mujoco.mjtCone.mjCONE_PYRAMIDAL),
                 mj_data.ncon,
                 mj_contact.frame,
                 mj_contact.dim,
                 mj_contact.geom,
                 mj_contact.efc_address,
                 mj_contact.worldid,
                 mj_data.efc.force,
             ],

If you prefer a rename for clarity (outside this hunk), consider updating the kernel signature and its uses:

- def convert_mjw_contact_to_warp_kernel(
-     contact_geom_mapping: wp.array2d(dtype=wp.int32),
+ def convert_mjw_contact_to_warp_kernel(
+     to_newton_shape_index: wp.array2d(dtype=wp.int32),
      pyramidal_cone: bool,
      ...
 )
@@
-    for i in range(2):
-        pair[i] = contact_geom_mapping[worldid, geoms_mjw[i]]
+    for i in range(2):
+        pair[i] = to_newton_shape_index[worldid, geoms_mjw[i]]

---

1979-1984: Avoid in-place mutation of shape_size when zero-filling.

size = shape_size[shape] returns a view; zero-filling will mutate model.shape_scale unexpectedly. Copy before patching zeros.

-                size = shape_size[shape]
+                # copy to avoid mutating model.shape_scale in-place
+                size = np.array(shape_size[shape], copy=True)
                 if np.any(size > 0.0):
                     # duplicate nonzero entries at places where size is 0
                     nonzero = size[size > 0.0][0]
                     size[size == 0.0] = nonzero
                     geom_params["size"] = size

Configuration used: Path: .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: preist-nvidia
PR: newton-physics/newton#579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, Example.use_mujoco is a high-level attribute that controls whether to use MuJoCo solver vs other solvers (like XPBD), while SolverMuJoCo.use_mujoco_cpu controls the backend within MuJoCo (CPU vs Warp). These are separate concepts serving different purposes - the PR rename only applies to the solver parameter, not the Example class attributes.

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

2014-2017: Recording geoms by name and resolving indices post-compile: solid improvement.

Storing shape -> geom_name at creation and resolving to indices post-compile removes fragility around assumed ordering and fixes the contact shape mapping issue targeted by this PR. The approach fits well with MuJoCo’s name lookup and your later two-way mapping builds.

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

2014-2017: Name-based geom bookkeeping is safer; mapping computed post-compile is correct.

Using body.add_geom(**geom_params) and storing the geom name, then resolving indices with mj_name2id post-compile, avoids brittle index assumptions during spec construction. The uniqueness guarantees above make this robust.

Minor nit:

• The local shape_mapping now holds names during spec-build and indices afterwards; consider renaming the temporary (e.g., shape_to_geom_name) to reduce cognitive load while reading this function. No functional change.

---

1616-1629: Good fix: use the canonical geom→shape map in contacts.

Switching the first kernel arg to self.to_newton_shape_index aligns contact-pair reporting with the mapping used elsewhere (built in update_incoming_shape_xform_kernel). Shapes/dtypes match (int32, [nworld, ngeom]), so this resolves the incorrect contact sensor mapping.

✅ Verified: no leftover references to the old mapping helpers in solver_mujoco.py. Searches for create_newton_contact_geom_mapping and mj_geoms returned no matches, and contact_geom_mapping only appears in the kernel’s parameter list and its use sites as intended. All shape_map hits are confined to the USD-import utility (import_usd.py) and are unrelated to this change.

Two small follow-ups:

• Consider renaming the kernel’s parameter from contact_geom_mapping to geom_to_shape_index for clarity, since it now directly receives to_newton_shape_index. This is a naming-only change.
• Optional guard: if any pair[i] ends up -1, you may want to skip/zero that contact to avoid downstream consumers tripping on invalid shape ids (rare, but can occur with visual-only or filtered geoms if assumptions drift).

Configuration used: Path: .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: preist-nvidia
PR: newton-physics/newton#579
File: newton/examples/example_mujoco.py:350-354
Timestamp: 2025-08-19T12:27:30.629Z
Learning: In Newton examples, Example.use_mujoco is a high-level attribute that controls whether to use MuJoCo solver vs other solvers (like XPBD), while SolverMuJoCo.use_mujoco_cpu controls the backend within MuJoCo (CPU vs Warp). These are separate concepts serving different purposes - the PR rename only applies to the solver parameter, not the Example class attributes.