From 1eb6ccd90c156ce35915472cf634f27ebab5a018 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 3 Aug 2026 23:44:05 -0700 Subject: [PATCH 01/13] Show transforms --- .../nav_3d/mls_planner/utils/plan_rrd.py | 67 ++-- .../nav_3d/mls_planner/utils/test_plan_rrd.py | 73 +++++ dimos/protocol/tf/tf.py | 18 ++ .../go2/blueprints/basic/unitree_go2_basic.py | 4 +- .../navigation/unitree_go2_nav_3d.py | 35 +- dimos/visualization/rerun/bridge.py | 45 ++- dimos/visualization/rerun/test_tf_tree.py | 302 ++++++++++++++++++ dimos/visualization/rerun/tf_tree.py | 257 +++++++++++++++ dimos/visualization/vis_module.py | 2 +- 9 files changed, 717 insertions(+), 86 deletions(-) create mode 100644 dimos/navigation/nav_3d/mls_planner/utils/test_plan_rrd.py create mode 100644 dimos/visualization/rerun/test_tf_tree.py create mode 100644 dimos/visualization/rerun/tf_tree.py diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index ffea4626b9..88b828faf1 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -21,7 +21,7 @@ from pathlib import Path as FsPath from time import perf_counter -from typing import TYPE_CHECKING, NamedTuple +from typing import TYPE_CHECKING, Any, NamedTuple import numpy as np from numpy.typing import NDArray @@ -37,19 +37,18 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation +from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.navigation.nav_3d.mls_planner.mls_planner import MLSPlanner from dimos.navigation.tf_pose import base_height_above_ground from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT, ROBOT_LENGTH, ROBOT_WIDTH from dimos.utils.data import resolve_named_path +from dimos.visualization.rerun.tf_tree import RerunTFTree if TYPE_CHECKING: import rerun.blueprint as rrb TIMELINE = "ts" -AXIS_LEN = 0.5 -AXIS_RADIUS_RATIO = 25 - # Mount frames as recorded on the tf stream. BASE_FRAME = "base_link" SENSOR_FRAME = "mid360_link" @@ -137,6 +136,33 @@ def _log_path_wp(waypoints: NDArray[np.float32] | None, entity: str, color: list rr.log(entity, rr.LineStrips3D([points], colors=[color], radii=0.05)) +def _window(stream: Any, from_time: float | None, to_time: float | None) -> Any: + """Clip a stream to the replay window, both bounds relative to its start.""" + if from_time is not None: + stream = stream.from_time(from_time) + if to_time is not None: + stream = stream.to_time(to_time) + return stream + + +def _tf_over(store: SqliteStore, window: Any) -> Any: + """The recorded tf stream clipped to another stream's time span. + + Absolute bounds, because the relative ones anchor on each stream's own + first observation and tf rarely starts on the same sample as the lidar. + Returns None when the recording has no tf, which must not be probed for + with ``store.stream``: that registers the stream and writes its tables. + """ + if "tf" not in store.list_streams(): + print("no tf stream in the recording; skipping the tf tree") + return None + try: + first, last = window.first().ts, window.last().ts + except LookupError: + return None + return store.stream("tf", TFMessage).order_by("ts").time_range(first, last) + + def _base_from_sensor(store: SqliteStore) -> Transform | None: """Sensor to robot base link transform from the recorded tf stream.""" tf = StreamTF.from_store(store) @@ -323,6 +349,12 @@ def _blueprint(crop: LocalCrop) -> rrb.Blueprint: origin="world", name="world", contents=["+ $origin/**", "- $origin/local/**"], + # The graph buries the map it was built from. Tick it back on in + # the viewer when the question is why a path went the way it did. + overrides={ + "world/nodes": rrb.EntityBehavior(visible=False), + "world/node_edges": rrb.EntityBehavior(visible=False), + }, ), rrb.Vertical( rrb.Spatial3DView( @@ -547,12 +579,9 @@ def main( store = SqliteStore(path=str(db_path)) with store: - lidar = store.stream(lidar_stream, PointCloud2).order_by("ts") - if from_time is not None: - lidar = lidar.from_time(from_time) - if to_time is not None: - lidar = lidar.to_time(to_time) + lidar = _window(store.stream(lidar_stream, PointCloud2).order_by("ts"), from_time, to_time) odom = store.stream(odom_stream, Odometry).order_by("ts") + tf = _tf_over(store, lidar) pose_tagged = lidar.align(odom, tolerance=align_tol).transform( FnTransformer(_attach_pose_from_odom) @@ -570,6 +599,8 @@ def main( support_min=support_min, ) ) + if tf is not None: + ray_pipeline = ray_pipeline.transform(RerunTFTree(tf)) configs = _parse_configs(config, wall_clearance, wall_buffer, wall_buffer_weight) ref_clearance = configs[0][0] @@ -592,24 +623,6 @@ def main( if base_from_sensor is not None else 0.0 ) - entities = ["world/mid360_link/axes"] + ( - ["world/base_link/axes"] if base_from_sensor else [] - ) - for entity in entities: - rr.log( - entity, - rr.Arrows3D( - origins=[[0.0, 0.0, 0.0]] * 3, - vectors=[ - [AXIS_LEN, 0.0, 0.0], - [0.0, AXIS_LEN, 0.0], - [0.0, 0.0, AXIS_LEN], - ], - colors=[[255, 0, 0], [0, 255, 0], [0, 0, 255]], - radii=AXIS_LEN / AXIS_RADIUS_RATIO, - ), - static=True, - ) if base_from_sensor is not None: rr.log( "world/base_link/outline", diff --git a/dimos/navigation/nav_3d/mls_planner/utils/test_plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/test_plan_rrd.py new file mode 100644 index 0000000000..876e20a0f1 --- /dev/null +++ b/dimos/navigation/nav_3d/mls_planner/utils/test_plan_rrd.py @@ -0,0 +1,73 @@ +# 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. + +"""Picking the tf stream out of a recording without disturbing it.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from dimos.memory2.store.sqlite import SqliteStore +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.navigation.nav_3d.mls_planner.utils.plan_rrd import _tf_over + +if TYPE_CHECKING: + from pathlib import Path + + +def _store(tmp_path: Path, tf_stamps: list[float] | None, lidar_stamps: list[float]) -> SqliteStore: + store = SqliteStore(path=str(tmp_path / "rec.db")) + lidar = store.stream("pointlio_lidar", str) + for ts in lidar_stamps: + lidar.append("cloud", ts=ts) + if tf_stamps is not None: + tf = store.stream("tf", TFMessage) + for ts in tf_stamps: + tf.append( + TFMessage(Transform(frame_id="odom", child_frame_id="base_link", ts=ts)), ts=ts + ) + return store + + +@pytest.mark.skipif_macos +@pytest.mark.skipif_aarch64 +def test_missing_tf_is_not_created_by_looking_for_it(tmp_path: Path) -> None: + """Probing with ``store.stream`` would register tf and write its tables.""" + with _store(tmp_path, tf_stamps=None, lidar_stamps=[1.0, 2.0]) as store: + assert _tf_over(store, store.stream("pointlio_lidar", str)) is None + assert "tf" not in store.list_streams() + + +@pytest.mark.skipif_macos +@pytest.mark.skipif_aarch64 +def test_tf_window_follows_the_lidar_window(tmp_path: Path) -> None: + """tf starts 10s before the lidar here, so a relative window would be shifted.""" + with _store( + tmp_path, tf_stamps=[90.0, 100.0, 105.0, 110.0, 115.0], lidar_stamps=[100.0, 110.0] + ) as store: + tf = _tf_over(store, store.stream("pointlio_lidar", str)) + + assert [obs.ts for obs in tf] == [100.0, 105.0, 110.0] + + +@pytest.mark.skipif_macos +@pytest.mark.skipif_aarch64 +def test_empty_lidar_window_has_no_tf(tmp_path: Path) -> None: + with _store(tmp_path, tf_stamps=[1.0], lidar_stamps=[1.0]) as store: + empty = store.stream("pointlio_lidar", str).after(500.0) + + assert _tf_over(store, empty) is None diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 2b0309e52e..9b9f0f5cc1 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -117,6 +117,24 @@ def get_frames(self) -> set[str]: frames.add(child) return frames + def latest_transforms(self) -> list[Transform]: + """Most recent transform on each edge of the tree.""" + with self._cv: + latest = [buffer.last() for buffer in self.buffers.values()] + return [transform for transform in latest if transform is not None] + + def get_parent(self, frame_id: str) -> str | None: + """Parent of a frame, or None if it is a root of the tree. + + A frame with more than one parent is a graph, not a tree. The first + parent seen wins so the answer stays stable once given. + """ + with self._cv: + for parent, child in self.buffers: + if child == frame_id: + return parent + return None + def get_connections(self, frame_id: str) -> set[str]: """Get all frames connected to the given frame (both as parent and child).""" connections = set() diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index c342c74f92..3540c49287 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -97,9 +97,9 @@ def _go2_rerun_blueprint() -> Any: "world/global_costmap": 0, # publishes at ~7.6 Hz "world/lidar": 1, # publishes at ~7.7 Hz; hidden by default in the blueprint }, - # slapping a go2 shaped box on top of tf/base_link + # slapping a go2 shaped box on the base_link frame "static": { - "world/tf/base_link": _static_base_link, + "world/robot_body": _static_base_link, }, } diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 2c03b51142..4f67155344 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -46,11 +46,6 @@ # Raise above 0 to draw what the planner searched over (surface, nodes, weighted edges). planner_viz_hz = 0.0 -# Body-frame axis-triad length (m). -_axis_len = 0.5 -# Arrow radius as a fraction of the triad length. -_AXIS_RADIUS_RATIO = 25 - class Go2Mid360Recorder(PointlioRecorder): lidar_l1: In[PointCloud2] @@ -104,30 +99,6 @@ def _static_robot_body(rr: Any) -> list[Any]: ] -def _axis_triad(rr: Any) -> Any: - """XYZ axis triad, red/green/blue for x/y/z.""" - return rr.Arrows3D( - origins=[[0.0, 0.0, 0.0]] * 3, - vectors=[ - [_axis_len, 0.0, 0.0], - [0.0, _axis_len, 0.0], - [0.0, 0.0, _axis_len], - ], - colors=[[255, 0, 0], [0, 255, 0], [0, 0, 255]], - radii=_axis_len / _AXIS_RADIUS_RATIO, - ) - - -def _static_body_axes(rr: Any) -> Any: - """XYZ triad on the robot body (child of the box).""" - return _axis_triad(rr) - - -def _static_sensor_axes(rr: Any) -> list[Any]: - """XYZ triad on pointlio's raw sensor frame, tilted by the lidar pitch.""" - return [_axis_triad(rr), rr.Transform3D(parent_frame="tf#/mid360_link")] - - _nav_rerun_config = { **rerun_config, "max_hz": { @@ -138,12 +109,10 @@ def _static_sensor_axes(rr: Any) -> list[Any]: }, # Ring buffer replayed to a connecting viewer. Small so connect catches up fast. "memory_limit": "64MB", - # The robot box hangs off base_link. It lives on its own entity: a static - # transform on world/tf/base_link would override the live tf. + # The robot box hangs off the base_link frame, on its own entity: world/tf + # holds the frame tree. "static": { "world/robot_body": _static_robot_body, - "world/robot_body/axes": _static_body_axes, - "world/sensor_axes": _static_sensor_axes, }, "visual_override": { **rerun_config["visual_override"], diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 25f16e0271..287132dc76 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -41,6 +41,7 @@ from dimos.core.core import rpc from dimos.core.global_config import global_config from dimos.core.module import Module, ModuleConfig +from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.protocol.pubsub.impl.zenohpubsub import Zenoh from dimos.protocol.pubsub.patterns import Glob, pattern_matches @@ -56,6 +57,7 @@ RerunOpenOption, ) from dimos.visualization.rerun.init import rerun_init +from dimos.visualization.rerun.tf_tree import TFTreeVis if TYPE_CHECKING: from rerun._baseclasses import Archetype @@ -79,29 +81,6 @@ # # as well as pubsubs={} to specify which protocols to listen to. -# TODO better TF processing -# -# this is rerun bridge specific, rerun has a specific (better) way of handling TFs -# using entity path conventions, each of these nodes in a path are TF frames: -# -# /world/robot1/base_link/camera/optical -# -# While here since we are just listening on TFMessage messages which optionally contain -# just a subset of full TF tree we don't know the full tree structure to build full entity -# path for a transform being published -# -# This is easy to reconstruct but a service/tf.py already does this so should be integrated here -# -# we have decoupled entity paths and actual transforms (like ROS TF frames) -# https://rerun.io/docs/concepts/logging-and-ingestion/transforms -# -# tf#/world -# tf#/base_link -# tf#/camera -# -# In order to solve this, bridge needs to own it's own tf service -# and render it's tf tree into correct rerun entity paths - logger = setup_logger() RerunMulti: TypeAlias = "list[tuple[str, Archetype]]" @@ -225,6 +204,10 @@ class Config(ModuleConfig): max_hz: dict[str, float] = field(default_factory=dict) entity_prefix: str = "world" + # Length in meters of the labeled triad drawn on every tf frame. 0 disables. + # Frames nest under `world/tf` mirroring the tree, so keep other entities out + # of that path and attach them to `tf#/` instead. + tf_axes: float = 0.5 topic_to_entity: Callable[[Any], str] | None = None connect_url: str | None = None memory_limit: str = "25%" @@ -264,6 +247,15 @@ def __init__(self, **kwargs: Any) -> None: self._last_log = {} self._override_cache: dict[str, Callable[[Any], RerunData | None]] = {} self._frame_attached: dict[str, str] = {} + self._tf_tree = self._new_tf_tree() + + def _new_tf_tree(self) -> TFTreeVis | None: + if self.config.tf_axes <= 0: + return None + return TFTreeVis( + axis_length=self.config.tf_axes, + root=f"{self.config.entity_prefix}/tf", + ) @property def host(self) -> str: @@ -351,6 +343,11 @@ def _on_message(self, msg: Any, topic: Any) -> None: # TFMessage for example returns list of (entity_path, archetype) tuples if is_rerun_multi(rerun_data): + # Bound locally: stop() clears the tree from another thread. + tf_tree = self._tf_tree + if tf_tree is not None and isinstance(msg, TFMessage): + tf_tree.log(msg) + return for path, archetype in rerun_data: rr.log(path, archetype) else: @@ -373,6 +370,7 @@ def start(self) -> None: self._last_log = {} self._frame_attached = {} + self._tf_tree = self._new_tf_tree() self._min_intervals: dict[str, float] = { entity: 1.0 / hz for entity, hz in self.config.max_hz.items() if hz > 0 } @@ -600,6 +598,7 @@ def log_blueprint_graph(self, dot_code: str, module_names: list[str]) -> None: def stop(self) -> None: self._override_cache.clear() self._frame_attached.clear() + self._tf_tree = None super().stop() diff --git a/dimos/visualization/rerun/test_tf_tree.py b/dimos/visualization/rerun/test_tf_tree.py new file mode 100644 index 0000000000..5b33d0c7a5 --- /dev/null +++ b/dimos/visualization/rerun/test_tf_tree.py @@ -0,0 +1,302 @@ +# 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. + +"""Entity paths the tf triads are drawn at. Rendering itself needs a human.""" + +from __future__ import annotations + +from dataclasses import dataclass +import logging +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any, cast +from unittest.mock import patch + +import pytest +import rerun as rr + +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.visualization.rerun.bridge import RerunBridgeModule +from dimos.visualization.rerun.tf_tree import RerunTFTree, TFTreeVis + +if TYPE_CHECKING: + from collections.abc import Iterator + +# tf_tree.py logs through this stdlib logger name (setup_logger() derives it +# from the module's file path). +_TF_TREE_LOGGER = "dimos/visualization/rerun/tf_tree.py" + + +def edge(parent: str, child: str, ts: float = 1.0) -> Transform: + return Transform(frame_id=parent, child_frame_id=child, ts=ts) + + +@dataclass +class Topic: + name: str + + +def _bridge_tf_paths(tf_axes: float) -> list[str]: + """Entity paths a bridge logs for the same two tf messages.""" + bridge = RerunBridgeModule(tf_axes=tf_axes) + bridge._min_intervals = {} + try: + with patch("rerun.log") as mock_log: + bridge._on_message(TFMessage(edge("odom", "base_link")), Topic("/tf")) + bridge._on_message(TFMessage(edge("odom", "base_link", ts=2.0)), Topic("/tf")) + finally: + bridge.stop() + return [call.args[0] for call in mock_log.call_args_list] + + +@pytest.fixture +def recording() -> Iterator[None]: + """Give ``rr.log`` somewhere to write, so nothing reaches a real viewer.""" + import rerun as rr + + with rr.RecordingStream("dimos_test_tf_tree"): + yield + + +@pytest.fixture +def tf_warnings(caplog: pytest.LogCaptureFixture) -> Iterator[pytest.LogCaptureFixture]: + """Capture tf_tree log lines via ``caplog``. + + The dimos logger is structlog over a stdlib logger with + ``propagate=False``, so caplog's root-level handler never sees it. + """ + lg = logging.getLogger(_TF_TREE_LOGGER) + lg.addHandler(caplog.handler) + caplog.set_level(logging.WARNING, logger=_TF_TREE_LOGGER) + try: + yield caplog + finally: + lg.removeHandler(caplog.handler) + + +def test_paths_mirror_the_tree() -> None: + vis = TFTreeVis() + vis.buffer.receive_tfmessage( + TFMessage(edge("odom", "base_link"), edge("base_link", "mid360_link")) + ) + + assert vis.path("mid360_link") == "world/tf/odom/base_link/mid360_link" + assert vis.path("base_link") == "world/tf/odom/base_link" + + +def test_root_frame_gets_a_path() -> None: + vis = TFTreeVis() + vis.buffer.receive_tfmessage(TFMessage(edge("odom", "base_link"))) + + assert vis.path("odom") == "world/tf/odom" + + +def test_root_honors_the_configured_prefix() -> None: + vis = TFTreeVis(root="scene/tf") + vis.buffer.receive_tfmessage(TFMessage(edge("odom", "base_link"))) + + assert vis.path("base_link") == "scene/tf/odom/base_link" + + +def test_frame_names_are_escaped() -> None: + vis = TFTreeVis() + vis.buffer.receive_tfmessage(TFMessage(edge("odom", "camera/optical"))) + + assert vis.path("camera/optical") == "world/tf/odom/camera\\/optical" + + +def test_late_reroot_leaves_paths_alone() -> None: + vis = TFTreeVis() + vis.buffer.receive_tfmessage(TFMessage(edge("odom", "base_link"))) + assert vis.path("base_link") == "world/tf/odom/base_link" + + vis.buffer.receive_tfmessage(TFMessage(edge("map", "odom"))) + + assert vis.path("odom") == "world/tf/odom" + assert vis.path("base_link") == "world/tf/odom/base_link" + assert vis.path("map") == "world/tf/map" + + +def test_settle_window_waits_for_the_whole_tree(recording: None) -> None: + """A tf tree arrives edge by edge, in whatever order its publishers start.""" + vis = TFTreeVis(settle=1.0) + vis.log(TFMessage(edge("base_link", "front_camera"))) + vis.log(TFMessage(edge("mid360_link", "base_link"))) + vis.log(TFMessage(edge("odom", "mid360_link"))) + assert vis.frame_paths() == {} + + vis.log(TFMessage(edge("odom", "mid360_link", ts=2.0))) + + assert vis.path("front_camera") == "world/tf/odom/mid360_link/base_link/front_camera" + + +def test_settle_window_does_not_swallow_a_one_shot_edge(recording: None) -> None: + """``world -> map`` shows up twice at startup in a real recording, then never again.""" + vis = TFTreeVis(settle=1.0) + vis.log(TFMessage(edge("world", "map"))) + vis.log(TFMessage(edge("map", "odom"), edge("odom", "base_link"))) + + vis.log(TFMessage(edge("odom", "base_link", ts=2.0))) + + assert vis._axes_logged == { + "world/tf/world", + "world/tf/world/map", + "world/tf/world/map/odom", + "world/tf/world/map/odom/base_link", + } + + +def test_slash_in_a_frame_name_is_not_a_reparent( + recording: None, tf_warnings: pytest.LogCaptureFixture +) -> None: + """The escaped name keeps its slash, so the path cannot be split to find the parent.""" + vis = TFTreeVis(settle=0.0) + vis.log(TFMessage(edge("odom", "camera/optical"))) + + assert [r for r in tf_warnings.records if "re-parented" in r.getMessage()] == [] + + +def test_reparent_of_a_slashed_frame_still_warns( + recording: None, tf_warnings: pytest.LogCaptureFixture +) -> None: + vis = TFTreeVis(settle=0.0) + vis.log(TFMessage(edge("odom", "camera/optical"))) + vis.log(TFMessage(edge("base_link", "camera/optical", ts=2.0))) + + assert len([r for r in tf_warnings.records if "re-parented" in r.getMessage()]) == 1 + + +def test_reparented_frame_warns_once( + recording: None, tf_warnings: pytest.LogCaptureFixture +) -> None: + vis = TFTreeVis(settle=0.0) + vis.log(TFMessage(edge("odom", "base_link"))) + vis.log(TFMessage(edge("chassis", "base_link", ts=2.0))) + vis.log(TFMessage(edge("chassis", "base_link", ts=3.0))) + + warnings = [r for r in tf_warnings.records if "re-parented" in r.getMessage()] + assert len(warnings) == 1 + assert vis.path("base_link") == "world/tf/odom/base_link" + + +def test_every_frame_gets_axes_once(recording: None) -> None: + vis = TFTreeVis(settle=0.0) + vis.log(TFMessage(edge("odom", "base_link"), edge("base_link", "mid360_link"))) + vis.log(TFMessage(edge("odom", "base_link", ts=2.0))) + + assert vis._axes_logged == { + "world/tf/odom", + "world/tf/odom/base_link", + "world/tf/odom/base_link/mid360_link", + } + + +def _arrow_length(arrows: rr.Arrows3D) -> float: + assert arrows.vectors is not None + return float(max(arrows.vectors.as_arrow_array().to_pylist()[0])) + + +def test_triads_shrink_with_depth() -> None: + vis = TFTreeVis(axis_length=1.0, settle=0.0) + with patch("rerun.log") as mock_log: + vis.log(TFMessage(edge("odom", "base_link"), edge("base_link", "mid360_link"))) + + lengths = { + call.args[0]: _arrow_length(arrows) + for call in mock_log.call_args_list + for arrows in call.args[1:] + if isinstance(arrows, rr.Arrows3D) + } + assert lengths == pytest.approx( + { + "world/tf/odom": 1.0, + "world/tf/odom/base_link": 0.8, + "world/tf/odom/base_link/mid360_link": 0.64, + } + ) + + +class FakeStream: + """Re-iterable stand-in for a memory2 stream of tf observations.""" + + def __init__(self, *stamped: tuple[float, TFMessage]) -> None: + self._stamped = stamped + + def __iter__(self) -> Iterator[SimpleNamespace]: + return iter([SimpleNamespace(ts=ts, data=msg) for ts, msg in self._stamped]) + + +def _drive(tf: FakeStream, stamps: list[float]) -> list[tuple[float, str]]: + """Replay ``stamps`` through the transformer, returning (time, entity) in log order.""" + upstream = iter([SimpleNamespace(ts=ts) for ts in stamps]) + events: list[tuple[float, str]] = [] + now = [0.0] + + def set_time(_timeline: str, *, timestamp: float) -> None: + now[0] = timestamp + + with ( + patch("rerun.set_time", side_effect=set_time), + patch("rerun.log", side_effect=lambda path, *a, **k: events.append((now[0], path))), + ): + list(cast("Any", RerunTFTree(cast("Any", tf)))(cast("Any", upstream))) + return events + + +def test_transformer_logs_tf_in_step_with_the_stream() -> None: + tf = FakeStream( + (1.0, TFMessage(edge("odom", "base_link"))), + (2.0, TFMessage(edge("odom", "base_link", ts=2.0))), + (3.0, TFMessage(edge("odom", "base_link", ts=3.0))), + ) + + stamps = [t for t, _ in _drive(tf, [1.5, 2.5, 3.5])] + + assert stamps == sorted(stamps) + assert set(stamps) == {1.0, 2.0, 3.0} + + +def test_transformer_does_not_run_ahead_of_the_stream() -> None: + tf = FakeStream( + (1.0, TFMessage(edge("odom", "base_link"))), + (9.0, TFMessage(edge("odom", "base_link", ts=9.0))), + ) + + # Upstream stops at 2.0, so the tf message at 9.0 is outside the replay. + assert {t for t, _ in _drive(tf, [2.0])} == {1.0} + + +def test_transformer_nests_from_the_whole_stream() -> None: + """Topology is read up front, so the first message already knows its parents.""" + tf = FakeStream( + (1.0, TFMessage(edge("base_link", "front_camera"))), + (1.0, TFMessage(edge("odom", "base_link", ts=1.0))), + ) + + paths = {path for _, path in _drive(tf, [5.0])} + + assert "world/tf/odom/base_link/front_camera" in paths + + +def test_bridge_nests_tf_when_axes_are_on() -> None: + assert _bridge_tf_paths(tf_axes=0.4) == [ + "world/tf/odom", + "world/tf/odom", + "world/tf/odom/base_link", + "world/tf/odom/base_link", + ] + + +def test_bridge_leaves_tf_flat_when_axes_are_off() -> None: + assert _bridge_tf_paths(tf_axes=0.0) == ["world/tf/base_link", "world/tf/base_link"] diff --git a/dimos/visualization/rerun/tf_tree.py b/dimos/visualization/rerun/tf_tree.py new file mode 100644 index 0000000000..1ab7c5ce8f --- /dev/null +++ b/dimos/visualization/rerun/tf_tree.py @@ -0,0 +1,257 @@ +# 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. + +"""Labeled axis triads for every frame of a tf tree.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, TypeVar + +from dimos.memory2.transform import Transformer +from dimos.protocol.tf.tf import MultiTBuffer +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + + import rerun as rr + + from dimos.memory2.stream import Stream + from dimos.memory2.type.observation import Observation + from dimos.msgs.geometry_msgs.Transform import Transform + from dimos.msgs.tf2_msgs.TFMessage import TFMessage + +T = TypeVar("T") + +logger = setup_logger() + +DEFAULT_TF_ROOT = "world/tf" +DEFAULT_AXIS_LENGTH = 0.5 +DEFAULT_TIMELINE = "ts" +# Seconds of tf to collect before handing out entity paths. Static mount trees +# publish at 5 Hz, so this covers several full cycles. +SETTLE_SECONDS = 1.0 +# Each level's triad relative to its parent's, so deeper frames read as smaller. +DEPTH_SCALE = 0.8 +# Arrow width in UI points. Rerun's own TransformAxes3D draws at 1.0. +AXIS_WIDTH_UI_POINTS = 2.0 +AXIS_COLORS = [[255, 0, 0], [0, 255, 0], [0, 0, 255]] + + +def triad(length: float) -> rr.Arrows3D: + """XYZ arrows for one frame, red/green/blue for x/y/z. + + Drawn by hand rather than with ``TransformAxes3D`` because that archetype + fixes its own width and carries a frame label that cannot be styled. + """ + import rerun as rr + + return rr.Arrows3D( + origins=[[0.0, 0.0, 0.0]] * 3, + vectors=[[length, 0.0, 0.0], [0.0, length, 0.0], [0.0, 0.0, length]], + colors=AXIS_COLORS, + radii=rr.components.Radius.ui_points(AXIS_WIDTH_UI_POINTS), + ) + + +class TFTreeVis: + """Draws each tf frame as a labeled triad, nested by entity path. + + Placement still comes from the tf graph: every ``Transform3D`` keeps its + explicit ``tf#/parent`` and ``tf#/child`` frames, so anything attached to a + named frame is unaffected. The entity path only mirrors the tree + (``world/tf/odom/base_link/mid360_link``) so the viewer's entity panel shows + its shape. + + The rerun bridge drives this off live tf. To replay a recorded stream, use + :class:`RerunTFTree` rather than driving it by hand. + """ + + def __init__( + self, + buffer: MultiTBuffer | None = None, + axis_length: float = DEFAULT_AXIS_LENGTH, + root: str = DEFAULT_TF_ROOT, + settle: float = SETTLE_SECONDS, + ) -> None: + self.buffer = buffer if buffer is not None else MultiTBuffer() + self.axis_length = axis_length + self.root = root + self.settle = settle + self._paths: dict[str, str] = {} + self._depths: dict[str, int] = {} + self._parents: dict[str, str | None] = {} + self._axes_logged: set[str] = set() + self._reparented: set[str] = set() + self._settle_deadline: float | None = None + self._flushed = False + + def log(self, msg: TFMessage) -> None: + """Feed a tf message into the buffer, then log its transforms.""" + if not msg.transforms: + return + self.buffer.receive_tfmessage(msg) + if not self._settled(msg.transforms): + return + transforms = msg.transforms + if self.settle > 0 and not self._flushed: + # An edge published only while the tree settled, like a root sent + # twice at startup, would otherwise never be drawn. + transforms = self.buffer.latest_transforms() + self._flushed = True + self._log_transforms(transforms) + + def _settled(self, transforms: Iterable[Transform]) -> bool: + """Whether the tree has had time to fill in. + + A path is frozen the first time its frame is seen, so a frame that gets + its path before its own parent arrives stays a root for the session. + Publishers put a full tree on the wire within a few messages, and the tf + that falls in this window is republished right after it. + """ + if self.settle <= 0: + return True + latest = max(transform.ts for transform in transforms) + if self._settle_deadline is None: + self._settle_deadline = latest + self.settle + return latest >= self._settle_deadline + + def _log_transforms(self, transforms: Iterable[Transform]) -> None: + import rerun as rr + + for transform in transforms: + parent_path = self.path(transform.frame_id) + child_path = self.path(transform.child_frame_id) + self._warn_on_reparent(transform) + self._log_axes(transform.frame_id, parent_path) + rr.log(child_path, transform.to_rerun()) + self._log_axes(transform.child_frame_id, child_path) + + def frame_paths(self) -> dict[str, str]: + """Entity path assigned to each frame seen so far.""" + return dict(self._paths) + + def path(self, frame: str) -> str: + """Entity path of a frame, assigned the first time the frame is seen. + + Rerun forbids a child frame's declaring entity from changing over time, + and tf trees re-root late, so a path never moves once handed out. + """ + import rerun as rr + + known = self._paths.get(frame) + if known is not None: + return known + + chain: list[str] = [] + base = self.root + depth = 0 + node: str | None = frame + visited: set[str] = set() + while node is not None and node not in visited: + visited.add(node) + if node in self._paths: + base = self._paths[node] + depth = self._depths[node] + 1 + break + chain.append(node) + node = self.buffer.get_parent(node) + + # Whatever the walk stopped on is the parent of the top of the chain. + parent = node + for name in reversed(chain): + base = f"{base}/{rr.escape_entity_path_part(name)}" + self._paths[name] = base + self._depths[name] = depth + self._parents[name] = parent + parent = name + depth += 1 + + return self._paths[frame] + + def _warn_on_reparent(self, transform: Transform) -> None: + child = transform.child_frame_id + if self._parents.get(child) == transform.frame_id or child in self._reparented: + return + self._reparented.add(child) + logger.warning( + "tf frame re-parented after its entity path was assigned, panel nesting is stale", + frame=child, + new_parent=transform.frame_id, + entity_path=self._paths[child], + ) + + def _log_axes(self, frame: str, path: str) -> None: + import rerun as rr + + if path in self._axes_logged: + return + self._axes_logged.add(path) + if self.buffer.get_parent(frame) is None: + # A root is never a child_frame_id, so nothing else declares it. + rr.log(path, rr.Transform3D(child_frame=f"tf#/{frame}")) + rr.log( + path, + # Without this the arrows sit in the entity path's implicit frame, + # which is pinned to the parent path and never moves. + rr.CoordinateFrame(f"tf#/{frame}"), + triad(self.axis_length * DEPTH_SCALE ** self._depths[frame]), + static=True, + ) + + +class RerunTFTree(Transformer[T, T]): + """Draw the tf tree's triads in step with the stream it passes through. + + Drop it into a replay pipeline and every tf frame gets its labeled triad, + each logged at its own place on the timeline rather than in one lump up + front:: + + pipeline = lidar.transform(RerunTFTree(store.stream("tf", TFMessage))) + + Window the tf stream the same way as the pipeline, or the tf that predates + the first observation all lands on that first frame. + """ + + def __init__( + self, + tf: Stream[TFMessage], + axis_length: float = DEFAULT_AXIS_LENGTH, + timeline: str = DEFAULT_TIMELINE, + root: str = DEFAULT_TF_ROOT, + ) -> None: + self._tf = tf + self._timeline = timeline + self._vis = TFTreeVis(axis_length=axis_length, root=root, settle=0.0) + + @property + def vis(self) -> TFTreeVis: + return self._vis + + def __call__(self, upstream: Iterator[Observation[T]]) -> Iterator[Observation[T]]: + import rerun as rr + + # Topology first, so no frame is given a path before its parent is known. + for tf_obs in self._tf: + self._vis.buffer.receive_tfmessage(tf_obs.data) + + pending = iter(self._tf) + head = next(pending, None) + for obs in upstream: + while head is not None and head.ts <= obs.ts: + rr.set_time(self._timeline, timestamp=head.ts) + self._vis.log(head.data) + head = next(pending, None) + rr.set_time(self._timeline, timestamp=obs.ts) + yield obs diff --git a/dimos/visualization/vis_module.py b/dimos/visualization/vis_module.py index 4731f6f88f..dd200310e4 100644 --- a/dimos/visualization/vis_module.py +++ b/dimos/visualization/vis_module.py @@ -43,7 +43,7 @@ def vis_module( "world/camera_info": lambda ci: ci.to_rerun(...), }, "static": { - "world/tf/base_link": lambda rr: [rr.Boxes3D(...)], + "world/robot_body": lambda rr: [rr.Boxes3D(...)], }, }, ) From 6b9cfaad3ad993d402a8d6201b070fbedb4ccfb4 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 4 Aug 2026 00:39:21 -0700 Subject: [PATCH 02/13] Fix parent issue when tree changes --- dimos/memory2/tf.py | 16 +- .../nav_3d/mls_planner/utils/plan_rrd.py | 37 +-- .../nav_3d/mls_planner/utils/test_plan_rrd.py | 73 ----- dimos/protocol/tf/tf.py | 18 -- .../go2/blueprints/basic/unitree_go2_basic.py | 4 +- dimos/visualization/rerun/bridge.py | 15 +- dimos/visualization/rerun/test_tf_tree.py | 295 ++---------------- dimos/visualization/rerun/tf_tree.py | 284 +++++++---------- 8 files changed, 183 insertions(+), 559 deletions(-) delete mode 100644 dimos/navigation/nav_3d/mls_planner/utils/test_plan_rrd.py diff --git a/dimos/memory2/tf.py b/dimos/memory2/tf.py index 27806dddf2..76254c91bb 100644 --- a/dimos/memory2/tf.py +++ b/dimos/memory2/tf.py @@ -28,6 +28,17 @@ from dimos.protocol.tf.tf import TFLookup +def tf_stream(store: Any, stream: str = "tf") -> Stream[TFMessage] | None: + """The recording's tf stream, or None if it has none. + + Asking a store for an absent stream registers it and writes its tables, so + the name has to be checked first. + """ + if stream not in store.list_streams(): + return None + return cast("Stream[TFMessage]", store.stream(stream, TFMessage)) + + class StreamTF(MultiTBuffer): def __init__( self, @@ -47,9 +58,8 @@ def __init__( @classmethod def from_store(cls, store: Any, stream: str = "tf") -> StreamTF | None: - if stream not in store.list_streams(): - return None - return cls(store.stream(stream, TFMessage)) + recorded = tf_stream(store, stream) + return None if recorded is None else cls(recorded) def publish(self, *args: Transform) -> None: raise NotImplementedError("StreamTF is a read-only replay service.") diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index 88b828faf1..594d35c507 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -29,7 +29,7 @@ from dimos.mapping.ray_tracing.transformer import RayTraceMap from dimos.memory2.store.sqlite import SqliteStore -from dimos.memory2.tf import StreamTF +from dimos.memory2.tf import StreamTF, tf_stream from dimos.memory2.transform import FnTransformer from dimos.memory2.type.observation import Observation from dimos.msgs.geometry_msgs.Quaternion import Quaternion @@ -47,6 +47,8 @@ if TYPE_CHECKING: import rerun.blueprint as rrb + from dimos.memory2.stream import Stream + TIMELINE = "ts" # Mount frames as recorded on the tf stream. @@ -136,31 +138,21 @@ def _log_path_wp(waypoints: NDArray[np.float32] | None, entity: str, color: list rr.log(entity, rr.LineStrips3D([points], colors=[color], radii=0.05)) -def _window(stream: Any, from_time: float | None, to_time: float | None) -> Any: - """Clip a stream to the replay window, both bounds relative to its start.""" - if from_time is not None: - stream = stream.from_time(from_time) - if to_time is not None: - stream = stream.to_time(to_time) - return stream - - -def _tf_over(store: SqliteStore, window: Any) -> Any: - """The recorded tf stream clipped to another stream's time span. +def _tf_over(store: SqliteStore, window: Stream[Any]) -> Stream[TFMessage] | None: + """The recorded tf stream clipped to another stream's span. - Absolute bounds, because the relative ones anchor on each stream's own - first observation and tf rarely starts on the same sample as the lidar. - Returns None when the recording has no tf, which must not be probed for - with ``store.stream``: that registers the stream and writes its tables. + Absolute bounds: the relative ones anchor on each stream's own first + observation, and tf rarely starts on the same sample as the lidar. """ - if "tf" not in store.list_streams(): + recorded = tf_stream(store) + if recorded is None: print("no tf stream in the recording; skipping the tf tree") return None try: first, last = window.first().ts, window.last().ts except LookupError: return None - return store.stream("tf", TFMessage).order_by("ts").time_range(first, last) + return recorded.order_by("ts").time_range(first, last) def _base_from_sensor(store: SqliteStore) -> Transform | None: @@ -349,8 +341,7 @@ def _blueprint(crop: LocalCrop) -> rrb.Blueprint: origin="world", name="world", contents=["+ $origin/**", "- $origin/local/**"], - # The graph buries the map it was built from. Tick it back on in - # the viewer when the question is why a path went the way it did. + # The graph buries the map it was built from. overrides={ "world/nodes": rrb.EntityBehavior(visible=False), "world/node_edges": rrb.EntityBehavior(visible=False), @@ -579,7 +570,11 @@ def main( store = SqliteStore(path=str(db_path)) with store: - lidar = _window(store.stream(lidar_stream, PointCloud2).order_by("ts"), from_time, to_time) + lidar = store.stream(lidar_stream, PointCloud2).order_by("ts") + if from_time is not None: + lidar = lidar.from_time(from_time) + if to_time is not None: + lidar = lidar.to_time(to_time) odom = store.stream(odom_stream, Odometry).order_by("ts") tf = _tf_over(store, lidar) diff --git a/dimos/navigation/nav_3d/mls_planner/utils/test_plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/test_plan_rrd.py deleted file mode 100644 index 876e20a0f1..0000000000 --- a/dimos/navigation/nav_3d/mls_planner/utils/test_plan_rrd.py +++ /dev/null @@ -1,73 +0,0 @@ -# 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. - -"""Picking the tf stream out of a recording without disturbing it.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import pytest - -from dimos.memory2.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.navigation.nav_3d.mls_planner.utils.plan_rrd import _tf_over - -if TYPE_CHECKING: - from pathlib import Path - - -def _store(tmp_path: Path, tf_stamps: list[float] | None, lidar_stamps: list[float]) -> SqliteStore: - store = SqliteStore(path=str(tmp_path / "rec.db")) - lidar = store.stream("pointlio_lidar", str) - for ts in lidar_stamps: - lidar.append("cloud", ts=ts) - if tf_stamps is not None: - tf = store.stream("tf", TFMessage) - for ts in tf_stamps: - tf.append( - TFMessage(Transform(frame_id="odom", child_frame_id="base_link", ts=ts)), ts=ts - ) - return store - - -@pytest.mark.skipif_macos -@pytest.mark.skipif_aarch64 -def test_missing_tf_is_not_created_by_looking_for_it(tmp_path: Path) -> None: - """Probing with ``store.stream`` would register tf and write its tables.""" - with _store(tmp_path, tf_stamps=None, lidar_stamps=[1.0, 2.0]) as store: - assert _tf_over(store, store.stream("pointlio_lidar", str)) is None - assert "tf" not in store.list_streams() - - -@pytest.mark.skipif_macos -@pytest.mark.skipif_aarch64 -def test_tf_window_follows_the_lidar_window(tmp_path: Path) -> None: - """tf starts 10s before the lidar here, so a relative window would be shifted.""" - with _store( - tmp_path, tf_stamps=[90.0, 100.0, 105.0, 110.0, 115.0], lidar_stamps=[100.0, 110.0] - ) as store: - tf = _tf_over(store, store.stream("pointlio_lidar", str)) - - assert [obs.ts for obs in tf] == [100.0, 105.0, 110.0] - - -@pytest.mark.skipif_macos -@pytest.mark.skipif_aarch64 -def test_empty_lidar_window_has_no_tf(tmp_path: Path) -> None: - with _store(tmp_path, tf_stamps=[1.0], lidar_stamps=[1.0]) as store: - empty = store.stream("pointlio_lidar", str).after(500.0) - - assert _tf_over(store, empty) is None diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 9b9f0f5cc1..2b0309e52e 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -117,24 +117,6 @@ def get_frames(self) -> set[str]: frames.add(child) return frames - def latest_transforms(self) -> list[Transform]: - """Most recent transform on each edge of the tree.""" - with self._cv: - latest = [buffer.last() for buffer in self.buffers.values()] - return [transform for transform in latest if transform is not None] - - def get_parent(self, frame_id: str) -> str | None: - """Parent of a frame, or None if it is a root of the tree. - - A frame with more than one parent is a graph, not a tree. The first - parent seen wins so the answer stays stable once given. - """ - with self._cv: - for parent, child in self.buffers: - if child == frame_id: - return parent - return None - def get_connections(self, frame_id: str) -> set[str]: """Get all frames connected to the given frame (both as parent and child).""" connections = set() diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index 3540c49287..dfdbd35593 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -42,7 +42,7 @@ def _convert_navigation_costmap(grid: Any) -> Any: ) -def _static_base_link(rr: Any) -> list[Any]: +def _static_robot_body(rr: Any) -> list[Any]: return [ rr.Boxes3D( half_sizes=[0.35, 0.155, 0.2], @@ -99,7 +99,7 @@ def _go2_rerun_blueprint() -> Any: }, # slapping a go2 shaped box on the base_link frame "static": { - "world/robot_body": _static_base_link, + "world/robot_body": _static_robot_body, }, } diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 287132dc76..99a66035b2 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -57,7 +57,7 @@ RerunOpenOption, ) from dimos.visualization.rerun.init import rerun_init -from dimos.visualization.rerun.tf_tree import TFTreeVis +from dimos.visualization.rerun.tf_tree import DEFAULT_AXIS_LENGTH, TFTreeVis if TYPE_CHECKING: from rerun._baseclasses import Archetype @@ -204,10 +204,8 @@ class Config(ModuleConfig): max_hz: dict[str, float] = field(default_factory=dict) entity_prefix: str = "world" - # Length in meters of the labeled triad drawn on every tf frame. 0 disables. - # Frames nest under `world/tf` mirroring the tree, so keep other entities out - # of that path and attach them to `tf#/` instead. - tf_axes: float = 0.5 + # Length in meters of the triad drawn on every tf frame, 0 to draw none. + tf_axes: float = DEFAULT_AXIS_LENGTH topic_to_entity: Callable[[Any], str] | None = None connect_url: str | None = None memory_limit: str = "25%" @@ -254,7 +252,7 @@ def _new_tf_tree(self) -> TFTreeVis | None: return None return TFTreeVis( axis_length=self.config.tf_axes, - root=f"{self.config.entity_prefix}/tf", + root=f"{self.config.entity_prefix}/frames", ) @property @@ -343,13 +341,12 @@ def _on_message(self, msg: Any, topic: Any) -> None: # TFMessage for example returns list of (entity_path, archetype) tuples if is_rerun_multi(rerun_data): + for path, archetype in rerun_data: + rr.log(path, archetype) # Bound locally: stop() clears the tree from another thread. tf_tree = self._tf_tree if tf_tree is not None and isinstance(msg, TFMessage): tf_tree.log(msg) - return - for path, archetype in rerun_data: - rr.log(path, archetype) else: rr.log(entity_path, cast("Archetype", rerun_data)) # if source msg carries a frame_id, attach the entity to that TF frame diff --git a/dimos/visualization/rerun/test_tf_tree.py b/dimos/visualization/rerun/test_tf_tree.py index 5b33d0c7a5..0477f52b68 100644 --- a/dimos/visualization/rerun/test_tf_tree.py +++ b/dimos/visualization/rerun/test_tf_tree.py @@ -12,291 +12,58 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Entity paths the tf triads are drawn at. Rendering itself needs a human.""" +"""Where the triads get drawn. How they look is for a human with the viewer open.""" from __future__ import annotations -from dataclasses import dataclass -import logging -from types import SimpleNamespace -from typing import TYPE_CHECKING, Any, cast -from unittest.mock import patch - -import pytest -import rerun as rr - from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.visualization.rerun.bridge import RerunBridgeModule -from dimos.visualization.rerun.tf_tree import RerunTFTree, TFTreeVis - -if TYPE_CHECKING: - from collections.abc import Iterator - -# tf_tree.py logs through this stdlib logger name (setup_logger() derives it -# from the module's file path). -_TF_TREE_LOGGER = "dimos/visualization/rerun/tf_tree.py" - - -def edge(parent: str, child: str, ts: float = 1.0) -> Transform: - return Transform(frame_id=parent, child_frame_id=child, ts=ts) - - -@dataclass -class Topic: - name: str - +from dimos.visualization.rerun.tf_tree import TFTreeVis -def _bridge_tf_paths(tf_axes: float) -> list[str]: - """Entity paths a bridge logs for the same two tf messages.""" - bridge = RerunBridgeModule(tf_axes=tf_axes) - bridge._min_intervals = {} - try: - with patch("rerun.log") as mock_log: - bridge._on_message(TFMessage(edge("odom", "base_link")), Topic("/tf")) - bridge._on_message(TFMessage(edge("odom", "base_link", ts=2.0)), Topic("/tf")) - finally: - bridge.stop() - return [call.args[0] for call in mock_log.call_args_list] +def edge(parent: str, child: str) -> Transform: + return Transform(frame_id=parent, child_frame_id=child, ts=1.0) -@pytest.fixture -def recording() -> Iterator[None]: - """Give ``rr.log`` somewhere to write, so nothing reaches a real viewer.""" - import rerun as rr - with rr.RecordingStream("dimos_test_tf_tree"): - yield +def paths(vis: TFTreeVis) -> dict[str, str]: + return {frame: spot.path for frame, spot in vis.placements().items()} -@pytest.fixture -def tf_warnings(caplog: pytest.LogCaptureFixture) -> Iterator[pytest.LogCaptureFixture]: - """Capture tf_tree log lines via ``caplog``. +def feed(vis: TFTreeVis, *messages: TFMessage) -> None: + """tf republishes, and the tree draws once a message adds nothing new.""" + for msg in messages: + vis.log(msg) + vis.log(messages[-1]) - The dimos logger is structlog over a stdlib logger with - ``propagate=False``, so caplog's root-level handler never sees it. - """ - lg = logging.getLogger(_TF_TREE_LOGGER) - lg.addHandler(caplog.handler) - caplog.set_level(logging.WARNING, logger=_TF_TREE_LOGGER) - try: - yield caplog - finally: - lg.removeHandler(caplog.handler) - -def test_paths_mirror_the_tree() -> None: - vis = TFTreeVis() - vis.buffer.receive_tfmessage( - TFMessage(edge("odom", "base_link"), edge("base_link", "mid360_link")) - ) - - assert vis.path("mid360_link") == "world/tf/odom/base_link/mid360_link" - assert vis.path("base_link") == "world/tf/odom/base_link" - - -def test_root_frame_gets_a_path() -> None: +def test_triads_nest_along_the_tree() -> None: vis = TFTreeVis() - vis.buffer.receive_tfmessage(TFMessage(edge("odom", "base_link"))) - - assert vis.path("odom") == "world/tf/odom" - - -def test_root_honors_the_configured_prefix() -> None: - vis = TFTreeVis(root="scene/tf") - vis.buffer.receive_tfmessage(TFMessage(edge("odom", "base_link"))) - - assert vis.path("base_link") == "scene/tf/odom/base_link" - - -def test_frame_names_are_escaped() -> None: - vis = TFTreeVis() - vis.buffer.receive_tfmessage(TFMessage(edge("odom", "camera/optical"))) - - assert vis.path("camera/optical") == "world/tf/odom/camera\\/optical" - - -def test_late_reroot_leaves_paths_alone() -> None: - vis = TFTreeVis() - vis.buffer.receive_tfmessage(TFMessage(edge("odom", "base_link"))) - assert vis.path("base_link") == "world/tf/odom/base_link" - - vis.buffer.receive_tfmessage(TFMessage(edge("map", "odom"))) - - assert vis.path("odom") == "world/tf/odom" - assert vis.path("base_link") == "world/tf/odom/base_link" - assert vis.path("map") == "world/tf/map" - - -def test_settle_window_waits_for_the_whole_tree(recording: None) -> None: - """A tf tree arrives edge by edge, in whatever order its publishers start.""" - vis = TFTreeVis(settle=1.0) - vis.log(TFMessage(edge("base_link", "front_camera"))) - vis.log(TFMessage(edge("mid360_link", "base_link"))) - vis.log(TFMessage(edge("odom", "mid360_link"))) - assert vis.frame_paths() == {} - - vis.log(TFMessage(edge("odom", "mid360_link", ts=2.0))) - - assert vis.path("front_camera") == "world/tf/odom/mid360_link/base_link/front_camera" - - -def test_settle_window_does_not_swallow_a_one_shot_edge(recording: None) -> None: - """``world -> map`` shows up twice at startup in a real recording, then never again.""" - vis = TFTreeVis(settle=1.0) - vis.log(TFMessage(edge("world", "map"))) - vis.log(TFMessage(edge("map", "odom"), edge("odom", "base_link"))) - - vis.log(TFMessage(edge("odom", "base_link", ts=2.0))) - - assert vis._axes_logged == { - "world/tf/world", - "world/tf/world/map", - "world/tf/world/map/odom", - "world/tf/world/map/odom/base_link", - } - - -def test_slash_in_a_frame_name_is_not_a_reparent( - recording: None, tf_warnings: pytest.LogCaptureFixture -) -> None: - """The escaped name keeps its slash, so the path cannot be split to find the parent.""" - vis = TFTreeVis(settle=0.0) - vis.log(TFMessage(edge("odom", "camera/optical"))) - - assert [r for r in tf_warnings.records if "re-parented" in r.getMessage()] == [] - - -def test_reparent_of_a_slashed_frame_still_warns( - recording: None, tf_warnings: pytest.LogCaptureFixture -) -> None: - vis = TFTreeVis(settle=0.0) - vis.log(TFMessage(edge("odom", "camera/optical"))) - vis.log(TFMessage(edge("base_link", "camera/optical", ts=2.0))) - - assert len([r for r in tf_warnings.records if "re-parented" in r.getMessage()]) == 1 - - -def test_reparented_frame_warns_once( - recording: None, tf_warnings: pytest.LogCaptureFixture -) -> None: - vis = TFTreeVis(settle=0.0) - vis.log(TFMessage(edge("odom", "base_link"))) - vis.log(TFMessage(edge("chassis", "base_link", ts=2.0))) - vis.log(TFMessage(edge("chassis", "base_link", ts=3.0))) + feed(vis, TFMessage(edge("odom", "base_link"), edge("base_link", "camera/optical"))) - warnings = [r for r in tf_warnings.records if "re-parented" in r.getMessage()] - assert len(warnings) == 1 - assert vis.path("base_link") == "world/tf/odom/base_link" - - -def test_every_frame_gets_axes_once(recording: None) -> None: - vis = TFTreeVis(settle=0.0) - vis.log(TFMessage(edge("odom", "base_link"), edge("base_link", "mid360_link"))) - vis.log(TFMessage(edge("odom", "base_link", ts=2.0))) - - assert vis._axes_logged == { - "world/tf/odom", - "world/tf/odom/base_link", - "world/tf/odom/base_link/mid360_link", + assert paths(vis) == { + "odom": "world/frames/odom", + "base_link": "world/frames/odom/base_link", + "camera/optical": "world/frames/odom/base_link/camera\\/optical", } -def _arrow_length(arrows: rr.Arrows3D) -> float: - assert arrows.vectors is not None - return float(max(arrows.vectors.as_arrow_array().to_pylist()[0])) - +def test_a_late_root_re_parents_the_tree() -> None: + """The mount tree is published seconds before the odometry that roots it.""" + vis = TFTreeVis() + feed(vis, TFMessage(edge("mid360_link", "base_link"))) + assert paths(vis)["base_link"] == "world/frames/mid360_link/base_link" -def test_triads_shrink_with_depth() -> None: - vis = TFTreeVis(axis_length=1.0, settle=0.0) - with patch("rerun.log") as mock_log: - vis.log(TFMessage(edge("odom", "base_link"), edge("base_link", "mid360_link"))) + feed(vis, TFMessage(edge("odom", "mid360_link"))) - lengths = { - call.args[0]: _arrow_length(arrows) - for call in mock_log.call_args_list - for arrows in call.args[1:] - if isinstance(arrows, rr.Arrows3D) + assert paths(vis) == { + "odom": "world/frames/odom", + "mid360_link": "world/frames/odom/mid360_link", + "base_link": "world/frames/odom/mid360_link/base_link", } - assert lengths == pytest.approx( - { - "world/tf/odom": 1.0, - "world/tf/odom/base_link": 0.8, - "world/tf/odom/base_link/mid360_link": 0.64, - } - ) - - -class FakeStream: - """Re-iterable stand-in for a memory2 stream of tf observations.""" - - def __init__(self, *stamped: tuple[float, TFMessage]) -> None: - self._stamped = stamped - - def __iter__(self) -> Iterator[SimpleNamespace]: - return iter([SimpleNamespace(ts=ts, data=msg) for ts, msg in self._stamped]) -def _drive(tf: FakeStream, stamps: list[float]) -> list[tuple[float, str]]: - """Replay ``stamps`` through the transformer, returning (time, entity) in log order.""" - upstream = iter([SimpleNamespace(ts=ts) for ts in stamps]) - events: list[tuple[float, str]] = [] - now = [0.0] - - def set_time(_timeline: str, *, timestamp: float) -> None: - now[0] = timestamp - - with ( - patch("rerun.set_time", side_effect=set_time), - patch("rerun.log", side_effect=lambda path, *a, **k: events.append((now[0], path))), - ): - list(cast("Any", RerunTFTree(cast("Any", tf)))(cast("Any", upstream))) - return events - - -def test_transformer_logs_tf_in_step_with_the_stream() -> None: - tf = FakeStream( - (1.0, TFMessage(edge("odom", "base_link"))), - (2.0, TFMessage(edge("odom", "base_link", ts=2.0))), - (3.0, TFMessage(edge("odom", "base_link", ts=3.0))), - ) - - stamps = [t for t, _ in _drive(tf, [1.5, 2.5, 3.5])] - - assert stamps == sorted(stamps) - assert set(stamps) == {1.0, 2.0, 3.0} - - -def test_transformer_does_not_run_ahead_of_the_stream() -> None: - tf = FakeStream( - (1.0, TFMessage(edge("odom", "base_link"))), - (9.0, TFMessage(edge("odom", "base_link", ts=9.0))), - ) - - # Upstream stops at 2.0, so the tf message at 9.0 is outside the replay. - assert {t for t, _ in _drive(tf, [2.0])} == {1.0} - - -def test_transformer_nests_from_the_whole_stream() -> None: - """Topology is read up front, so the first message already knows its parents.""" - tf = FakeStream( - (1.0, TFMessage(edge("base_link", "front_camera"))), - (1.0, TFMessage(edge("odom", "base_link", ts=1.0))), - ) - - paths = {path for _, path in _drive(tf, [5.0])} - - assert "world/tf/odom/base_link/front_camera" in paths - - -def test_bridge_nests_tf_when_axes_are_on() -> None: - assert _bridge_tf_paths(tf_axes=0.4) == [ - "world/tf/odom", - "world/tf/odom", - "world/tf/odom/base_link", - "world/tf/odom/base_link", - ] - +def test_a_cycle_does_not_hang() -> None: + vis = TFTreeVis() + feed(vis, TFMessage(edge("a", "b"), edge("b", "a"))) -def test_bridge_leaves_tf_flat_when_axes_are_off() -> None: - assert _bridge_tf_paths(tf_axes=0.0) == ["world/tf/base_link", "world/tf/base_link"] + assert set(paths(vis)) == {"a", "b"} diff --git a/dimos/visualization/rerun/tf_tree.py b/dimos/visualization/rerun/tf_tree.py index 1ab7c5ce8f..468eb3a54b 100644 --- a/dimos/visualization/rerun/tf_tree.py +++ b/dimos/visualization/rerun/tf_tree.py @@ -16,11 +16,11 @@ from __future__ import annotations +from dataclasses import dataclass +import threading from typing import TYPE_CHECKING, TypeVar from dimos.memory2.transform import Transformer -from dimos.protocol.tf.tf import MultiTBuffer -from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from collections.abc import Iterable, Iterator @@ -34,27 +34,18 @@ T = TypeVar("T") -logger = setup_logger() - -DEFAULT_TF_ROOT = "world/tf" +DEFAULT_FRAMES_ROOT = "world/frames" DEFAULT_AXIS_LENGTH = 0.5 DEFAULT_TIMELINE = "ts" -# Seconds of tf to collect before handing out entity paths. Static mount trees -# publish at 5 Hz, so this covers several full cycles. -SETTLE_SECONDS = 1.0 -# Each level's triad relative to its parent's, so deeper frames read as smaller. +# Each level's triad relative to its parent's. DEPTH_SCALE = 0.8 -# Arrow width in UI points. Rerun's own TransformAxes3D draws at 1.0. +# Rerun's own TransformAxes3D draws at 1.0. AXIS_WIDTH_UI_POINTS = 2.0 AXIS_COLORS = [[255, 0, 0], [0, 255, 0], [0, 0, 255]] def triad(length: float) -> rr.Arrows3D: - """XYZ arrows for one frame, red/green/blue for x/y/z. - - Drawn by hand rather than with ``TransformAxes3D`` because that archetype - fixes its own width and carries a frame label that cannot be styled. - """ + """XYZ arrows, red green blue for x y z.""" import rerun as rr return rr.Arrows3D( @@ -65,193 +56,148 @@ def triad(length: float) -> rr.Arrows3D: ) -class TFTreeVis: - """Draws each tf frame as a labeled triad, nested by entity path. +@dataclass(frozen=True) +class Placement: + """Where a frame's triad is drawn, and how big.""" - Placement still comes from the tf graph: every ``Transform3D`` keeps its - explicit ``tf#/parent`` and ``tf#/child`` frames, so anything attached to a - named frame is unaffected. The entity path only mirrors the tree - (``world/tf/odom/base_link/mid360_link``) so the viewer's entity panel shows - its shape. + path: str + depth: int - The rerun bridge drives this off live tf. To replay a recorded stream, use - :class:`RerunTFTree` rather than driving it by hand. + +class TFTreeVis: + """Draws a labeled triad per tf frame, nested by entity path. + + Draws markers only. The transforms themselves are logged by whoever owns the + tf stream, at the flat paths TFMessage.to_rerun assigns. """ def __init__( self, - buffer: MultiTBuffer | None = None, axis_length: float = DEFAULT_AXIS_LENGTH, - root: str = DEFAULT_TF_ROOT, - settle: float = SETTLE_SECONDS, + root: str = DEFAULT_FRAMES_ROOT, ) -> None: - self.buffer = buffer if buffer is not None else MultiTBuffer() self.axis_length = axis_length self.root = root - self.settle = settle - self._paths: dict[str, str] = {} - self._depths: dict[str, int] = {} - self._parents: dict[str, str | None] = {} - self._axes_logged: set[str] = set() - self._reparented: set[str] = set() - self._settle_deadline: float | None = None - self._flushed = False + self._lock = threading.Lock() + self._parents: dict[str, str] = {} + self._drawn: dict[str, Placement] = {} + self._pending = False def log(self, msg: TFMessage) -> None: - """Feed a tf message into the buffer, then log its transforms.""" + """Redraw once a message arrives that adds nothing new. + + Publishers split one tree across several messages, and drawing each of + them walks the tree through shapes it never really had, leaving a stale + entity behind every time. + """ if not msg.transforms: return - self.buffer.receive_tfmessage(msg) - if not self._settled(msg.transforms): - return - transforms = msg.transforms - if self.settle > 0 and not self._flushed: - # An edge published only while the tree settled, like a root sent - # twice at startup, would otherwise never be drawn. - transforms = self.buffer.latest_transforms() - self._flushed = True - self._log_transforms(transforms) - - def _settled(self, transforms: Iterable[Transform]) -> bool: - """Whether the tree has had time to fill in. - - A path is frozen the first time its frame is seen, so a frame that gets - its path before its own parent arrives stays a root for the session. - Publishers put a full tree on the wire within a few messages, and the tf - that falls in this window is republished right after it. - """ - if self.settle <= 0: - return True - latest = max(transform.ts for transform in transforms) - if self._settle_deadline is None: - self._settle_deadline = latest + self.settle - return latest >= self._settle_deadline - - def _log_transforms(self, transforms: Iterable[Transform]) -> None: - import rerun as rr - + with self._lock: + if self._learn(msg.transforms): + self._pending = True + elif self._pending: + self._pending = False + self._redraw() + + def flush(self) -> None: + """Draw a pending change that no later message arrived to trigger.""" + with self._lock: + if self._pending: + self._pending = False + self._redraw() + + def placements(self) -> dict[str, Placement]: + with self._lock: + return dict(self._drawn) + + def _learn(self, transforms: Iterable[Transform]) -> bool: + changed = False for transform in transforms: - parent_path = self.path(transform.frame_id) - child_path = self.path(transform.child_frame_id) - self._warn_on_reparent(transform) - self._log_axes(transform.frame_id, parent_path) - rr.log(child_path, transform.to_rerun()) - self._log_axes(transform.child_frame_id, child_path) - - def frame_paths(self) -> dict[str, str]: - """Entity path assigned to each frame seen so far.""" - return dict(self._paths) - - def path(self, frame: str) -> str: - """Entity path of a frame, assigned the first time the frame is seen. - - Rerun forbids a child frame's declaring entity from changing over time, - and tf trees re-root late, so a path never moves once handed out. - """ + if self._parents.get(transform.child_frame_id) != transform.frame_id: + self._parents[transform.child_frame_id] = transform.frame_id + changed = True + return changed + + def _layout(self) -> dict[str, Placement]: import rerun as rr - known = self._paths.get(frame) - if known is not None: - return known - - chain: list[str] = [] - base = self.root - depth = 0 - node: str | None = frame - visited: set[str] = set() - while node is not None and node not in visited: - visited.add(node) - if node in self._paths: - base = self._paths[node] - depth = self._depths[node] + 1 - break - chain.append(node) - node = self.buffer.get_parent(node) - - # Whatever the walk stopped on is the parent of the top of the chain. - parent = node - for name in reversed(chain): - base = f"{base}/{rr.escape_entity_path_part(name)}" - self._paths[name] = base - self._depths[name] = depth - self._parents[name] = parent - parent = name - depth += 1 - - return self._paths[frame] - - def _warn_on_reparent(self, transform: Transform) -> None: - child = transform.child_frame_id - if self._parents.get(child) == transform.frame_id or child in self._reparented: - return - self._reparented.add(child) - logger.warning( - "tf frame re-parented after its entity path was assigned, panel nesting is stale", - frame=child, - new_parent=transform.frame_id, - entity_path=self._paths[child], - ) - - def _log_axes(self, frame: str, path: str) -> None: + placed: dict[str, Placement] = {} + + def place(frame: str, walked: frozenset[str]) -> Placement: + known = placed.get(frame) + if known is not None: + return known + parent = self._parents.get(frame) + if parent is None or parent in walked: + spot = Placement(f"{self.root}/{rr.escape_entity_path_part(frame)}", 0) + else: + above = place(parent, walked | {frame}) + spot = Placement( + f"{above.path}/{rr.escape_entity_path_part(frame)}", above.depth + 1 + ) + placed[frame] = spot + return spot + + for frame in (*self._parents, *self._parents.values()): + place(frame, frozenset()) + return placed + + def _redraw(self) -> None: + """Move triads to match the tree as it is now. + + Rerun refuses to let the entity declaring a frame move, which is why the + triads carry a CoordinateFrame instead and declare nothing. + """ import rerun as rr - if path in self._axes_logged: - return - self._axes_logged.add(path) - if self.buffer.get_parent(frame) is None: - # A root is never a child_frame_id, so nothing else declares it. - rr.log(path, rr.Transform3D(child_frame=f"tf#/{frame}")) - rr.log( - path, - # Without this the arrows sit in the entity path's implicit frame, - # which is pinned to the parent path and never moves. - rr.CoordinateFrame(f"tf#/{frame}"), - triad(self.axis_length * DEPTH_SCALE ** self._depths[frame]), - static=True, - ) + layout = self._layout() + for frame, was in self._drawn.items(): + now = layout.get(frame) + if now is None or now.path != was.path: + rr.log(was.path, rr.Arrows3D(origins=[], vectors=[]), static=True) -class RerunTFTree(Transformer[T, T]): - """Draw the tf tree's triads in step with the stream it passes through. + for frame, spot in layout.items(): + if self._drawn.get(frame) != spot: + rr.log( + spot.path, + rr.CoordinateFrame(f"tf#/{frame}"), + triad(self.axis_length * DEPTH_SCALE**spot.depth), + static=True, + ) - Drop it into a replay pipeline and every tf frame gets its labeled triad, - each logged at its own place on the timeline rather than in one lump up - front:: + self._drawn = layout - pipeline = lidar.transform(RerunTFTree(store.stream("tf", TFMessage))) - Window the tf stream the same way as the pipeline, or the tf that predates - the first observation all lands on that first frame. - """ +class RerunTFTree(Transformer[T, T]): + """Logs a recorded tf stream in step with the stream it passes through.""" - def __init__( - self, - tf: Stream[TFMessage], - axis_length: float = DEFAULT_AXIS_LENGTH, - timeline: str = DEFAULT_TIMELINE, - root: str = DEFAULT_TF_ROOT, - ) -> None: + def __init__(self, tf: Stream[TFMessage]) -> None: self._tf = tf - self._timeline = timeline - self._vis = TFTreeVis(axis_length=axis_length, root=root, settle=0.0) - - @property - def vis(self) -> TFTreeVis: - return self._vis + self._vis = TFTreeVis() def __call__(self, upstream: Iterator[Observation[T]]) -> Iterator[Observation[T]]: import rerun as rr - # Topology first, so no frame is given a path before its parent is known. - for tf_obs in self._tf: - self._vis.buffer.receive_tfmessage(tf_obs.data) - pending = iter(self._tf) head = next(pending, None) + floor: float | None = None for obs in upstream: + if floor is None: + # tf older than the replay would otherwise all land on frame one. + floor = obs.ts while head is not None and head.ts <= obs.ts: - rr.set_time(self._timeline, timestamp=head.ts) - self._vis.log(head.data) + if head.ts >= floor: + self._log(head) head = next(pending, None) - rr.set_time(self._timeline, timestamp=obs.ts) + rr.set_time(DEFAULT_TIMELINE, timestamp=obs.ts) yield obs + self._vis.flush() + + def _log(self, tf_obs: Observation[TFMessage]) -> None: + import rerun as rr + + rr.set_time(DEFAULT_TIMELINE, timestamp=tf_obs.ts) + for path, archetype in tf_obs.data.to_rerun(): + rr.log(path, archetype) + self._vis.log(tf_obs.data) From 8458db0623f103aabd3509765c1ddb81121f8e81 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 4 Aug 2026 12:08:20 -0700 Subject: [PATCH 03/13] Fix tree vis --- .../nav_3d/mls_planner/utils/plan_rrd.py | 51 ++++++++----------- .../navigation/unitree_go2_nav_3d.py | 3 +- dimos/robot/unitree/go2/zenoh/blueprints.py | 2 + dimos/visualization/rerun/bridge.py | 16 +++--- dimos/visualization/rerun/test_tf_tree.py | 19 ++++--- dimos/visualization/rerun/tf_tree.py | 37 ++++++++------ 6 files changed, 64 insertions(+), 64 deletions(-) diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index 594d35c507..9ae0799e49 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -183,54 +183,38 @@ def _plan_start( base_from_sensor: Transform | None, base_height: float, robot_height: float, -) -> tuple[tuple[float, float, float], Transform | None]: - """Ground-projected planner start, plus the base pose when tf has the mount. +) -> tuple[float, float, float]: + """Ground-projected planner start. - Without a tf stream the start is the sensor pose dropped by the robot height. + Without a tf stream this is the sensor pose dropped by the robot height. """ px, py, pz, *_ = pose if base_from_sensor is None: - return (float(px), float(py), float(pz) - robot_height), None + return (float(px), float(py), float(pz) - robot_height) base = _base_pose(pose, ts, base_from_sensor) - start = ( + return ( float(base.translation.x), float(base.translation.y), float(base.translation.z) - base_height, ) - return start, base def _log_odometry( - pose: tuple[float, ...], - ts: float, - trail: list[tuple[float, float, float]], - base: Transform | None, + pose: tuple[float, ...], ts: float, trail: list[tuple[float, float, float]] ) -> None: - """Trace the sensor moving throughout the scene.""" + """Trace the sensor moving throughout the scene. + + The pose itself is drawn by the tf tree, off the recorded tf stream. + """ import rerun as rr - px, py, pz, qx, qy, qz, qw = pose + px, py, pz, *_ = pose rr.set_time(TIMELINE, timestamp=ts) - rr.log( - "world/mid360_link", - rr.Transform3D(translation=[px, py, pz], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw])), - ) trail.append((px, py, pz)) if len(trail) > 1: rr.log( "world/mid360_path", rr.LineStrips3D([trail], colors=[SENSOR_PATH_COLOR], radii=0.015) ) - if base is None: - return - rr.log( - "world/base_link", - rr.Transform3D( - translation=[base.translation.x, base.translation.y, base.translation.z], - quaternion=rr.Quaternion( - xyzw=[base.rotation.x, base.rotation.y, base.rotation.z, base.rotation.w] - ), - ), - ) def _clearance_colors(clearance: NDArray[np.float32], clamp_m: float) -> NDArray[np.uint8]: @@ -619,8 +603,13 @@ def main( else 0.0 ) if base_from_sensor is not None: + # Rides the tf frame rather than a pose of its own, so the box cannot + # disagree with the tree. + rr.log( + "world/robot_body", rr.Transform3D(parent_frame=f"tf#/{BASE_FRAME}"), static=True + ) rr.log( - "world/base_link/outline", + "world/robot_body/outline", rr.Boxes3D( half_sizes=[ROBOT_LENGTH / 2, ROBOT_WIDTH / 2, robot_height / 2], colors=[(0, 255, 127)], @@ -629,7 +618,7 @@ def main( ) # wall_clearance is the planner's proxy for the robot radius. rr.log( - "world/base_link/clearance", + "world/robot_body/clearance", rr.Cylinders3D( lengths=[robot_height], radii=[wall_clearance], @@ -645,7 +634,7 @@ def main( for ray_obs in ray_pipeline: if ray_obs.pose_tuple is None: continue - start, base = _plan_start( + start = _plan_start( ray_obs.pose_tuple, ray_obs.ts, base_from_sensor, base_height, robot_height ) ref_timing = _process_frame( @@ -658,7 +647,7 @@ def main( ref_clearance, crop, ) - _log_odometry(ray_obs.pose_tuple, ray_obs.ts, sensor_trail, base) + _log_odometry(ray_obs.pose_tuple, ray_obs.ts, sensor_trail) frame += 1 print( f"frame={frame} configs={len(planners)} " diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 4f67155344..cd5146d829 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -109,8 +109,9 @@ def _static_robot_body(rr: Any) -> list[Any]: }, # Ring buffer replayed to a connecting viewer. Small so connect catches up fast. "memory_limit": "64MB", + "tf_axes": 0.5, # The robot box hangs off the base_link frame, on its own entity: world/tf - # holds the frame tree. + # holds the frame edges. "static": { "world/robot_body": _static_robot_body, }, diff --git a/dimos/robot/unitree/go2/zenoh/blueprints.py b/dimos/robot/unitree/go2/zenoh/blueprints.py index 708c400023..d958e04c08 100644 --- a/dimos/robot/unitree/go2/zenoh/blueprints.py +++ b/dimos/robot/unitree/go2/zenoh/blueprints.py @@ -110,6 +110,8 @@ def _rerun_config(visual_override: dict[str, Any] | None = None) -> dict[str, An """The bridge's own view, plus whatever the layer above it adds.""" return { "blueprint": _rerun_blueprint, + # Triad length in meters, on every frame that reaches the tf topic. + "tf_axes": 0.5, "visual_override": { "world/camera_info": _camera_info_to_pinhole, "world/pointlio_map": _render_map, diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 99a66035b2..e17e16606e 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -57,7 +57,7 @@ RerunOpenOption, ) from dimos.visualization.rerun.init import rerun_init -from dimos.visualization.rerun.tf_tree import DEFAULT_AXIS_LENGTH, TFTreeVis +from dimos.visualization.rerun.tf_tree import TFTreeVis if TYPE_CHECKING: from rerun._baseclasses import Archetype @@ -204,8 +204,8 @@ class Config(ModuleConfig): max_hz: dict[str, float] = field(default_factory=dict) entity_prefix: str = "world" - # Length in meters of the triad drawn on every tf frame, 0 to draw none. - tf_axes: float = DEFAULT_AXIS_LENGTH + # Length of the triads to draw + tf_axes: float = 0.0 topic_to_entity: Callable[[Any], str] | None = None connect_url: str | None = None memory_limit: str = "25%" @@ -252,7 +252,7 @@ def _new_tf_tree(self) -> TFTreeVis | None: return None return TFTreeVis( axis_length=self.config.tf_axes, - root=f"{self.config.entity_prefix}/frames", + root=f"{self.config.entity_prefix}/tf", ) @property @@ -341,12 +341,14 @@ def _on_message(self, msg: Any, topic: Any) -> None: # TFMessage for example returns list of (entity_path, archetype) tuples if is_rerun_multi(rerun_data): - for path, archetype in rerun_data: - rr.log(path, archetype) # Bound locally: stop() clears the tree from another thread. tf_tree = self._tf_tree if tf_tree is not None and isinstance(msg, TFMessage): - tf_tree.log(msg) + # The tree re-paths these: its own layout, not the flat one. + tf_tree.log(msg, [archetype for _, archetype in rerun_data]) + return + for path, archetype in rerun_data: + rr.log(path, archetype) else: rr.log(entity_path, cast("Archetype", rerun_data)) # if source msg carries a frame_id, attach the entity to that TF frame diff --git a/dimos/visualization/rerun/test_tf_tree.py b/dimos/visualization/rerun/test_tf_tree.py index 0477f52b68..737f478650 100644 --- a/dimos/visualization/rerun/test_tf_tree.py +++ b/dimos/visualization/rerun/test_tf_tree.py @@ -31,9 +31,8 @@ def paths(vis: TFTreeVis) -> dict[str, str]: def feed(vis: TFTreeVis, *messages: TFMessage) -> None: """tf republishes, and the tree draws once a message adds nothing new.""" - for msg in messages: - vis.log(msg) - vis.log(messages[-1]) + for msg in (*messages, messages[-1]): + vis.log(msg, [archetype for _, archetype in msg.to_rerun()]) def test_triads_nest_along_the_tree() -> None: @@ -41,9 +40,9 @@ def test_triads_nest_along_the_tree() -> None: feed(vis, TFMessage(edge("odom", "base_link"), edge("base_link", "camera/optical"))) assert paths(vis) == { - "odom": "world/frames/odom", - "base_link": "world/frames/odom/base_link", - "camera/optical": "world/frames/odom/base_link/camera\\/optical", + "odom": "world/tf/odom", + "base_link": "world/tf/odom/base_link", + "camera/optical": "world/tf/odom/base_link/camera\\/optical", } @@ -51,14 +50,14 @@ def test_a_late_root_re_parents_the_tree() -> None: """The mount tree is published seconds before the odometry that roots it.""" vis = TFTreeVis() feed(vis, TFMessage(edge("mid360_link", "base_link"))) - assert paths(vis)["base_link"] == "world/frames/mid360_link/base_link" + assert paths(vis)["base_link"] == "world/tf/mid360_link/base_link" feed(vis, TFMessage(edge("odom", "mid360_link"))) assert paths(vis) == { - "odom": "world/frames/odom", - "mid360_link": "world/frames/odom/mid360_link", - "base_link": "world/frames/odom/mid360_link/base_link", + "odom": "world/tf/odom", + "mid360_link": "world/tf/odom/mid360_link", + "base_link": "world/tf/odom/mid360_link/base_link", } diff --git a/dimos/visualization/rerun/tf_tree.py b/dimos/visualization/rerun/tf_tree.py index 468eb3a54b..57f211c6ca 100644 --- a/dimos/visualization/rerun/tf_tree.py +++ b/dimos/visualization/rerun/tf_tree.py @@ -26,6 +26,7 @@ from collections.abc import Iterable, Iterator import rerun as rr + from rerun._baseclasses import Archetype from dimos.memory2.stream import Stream from dimos.memory2.type.observation import Observation @@ -34,7 +35,11 @@ T = TypeVar("T") -DEFAULT_FRAMES_ROOT = "world/frames" +DEFAULT_TF_ROOT = "world/tf" +# Where the transforms themselves are declared. Off to the side and free of +# geometry: rerun pins a frame to the entity that declares it for the life of a +# recording, so these cannot follow the tree when it re-roots. +DEFAULT_LINKS_ROOT = "tf_links" DEFAULT_AXIS_LENGTH = 0.5 DEFAULT_TIMELINE = "ts" # Each level's triad relative to its parent's. @@ -65,34 +70,42 @@ class Placement: class TFTreeVis: - """Draws a labeled triad per tf frame, nested by entity path. + """Draws the tf tree, one nested entity per frame carrying a labeled triad. - Draws markers only. The transforms themselves are logged by whoever owns the - tf stream, at the flat paths TFMessage.to_rerun assigns. + Under root sits nothing but the tree. The transforms are declared apart from + it, under links, because a frame is pinned to its declaring entity while the + tree has to re-parent when a publisher starts late. """ def __init__( self, axis_length: float = DEFAULT_AXIS_LENGTH, - root: str = DEFAULT_FRAMES_ROOT, + root: str = DEFAULT_TF_ROOT, + links: str = DEFAULT_LINKS_ROOT, ) -> None: self.axis_length = axis_length self.root = root + self.links = links self._lock = threading.Lock() self._parents: dict[str, str] = {} self._drawn: dict[str, Placement] = {} self._pending = False - def log(self, msg: TFMessage) -> None: - """Redraw once a message arrives that adds nothing new. + def log(self, msg: TFMessage, archetypes: Iterable[Archetype]) -> None: + """Declare the transforms, then redraw once a message adds nothing new. Publishers split one tree across several messages, and drawing each of them walks the tree through shapes it never really had, leaving a stale entity behind every time. """ + import rerun as rr + if not msg.transforms: return with self._lock: + for transform, archetype in zip(msg.transforms, archetypes, strict=True): + child = rr.escape_entity_path_part(transform.child_frame_id) + rr.log(f"{self.links}/{child}", archetype) if self._learn(msg.transforms): self._pending = True elif self._pending: @@ -143,11 +156,7 @@ def place(frame: str, walked: frozenset[str]) -> Placement: return placed def _redraw(self) -> None: - """Move triads to match the tree as it is now. - - Rerun refuses to let the entity declaring a frame move, which is why the - triads carry a CoordinateFrame instead and declare nothing. - """ + """Move the tree to match the shape tf has now.""" import rerun as rr layout = self._layout() @@ -198,6 +207,4 @@ def _log(self, tf_obs: Observation[TFMessage]) -> None: import rerun as rr rr.set_time(DEFAULT_TIMELINE, timestamp=tf_obs.ts) - for path, archetype in tf_obs.data.to_rerun(): - rr.log(path, archetype) - self._vis.log(tf_obs.data) + self._vis.log(tf_obs.data, [archetype for _, archetype in tf_obs.data.to_rerun()]) From 907b663210ad98336c02566d226d0bf4eb94030c Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 4 Aug 2026 12:42:19 -0700 Subject: [PATCH 04/13] Always plan from base_link tf --- .../nav_3d/mls_planner/utils/plan_rrd.py | 143 +++++++++++------- dimos/protocol/tf/tf.py | 10 ++ 2 files changed, 95 insertions(+), 58 deletions(-) diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index 9ae0799e49..fa8e92d6b2 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -39,7 +39,7 @@ from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.navigation.nav_3d.mls_planner.mls_planner import MLSPlanner -from dimos.navigation.tf_pose import base_height_above_ground +from dimos.navigation.tf_pose import OdomBasePose, base_height_above_ground from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT, ROBOT_LENGTH, ROBOT_WIDTH from dimos.utils.data import resolve_named_path from dimos.visualization.rerun.tf_tree import RerunTFTree @@ -155,13 +155,51 @@ def _tf_over(store: SqliteStore, window: Stream[Any]) -> Stream[TFMessage] | Non return recorded.order_by("ts").time_range(first, last) -def _base_from_sensor(store: SqliteStore) -> Transform | None: - """Sensor to robot base link transform from the recorded tf stream.""" +class BaseSource: + """Where the body pose comes from. The planner never starts from the sensor. + + The mount leg is what a healthy recording carries, and it rides the replayed + odometry rather than the recorded one. Without it the body is read straight + off tf, which on a recording whose base_link hangs under a second odometry + is the wrong spot, visibly so. + """ + + def __init__(self, tf: StreamTF) -> None: + self._tf = tf + # Also primes the buffer: a replay tf loads nothing until it is asked. + self.leg = OdomBasePose(tf, BASE_FRAME).sensor_to_base(SENSOR_FRAME) + if BASE_FRAME not in tf.get_frames(): + raise typer.BadParameter( + f"recording has no {BASE_FRAME} on tf, so there is no body pose to plan from. " + "Record with Go2Mid360StaticTf running, or repair the tf stream." + ) + self.root = tf.get_root(BASE_FRAME) + if self.leg is None: + print( + f"no {SENSOR_FRAME} -> {BASE_FRAME} on tf: reading the body off " + f"{self.root} -> {BASE_FRAME} instead, which the replayed odometry " + "does not correct" + ) + + def base_height(self, robot_height: float) -> float: + """How far to drop the body pose to the ground, 0 without the mount.""" + return ( + 0.0 if self.leg is None else base_height_above_ground(robot_height, self.leg.inverse()) + ) + + def pose(self, sensor_pose: tuple[float, ...], ts: float) -> Transform | None: + if self.leg is not None: + return _base_pose(sensor_pose, ts, self.leg) + return self._tf.get(self.root, BASE_FRAME, time_point=ts) + + +def _base_source(store: SqliteStore) -> BaseSource: tf = StreamTF.from_store(store) if tf is None: - print("no tf stream in the recording; skipping the base_link triad") - return None - return tf.get(SENSOR_FRAME, BASE_FRAME) + raise typer.BadParameter( + "recording has no tf stream, so there is no body pose to plan from" + ) + return BaseSource(tf) def _base_pose(pose: tuple[float, ...], ts: float, base_from_sensor: Transform) -> Transform: @@ -177,21 +215,8 @@ def _base_pose(pose: tuple[float, ...], ts: float, base_from_sensor: Transform) return sensor + base_from_sensor -def _plan_start( - pose: tuple[float, ...], - ts: float, - base_from_sensor: Transform | None, - base_height: float, - robot_height: float, -) -> tuple[float, float, float]: - """Ground-projected planner start. - - Without a tf stream this is the sensor pose dropped by the robot height. - """ - px, py, pz, *_ = pose - if base_from_sensor is None: - return (float(px), float(py), float(pz) - robot_height) - base = _base_pose(pose, ts, base_from_sensor) +def _plan_start(base: Transform, base_height: float) -> tuple[float, float, float]: + """The body pose dropped to the ground, which is where the planner starts.""" return ( float(base.translation.x), float(base.translation.y), @@ -200,11 +225,13 @@ def _plan_start( def _log_odometry( - pose: tuple[float, ...], ts: float, trail: list[tuple[float, float, float]] + pose: tuple[float, ...], ts: float, trail: list[tuple[float, float, float]], base: Transform ) -> None: - """Trace the sensor moving throughout the scene. + """Trace the sensor through the scene, and put the body where the planner has it. - The pose itself is drawn by the tf tree, off the recorded tf stream. + The body is drawn on its own entity rather than a tf frame: on a recording + whose tf carries a second odometry estimate, the frame and the pose the + planner used are not the same place. """ import rerun as rr @@ -215,6 +242,15 @@ def _log_odometry( rr.log( "world/mid360_path", rr.LineStrips3D([trail], colors=[SENSOR_PATH_COLOR], radii=0.015) ) + rr.log( + "world/robot_body", + rr.Transform3D( + translation=[base.translation.x, base.translation.y, base.translation.z], + quaternion=rr.Quaternion( + xyzw=[base.rotation.x, base.rotation.y, base.rotation.z, base.rotation.w] + ), + ), + ) def _clearance_colors(clearance: NDArray[np.float32], clamp_m: float) -> NDArray[np.uint8]: @@ -596,37 +632,27 @@ def main( rr.log("world/goal", rr.Points3D([goal], colors=[[255, 0, 0]], radii=0.1), static=True) - base_from_sensor = _base_from_sensor(store) - base_height = ( - base_height_above_ground(robot_height, base_from_sensor.inverse()) - if base_from_sensor is not None - else 0.0 + base_source = _base_source(store) + base_height = base_source.base_height(robot_height) + rr.log( + "world/robot_body/outline", + rr.Boxes3D( + half_sizes=[ROBOT_LENGTH / 2, ROBOT_WIDTH / 2, robot_height / 2], + colors=[(0, 255, 127)], + ), + static=True, + ) + # wall_clearance is the planner's proxy for the robot radius. + rr.log( + "world/robot_body/clearance", + rr.Cylinders3D( + lengths=[robot_height], + radii=[wall_clearance], + colors=[(255, 120, 120, 80)], + fill_mode="solid", + ), + static=True, ) - if base_from_sensor is not None: - # Rides the tf frame rather than a pose of its own, so the box cannot - # disagree with the tree. - rr.log( - "world/robot_body", rr.Transform3D(parent_frame=f"tf#/{BASE_FRAME}"), static=True - ) - rr.log( - "world/robot_body/outline", - rr.Boxes3D( - half_sizes=[ROBOT_LENGTH / 2, ROBOT_WIDTH / 2, robot_height / 2], - colors=[(0, 255, 127)], - ), - static=True, - ) - # wall_clearance is the planner's proxy for the robot radius. - rr.log( - "world/robot_body/clearance", - rr.Cylinders3D( - lengths=[robot_height], - radii=[wall_clearance], - colors=[(255, 120, 120, 80)], - fill_mode="solid", - ), - static=True, - ) sensor_trail: list[tuple[float, float, float]] = [] try: @@ -634,9 +660,10 @@ def main( for ray_obs in ray_pipeline: if ray_obs.pose_tuple is None: continue - start = _plan_start( - ray_obs.pose_tuple, ray_obs.ts, base_from_sensor, base_height, robot_height - ) + base = base_source.pose(ray_obs.pose_tuple, ray_obs.ts) + if base is None: + continue + start = _plan_start(base, base_height) ref_timing = _process_frame( ray_obs, planners, @@ -647,7 +674,7 @@ def main( ref_clearance, crop, ) - _log_odometry(ray_obs.pose_tuple, ray_obs.ts, sensor_trail) + _log_odometry(ray_obs.pose_tuple, ray_obs.ts, sensor_trail, base) frame += 1 print( f"frame={frame} configs={len(planners)} " diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 2b0309e52e..95b165a79a 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -117,6 +117,16 @@ def get_frames(self) -> set[str]: frames.add(child) return frames + def get_root(self, frame_id: str) -> str: + """The frame at the top of frame_id's parent chain, frame_id if it has none.""" + with self._cv: + parents = {child: parent for parent, child in self.buffers} + walked: set[str] = set() + while frame_id in parents and frame_id not in walked: + walked.add(frame_id) + frame_id = parents[frame_id] + return frame_id + def get_connections(self, frame_id: str) -> set[str]: """Get all frames connected to the given frame (both as parent and child).""" connections = set() From 3a55299fe301cec77acf3a46de76a478f5d2a8ea Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 4 Aug 2026 14:20:59 -0700 Subject: [PATCH 05/13] Bug fix --- dimos/memory2/tf.py | 3 +- .../nav_3d/mls_planner/utils/plan_rrd.py | 37 +++++++-------- .../navigation/unitree_go2_nav_3d.py | 4 +- dimos/visualization/rerun/bridge.py | 2 + dimos/visualization/rerun/test_tf_tree.py | 2 +- dimos/visualization/rerun/tf_tree.py | 47 +++++++++---------- 6 files changed, 46 insertions(+), 49 deletions(-) diff --git a/dimos/memory2/tf.py b/dimos/memory2/tf.py index 76254c91bb..5df8f45888 100644 --- a/dimos/memory2/tf.py +++ b/dimos/memory2/tf.py @@ -31,8 +31,7 @@ def tf_stream(store: Any, stream: str = "tf") -> Stream[TFMessage] | None: """The recording's tf stream, or None if it has none. - Asking a store for an absent stream registers it and writes its tables, so - the name has to be checked first. + Asking a store for an absent stream creates it, so the name is checked first. """ if stream not in store.list_streams(): return None diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index fa8e92d6b2..aa0f3dfb4c 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -39,7 +39,7 @@ from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.navigation.nav_3d.mls_planner.mls_planner import MLSPlanner -from dimos.navigation.tf_pose import OdomBasePose, base_height_above_ground +from dimos.navigation.tf_pose import base_height_above_ground from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT, ROBOT_LENGTH, ROBOT_WIDTH from dimos.utils.data import resolve_named_path from dimos.visualization.rerun.tf_tree import RerunTFTree @@ -141,8 +141,7 @@ def _log_path_wp(waypoints: NDArray[np.float32] | None, entity: str, color: list def _tf_over(store: SqliteStore, window: Stream[Any]) -> Stream[TFMessage] | None: """The recorded tf stream clipped to another stream's span. - Absolute bounds: the relative ones anchor on each stream's own first - observation, and tf rarely starts on the same sample as the lidar. + Absolute bounds: relative ones anchor on each stream's own first observation. """ recorded = tf_stream(store) if recorded is None: @@ -158,16 +157,15 @@ def _tf_over(store: SqliteStore, window: Stream[Any]) -> Stream[TFMessage] | Non class BaseSource: """Where the body pose comes from. The planner never starts from the sensor. - The mount leg is what a healthy recording carries, and it rides the replayed - odometry rather than the recorded one. Without it the body is read straight - off tf, which on a recording whose base_link hangs under a second odometry - is the wrong spot, visibly so. + Prefers the mount leg, which rides the replayed odometry rather than the + recorded one, and falls back to reading the body straight off tf. """ - def __init__(self, tf: StreamTF) -> None: + def __init__(self, tf: StreamTF, start: float) -> None: self._tf = tf - # Also primes the buffer: a replay tf loads nothing until it is asked. - self.leg = OdomBasePose(tf, BASE_FRAME).sensor_to_base(SENSOR_FRAME) + # A replay tf loads nothing until asked, and an untimed lookup anchors on + # the end of the stream. Ask at the window this run replays. + self.leg = tf.get(SENSOR_FRAME, BASE_FRAME, time_point=start) if BASE_FRAME not in tf.get_frames(): raise typer.BadParameter( f"recording has no {BASE_FRAME} on tf, so there is no body pose to plan from. " @@ -193,13 +191,17 @@ def pose(self, sensor_pose: tuple[float, ...], ts: float) -> Transform | None: return self._tf.get(self.root, BASE_FRAME, time_point=ts) -def _base_source(store: SqliteStore) -> BaseSource: +def _base_source(store: SqliteStore, window: Stream[Any]) -> BaseSource: tf = StreamTF.from_store(store) if tf is None: raise typer.BadParameter( "recording has no tf stream, so there is no body pose to plan from" ) - return BaseSource(tf) + try: + start = window.first().ts + except LookupError: + raise typer.BadParameter("no data in the requested window") from None + return BaseSource(tf, start) def _base_pose(pose: tuple[float, ...], ts: float, base_from_sensor: Transform) -> Transform: @@ -216,7 +218,7 @@ def _base_pose(pose: tuple[float, ...], ts: float, base_from_sensor: Transform) def _plan_start(base: Transform, base_height: float) -> tuple[float, float, float]: - """The body pose dropped to the ground, which is where the planner starts.""" + """The body pose dropped to the ground, where the planner starts.""" return ( float(base.translation.x), float(base.translation.y), @@ -227,12 +229,7 @@ def _plan_start(base: Transform, base_height: float) -> tuple[float, float, floa def _log_odometry( pose: tuple[float, ...], ts: float, trail: list[tuple[float, float, float]], base: Transform ) -> None: - """Trace the sensor through the scene, and put the body where the planner has it. - - The body is drawn on its own entity rather than a tf frame: on a recording - whose tf carries a second odometry estimate, the frame and the pose the - planner used are not the same place. - """ + """Trace the sensor through the scene, and put the body where the planner has it.""" import rerun as rr px, py, pz, *_ = pose @@ -632,7 +629,7 @@ def main( rr.log("world/goal", rr.Points3D([goal], colors=[[255, 0, 0]], radii=0.1), static=True) - base_source = _base_source(store) + base_source = _base_source(store, lidar) base_height = base_source.base_height(robot_height) rr.log( "world/robot_body/outline", diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index cd5146d829..fb64583258 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -110,8 +110,8 @@ def _static_robot_body(rr: Any) -> list[Any]: # Ring buffer replayed to a connecting viewer. Small so connect catches up fast. "memory_limit": "64MB", "tf_axes": 0.5, - # The robot box hangs off the base_link frame, on its own entity: world/tf - # holds the frame edges. + # The robot box hangs off base_link on its own entity: a static transform + # under world/tf would override the live one. "static": { "world/robot_body": _static_robot_body, }, diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index e17e16606e..acdc85f1a0 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -597,6 +597,8 @@ def log_blueprint_graph(self, dot_code: str, module_names: list[str]) -> None: def stop(self) -> None: self._override_cache.clear() self._frame_attached.clear() + if self._tf_tree is not None: + self._tf_tree.flush() self._tf_tree = None super().stop() diff --git a/dimos/visualization/rerun/test_tf_tree.py b/dimos/visualization/rerun/test_tf_tree.py index 737f478650..fa0f5bc7bc 100644 --- a/dimos/visualization/rerun/test_tf_tree.py +++ b/dimos/visualization/rerun/test_tf_tree.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Where the triads get drawn. How they look is for a human with the viewer open.""" +"""Where the triads get drawn, not how they look.""" from __future__ import annotations diff --git a/dimos/visualization/rerun/tf_tree.py b/dimos/visualization/rerun/tf_tree.py index 57f211c6ca..00e77d39d4 100644 --- a/dimos/visualization/rerun/tf_tree.py +++ b/dimos/visualization/rerun/tf_tree.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Labeled axis triads for every frame of a tf tree.""" +"""Axis triads for every frame of a tf tree.""" from __future__ import annotations @@ -36,9 +36,8 @@ T = TypeVar("T") DEFAULT_TF_ROOT = "world/tf" -# Where the transforms themselves are declared. Off to the side and free of -# geometry: rerun pins a frame to the entity that declares it for the life of a -# recording, so these cannot follow the tree when it re-roots. +# Where the transforms are declared. Rerun pins a frame to its declaring entity +# for the life of a recording, so these cannot move with the tree. DEFAULT_LINKS_ROOT = "tf_links" DEFAULT_AXIS_LENGTH = 0.5 DEFAULT_TIMELINE = "ts" @@ -50,7 +49,7 @@ def triad(length: float) -> rr.Arrows3D: - """XYZ arrows, red green blue for x y z.""" + """XYZ arrows, red green blue.""" import rerun as rr return rr.Arrows3D( @@ -70,11 +69,10 @@ class Placement: class TFTreeVis: - """Draws the tf tree, one nested entity per frame carrying a labeled triad. + """Draws the tf tree, one nested entity per frame carrying a triad. - Under root sits nothing but the tree. The transforms are declared apart from - it, under links, because a frame is pinned to its declaring entity while the - tree has to re-parent when a publisher starts late. + Under root sits nothing but the tree. The transforms are declared apart + from it, under links. """ def __init__( @@ -94,9 +92,7 @@ def __init__( def log(self, msg: TFMessage, archetypes: Iterable[Archetype]) -> None: """Declare the transforms, then redraw once a message adds nothing new. - Publishers split one tree across several messages, and drawing each of - them walks the tree through shapes it never really had, leaving a stale - entity behind every time. + Publishers split one tree across several messages. """ import rerun as rr @@ -113,7 +109,7 @@ def log(self, msg: TFMessage, archetypes: Iterable[Archetype]) -> None: self._redraw() def flush(self) -> None: - """Draw a pending change that no later message arrived to trigger.""" + """Draw a pending change nothing else triggered.""" with self._lock: if self._pending: self._pending = False @@ -191,17 +187,20 @@ def __call__(self, upstream: Iterator[Observation[T]]) -> Iterator[Observation[T pending = iter(self._tf) head = next(pending, None) floor: float | None = None - for obs in upstream: - if floor is None: - # tf older than the replay would otherwise all land on frame one. - floor = obs.ts - while head is not None and head.ts <= obs.ts: - if head.ts >= floor: - self._log(head) - head = next(pending, None) - rr.set_time(DEFAULT_TIMELINE, timestamp=obs.ts) - yield obs - self._vis.flush() + try: + for obs in upstream: + if floor is None: + # tf older than the replay would otherwise all land on frame one. + floor = obs.ts + while head is not None and head.ts <= obs.ts: + if head.ts >= floor: + self._log(head) + head = next(pending, None) + rr.set_time(DEFAULT_TIMELINE, timestamp=obs.ts) + yield obs + finally: + # Ctrl+C abandons the generator rather than exhausting it. + self._vis.flush() def _log(self, tf_obs: Observation[TFMessage]) -> None: import rerun as rr From 0fafe5b37b7143e500d4ff222e351d364fd51f13 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Wed, 5 Aug 2026 11:09:56 -0700 Subject: [PATCH 06/13] Revert the planning changes so now this is just vis pr --- .../nav_3d/mls_planner/utils/plan_rrd.py | 131 ++++++++---------- dimos/protocol/tf/tf.py | 10 -- 2 files changed, 56 insertions(+), 85 deletions(-) diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index aa0f3dfb4c..5e401b96eb 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -154,54 +154,12 @@ def _tf_over(store: SqliteStore, window: Stream[Any]) -> Stream[TFMessage] | Non return recorded.order_by("ts").time_range(first, last) -class BaseSource: - """Where the body pose comes from. The planner never starts from the sensor. - - Prefers the mount leg, which rides the replayed odometry rather than the - recorded one, and falls back to reading the body straight off tf. - """ - - def __init__(self, tf: StreamTF, start: float) -> None: - self._tf = tf - # A replay tf loads nothing until asked, and an untimed lookup anchors on - # the end of the stream. Ask at the window this run replays. - self.leg = tf.get(SENSOR_FRAME, BASE_FRAME, time_point=start) - if BASE_FRAME not in tf.get_frames(): - raise typer.BadParameter( - f"recording has no {BASE_FRAME} on tf, so there is no body pose to plan from. " - "Record with Go2Mid360StaticTf running, or repair the tf stream." - ) - self.root = tf.get_root(BASE_FRAME) - if self.leg is None: - print( - f"no {SENSOR_FRAME} -> {BASE_FRAME} on tf: reading the body off " - f"{self.root} -> {BASE_FRAME} instead, which the replayed odometry " - "does not correct" - ) - - def base_height(self, robot_height: float) -> float: - """How far to drop the body pose to the ground, 0 without the mount.""" - return ( - 0.0 if self.leg is None else base_height_above_ground(robot_height, self.leg.inverse()) - ) - - def pose(self, sensor_pose: tuple[float, ...], ts: float) -> Transform | None: - if self.leg is not None: - return _base_pose(sensor_pose, ts, self.leg) - return self._tf.get(self.root, BASE_FRAME, time_point=ts) - - -def _base_source(store: SqliteStore, window: Stream[Any]) -> BaseSource: +def _base_from_sensor(store: SqliteStore) -> Transform | None: + """Sensor to robot base link transform from the recorded tf stream.""" tf = StreamTF.from_store(store) if tf is None: - raise typer.BadParameter( - "recording has no tf stream, so there is no body pose to plan from" - ) - try: - start = window.first().ts - except LookupError: - raise typer.BadParameter("no data in the requested window") from None - return BaseSource(tf, start) + return None + return tf.get(SENSOR_FRAME, BASE_FRAME) def _base_pose(pose: tuple[float, ...], ts: float, base_from_sensor: Transform) -> Transform: @@ -217,19 +175,36 @@ def _base_pose(pose: tuple[float, ...], ts: float, base_from_sensor: Transform) return sensor + base_from_sensor -def _plan_start(base: Transform, base_height: float) -> tuple[float, float, float]: - """The body pose dropped to the ground, where the planner starts.""" - return ( +def _plan_start( + pose: tuple[float, ...], + ts: float, + base_from_sensor: Transform | None, + base_height: float, + robot_height: float, +) -> tuple[tuple[float, float, float], Transform | None]: + """Ground-projected planner start, plus the base pose when tf has the mount. + + Without a tf stream the start is the sensor pose dropped by the robot height. + """ + px, py, pz, *_ = pose + if base_from_sensor is None: + return (float(px), float(py), float(pz) - robot_height), None + base = _base_pose(pose, ts, base_from_sensor) + start = ( float(base.translation.x), float(base.translation.y), float(base.translation.z) - base_height, ) + return start, base def _log_odometry( - pose: tuple[float, ...], ts: float, trail: list[tuple[float, float, float]], base: Transform + pose: tuple[float, ...], + ts: float, + trail: list[tuple[float, float, float]], + base: Transform | None, ) -> None: - """Trace the sensor through the scene, and put the body where the planner has it.""" + """Trace the sensor moving throughout the scene.""" import rerun as rr px, py, pz, *_ = pose @@ -239,6 +214,8 @@ def _log_odometry( rr.log( "world/mid360_path", rr.LineStrips3D([trail], colors=[SENSOR_PATH_COLOR], radii=0.015) ) + if base is None: + return rr.log( "world/robot_body", rr.Transform3D( @@ -629,27 +606,32 @@ def main( rr.log("world/goal", rr.Points3D([goal], colors=[[255, 0, 0]], radii=0.1), static=True) - base_source = _base_source(store, lidar) - base_height = base_source.base_height(robot_height) - rr.log( - "world/robot_body/outline", - rr.Boxes3D( - half_sizes=[ROBOT_LENGTH / 2, ROBOT_WIDTH / 2, robot_height / 2], - colors=[(0, 255, 127)], - ), - static=True, - ) - # wall_clearance is the planner's proxy for the robot radius. - rr.log( - "world/robot_body/clearance", - rr.Cylinders3D( - lengths=[robot_height], - radii=[wall_clearance], - colors=[(255, 120, 120, 80)], - fill_mode="solid", - ), - static=True, + base_from_sensor = _base_from_sensor(store) + base_height = ( + base_height_above_ground(robot_height, base_from_sensor.inverse()) + if base_from_sensor is not None + else 0.0 ) + if base_from_sensor is not None: + rr.log( + "world/robot_body/outline", + rr.Boxes3D( + half_sizes=[ROBOT_LENGTH / 2, ROBOT_WIDTH / 2, robot_height / 2], + colors=[(0, 255, 127)], + ), + static=True, + ) + # wall_clearance is the planner's proxy for the robot radius. + rr.log( + "world/robot_body/clearance", + rr.Cylinders3D( + lengths=[robot_height], + radii=[wall_clearance], + colors=[(255, 120, 120, 80)], + fill_mode="solid", + ), + static=True, + ) sensor_trail: list[tuple[float, float, float]] = [] try: @@ -657,10 +639,9 @@ def main( for ray_obs in ray_pipeline: if ray_obs.pose_tuple is None: continue - base = base_source.pose(ray_obs.pose_tuple, ray_obs.ts) - if base is None: - continue - start = _plan_start(base, base_height) + start, base = _plan_start( + ray_obs.pose_tuple, ray_obs.ts, base_from_sensor, base_height, robot_height + ) ref_timing = _process_frame( ray_obs, planners, diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 95b165a79a..2b0309e52e 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -117,16 +117,6 @@ def get_frames(self) -> set[str]: frames.add(child) return frames - def get_root(self, frame_id: str) -> str: - """The frame at the top of frame_id's parent chain, frame_id if it has none.""" - with self._cv: - parents = {child: parent for parent, child in self.buffers} - walked: set[str] = set() - while frame_id in parents and frame_id not in walked: - walked.add(frame_id) - frame_id = parents[frame_id] - return frame_id - def get_connections(self, frame_id: str) -> set[str]: """Get all frames connected to the given frame (both as parent and child).""" connections = set() From 6eb120e3c5f90fa1be1ffc95cef6b7f261b66b5e Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Wed, 5 Aug 2026 11:11:45 -0700 Subject: [PATCH 07/13] Increase test time --- native/rust/dimos-module/src/module.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index dc953bd24b..c59bc614ad 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -659,7 +659,7 @@ mod tests { } async fn wait_for(what: &str, mut cond: impl FnMut() -> bool) { - let deadline = Instant::now() + Duration::from_secs(1); + let deadline = Instant::now() + Duration::from_secs(5); while !cond() { assert!(Instant::now() < deadline, "timed out waiting for {what}"); tokio::time::sleep(Duration::from_millis(10)).await; From d2fc831f5fb12281f0cb3ebfd2f36dbd4403227c Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Wed, 5 Aug 2026 11:22:12 -0700 Subject: [PATCH 08/13] Clean up some stuff --- dimos/memory2/tf.py | 6 ++---- dimos/robot/unitree/go2/zenoh/blueprints.py | 1 - dimos/visualization/rerun/bridge.py | 2 -- dimos/visualization/rerun/tf_tree.py | 7 +------ 4 files changed, 3 insertions(+), 13 deletions(-) diff --git a/dimos/memory2/tf.py b/dimos/memory2/tf.py index 5df8f45888..e9d0041d44 100644 --- a/dimos/memory2/tf.py +++ b/dimos/memory2/tf.py @@ -29,10 +29,8 @@ def tf_stream(store: Any, stream: str = "tf") -> Stream[TFMessage] | None: - """The recording's tf stream, or None if it has none. - - Asking a store for an absent stream creates it, so the name is checked first. - """ + """The recording's tf stream, or None if it has none.""" + # check if it's there first so we don't create one if stream not in store.list_streams(): return None return cast("Stream[TFMessage]", store.stream(stream, TFMessage)) diff --git a/dimos/robot/unitree/go2/zenoh/blueprints.py b/dimos/robot/unitree/go2/zenoh/blueprints.py index d958e04c08..ffd3ad4e2c 100644 --- a/dimos/robot/unitree/go2/zenoh/blueprints.py +++ b/dimos/robot/unitree/go2/zenoh/blueprints.py @@ -110,7 +110,6 @@ def _rerun_config(visual_override: dict[str, Any] | None = None) -> dict[str, An """The bridge's own view, plus whatever the layer above it adds.""" return { "blueprint": _rerun_blueprint, - # Triad length in meters, on every frame that reaches the tf topic. "tf_axes": 0.5, "visual_override": { "world/camera_info": _camera_info_to_pinhole, diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index acdc85f1a0..6b21306dbb 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -341,10 +341,8 @@ def _on_message(self, msg: Any, topic: Any) -> None: # TFMessage for example returns list of (entity_path, archetype) tuples if is_rerun_multi(rerun_data): - # Bound locally: stop() clears the tree from another thread. tf_tree = self._tf_tree if tf_tree is not None and isinstance(msg, TFMessage): - # The tree re-paths these: its own layout, not the flat one. tf_tree.log(msg, [archetype for _, archetype in rerun_data]) return for path, archetype in rerun_data: diff --git a/dimos/visualization/rerun/tf_tree.py b/dimos/visualization/rerun/tf_tree.py index 00e77d39d4..745f0e40d4 100644 --- a/dimos/visualization/rerun/tf_tree.py +++ b/dimos/visualization/rerun/tf_tree.py @@ -69,11 +69,7 @@ class Placement: class TFTreeVis: - """Draws the tf tree, one nested entity per frame carrying a triad. - - Under root sits nothing but the tree. The transforms are declared apart - from it, under links. - """ + """Draws the tf tree, one nested entity per frame carrying a triad.""" def __init__( self, @@ -199,7 +195,6 @@ def __call__(self, upstream: Iterator[Observation[T]]) -> Iterator[Observation[T rr.set_time(DEFAULT_TIMELINE, timestamp=obs.ts) yield obs finally: - # Ctrl+C abandons the generator rather than exhausting it. self._vis.flush() def _log(self, tf_obs: Observation[TFMessage]) -> None: From 7799f21ed54627d88870e47900083e12453141d3 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Wed, 5 Aug 2026 11:47:19 -0700 Subject: [PATCH 09/13] Add vis by default --- dimos/mapping/ray_tracing/rust/flake.nix | 2 +- dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py | 1 + .../unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py | 1 - 3 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dimos/mapping/ray_tracing/rust/flake.nix b/dimos/mapping/ray_tracing/rust/flake.nix index 15c40909f3..1c37037eee 100644 --- a/dimos/mapping/ray_tracing/rust/flake.nix +++ b/dimos/mapping/ray_tracing/rust/flake.nix @@ -34,7 +34,7 @@ cargoRoot = "dimos/mapping/ray_tracing/rust"; buildAndTestSubdir = "dimos/mapping/ray_tracing/rust"; - cargoHash = "sha256-6a8GHRSKI6mjg9HNbrestCud8xZtF8HaD0bWVMbl7N8="; + cargoHash = "sha256-0xv5Hb9q0goNiFUU2FaTg7NhUFD0eZTobcR4ssXVDNg="; meta.mainProgram = "voxel_ray_tracing"; }; diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index dfdbd35593..c55c6e1b9c 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -97,6 +97,7 @@ def _go2_rerun_blueprint() -> Any: "world/global_costmap": 0, # publishes at ~7.6 Hz "world/lidar": 1, # publishes at ~7.7 Hz; hidden by default in the blueprint }, + "tf_axes": 0.5, # slapping a go2 shaped box on the base_link frame "static": { "world/robot_body": _static_robot_body, diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index fb64583258..6b676f587a 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -109,7 +109,6 @@ def _static_robot_body(rr: Any) -> list[Any]: }, # Ring buffer replayed to a connecting viewer. Small so connect catches up fast. "memory_limit": "64MB", - "tf_axes": 0.5, # The robot box hangs off base_link on its own entity: a static transform # under world/tf would override the live one. "static": { From bd0f15cc351293e2103b01efcfe1c3eec1f9264a Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Wed, 5 Aug 2026 12:02:37 -0700 Subject: [PATCH 10/13] Show robot frame --- dimos/robot/unitree/go2/zenoh/blueprints.py | 20 +++++++++++++++++-- .../unitree/go2/zenoh/zenohconnection.py | 7 +++---- 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/dimos/robot/unitree/go2/zenoh/blueprints.py b/dimos/robot/unitree/go2/zenoh/blueprints.py index ffd3ad4e2c..f7166f0050 100644 --- a/dimos/robot/unitree/go2/zenoh/blueprints.py +++ b/dimos/robot/unitree/go2/zenoh/blueprints.py @@ -38,7 +38,7 @@ from dimos.navigation.nav_3d.mls_planner.goal_relay import GoalRelay from dimos.navigation.nav_3d.mls_planner.mls_planner_native import MLSPlannerNative from dimos.navigation.nav_3d.mls_planner.viz import planner_visual_override -from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT +from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT, ROBOT_LENGTH, ROBOT_WIDTH from dimos.robot.unitree.go2.zenoh.zenohconnection import GO2Zenoh from dimos.visualization.vis_module import vis_module @@ -48,7 +48,18 @@ planner_viz_hz = 2.0 # GO2Zenoh publishes this mount onto tf, where nav reads its odometry corrections. -MID360_MOUNT_RPY_DEG = (-60.0, 0.0, -90.0) +MID360_MOUNT_RPY_DEG = (0.0, 60.0, 0.0) + + +def _static_robot_body(rr: Any) -> list[Any]: + """Go2-shaped box on the body frame.""" + return [ + rr.Boxes3D( + half_sizes=[ROBOT_LENGTH / 2, ROBOT_WIDTH / 2, ROBOT_HEIGHT / 2], + colors=[(0, 255, 127)], + ), + rr.Transform3D(parent_frame="tf#/base_link"), + ] def _camera_info_to_pinhole(camera_info: Any) -> Any: @@ -111,6 +122,11 @@ def _rerun_config(visual_override: dict[str, Any] | None = None) -> dict[str, An return { "blueprint": _rerun_blueprint, "tf_axes": 0.5, + # The robot box hangs off base_link on its own entity: a static transform + # under world/tf would override the live one. + "static": { + "world/robot_body": _static_robot_body, + }, "visual_override": { "world/camera_info": _camera_info_to_pinhole, "world/pointlio_map": _render_map, diff --git a/dimos/robot/unitree/go2/zenoh/zenohconnection.py b/dimos/robot/unitree/go2/zenoh/zenohconnection.py index 02c3a189ce..f109d94d1a 100644 --- a/dimos/robot/unitree/go2/zenoh/zenohconnection.py +++ b/dimos/robot/unitree/go2/zenoh/zenohconnection.py @@ -57,10 +57,9 @@ class GO2ZenohConfig(StaticTfPublisherConfig): - # front_camera -> mid360_link, fixed-axis rpy in degrees. The 60 deg tilt lands on - # roll because the lidar sits yawed 90 deg on its bracket. Both yaw signs level the - # body but differ by 180 deg of heading — flip it if the camera looks backwards. - mid360_mount_rpy_deg: tuple[float, float, float] = (-60.0, 0.0, -90.0) + # front_camera -> mid360_link, fixed-axis rpy in degrees: pointing straight ahead, + # pitched 60 deg down. + mid360_mount_rpy_deg: tuple[float, float, float] = (0.0, 60.0, 0.0) camera_info_hz: float = Field(default=1.0, gt=0.0) From 92ac379f4a1aa7598808ad015080616591705a10 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Wed, 5 Aug 2026 12:19:25 -0700 Subject: [PATCH 11/13] Add lidar presets --- dimos/robot/unitree/go2/zenoh/blueprints.py | 5 +-- .../unitree/go2/zenoh/zenohconnection.py | 32 ++++++++++++++++--- 2 files changed, 30 insertions(+), 7 deletions(-) diff --git a/dimos/robot/unitree/go2/zenoh/blueprints.py b/dimos/robot/unitree/go2/zenoh/blueprints.py index f7166f0050..bada519f0d 100644 --- a/dimos/robot/unitree/go2/zenoh/blueprints.py +++ b/dimos/robot/unitree/go2/zenoh/blueprints.py @@ -48,7 +48,8 @@ planner_viz_hz = 2.0 # GO2Zenoh publishes this mount onto tf, where nav reads its odometry corrections. -MID360_MOUNT_RPY_DEG = (0.0, 60.0, 0.0) +# Either a raw (roll, pitch, yaw) tuple in degrees or a GO2ZenohConfig.mid360_mount preset. +MID360_MOUNT = "SF" def _static_robot_body(rr: Any) -> list[Any]: @@ -144,7 +145,7 @@ def _rerun_config(visual_override: dict[str, Any] | None = None) -> dict[str, An # is the layer to drive from when something upstream is suspect. go2_zenoh_basic = autoconnect( vis_module(viewer_backend=global_config.viewer, rerun_config=_rerun_config()), - GO2Zenoh.blueprint(mid360_mount_rpy_deg=MID360_MOUNT_RPY_DEG), + GO2Zenoh.blueprint(mid360_mount=MID360_MOUNT), MovementManager.blueprint(), ).global_config(transport="zenoh", n_workers=4, robot_model="unitree_go2") diff --git a/dimos/robot/unitree/go2/zenoh/zenohconnection.py b/dimos/robot/unitree/go2/zenoh/zenohconnection.py index f109d94d1a..9c194b9acf 100644 --- a/dimos/robot/unitree/go2/zenoh/zenohconnection.py +++ b/dimos/robot/unitree/go2/zenoh/zenohconnection.py @@ -29,8 +29,9 @@ import math import threading import time +from typing import Any -from pydantic import Field +from pydantic import Field, field_validator from reactivex.disposable import Disposable from dimos.core.core import rpc @@ -55,13 +56,34 @@ # rpy mapping a sensor frame to its optical frame (x-right, y-down, z-forward) OPTICAL_RPY = Vector3(-math.pi / 2, 0.0, -math.pi / 2) +# front_camera -> mid360_link, fixed-axis rpy in degrees, by rig. +MID360_MOUNT_PRESETS: dict[str, tuple[float, float, float]] = { + # Pointing straight ahead, pitched 60 deg down. + "SF": (0.0, 60.0, 0.0), + # The 60 deg tilt lands on roll because this lidar sits yawed 90 deg on its bracket. + "ATHENS": (-60.0, 0.0, -90.0), +} + class GO2ZenohConfig(StaticTfPublisherConfig): - # front_camera -> mid360_link, fixed-axis rpy in degrees: pointing straight ahead, - # pitched 60 deg down. - mid360_mount_rpy_deg: tuple[float, float, float] = (0.0, 60.0, 0.0) + # front_camera -> mid360_link, fixed-axis rpy in degrees. Either a raw (roll, pitch, + # yaw) tuple or a name from MID360_MOUNT_PRESETS. + mid360_mount: tuple[float, float, float] | str = MID360_MOUNT_PRESETS["SF"] camera_info_hz: float = Field(default=1.0, gt=0.0) + @field_validator("mid360_mount", mode="before") + @classmethod + def _resolve_mid360_mount(cls, value: Any) -> Any: + if isinstance(value, str): + try: + return MID360_MOUNT_PRESETS[value] + except KeyError: + raise ValueError( + f"unknown mid360_mount preset {value!r}; " + f"expected one of {sorted(MID360_MOUNT_PRESETS)}" + ) from None + return value + class GO2Zenoh(StaticTfPublisher): """The go2's zenoh-side streams, plus the static data the robot doesn't send.""" @@ -161,7 +183,7 @@ def transforms(self) -> list[Transform]: camera_to_mid360 = Transform( translation=MID360_XYZ, rotation=Quaternion.from_euler( - Vector3(*(math.radians(d) for d in self.config.mid360_mount_rpy_deg)) + Vector3(*(math.radians(float(d)) for d in self.config.mid360_mount)) ), frame_id="front_camera", child_frame_id="mid360_link", From 1938bd5c2018e55ec01c3afe14802c6a1a6c8979 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Wed, 5 Aug 2026 13:31:29 -0700 Subject: [PATCH 12/13] Do everything in the tf messages --- dimos/msgs/tf2_msgs/TFMessage.py | 135 ++++++++---- .../nav_3d/mls_planner/utils/plan_rrd.py | 25 ++- dimos/visualization/rerun/bridge.py | 21 +- dimos/visualization/rerun/test_tf_tree.py | 68 ------ dimos/visualization/rerun/tf_tree.py | 204 ------------------ 5 files changed, 126 insertions(+), 327 deletions(-) delete mode 100644 dimos/visualization/rerun/test_tf_tree.py delete mode 100644 dimos/visualization/rerun/tf_tree.py diff --git a/dimos/msgs/tf2_msgs/TFMessage.py b/dimos/msgs/tf2_msgs/TFMessage.py index 7a47a96e6d..441e1772b6 100644 --- a/dimos/msgs/tf2_msgs/TFMessage.py +++ b/dimos/msgs/tf2_msgs/TFMessage.py @@ -12,19 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -# 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 __future__ import annotations from typing import TYPE_CHECKING, BinaryIO @@ -36,10 +23,90 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Iterable, Iterator from dimos.visualization.rerun.bridge import RerunMulti +DEFAULT_LINKS_ROOT = "tf_links" +DEFAULT_TF_ROOT = "world/tf" +DEFAULT_AXIS_LENGTH = 0.5 +DEPTH_SCALE = 0.8 +AXIS_WIDTH_UI_POINTS = 2.0 +AXIS_COLORS = [[255, 0, 0], [0, 255, 0], [0, 0, 255]] + + +def _triad(length: float): # type: ignore[no-untyped-def] + """XYZ arrows, red green blue.""" + import rerun as rr + + return rr.Arrows3D( + origins=[[0.0, 0.0, 0.0]] * 3, + vectors=[[length, 0.0, 0.0], [0.0, length, 0.0], [0.0, 0.0, length]], + colors=AXIS_COLORS, + radii=rr.components.Radius.ui_points(AXIS_WIDTH_UI_POINTS), + ) + + +class TfFrameTree: + """Nests each frame's triad under its parent, mirroring the tf tree in the panel. + + Rerun never lets an entity move, so a frame placed as a root before its + real parent arrives gets its old entity blanked and a new one declared + once that parent is known. + """ + + def __init__( + self, axis_length: float = DEFAULT_AXIS_LENGTH, root: str = DEFAULT_TF_ROOT + ) -> None: + self.axis_length = axis_length + self.root = root + self._parents: dict[str, str] = {} + self._placed: dict[str, tuple[str, int]] = {} # frame -> (path, depth) + + def placements(self) -> dict[str, str]: + return {frame: path for frame, (path, _depth) in self._placed.items()} + + def update(self, transforms: Iterable[Transform]) -> None: + learned = False + for transform in transforms: + if self._parents.get(transform.child_frame_id) != transform.frame_id: + self._parents[transform.child_frame_id] = transform.frame_id + learned = True + if learned: + self._redraw() + + def _place(self, frame: str, walked: frozenset[str]) -> tuple[str, int]: + import rerun as rr + + part = rr.escape_entity_path_part(frame) + parent = self._parents.get(frame) + if parent is None or parent in walked: + return f"{self.root}/{part}", 0 + parent_path, parent_depth = self._place(parent, walked | {frame}) + return f"{parent_path}/{part}", parent_depth + 1 + + def _redraw(self) -> None: + import rerun as rr + + frames = {*self._parents, *self._parents.values()} + placed = {frame: self._place(frame, frozenset()) for frame in frames} + + for frame, (old_path, _depth) in self._placed.items(): + new = placed.get(frame) + if new is None or new[0] != old_path: + rr.log(old_path, rr.Arrows3D(origins=[], vectors=[]), static=True) + + for frame, (path, depth) in placed.items(): + if self._placed.get(frame) != (path, depth): + rr.log( + path, + rr.CoordinateFrame(f"tf#/{frame}"), + _triad(self.axis_length * DEPTH_SCALE**depth), + static=True, + ) + + self._placed = placed + class TFMessage: """TFMessage that accepts Transform objects and encodes to LCM format.""" @@ -56,14 +123,8 @@ def add_transform(self, transform: Transform, child_frame_id: str = "base_link") self.transforms_length = len(self.transforms) def lcm_encode(self) -> bytes: - """Encode as LCM TFMessage. - - Args: - child_frame_ids: Optional list of child frame IDs for each transform. - If not provided, defaults to "base_link" for all. - """ - - res = list(map(lambda t: t.lcm_transform(), self.transforms)) + """Encode as LCM TFMessage.""" + res = [t.lcm_transform() for t in self.transforms] lcm_msg = LCMTFMessage( transforms_length=len(self.transforms), @@ -77,15 +138,12 @@ def lcm_decode(cls, data: bytes | BinaryIO) -> TFMessage: """Decode from LCM TFMessage bytes.""" lcm_msg = LCMTFMessage.lcm_decode(data) - # Convert LCM TransformStamped objects to Transform objects transforms = [] for lcm_transform_stamped in lcm_msg.transforms: - # Extract timestamp ts = lcm_transform_stamped.header.stamp.sec + ( lcm_transform_stamped.header.stamp.nsec / 1_000_000_000 ) - # Create Transform with our custom types lcm_trans = lcm_transform_stamped.transform.translation lcm_rot = lcm_transform_stamped.transform.rotation @@ -123,24 +181,19 @@ def __str__(self) -> str: ) return "\n".join(lines) - def to_rerun(self) -> RerunMulti: - """Convert to a list of rerun Transform3D archetypes. - - Returns a list of tuples (entity_path, Transform3D) for each transform - in the message. The entity_path is derived from the child_frame_id and - logged under `world/tf/...` so it is visible under the default `world` - origin while keeping TF visualization isolated from semantic entities - like `world/robot/...`. - - Returns: - List of (entity_path, rr.Transform3D) tuples + def to_rerun(self, tree: TfFrameTree | None = None) -> RerunMulti: + """Convert to (entity_path, archetype) pairs to log to rerun. - Example: - for path, transform in tf_msg.to_rerun(): - rr.log(path, transform) + Pass a TfFrameTree to also nest a triad per frame under its ancestors, + matching the tf tree's shape in the entity panel. """ + import rerun as rr + results: RerunMulti = [] for transform in self.transforms: - entity_path = f"world/tf/{transform.child_frame_id}" - results.append((entity_path, transform.to_rerun())) + path = f"{DEFAULT_LINKS_ROOT}/{rr.escape_entity_path_part(transform.child_frame_id)}" + results.append((path, transform.to_rerun())) + + if tree is not None: + tree.update(self.transforms) return results diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index 5e401b96eb..8e3261926d 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -37,12 +37,11 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.tf2_msgs.TFMessage import TfFrameTree, TFMessage from dimos.navigation.nav_3d.mls_planner.mls_planner import MLSPlanner from dimos.navigation.tf_pose import base_height_above_ground from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT, ROBOT_LENGTH, ROBOT_WIDTH from dimos.utils.data import resolve_named_path -from dimos.visualization.rerun.tf_tree import RerunTFTree if TYPE_CHECKING: import rerun.blueprint as rrb @@ -154,6 +153,24 @@ def _tf_over(store: SqliteStore, window: Stream[Any]) -> Stream[TFMessage] | Non return recorded.order_by("ts").time_range(first, last) +class _TfSync: + """Logs a recorded tf stream up to a given timestamp, in step with another stream.""" + + def __init__(self, tf: Stream[TFMessage] | None) -> None: + self._pending = iter(tf) if tf is not None else iter(()) + self._next = next(self._pending, None) + self._tree = TfFrameTree() + + def up_to(self, ts: float) -> None: + import rerun as rr + + while self._next is not None and self._next.ts <= ts: + rr.set_time(TIMELINE, timestamp=self._next.ts) + for path, archetype in self._next.data.to_rerun(self._tree): + rr.log(path, archetype) + self._next = next(self._pending, None) + + def _base_from_sensor(store: SqliteStore) -> Transform | None: """Sensor to robot base link transform from the recorded tf stream.""" tf = StreamTF.from_store(store) @@ -588,8 +605,7 @@ def main( support_min=support_min, ) ) - if tf is not None: - ray_pipeline = ray_pipeline.transform(RerunTFTree(tf)) + tf_sync = _TfSync(tf) configs = _parse_configs(config, wall_clearance, wall_buffer, wall_buffer_weight) ref_clearance = configs[0][0] @@ -637,6 +653,7 @@ def main( try: frame = 0 for ray_obs in ray_pipeline: + tf_sync.up_to(ray_obs.ts) if ray_obs.pose_tuple is None: continue start, base = _plan_start( diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 6b21306dbb..e66517015f 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -22,6 +22,7 @@ import socket import subprocess import sys +import threading import time from typing import ( TYPE_CHECKING, @@ -41,7 +42,7 @@ from dimos.core.core import rpc from dimos.core.global_config import global_config from dimos.core.module import Module, ModuleConfig -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.msgs.tf2_msgs.TFMessage import TfFrameTree, TFMessage from dimos.protocol.pubsub.impl.lcmpubsub import LCM from dimos.protocol.pubsub.impl.zenohpubsub import Zenoh from dimos.protocol.pubsub.patterns import Glob, pattern_matches @@ -57,7 +58,6 @@ RerunOpenOption, ) from dimos.visualization.rerun.init import rerun_init -from dimos.visualization.rerun.tf_tree import TFTreeVis if TYPE_CHECKING: from rerun._baseclasses import Archetype @@ -245,12 +245,13 @@ def __init__(self, **kwargs: Any) -> None: self._last_log = {} self._override_cache: dict[str, Callable[[Any], RerunData | None]] = {} self._frame_attached: dict[str, str] = {} + self._tf_lock = threading.Lock() self._tf_tree = self._new_tf_tree() - def _new_tf_tree(self) -> TFTreeVis | None: + def _new_tf_tree(self) -> TfFrameTree | None: if self.config.tf_axes <= 0: return None - return TFTreeVis( + return TfFrameTree( axis_length=self.config.tf_axes, root=f"{self.config.entity_prefix}/tf", ) @@ -334,6 +335,12 @@ def _on_message(self, msg: Any, topic: Any) -> None: return self._last_log[entity_path] = now + if self._tf_tree is not None and isinstance(msg, TFMessage): + with self._tf_lock: + for path, archetype in msg.to_rerun(self._tf_tree): + rr.log(path, archetype) + return + rerun_data: RerunData | None = self._visual_override_for_entity_path(entity_path)(msg) if not rerun_data: @@ -341,10 +348,6 @@ def _on_message(self, msg: Any, topic: Any) -> None: # TFMessage for example returns list of (entity_path, archetype) tuples if is_rerun_multi(rerun_data): - tf_tree = self._tf_tree - if tf_tree is not None and isinstance(msg, TFMessage): - tf_tree.log(msg, [archetype for _, archetype in rerun_data]) - return for path, archetype in rerun_data: rr.log(path, archetype) else: @@ -595,8 +598,6 @@ def log_blueprint_graph(self, dot_code: str, module_names: list[str]) -> None: def stop(self) -> None: self._override_cache.clear() self._frame_attached.clear() - if self._tf_tree is not None: - self._tf_tree.flush() self._tf_tree = None super().stop() diff --git a/dimos/visualization/rerun/test_tf_tree.py b/dimos/visualization/rerun/test_tf_tree.py deleted file mode 100644 index fa0f5bc7bc..0000000000 --- a/dimos/visualization/rerun/test_tf_tree.py +++ /dev/null @@ -1,68 +0,0 @@ -# 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. - -"""Where the triads get drawn, not how they look.""" - -from __future__ import annotations - -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.visualization.rerun.tf_tree import TFTreeVis - - -def edge(parent: str, child: str) -> Transform: - return Transform(frame_id=parent, child_frame_id=child, ts=1.0) - - -def paths(vis: TFTreeVis) -> dict[str, str]: - return {frame: spot.path for frame, spot in vis.placements().items()} - - -def feed(vis: TFTreeVis, *messages: TFMessage) -> None: - """tf republishes, and the tree draws once a message adds nothing new.""" - for msg in (*messages, messages[-1]): - vis.log(msg, [archetype for _, archetype in msg.to_rerun()]) - - -def test_triads_nest_along_the_tree() -> None: - vis = TFTreeVis() - feed(vis, TFMessage(edge("odom", "base_link"), edge("base_link", "camera/optical"))) - - assert paths(vis) == { - "odom": "world/tf/odom", - "base_link": "world/tf/odom/base_link", - "camera/optical": "world/tf/odom/base_link/camera\\/optical", - } - - -def test_a_late_root_re_parents_the_tree() -> None: - """The mount tree is published seconds before the odometry that roots it.""" - vis = TFTreeVis() - feed(vis, TFMessage(edge("mid360_link", "base_link"))) - assert paths(vis)["base_link"] == "world/tf/mid360_link/base_link" - - feed(vis, TFMessage(edge("odom", "mid360_link"))) - - assert paths(vis) == { - "odom": "world/tf/odom", - "mid360_link": "world/tf/odom/mid360_link", - "base_link": "world/tf/odom/mid360_link/base_link", - } - - -def test_a_cycle_does_not_hang() -> None: - vis = TFTreeVis() - feed(vis, TFMessage(edge("a", "b"), edge("b", "a"))) - - assert set(paths(vis)) == {"a", "b"} diff --git a/dimos/visualization/rerun/tf_tree.py b/dimos/visualization/rerun/tf_tree.py deleted file mode 100644 index 745f0e40d4..0000000000 --- a/dimos/visualization/rerun/tf_tree.py +++ /dev/null @@ -1,204 +0,0 @@ -# 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. - -"""Axis triads for every frame of a tf tree.""" - -from __future__ import annotations - -from dataclasses import dataclass -import threading -from typing import TYPE_CHECKING, TypeVar - -from dimos.memory2.transform import Transformer - -if TYPE_CHECKING: - from collections.abc import Iterable, Iterator - - import rerun as rr - from rerun._baseclasses import Archetype - - from dimos.memory2.stream import Stream - from dimos.memory2.type.observation import Observation - from dimos.msgs.geometry_msgs.Transform import Transform - from dimos.msgs.tf2_msgs.TFMessage import TFMessage - -T = TypeVar("T") - -DEFAULT_TF_ROOT = "world/tf" -# Where the transforms are declared. Rerun pins a frame to its declaring entity -# for the life of a recording, so these cannot move with the tree. -DEFAULT_LINKS_ROOT = "tf_links" -DEFAULT_AXIS_LENGTH = 0.5 -DEFAULT_TIMELINE = "ts" -# Each level's triad relative to its parent's. -DEPTH_SCALE = 0.8 -# Rerun's own TransformAxes3D draws at 1.0. -AXIS_WIDTH_UI_POINTS = 2.0 -AXIS_COLORS = [[255, 0, 0], [0, 255, 0], [0, 0, 255]] - - -def triad(length: float) -> rr.Arrows3D: - """XYZ arrows, red green blue.""" - import rerun as rr - - return rr.Arrows3D( - origins=[[0.0, 0.0, 0.0]] * 3, - vectors=[[length, 0.0, 0.0], [0.0, length, 0.0], [0.0, 0.0, length]], - colors=AXIS_COLORS, - radii=rr.components.Radius.ui_points(AXIS_WIDTH_UI_POINTS), - ) - - -@dataclass(frozen=True) -class Placement: - """Where a frame's triad is drawn, and how big.""" - - path: str - depth: int - - -class TFTreeVis: - """Draws the tf tree, one nested entity per frame carrying a triad.""" - - def __init__( - self, - axis_length: float = DEFAULT_AXIS_LENGTH, - root: str = DEFAULT_TF_ROOT, - links: str = DEFAULT_LINKS_ROOT, - ) -> None: - self.axis_length = axis_length - self.root = root - self.links = links - self._lock = threading.Lock() - self._parents: dict[str, str] = {} - self._drawn: dict[str, Placement] = {} - self._pending = False - - def log(self, msg: TFMessage, archetypes: Iterable[Archetype]) -> None: - """Declare the transforms, then redraw once a message adds nothing new. - - Publishers split one tree across several messages. - """ - import rerun as rr - - if not msg.transforms: - return - with self._lock: - for transform, archetype in zip(msg.transforms, archetypes, strict=True): - child = rr.escape_entity_path_part(transform.child_frame_id) - rr.log(f"{self.links}/{child}", archetype) - if self._learn(msg.transforms): - self._pending = True - elif self._pending: - self._pending = False - self._redraw() - - def flush(self) -> None: - """Draw a pending change nothing else triggered.""" - with self._lock: - if self._pending: - self._pending = False - self._redraw() - - def placements(self) -> dict[str, Placement]: - with self._lock: - return dict(self._drawn) - - def _learn(self, transforms: Iterable[Transform]) -> bool: - changed = False - for transform in transforms: - if self._parents.get(transform.child_frame_id) != transform.frame_id: - self._parents[transform.child_frame_id] = transform.frame_id - changed = True - return changed - - def _layout(self) -> dict[str, Placement]: - import rerun as rr - - placed: dict[str, Placement] = {} - - def place(frame: str, walked: frozenset[str]) -> Placement: - known = placed.get(frame) - if known is not None: - return known - parent = self._parents.get(frame) - if parent is None or parent in walked: - spot = Placement(f"{self.root}/{rr.escape_entity_path_part(frame)}", 0) - else: - above = place(parent, walked | {frame}) - spot = Placement( - f"{above.path}/{rr.escape_entity_path_part(frame)}", above.depth + 1 - ) - placed[frame] = spot - return spot - - for frame in (*self._parents, *self._parents.values()): - place(frame, frozenset()) - return placed - - def _redraw(self) -> None: - """Move the tree to match the shape tf has now.""" - import rerun as rr - - layout = self._layout() - - for frame, was in self._drawn.items(): - now = layout.get(frame) - if now is None or now.path != was.path: - rr.log(was.path, rr.Arrows3D(origins=[], vectors=[]), static=True) - - for frame, spot in layout.items(): - if self._drawn.get(frame) != spot: - rr.log( - spot.path, - rr.CoordinateFrame(f"tf#/{frame}"), - triad(self.axis_length * DEPTH_SCALE**spot.depth), - static=True, - ) - - self._drawn = layout - - -class RerunTFTree(Transformer[T, T]): - """Logs a recorded tf stream in step with the stream it passes through.""" - - def __init__(self, tf: Stream[TFMessage]) -> None: - self._tf = tf - self._vis = TFTreeVis() - - def __call__(self, upstream: Iterator[Observation[T]]) -> Iterator[Observation[T]]: - import rerun as rr - - pending = iter(self._tf) - head = next(pending, None) - floor: float | None = None - try: - for obs in upstream: - if floor is None: - # tf older than the replay would otherwise all land on frame one. - floor = obs.ts - while head is not None and head.ts <= obs.ts: - if head.ts >= floor: - self._log(head) - head = next(pending, None) - rr.set_time(DEFAULT_TIMELINE, timestamp=obs.ts) - yield obs - finally: - self._vis.flush() - - def _log(self, tf_obs: Observation[TFMessage]) -> None: - import rerun as rr - - rr.set_time(DEFAULT_TIMELINE, timestamp=tf_obs.ts) - self._vis.log(tf_obs.data, [archetype for _, archetype in tf_obs.data.to_rerun()]) From 7a5e6b2a7fec7c31d19fb05c43e78932122a2ceb Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Wed, 5 Aug 2026 13:40:46 -0700 Subject: [PATCH 13/13] Nit --- dimos/msgs/tf2_msgs/TFMessage.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/dimos/msgs/tf2_msgs/TFMessage.py b/dimos/msgs/tf2_msgs/TFMessage.py index 441e1772b6..5e44b7762d 100644 --- a/dimos/msgs/tf2_msgs/TFMessage.py +++ b/dimos/msgs/tf2_msgs/TFMessage.py @@ -48,12 +48,7 @@ def _triad(length: float): # type: ignore[no-untyped-def] class TfFrameTree: - """Nests each frame's triad under its parent, mirroring the tf tree in the panel. - - Rerun never lets an entity move, so a frame placed as a root before its - real parent arrives gets its old entity blanked and a new one declared - once that parent is known. - """ + """Store each frame under its parents. This lets us view the tree in the left panel.""" def __init__( self, axis_length: float = DEFAULT_AXIS_LENGTH, root: str = DEFAULT_TF_ROOT