Raise on ambiguous heterogeneous morph access

Author: kunni918

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -1951,6 +1951,16 @@ def n_links(self):
         """The number of `RigidLink` in the entity."""
         return len(self._links)
 
+    @property
+    def morph(self):
+        """The morph of the entity."""
+        if self._enable_heterogeneous:
+            gs.raise_exception(
+                "Heterogeneous entities have multiple morph variants. Use `.morphs` for all variants, "
+                "or `.main_morph` only when explicitly using the first variant."
+            )
+        return self._morph
+
     @property
     def main_morph(self):
         """The main morph of the entity (first morph for heterogeneous entities)."""

genesis/ext/pyrender/overlay/plugin.py

@@ -373,7 +373,12 @@ def _capture_pending_entities_kwargs(self):
         editor panel. Keyed by entity name; values are the kwargs forwarded to ``scene.add_entity``."""
         self._pending_entities_kwargs = {}
         for entity in self.scene.entities:
-            kwargs: dict[str, Any] = {"morph": entity.morph}
+            # Heterogeneous entities raise on .morph; capture all variants so the rebuild reproduces them.
+            if isinstance(entity, gs.engine.entities.KinematicEntity) and entity._enable_heterogeneous:
+                morph = tuple(entity.morphs)
+            else:
+                morph = entity.morph
+            kwargs: dict[str, Any] = {"morph": morph}
             # Carry the material and surface so a rebuild preserves the entity's solver (e.g. a Kinematic entity must
             # not silently become Rigid, the add_entity default). visualize_contact is rigid-only.
             if isinstance(entity, gs.engine.entities.KinematicEntity):

genesis/vis/rasterizer_context.py

@@ -446,8 +446,8 @@ def on_rigid(self):
                     # For z-axis normal planes, render a single instance shared across all envs to avoid z-fighting,
                     # unless they do not overlap.
                     env_shared = not self.env_separate_rigid
-                    if not env_shared and isinstance(entity.morph, gs.morphs.Plane):
-                        plane_normal, plane_size = entity.morph.normal, entity.morph.plane_size
+                    if not env_shared and isinstance(entity.main_morph, gs.morphs.Plane):
+                        plane_normal, plane_size = entity.main_morph.normal, entity.main_morph.plane_size
                         if (
                             abs(plane_normal[0]) < gs.EPS
                             and abs(plane_normal[1]) < gs.EPS
@@ -565,8 +565,8 @@ def update_rigid(self):
                     geom_T = geoms_T[geom.idx][geom_envs_idx]
 
                     # Keep single-instance for z-axis normal planes (see on_rigid)
-                    if isinstance(entity.morph, gs.morphs.Plane):
-                        plane_normal, plane_size = entity.morph.normal, entity.morph.plane_size
+                    if isinstance(entity.main_morph, gs.morphs.Plane):
+                        plane_normal, plane_size = entity.main_morph.normal, entity.main_morph.plane_size
                         if (
                             abs(plane_normal[0]) < gs.EPS
                             and abs(plane_normal[1]) < gs.EPS

tests/test_imgui_overlay.py

@@ -20,6 +20,45 @@
     _IMGUI_BUNDLE_AVAILABLE = False
 
 
+@pytest.mark.required
+def test_imgui_overlay_capture_pending_entities_preserves_heterogeneous_morphs():
+    scene = gs.Scene(show_viewer=False)
+    single_morph = gs.morphs.Box(size=(0.1, 0.1, 0.1))
+    single_entity = scene.add_entity(
+        morph=single_morph,
+        visualize_contact=True,
+        name="single",
+    )
+    heterogeneous_morphs = (
+        gs.morphs.Box(size=(0.2, 0.2, 0.2)),
+        gs.morphs.Cylinder(radius=0.05, height=0.2),
+    )
+    heterogeneous_entity = scene.add_entity(
+        morph=heterogeneous_morphs,
+        visualize_contact=True,
+        name="heterogeneous",
+    )
+
+    plugin = ImGuiOverlayPlugin.__new__(ImGuiOverlayPlugin)
+    plugin.scene = scene
+
+    plugin._capture_pending_entities_kwargs()
+
+    single_kwargs = plugin._pending_entities_kwargs[single_entity.name]
+    assert single_kwargs["morph"] is single_morph
+    assert single_kwargs["material"] is single_entity.material
+    assert single_kwargs["surface"] is single_entity.surface
+    assert single_kwargs["visualize_contact"] is True
+
+    heterogeneous_kwargs = plugin._pending_entities_kwargs[heterogeneous_entity.name]
+    assert heterogeneous_kwargs["morph"] == heterogeneous_morphs
+    assert heterogeneous_kwargs["morph"][0] is heterogeneous_morphs[0]
+    assert heterogeneous_kwargs["morph"][1] is heterogeneous_morphs[1]
+    assert heterogeneous_kwargs["material"] is heterogeneous_entity.material
+    assert heterogeneous_kwargs["surface"] is heterogeneous_entity.surface
+    assert heterogeneous_kwargs["visualize_contact"] is True
+
+
 def _apply_deterministic_imgui_overrides(monkeypatch):
     """Make ImGui rendering and timing pixel-identical across renderers for snapshot tests."""
     from imgui_bundle import imgui

tests/test_render.py

@@ -79,6 +79,42 @@ def skip_if_not_installed(renderer_type):
             pytest.skip(SKIP_NO_LUISA)
 
 
+@pytest.mark.required
+@pytest.mark.parametrize("renderer_type", [RENDERER_TYPE.RASTERIZER])
+def test_rasterizer_context_heterogeneous_main_morph_no_throw(monkeypatch, renderer):
+    from genesis.vis.rasterizer import Rasterizer
+
+    def _skip_offscreen_renderer_build(self):
+        self.visualizer = self._context.visualizer
+
+    monkeypatch.setattr(Rasterizer, "build", _skip_offscreen_renderer_build)
+
+    scene = gs.Scene(
+        renderer=renderer,
+        vis_options=gs.options.VisOptions(
+            env_separate_rigid=True,
+            shadow=False,
+        ),
+        show_viewer=False,
+    )
+    scene.add_entity(
+        morph=(
+            gs.morphs.Box(size=(0.1, 0.1, 0.1)),
+            gs.morphs.Cylinder(radius=0.05, height=0.2),
+        ),
+    )
+    scene.add_entity(
+        morph=(
+            gs.morphs.Box(size=(0.2, 0.2, 0.2)),
+            gs.morphs.Sphere(radius=0.1),
+        ),
+        material=gs.materials.Kinematic(),
+    )
+
+    scene.build(n_envs=2)
+    scene.visualizer.context.update(force_render=True)
+
+
 @pytest.mark.required
 @pytest.mark.parametrize(
     "renderer_type",

tests/test_rigid_physics.py

@@ -6998,6 +6998,45 @@ def test_heterogeneous_invalid_material_raises():
 
 
 @pytest.mark.slow  # ~200s
+@pytest.mark.required
+def test_heterogeneous_morph_property_raises():
+    """Test that heterogeneous entities require explicit access through morphs."""
+    scene = gs.Scene(show_viewer=False)
+
+    single_morph = gs.morphs.Box(size=(0.1, 0.1, 0.1))
+    single_obj = scene.add_entity(morph=single_morph)
+
+    rigid_morphs_heterogeneous = (
+        gs.morphs.Box(size=(0.1, 0.1, 0.1)),
+        gs.morphs.Cylinder(radius=0.05, height=0.2),
+    )
+    rigid_obj = scene.add_entity(morph=rigid_morphs_heterogeneous)
+    kinematic_morphs_heterogeneous = (
+        gs.morphs.Box(size=(0.2, 0.2, 0.2)),
+        gs.morphs.Sphere(radius=0.1),
+    )
+    kinematic_obj = scene.add_entity(
+        morph=kinematic_morphs_heterogeneous,
+        material=gs.materials.Kinematic(),
+    )
+
+    assert single_obj.morph is single_morph
+    assert rigid_obj.main_morph is rigid_morphs_heterogeneous[0]
+    assert list(rigid_obj.morphs) == list(rigid_morphs_heterogeneous)
+    with pytest.raises(gs.GenesisException, match=r"Heterogeneous.*\.morphs") as exc_info:
+        _ = rigid_obj.morph
+    assert ".main_morph" in str(exc_info.value)
+
+    assert kinematic_obj.main_morph is kinematic_morphs_heterogeneous[0]
+    assert list(kinematic_obj.morphs) == list(kinematic_morphs_heterogeneous)
+    with pytest.raises(gs.GenesisException, match=r"Heterogeneous.*\.morphs"):
+        _ = kinematic_obj.morph
+
+    repr(rigid_obj)
+    repr(kinematic_obj)
+    repr(scene.entities)
+
+
 @pytest.mark.required
 def test_heterogeneous_fewer_envs_than_variants():
     """Test that having fewer environments than variants works correctly.