Fix Viser rendering for simulation axis markers

embodichain/lab/sim/sim_manager.py

@@ -213,6 +213,16 @@ class _WindowRecordState:
     loop_handle: object | None = None
 
 
+@dataclass(frozen=True)
+class _AxisMarkerGroup:
+    """Native axis handles and their backend-neutral display dimensions."""
+
+    handles: tuple[MeshObject, ...]
+    arena_index: int
+    axis_length: float
+    axis_radius: float
+
+
 class SimulationManager:
     r"""Global Embodied AI simulation manager.
 
@@ -326,7 +336,7 @@ def __init__(
         self._gizmos: Dict[str, object] = dict()  # Store active gizmos
 
         # marker management
-        self._markers: Dict[str, MeshObject] = dict()
+        self._markers: dict[str, _AxisMarkerGroup] = {}
 
         self._rigid_objects: Dict[str, RigidObject] = dict()
         self._constraints: Dict[str, RigidConstraint] = dict()
@@ -2369,7 +2379,12 @@ def draw_marker(
         #     # Create point markers
         #     pass
 
-        self._markers[name] = (marker_handles, cfg.arena_index)
+        self._markers[name] = _AxisMarkerGroup(
+            handles=tuple(marker_handles),
+            arena_index=cfg.arena_index,
+            axis_length=cfg.axis_len,
+            axis_radius=cfg.axis_size,
+        )
 
         if self.is_physics_manually_update:
             self.update(step=1)
@@ -2388,9 +2403,9 @@ def remove_marker(self, name: str) -> bool:
             logger.log_warning(f"Marker {name} not found.")
             return False
         try:
-            env = self.get_env(self._markers[name][1])
-            marker_handles, arena_index = self._markers[name]
-            for marker_handle in marker_handles:
+            marker_group = self._markers[name]
+            env = self.get_env(marker_group.arena_index)
+            for marker_handle in marker_group.handles:
                 if marker_handle is not None:
                     env.remove_actor(marker_handle.get_name())
             self._markers.pop(name)
@@ -2399,6 +2414,25 @@ def remove_marker(self, name: str) -> bool:
             logger.log_warning(f"Failed to remove marker {name}: {str(e)}")
             return False
 
+    def get_axis_marker_items(
+        self,
+    ) -> tuple[tuple[str, tuple[MeshObject, ...], float, float], ...]:
+        """Return active axes for backend-neutral visualization.
+
+        Returns:
+            Tuples containing the marker name, native handles, axis length, and
+            axis radius for each active marker group.
+        """
+        return tuple(
+            (
+                name,
+                group.handles,
+                group.axis_length,
+                group.axis_radius,
+            )
+            for name, group in self._markers.items()
+        )
+
     def add_custom_window_control(self, controls: list[ObjectManipulator]) -> None:
         """Add one or more custom window input controls.
 

embodichain/lab/visualization/scene_exporter.py

@@ -32,6 +32,7 @@
     CameraImageFrame,
     CameraSpec,
     DynamicMeshUpdate,
+    FrameOverlay,
     GizmoSpec,
     GizmoState,
     JointControlProvider,
@@ -702,9 +703,54 @@ def _append_articulations(
                         )
                     )
 
+    def _capture_axis_marker_overlays(
+        self, reserved_frame_ids: set[str] | None = None
+    ) -> tuple[FrameOverlay, ...]:
+        """Capture native simulation axes as Viser coordinate-frame overlays.
+
+        Args:
+            reserved_frame_ids: Caller-owned frame IDs that generated markers
+                must not replace.
+        """
+        get_axis_marker_items = getattr(self._sim, "get_axis_marker_items", None)
+        if get_axis_marker_items is None:
+            return ()
+
+        frames: list[FrameOverlay] = []
+        used_frame_ids = set(reserved_frame_ids or ())
+        for marker_name, handles, axis_length, axis_radius in get_axis_marker_items():
+            for index, handle in enumerate(handles):
+                position, wxyz = pose_to_position_wxyz(handle.get_world_pose())
+                base_id = f"marker:{marker_name}:{index}"
+                overlay_id = base_id
+                suffix = 1
+                while overlay_id in used_frame_ids:
+                    overlay_id = f"{base_id}#{suffix}"
+                    suffix += 1
+                used_frame_ids.add(overlay_id)
+                frames.append(
+                    FrameOverlay(
+                        overlay_id=overlay_id,
+                        position=position,
+                        wxyz=wxyz,
+                        axes_length=axis_length,
+                        axes_radius=axis_radius,
+                        # Native handles report hidden in headless mode even
+                        # though draw_marker() requested a visible marker.
+                        visible=True,
+                    )
+                )
+        return tuple(frames)
+
     def _prepare_overlays(self, overlays: SceneOverlays | None) -> SceneOverlays:
