diff --git a/dimos/perception/test_worldbelief_replay.py b/dimos/perception/test_worldbelief_replay.py new file mode 100644 index 0000000000..8e1ae5a6fe --- /dev/null +++ b/dimos/perception/test_worldbelief_replay.py @@ -0,0 +1,50 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +import pytest + +from dimos.perception.worldbelief_replay import _resolve_recording_path + + +def test_resolve_explicit_worldbelief_db(tmp_path: Path) -> None: + recording = tmp_path / "recording.db" + recording.touch() + + assert _resolve_recording_path(recording) == recording + + +def test_resolve_worldbelief_dataset_directory(tmp_path: Path) -> None: + recording = tmp_path / "session.db" + recording.touch() + + assert _resolve_recording_path(tmp_path) == recording + + +@pytest.mark.parametrize("count", [0, 2]) +def test_worldbelief_dataset_directory_requires_one_db(tmp_path: Path, count: int) -> None: + for index in range(count): + (tmp_path / f"session_{index}.db").touch() + + with pytest.raises(ValueError, match="exactly one .db"): + _resolve_recording_path(tmp_path) + + +def test_worldbelief_replay_rejects_non_db_file(tmp_path: Path) -> None: + recording = tmp_path / "recording.mcap" + recording.touch() + + with pytest.raises(ValueError, match=r"expects a Memory2 \.db recording"): + _resolve_recording_path(recording) diff --git a/dimos/perception/worldbelief_replay.py b/dimos/perception/worldbelief_replay.py new file mode 100644 index 0000000000..638d4e62e5 --- /dev/null +++ b/dimos/perception/worldbelief_replay.py @@ -0,0 +1,326 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Replay a WorldBelief recording through the same live streams as the physical stack.""" + +from __future__ import annotations + +from pathlib import Path +import threading +from typing import Any, ClassVar + +from pydantic import Field +from reactivex.disposable import CompositeDisposable + +from dimos.agents.annotation import skill +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import Out +from dimos.memory2.store.sqlite import SqliteStore +from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.sensor_msgs.JointState import JointState +from dimos.utils.data import get_data +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +_REQUIRED_STREAMS = frozenset({"color_image", "depth_image", "camera_info", "tf"}) +_REPLAY_ORDER = ( + "tf", + "camera_info", + "depth_camera_info", + "coordinator_joint_state", + "depth_image", + "color_image", +) +_RESUME_EPSILON_S = 1e-6 + + +def _resolve_recording_path(dataset: str | Path) -> Path: + """Resolve an explicit DB or a named LFS dataset containing exactly one DB.""" + requested = Path(dataset).expanduser() + resolved = requested if requested.is_absolute() or requested.exists() else get_data(requested) + if resolved.is_dir(): + databases = sorted(resolved.glob("*.db")) + if len(databases) != 1: + raise ValueError( + f"WorldBelief replay directory must contain exactly one .db file: {resolved} " + f"(found {len(databases)})" + ) + resolved = databases[0] + if resolved.suffix != ".db": + raise ValueError(f"WorldBelief replay expects a Memory2 .db recording, got: {resolved}") + return resolved + + +class _WorldBeliefReplaySourceConfig(ModuleConfig): + dataset: str | Path = Field(default_factory=lambda data: data["g"].replay_db) + speed: float = Field(default=1.0, gt=0.0) + seek: float = Field(default=0.0, ge=0.0) + duration: float | None = Field(default=None, gt=0.0) + autoplay: bool = True + + +class _WorldBeliefReplaySource(Module): + """Publish a recorded RGB-D/TF session with its original relative timing.""" + + config: _WorldBeliefReplaySourceConfig + dedicated_worker: ClassVar[bool] = True + + color_image: Out[Image] + depth_image: Out[Image] + camera_info: Out[CameraInfo] + depth_camera_info: Out[CameraInfo] + coordinator_joint_state: Out[JointState] + + @rpc + def build(self) -> None: + self._recording_path = _resolve_recording_path(self.config.dataset) + logger.info("WorldBelief replay source: %s", self._recording_path) + + @rpc + def start(self) -> None: + super().start() + self._lock = threading.RLock() + self._playback: CompositeDisposable | None = None + self._playback_generation = 0 + self._started_once = False + self._error: str | None = None + + path = getattr(self, "_recording_path", None) + if path is None: + path = _resolve_recording_path(self.config.dataset) + self._recording_path = path + + self._store = self.register_disposable( + SqliteStore(path=str(path), must_exist=True), + ) + self._store.start() + available = set(self._store.list_streams()) + missing = sorted(_REQUIRED_STREAMS - available) + if missing: + raise RuntimeError( + f"WorldBelief replay recording is missing required stream(s): {missing}; " + f"available: {sorted(available)}" + ) + stream_names: list[str] = [] + first_timestamps: list[float] = [] + last_timestamps: list[float] = [] + for name in _REPLAY_ORDER: + if name not in available: + continue + stream = self._store.stream(name) + try: + first_timestamps.append(float(stream.first().ts)) + last_timestamps.append(float(stream.last().ts)) + except LookupError: + if name in _REQUIRED_STREAMS: + raise RuntimeError( + f"WorldBelief replay required stream is empty: {name}" + ) from None + logger.warning("Skipping empty optional replay stream: %s", name) + continue + stream_names.append(name) + self._stream_names = tuple(stream_names) + + self._timeline_start_ts = min(first_timestamps) + self._timeline_end_ts = max(last_timestamps) + self._window_start_ts = self._timeline_start_ts + self.config.seek + if self._window_start_ts >= self._timeline_end_ts: + raise ValueError( + f"Replay seek {self.config.seek:.3f}s is outside the " + f"{self._timeline_end_ts - self._timeline_start_ts:.3f}s recording" + ) + self._window_end_ts = self._timeline_end_ts + if self.config.duration is not None: + self._window_end_ts = min( + self._window_end_ts, + self._window_start_ts + self.config.duration, + ) + self._cursor_ts = self._window_start_ts + self._active_streams: set[str] = set() + self._state = "ready" + logger.info( + "WorldBelief replay ready: %.3fs..%.3fs at %.2fx", + self._window_start_ts - self._timeline_start_ts, + self._window_end_ts - self._timeline_start_ts, + self.config.speed, + ) + + @rpc + def on_system_modules(self, _modules: list[Any]) -> None: + """Autoplay only after Recorder, WorldBelief, Rerun, and MCP have started.""" + if self.config.autoplay: + self._begin_playback() + + def _message_timestamp(self, msg: Any) -> float | None: + ts = getattr(msg, "ts", None) + if ts is not None: + return float(ts) + transforms = getattr(msg, "transforms", None) + if transforms: + return max(float(transform.ts) for transform in transforms) + return None + + def _publish(self, name: str, msg: Any, generation: int) -> None: + with self._lock: + if generation != self._playback_generation or self._state != "playing": + return + ts = self._message_timestamp(msg) + if ts is not None: + self._cursor_ts = max(self._cursor_ts, min(ts, self._window_end_ts)) + + if name == "tf": + self.tf.publish(*msg.transforms) + else: + getattr(self, name).publish(msg) + + def _stream_completed(self, name: str, generation: int) -> None: + with self._lock: + if generation != self._playback_generation: + return + self._active_streams.discard(name) + if self._active_streams: + return + self._cursor_ts = self._window_end_ts + self._state = "ended" + self._playback = None + logger.info("WorldBelief replay ended") + + def _stream_failed(self, name: str, error: Exception, generation: int) -> None: + with self._lock: + if generation != self._playback_generation: + return + self._error = f"{name}: {error}" + self._state = "error" + self._active_streams.clear() + playback = self._playback + self._playback = None + if playback is not None: + playback.dispose() + logger.error("WorldBelief replay failed on %s: %s", name, error) + + def _begin_playback(self) -> None: + with self._lock: + if self._state == "playing": + return + if self._state == "error": + raise RuntimeError(f"WorldBelief replay is in error: {self._error}") + start_ts = self._cursor_ts + if self._started_once: + start_ts += _RESUME_EPSILON_S + if start_ts >= self._window_end_ts: + self._cursor_ts = self._window_end_ts + self._state = "ended" + return + + duration = self._window_end_ts - start_ts + replay = self._store.replay( + speed=self.config.speed, + from_timestamp=start_ts, + duration=duration, + ) + self._playback_generation += 1 + generation = self._playback_generation + self._active_streams = set(self._stream_names) + self._state = "playing" + self._error = None + self._started_once = True + playback = CompositeDisposable() + self._playback = playback + + for name in self._stream_names: + disposable = ( + replay.stream(name) + .observable() + .subscribe( + on_next=lambda msg, stream=name: self._publish(stream, msg, generation), + on_error=lambda error, stream=name: self._stream_failed( + stream, error, generation + ), + on_completed=lambda stream=name: self._stream_completed(stream, generation), + ) + ) + playback.add(disposable) + + def _status_locked(self) -> dict[str, Any]: + timeline_start = getattr(self, "_timeline_start_ts", 0.0) + timeline_end = getattr(self, "_timeline_end_ts", timeline_start) + window_start = getattr(self, "_window_start_ts", timeline_start) + window_end = getattr(self, "_window_end_ts", timeline_end) + cursor = getattr(self, "_cursor_ts", window_start) + window_duration = max(0.0, window_end - window_start) + window_progress = 0.0 + if window_duration > 0.0: + window_progress = min(1.0, max(0.0, (cursor - window_start) / window_duration)) + return { + "state": getattr(self, "_state", "not_started"), + "dataset": str(self.config.dataset), + "recording": str(getattr(self, "_recording_path", "")), + "position_s": round(max(0.0, cursor - timeline_start), 3), + "recording_duration_s": round(max(0.0, timeline_end - timeline_start), 3), + "window_start_s": round(max(0.0, window_start - timeline_start), 3), + "window_end_s": round(max(0.0, window_end - timeline_start), 3), + "window_progress": round(window_progress, 4), + "speed": self.config.speed, + "autoplay": self.config.autoplay, + "error": getattr(self, "_error", None), + } + + @skill + def replay_status(self) -> dict[str, Any]: + """Report replay position; WorldBelief scan sees evidence through this point.""" + with self._lock: + return self._status_locked() + + @skill + def replay_pause(self) -> dict[str, Any]: + """Pause replay without resetting Recorder or WorldBelief state.""" + with self._lock: + if self._state == "playing": + self._playback_generation += 1 + playback = self._playback + self._playback = None + self._active_streams.clear() + self._state = "paused" + else: + playback = None + if playback is not None: + playback.dispose() + with self._lock: + return self._status_locked() + + @skill + def replay_resume(self) -> dict[str, Any]: + """Resume replay from the last published timestamp.""" + self._begin_playback() + with self._lock: + return self._status_locked() + + @rpc + def stop(self) -> None: + lock = getattr(self, "_lock", None) + if lock is not None: + with lock: + self._playback_generation += 1 + playback = self._playback + self._playback = None + getattr(self, "_active_streams", set()).clear() + self._state = "stopped" + else: + playback = None + if playback is not None: + playback.dispose() + super().stop() diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 05702246a7..38d06bf0f5 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -156,6 +156,9 @@ "xarm-perception-sim-agent": "dimos.robot.manipulators.xarm.blueprints.agentic:xarm_perception_sim_agent", "xarm6-planner-only": "dimos.robot.manipulators.xarm.blueprints.basic:xarm6_planner_only", "xarm6-worldbelief": "dimos.robot.manipulators.xarm.blueprints.worldbelief:xarm6_worldbelief", + "xarm6-worldbelief-replay": "dimos.robot.manipulators.xarm.blueprints.worldbelief:xarm6_worldbelief_replay", + "xarm6-worldbelief-replay-kitchen": "dimos.robot.manipulators.xarm.blueprints.worldbelief:xarm6_worldbelief_replay_kitchen", + "xarm6-worldbelief-replay-stationery": "dimos.robot.manipulators.xarm.blueprints.worldbelief:xarm6_worldbelief_replay_stationery", "xarm7-planner-coordinator": "dimos.robot.manipulators.xarm.blueprints.basic:xarm7_planner_coordinator", "xarm7-planner-coordinator-agent": "dimos.robot.manipulators.xarm.blueprints.agentic:xarm7_planner_coordinator_agent", } diff --git a/dimos/robot/manipulators/xarm/blueprints/worldbelief.py b/dimos/robot/manipulators/xarm/blueprints/worldbelief.py index 73d8d8f8fa..83647e4e04 100644 --- a/dimos/robot/manipulators/xarm/blueprints/worldbelief.py +++ b/dimos/robot/manipulators/xarm/blueprints/worldbelief.py @@ -12,18 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""xArm6 WorldBelief perception stack.""" +"""xArm6 WorldBelief perception and recording-replay stacks.""" from __future__ import annotations from functools import partial +from pathlib import Path from typing import Any, cast import rerun.blueprint as rrb from dimos.agents.mcp.mcp_server import McpServer from dimos.constants import STATE_DIR -from dimos.core.coordination.blueprints import autoconnect +from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.hardware.sensors.camera.realsense.camera import RealSenseCamera from dimos.manipulation.manipulation_module import ManipulationModule from dimos.msgs.geometry_msgs.Quaternion import Quaternion @@ -31,6 +32,7 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.perception.worldbelief_module import WorldBeliefModule from dimos.perception.worldbelief_recorder import WorldBeliefRecorder +from dimos.perception.worldbelief_replay import _WorldBeliefReplaySource from dimos.robot.manipulators.common.blueprints import coordinator, trajectory_task from dimos.robot.manipulators.xarm.config import make_xarm6_model_config, xarm6_hardware from dimos.visualization.rerun.bridge import RerunBridgeModule @@ -69,29 +71,8 @@ def _rerun_blueprint() -> rrb.Blueprint: ) -_hw = xarm6_hardware("arm") -_hw.auto_enable = True - -xarm6_worldbelief = autoconnect( - # Provides wrist-camera FK/TF. - ManipulationModule.blueprint( - robots=[ - make_xarm6_model_config( - name="arm", - add_gripper=False, - # Enables TF publication. - tf_extra_links=["link_base"], - ), - ], - ), - RealSenseCamera.blueprint( - width=640, - height=480, - fps=15, - base_frame_id="link6", - base_transform=XARM6_WORLDBELIEF_CAMERA_TRANSFORM, - ), - RerunBridgeModule.blueprint( +def _rerun_bridge() -> Blueprint: + return RerunBridgeModule.blueprint( blueprint=_rerun_blueprint, topic_to_entity=_topic_to_entity, visual_override={ @@ -108,19 +89,52 @@ def _rerun_blueprint() -> rrb.Blueprint: "world/detections_3d": 10.0, "world/pointcloud": 5.0, }, - ), - WorldBeliefRecorder.blueprint( - db_path=STATE_DIR / "worldbelief" / "xarm6" / "recordings" / "xarm6_worldbelief.db", - ), - WorldBeliefModule.blueprint( - db_path=STATE_DIR / "worldbelief" / "xarm6" / "recordings" / "xarm6_worldbelief.db", - history_path=STATE_DIR / "worldbelief" / "xarm6" / "worldbelief_history.db", + ) + + +def _worldbelief_module(db_path: Path, history_path: Path) -> Blueprint: + return WorldBeliefModule.blueprint( + db_path=db_path, + history_path=history_path, scan_prompts=[], depth_tolerance_s=0.1, stationary_hz=4.0, yoloe_model_name="yoloe-11l-seg.pt", dino_model_name="facebook/dinov2-base", clip_model_name="openai/clip-vit-base-patch32", + ) + + +_hw = xarm6_hardware("arm") +_hw.auto_enable = True + +_xarm6_state = STATE_DIR / "worldbelief" / "xarm6" +_xarm6_recording_base = _xarm6_state / "recordings" / "xarm6_worldbelief.db" + +xarm6_worldbelief = autoconnect( + # Provides wrist-camera FK/TF. + ManipulationModule.blueprint( + robots=[ + make_xarm6_model_config( + name="arm", + add_gripper=False, + # Enables TF publication. + tf_extra_links=["link_base"], + ), + ], + ), + RealSenseCamera.blueprint( + width=640, + height=480, + fps=15, + base_frame_id="link6", + base_transform=XARM6_WORLDBELIEF_CAMERA_TRANSFORM, + ), + _rerun_bridge(), + WorldBeliefRecorder.blueprint(db_path=_xarm6_recording_base), + _worldbelief_module( + _xarm6_recording_base, + _xarm6_state / "worldbelief_history.db", ), McpServer.blueprint(), coordinator( @@ -128,3 +142,36 @@ def _rerun_blueprint() -> rrb.Blueprint: tasks=[trajectory_task(_hw)], ), ).global_config(n_workers=8) + + +def _xarm6_worldbelief_replay(dataset: str | Path | None, state_name: str) -> Blueprint: + """Replace only the physical sensor/TF source; keep the live belief stack.""" + state = _xarm6_state / "replay" / state_name + recording_base = state / "recordings" / "xarm6_worldbelief.db" + source_kwargs: dict[str, Any] = {"instance_name": "worldbelief_replay_source"} + if dataset is not None: + source_kwargs["dataset"] = dataset + + return autoconnect( + _rerun_bridge(), + WorldBeliefRecorder.blueprint(db_path=recording_base), + _worldbelief_module(recording_base, state / "worldbelief_history.db"), + McpServer.blueprint(), + _WorldBeliefReplaySource.blueprint(**source_kwargs), + ) + + +xarm6_worldbelief_replay = _xarm6_worldbelief_replay( + None, + "custom", +).global_config(n_workers=8) + +xarm6_worldbelief_replay_kitchen = _xarm6_worldbelief_replay( + "xarm6_worldbelief_realsense_d435i_kitchen", + "kitchen", +).global_config(n_workers=8) + +xarm6_worldbelief_replay_stationery = _xarm6_worldbelief_replay( + "xarm6_worldbelief_realsense_d435i_stationery", + "stationery", +).global_config(n_workers=8)