diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 6b71f2adf..c8d1cf875 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/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. diff --git a/embodichain/lab/visualization/scene_exporter.py b/embodichain/lab/visualization/scene_exporter.py index cf77446b8..31d12c027 100644 --- a/embodichain/lab/visualization/scene_exporter.py +++ b/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), diff --git a/tests/visualization/test_scene_exporter.py b/tests/visualization/test_scene_exporter.py index 8e295fffb..ef94f03d5 100644 --- a/tests/visualization/test_scene_exporter.py +++ b/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( diff --git a/tests/visualization/test_viser_backend.py b/tests/visualization/test_viser_backend.py index b8af652f3..aac24485c 100644 --- a/tests/visualization/test_viser_backend.py +++ b/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)