+        reserved_frame_ids = (
+            {frame.overlay_id for frame in overlays.frames}
+            if overlays is not None
+            else None
+        )
+        marker_frames = self._capture_axis_marker_overlays(reserved_frame_ids)
         if overlays is None:
-            return SceneOverlays()
+            return SceneOverlays(frames=marker_frames)
         point_clouds: list[PointCloudOverlay] = []
         for point_cloud in overlays.point_clouds:
             point_count = point_cloud.points.shape[0]
@@ -732,7 +778,7 @@ def _prepare_overlays(self, overlays: SceneOverlays | None) -> SceneOverlays:
                 )
             )
         return SceneOverlays(
-            frames=overlays.frames,
+            frames=marker_frames + overlays.frames,
             trajectories=overlays.trajectories,
             targets=overlays.targets,
             point_clouds=tuple(point_clouds),

tests/visualization/test_scene_exporter.py

@@ -21,6 +21,7 @@
 import numpy as np
 
 from embodichain.lab.visualization import (
+    FrameOverlay,
     JointControlSpec,
     JointControlState,
     PointCloudOverlay,
@@ -342,6 +343,30 @@ def get_gizmo_items(self) -> tuple[tuple[str, _Gizmo], ...]:
         return (("cube", _Gizmo()),)
 
 
+class _AxisMarkerHandle:
+    def __init__(self, pose: np.ndarray, visible: bool = False) -> None:
+        self._pose = pose
+        self._visible = visible
+
+    def get_world_pose(self) -> np.ndarray:
+        return self._pose
+
+    def is_visible(self) -> bool:
+        return self._visible
+
+
+class _AxisMarkerSimulation(_EmptySimulation):
+    def __init__(self) -> None:
+        pose = np.eye(4, dtype=np.float32)
+        pose[:3, 3] = [-0.4, 0.48, 0.1]
+        self.marker = _AxisMarkerHandle(pose)
+
+    def get_axis_marker_items(
+        self,
+    ) -> tuple[tuple[str, tuple[_AxisMarkerHandle, ...], float, float], ...]:
+        return (("place_target_axis", (self.marker,), 0.2, 0.01),)
+
+
 class _Camera:
     def __init__(self, visualization_role: str = "sensor") -> None:
         self.cfg = SimpleNamespace(
@@ -567,6 +592,54 @@ def test_gizmo_manifest_and_authoritative_pose_are_exported() -> None:
     np.testing.assert_allclose(result.frame.gizmos[0].position, [0.2, 0.3, 0.4])
 
 
+def test_axis_markers_ignore_headless_native_visibility() -> None:
+    exporter = SceneExporter(
+        _AxisMarkerSimulation(),
+        VisualizationCfg(backend="viser"),
+        run_id="axis-marker-run",
+    )
+
+    exporter.build_manifest()
+    result = exporter.capture(sim_step=1, sim_time=0.01)
+
+    assert len(result.frame.overlays.frames) == 1
+    overlay = result.frame.overlays.frames[0]
+    assert overlay.overlay_id == "marker:place_target_axis:0"
+    np.testing.assert_allclose(overlay.position, [-0.4, 0.48, 0.1])
+    np.testing.assert_allclose(overlay.wxyz, [1.0, 0.0, 0.0, 0.0])
+    assert overlay.axes_length == 0.2
+    assert overlay.axes_radius == 0.01
+    assert overlay.visible
+
+
+def test_axis_marker_id_does_not_collide_with_caller_frame() -> None:
+    exporter = SceneExporter(
+        _AxisMarkerSimulation(),
+        VisualizationCfg(backend="viser"),
+        run_id="axis-marker-collision-run",
+    )
+    caller_frame = FrameOverlay(
+        overlay_id="marker:place_target_axis:0",
+        position=np.array([0.1, 0.2, 0.3], dtype=np.float32),
+        wxyz=np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+    )
+
+    exporter.build_manifest()
+    result = exporter.capture(
+        sim_step=1,
+        sim_time=0.01,
+        overlays=SceneOverlays(frames=(caller_frame,)),
+    )
+
+    frames = result.frame.overlays.frames
+    assert [frame.overlay_id for frame in frames] == [
+        "marker:place_target_axis:0#1",
+        "marker:place_target_axis:0",
+    ]
+    np.testing.assert_allclose(frames[0].position, [-0.4, 0.48, 0.1])
+    np.testing.assert_allclose(frames[1].position, caller_frame.position)
+
+
 def test_camera_frustum_pose_and_low_frequency_rgb_are_exported() -> None:
     simulation = _CameraSimulation()
     exporter = SceneExporter(

tests/visualization/test_viser_backend.py

@@ -25,6 +25,7 @@
     CameraImageFrame,
     CameraSpec,
     DynamicMeshUpdate,
+    FrameOverlay,
     GizmoSpec,
     GizmoState,
     JointControlSpec,
@@ -33,6 +34,7 @@
     SceneFrame,
     SceneManifest,
     SceneNode,
+    SceneOverlays,
     ViserServerCfg,
 )
 from embodichain.lab.visualization.backends.viser import ViserBackend
@@ -222,6 +224,7 @@ def __init__(self) -> None:
         self.mesh_handles: list[_Handle] = []
         self.dynamic_mesh_handles: list[_Handle] = []
         self.camera_handles: list[_Handle] = []
+        self.frame_handles: list[_Handle] = []
         self.grid_handles: list[_Handle] = []
         self.transform_controls: list[_TransformControls] = []
 
@@ -233,7 +236,9 @@ def set_up_direction(self, direction: str) -> None:
 
     def add_frame(self, name: str, **kwargs: object) -> _Handle:
         kwargs.setdefault("visible", True)
-        return _Handle(name=name, **kwargs)
+        handle = _Handle(name=name, removed=False, **kwargs)
+        self.frame_handles.append(handle)
+        return handle
 
     def add_batched_meshes_simple(self, name: str, **kwargs: object) -> _Handle:
         self.mesh_uploads += 1
@@ -321,6 +326,47 @@ def test_viser_backend_adds_one_meter_default_ground_grid() -> None:
     backend.stop()
 
 
+def test_viser_backend_renders_axis_marker_frame_overlay() -> None:
+    server = _Server()
+    backend = ViserBackend(ViserServerCfg(port=8765), server_factory=lambda **_: server)
+    marker = FrameOverlay(
+        overlay_id="marker:place_target_axis:0",
+        position=np.array([-0.4, 0.48, 0.1], dtype=np.float32),
+        wxyz=np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32),
+        axes_length=0.2,
+        axes_radius=0.01,
+    )
+    frame = SceneFrame(
+        run_id="run",
+        scene_revision=1,
+        sequence=1,
+        sim_step=1,
+        sim_time=0.01,
+        node_ids=(),
+        positions=np.empty((0, 3), dtype=np.float32),
+        wxyz=np.empty((0, 4), dtype=np.float32),
+        visible=np.empty((0,), dtype=np.bool_),
+        overlays=SceneOverlays(frames=(marker,)),
+    )
+
+    backend.start()
+    backend.publish_manifest(SceneManifest("run", 1, (), ()))
+    assert backend.publish_frame(frame)
+
+    handle = next(
+        handle
+        for handle in server.scene.frame_handles
+        if handle.name.startswith("/overlays/frames/")
+    )
+    assert handle.name == "/overlays/frames/marker%3Aplace_target_axis%3A0"
+    assert handle.axes_length == 0.2
+    assert handle.axes_radius == 0.01
+    np.testing.assert_allclose(handle.position, marker.position)
+    np.testing.assert_allclose(handle.wxyz, marker.wxyz)
+    assert handle.visible
+    backend.stop()
+
+
 def test_viser_backend_uploads_static_mesh_once_and_updates_only_poses() -> None:
     server = _Server()
     backend = ViserBackend(ViserServerCfg(port=8765), server_factory=lambda **_: server)