diff --git a/dimos/experimental/robot/bosdyn/spot/blueprints/spot.py b/dimos/experimental/robot/bosdyn/spot/blueprints/spot.py index f17478613e..bcb44de199 100644 --- a/dimos/experimental/robot/bosdyn/spot/blueprints/spot.py +++ b/dimos/experimental/robot/bosdyn/spot/blueprints/spot.py @@ -37,14 +37,23 @@ from dimos.core.coordination.blueprints import autoconnect from dimos.core.global_config import global_config from dimos.experimental.robot.bosdyn.spot.effectors.high_level import SpotHighLevel -from dimos.experimental.robot.bosdyn.spot.rerun import spot_camera_layout +from dimos.experimental.robot.bosdyn.spot.rerun import ( + spot_camera_layout, + spot_camera_visual_overrides, +) from dimos.navigation.movement_manager.movement_manager import MovementManager from dimos.visualization.vis_module import vis_module spot = autoconnect( SpotHighLevel.blueprint(), MovementManager.blueprint(), - vis_module(global_config.viewer, rerun_config={"blueprint": spot_camera_layout}), + vis_module( + global_config.viewer, + rerun_config={ + "blueprint": spot_camera_layout, + "visual_override": spot_camera_visual_overrides(), + }, + ), ).remappings( [ # No nav stack here, so MovementManager's goal/way_point/stop_movement diff --git a/dimos/experimental/robot/bosdyn/spot/blueprints/spot_replay.py b/dimos/experimental/robot/bosdyn/spot/blueprints/spot_replay.py new file mode 100644 index 0000000000..3f9fd629a0 --- /dev/null +++ b/dimos/experimental/robot/bosdyn/spot/blueprints/spot_replay.py @@ -0,0 +1,62 @@ +# 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 the latest Spot recording to Rerun — no robot, runs anywhere. + +`SpotReplay` plays a memory2 recording (from `spot-record`) back onto the same +stream names `SpotHighLevel` uses, so the Spot camera layout and per-camera +frustums light up in the Rerun 3D view. The `visual_override` routes the two +shared CameraInfo streams onto each camera's image entity so every frustum +anchors to its optical frame. + +Usage: + # the shared recording, pulled from LFS on first run: + dimos run spot-replay --db-path=spot_small_loop.db + # newest *.db under ~/datasets/spot: + dimos run spot-replay + # a specific recording: + dimos run spot-replay --db-path=/path/to/spot.db +""" + +from __future__ import annotations + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.global_config import global_config +from dimos.experimental.robot.bosdyn.spot.replay import SpotReplay +from dimos.experimental.robot.bosdyn.spot.rerun import ( + spot_body_static_overrides, + spot_camera_layout, + spot_camera_visual_overrides, +) +from dimos.mapping.odometry_path import OdometryPath +from dimos.protocol.pubsub.impl.lcmpubsub import LCM +from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.rerun.websocket_server import RerunWebSocketServer + +# Compose only the Rerun bridge (+ its websocket server) directly instead of the +# shared `vis_module`: replay just needs the 3D viewer, and `vis_module` also +# bundles the WebsocketVisModule, which auto-opens the 7779 Command Center tab. +spot_replay = autoconnect( + SpotReplay.blueprint(), + OdometryPath.blueprint(), + RerunBridgeModule.blueprint( + pubsubs=[LCM()], + rerun_open=global_config.rerun_open, + rerun_web=global_config.rerun_web, + blueprint=spot_camera_layout, + visual_override=spot_camera_visual_overrides(), + static=spot_body_static_overrides(), + ), + RerunWebSocketServer.blueprint(), +).remappings([(OdometryPath, "path", "odom_path")]) diff --git a/dimos/experimental/robot/bosdyn/spot/replay.py b/dimos/experimental/robot/bosdyn/spot/replay.py new file mode 100644 index 0000000000..5085640f3a --- /dev/null +++ b/dimos/experimental/robot/bosdyn/spot/replay.py @@ -0,0 +1,171 @@ +# 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 recorded Spot session to streams, mirroring `SpotHighLevel`'s outputs. + +Opens a memory2 SQLite recording (written by `SpotRecorder`) and replays every +camera, depth, and odometry stream onto Out ports named exactly like +`SpotHighLevel`'s, so the same Rerun visualization wires up by name — no robot +required. The recorded ``tf`` tree (odom->base_link plus the base_link->camera +mounts) is republished so every frame stays spatially anchored in 3D. + +The travelled trail is not built here: ``odometry`` feeds ``OdometryPath``, +which accumulates it and is equally happy on a live robot. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +import math +from pathlib import Path + +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import Out +from dimos.experimental.robot.bosdyn.spot.config import ( + CAMERA_STREAM_SUFFIXES, + FRONT_CAMERA_ROTATE_UPRIGHT, +) +from dimos.memory2.replay import resolve_db_path +from dimos.memory2.store.sqlite import SqliteStore +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo +from dimos.msgs.sensor_msgs.Image import Image +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +# Image Out ports to replay, matching SpotHighLevel / SpotRecorder stream names. +_IMAGE_STREAMS = [ + f"{kind}_image_{suffix}" for kind in ("grayscale", "depth") for suffix in CAMERA_STREAM_SUFFIXES +] +_PLAYBACK_STREAMS = [*_IMAGE_STREAMS, "grayscale_info", "depth_info", "odometry"] + +# SpotHighLevel rights the sideways front camera images but leaves their optical +# tf frames at the raw mount orientation, so the 3D frustum + depth +# back-projection land rotated. Roll each front optical frame about its viewing +# (z) axis to bring it back upright. The two front cameras mount mirror-imaged, +# so frontright sits a half turn (2 quarter turns) past frontleft. +_HALF_TURN_QUARTERS = 2 +_OPTICAL_FRAME_ROLL_TURNS = { + "frontleft_camera_optical": FRONT_CAMERA_ROTATE_UPRIGHT, + "frontright_camera_optical": FRONT_CAMERA_ROTATE_UPRIGHT + _HALF_TURN_QUARTERS, +} + + +class SpotReplayConfig(ModuleConfig): + """Where to read the recording from and how to play it back.""" + + # An on-disk path, or the name of a dataset to pull from LFS (e.g. + # ``spot_small_loop.db``). Empty -> newest ``*.db`` in ``dataset_dir``. + db_path: str = "" + dataset_dir: str = "~/datasets/spot" + + speed: float = 1.0 + loop: bool = True + seek: float | None = None + duration: float | None = None + + # Roll the front optical tf frames upright at replay time. Off by default: + # recordings are expected to already store upright front frames. Turn on to + # view an old recording whose tf still holds the raw sideways mount. + roll_front_frames: bool = False + + +class SpotReplay(Module): + """Replays Spot's fisheye + depth cameras, odometry, and tf from a recording.""" + + config: SpotReplayConfig + dedicated_worker = True + + grayscale_image_front_left: Out[Image] + grayscale_image_front_right: Out[Image] + grayscale_image_left: Out[Image] + grayscale_image_right: Out[Image] + grayscale_image_back: Out[Image] + + depth_image_front_left: Out[Image] + depth_image_front_right: Out[Image] + depth_image_left: Out[Image] + depth_image_right: Out[Image] + depth_image_back: Out[Image] + + grayscale_info: Out[CameraInfo] + depth_info: Out[CameraInfo] + + odometry: Out[Odometry] + tf: Out[TFMessage] + + def _resolve_db_path(self) -> Path: + if self.config.db_path: + return resolve_db_path(Path(self.config.db_path).expanduser()) + directory = Path(self.config.dataset_dir).expanduser() + recordings = sorted(directory.glob("*.db"), key=lambda path: path.stat().st_mtime) + if not recordings: + raise FileNotFoundError(f"No .db recordings found in {directory}") + return recordings[-1] + + def _republish_tf(self, message: TFMessage) -> None: + self.tf.publish( + TFMessage(*(self._roll_optical_frame(transform) for transform in message.transforms)) + ) + + def _roll_optical_frame(self, transform: Transform) -> Transform: + if not self.config.roll_front_frames: + return transform + turns = _OPTICAL_FRAME_ROLL_TURNS.get(transform.child_frame_id) + if not turns: + return transform + roll = Quaternion.from_euler(Vector3(0.0, 0.0, turns * math.pi / 2)) + return Transform( + translation=transform.translation, + rotation=transform.rotation * roll, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ts=transform.ts, + ) + + async def main(self) -> AsyncIterator[None]: + db_path = self._resolve_db_path() + logger.info(f"Replaying Spot recording from {db_path}") + + store = SqliteStore(path=str(db_path), must_exist=True) + store.start() + self.register_disposable(store) + + replay = store.replay( + speed=self.config.speed, + loop=self.config.loop, + seek=self.config.seek, + duration=self.config.duration, + ) + available = set(replay.list_streams()) + + for name in _PLAYBACK_STREAMS: + if name not in available: + logger.warning(f"Spot replay: stream {name!r} missing from recording; skipping") + continue + self.register_disposable( + replay.stream(name).observable().subscribe(getattr(self, name).publish) + ) + + if "tf" in available: + self.register_disposable(replay.stream("tf").observable().subscribe(self._republish_tf)) + else: + logger.warning("Spot replay: no tf stream in recording; 3D frames will be missing") + + yield diff --git a/dimos/experimental/robot/bosdyn/spot/rerun.py b/dimos/experimental/robot/bosdyn/spot/rerun.py index a7ef8ac80f..0f028b571a 100644 --- a/dimos/experimental/robot/bosdyn/spot/rerun.py +++ b/dimos/experimental/robot/bosdyn/spot/rerun.py @@ -16,11 +16,56 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any + +from dimos.protocol.tf.tf import MultiTBuffer if TYPE_CHECKING: + from collections.abc import Callable + import rerun.blueprint as rrb + from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo + from dimos.msgs.tf2_msgs.TFMessage import TFMessage + from dimos.visualization.rerun.bridge import RerunData + +# Optical tf frame_id (SpotHighLevelConfig defaults) -> stream-name suffix. +_OPTICAL_FRAME_TO_SUFFIX = { + "frontleft_camera_optical": "front_left", + "frontright_camera_optical": "front_right", + "left_camera_optical": "left", + "right_camera_optical": "right", + "back_camera_optical": "back", +} + +# The five cameras mount ~0.3 m apart on Spot's body; a full 1.0 m image plane +# makes neighbouring frustums overlap in the 3D view, so draw them shorter. +_FRUSTUM_PLANE_DISTANCE = 0.3 + +# Spot's body is ~1.1 m long, 0.5 m wide, 0.19 m tall and base_link sits at its +# center, so a box that size on the frame stands in for the robot. +_SPOT_BODY_SIZE = (1.1, 0.5, 0.19) +_SPOT_GREEN = (0, 255, 0) + +# Mirrors the tf stream so a camera pose can be looked up at the moment its +# image was captured. Module-level, not an attribute of the override objects: +# those get pickled out to the bridge's worker and would each carry a copy. +_tf_buffer = MultiTBuffer() + + +def _anchor_frame() -> str | None: + """The frame that is a parent but never a child - the one that never moves. + + Everything else in the tf tree hangs off it, so a pose expressed against it + stays put no matter what the robot does afterwards. Discovered rather than + named because the odometry frame is configurable. None while the tree is + still filling in, or if it has more than one root. + """ + parents = {parent for parent, _ in _tf_buffer.buffers} + children = {child for _, child in _tf_buffer.buffers} + roots = parents - children + return roots.pop() if len(roots) == 1 else None + def _grayscale_origin(suffix: str) -> str: return f"world/grayscale_image_{suffix}" @@ -30,10 +75,85 @@ def _depth_origin(suffix: str) -> str: return f"world/depth_image_{suffix}" +def _camera_entity(origin: str) -> str: + """Child entity holding the projection and the pixels, under the camera's pose. + + A Pinhole is a parent link, and rerun gives an entity exactly one — so an + entity carrying a projection cannot also carry the camera's pose. One level + down the pixels inherit the pose logged above them. + """ + return f"{origin}/camera" + + +def _tf_to_rerun(tf_message: TFMessage) -> RerunData: + """Log tf as usual, and keep a copy for the camera pose lookups.""" + _tf_buffer.receive_tfmessage(tf_message) + return tf_message.to_rerun() + + +def _camera_info_pinhole(camera_info: CameraInfo, origin: Callable[[str], str]) -> RerunData | None: + """Re-emit a shared CameraInfo onto its camera's image entity as a Pinhole.""" + suffix = _OPTICAL_FRAME_TO_SUFFIX.get(camera_info.frame_id) + if suffix is None: + return None + return camera_info.to_rerun( + image_plane_distance=_FRUSTUM_PLANE_DISTANCE, + image_topic=_camera_entity(origin(suffix)), + ) + + +# Module-level (not closures) so the RerunBridgeModule config stays picklable +# when it is shipped to its worker process. +def _grayscale_info_to_pinhole(camera_info: CameraInfo) -> RerunData | None: + return _camera_info_pinhole(camera_info, _grayscale_origin) + + +def _depth_info_to_pinhole(camera_info: CameraInfo) -> RerunData | None: + return _camera_info_pinhole(camera_info, _depth_origin) + + +def spot_camera_visual_overrides() -> dict[str, Callable[[Any], RerunData | None]]: + """Give each camera its own frustum, posed where it was at capture time. + + ``grayscale_info``/``depth_info`` are shared across all five cameras, so + their default Pinhole lands on one throwaway entity and no camera renders a + frustum. Route each message onto its camera's entity (picked by the + CameraInfo's optical ``frame_id``) so every image plane gets a projection, + and hand the images a pose of their own — see :class:`_ImageBakedIntoAnchor`. + """ + overrides: dict[str, Callable[[Any], RerunData | None]] = { + "world/tf": _tf_to_rerun, + "world/grayscale_info": _grayscale_info_to_pinhole, + "world/depth_info": _depth_info_to_pinhole, + } + for suffix in _OPTICAL_FRAME_TO_SUFFIX.values(): + for origin in (_grayscale_origin, _depth_origin): + overrides[origin(suffix)] = _ImageBakedIntoAnchor(origin(suffix)) + return overrides + + +def _spot_body_box(rerun_module: Any) -> list[Any]: + """A green box at base_link standing in for Spot's body.""" + return [ + rerun_module.Transform3D(parent_frame="tf#/base_link"), + rerun_module.Boxes3D( + centers=[(0.0, 0.0, 0.0)], + sizes=[_SPOT_BODY_SIZE], + colors=[_SPOT_GREEN], + fill_mode="solid", + ), + ] + + +def spot_body_static_overrides() -> dict[str, Callable[[Any], Any]]: + """Draw a green body box anchored to the moving base_link frame.""" + return {"world/spot_body": _spot_body_box} + + def _camera_view(origin: str, name: str) -> rrb.Spatial2DView: import rerun.blueprint as rrb - return rrb.Spatial2DView(origin=origin, name=name) + return rrb.Spatial2DView(origin=_camera_entity(origin), name=name) def spot_camera_layout() -> rrb.Blueprint: @@ -96,3 +216,35 @@ def spot_camera_layout() -> rrb.Blueprint: ), collapse_panels=True, ) + + +# This is a correction to depth jitter +# Depth starts in the right location then immediately moves slightly to the wrong location because of the odom updates +# Odom is faster than depth hz, but depth is anchored dynamically to odom +# The solution is to anchor the image to the frame *at a specific time* +# Probably hasn't been noticed earlier because its less obvious for color image and for 60/30zh +# we should fix this in our rerun bridge but I'm not going to do that in this spot PR +class _ImageBakedIntoAnchor: + def __init__(self, origin: str) -> None: + self.origin = origin + + def __call__(self, image: Any) -> RerunData | None: + anchor = _anchor_frame() + if anchor is None: + return None + pose = _tf_buffer.get(anchor, image.frame_id, image.ts) + if pose is None: + return None + import rerun as rr + + return [ + ( + self.origin, + rr.Transform3D( + translation=[pose.translation.x, pose.translation.y, pose.translation.z], + rotation=pose.rotation.to_rerun(), + parent_frame=f"tf#/{anchor}", + ), + ), + (_camera_entity(self.origin), image.to_rerun()), + ] diff --git a/dimos/mapping/odometry_path.py b/dimos/mapping/odometry_path.py new file mode 100644 index 0000000000..da71de1eaf --- /dev/null +++ b/dimos/mapping/odometry_path.py @@ -0,0 +1,74 @@ +# Copyright 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. + +"""Accumulate an odometry stream into the path it has travelled.""" + +from __future__ import annotations + +from collections import deque +import math +from typing import Any + +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.nav_msgs.Path import Path + + +class OdometryPathConfig(ModuleConfig): + # Empty follows the odometry's own frame_id. + frame_id: str = "" + min_step_meters: float = 0.02 + max_poses: int = 20000 + min_publish_interval_seconds: float = 0.1 + + +class OdometryPath(Module): + """``odometry`` in, the trail it has drawn out, as a ``nav_msgs/Path``.""" + + config: OdometryPathConfig + + odometry: In[Odometry] + + path: Out[Path] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._poses: deque[PoseStamped] = deque(maxlen=self.config.max_poses) + self._last_publish_ts = 0.0 + + async def handle_odometry(self, msg: Odometry) -> None: + position = msg.pose.position + point = (position.x, position.y, position.z) + if self._poses: + previous = self._poses[-1] + if math.dist((previous.x, previous.y, previous.z), point) < self.config.min_step_meters: + return + + orientation = msg.pose.orientation + frame_id = self.config.frame_id or msg.frame_id + self._poses.append( + PoseStamped( + ts=msg.ts, + frame_id=frame_id, + position=list(point), + orientation=[orientation.x, orientation.y, orientation.z, orientation.w], + ) + ) + if msg.ts - self._last_publish_ts < self.config.min_publish_interval_seconds: + return + self._last_publish_ts = msg.ts + # A copy: Path holds the list by reference. + self.path.publish(Path(ts=msg.ts, frame_id=frame_id, poses=list(self._poses))) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index b8693374a0..621ab0a888 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -92,6 +92,7 @@ "openyam-planner-coordinator": "dimos.robot.manipulators.openyam.blueprints.basic:openyam_planner_coordinator", "spot": "dimos.experimental.robot.bosdyn.spot.blueprints.spot:spot", "spot-record": "dimos.experimental.robot.bosdyn.spot.blueprints.spot_record:spot_record", + "spot-replay": "dimos.experimental.robot.bosdyn.spot.blueprints.spot_replay:spot_replay", "teleop-hosted-go2-multicam": "dimos.teleop.hosted.blueprints.cloudflare:teleop_hosted_go2_multicam", "teleop-hosted-go2-transport": "dimos.teleop.hosted.blueprints.cloudflare:teleop_hosted_go2_transport", "teleop-hosted-xarm6": "dimos.teleop.hosted.blueprints.cloudflare:teleop_hosted_xarm6", @@ -249,6 +250,7 @@ "object-tracker2-d": "dimos.perception.experimental.object_tracker_2d.ObjectTracker2D", "object-tracker3-d": "dimos.perception.experimental.object_tracker_3d.ObjectTracker3D", "object-tracking": "dimos.perception.experimental.object_tracker.ObjectTracking", + "odometry-path": "dimos.mapping.odometry_path.OdometryPath", "osm-skill": "dimos.agents.skills.osm.OsmSkill", "path-follower": "dimos.navigation.cmu_nav.modules.path_follower.path_follower.PathFollower", "path-following-coordinator": "dimos.control.path_following_coordinator.PathFollowingCoordinator",