feat: rerun tf tree visualization

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";
         };

dimos/memory2/tf.py

@@ -28,6 +28,14 @@
     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."""
+    # 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))
+
+
 class StreamTF(MultiTBuffer):
     def __init__(
         self,
@@ -47,9 +55,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.")

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,85 @@
 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:
+    """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
+    ) -> 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 +118,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 +133,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 +176,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

dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py

@@ -21,22 +21,23 @@
 
 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
 import typer
 
 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
 from dimos.msgs.geometry_msgs.Transform import Transform
 from dimos.msgs.geometry_msgs.Vector3 import Vector3
 from dimos.msgs.nav_msgs.Odometry import Odometry
 from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation
+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
@@ -45,10 +46,9 @@
 if TYPE_CHECKING:
     import rerun.blueprint as rrb
 
-TIMELINE = "ts"
+    from dimos.memory2.stream import Stream
 
-AXIS_LEN = 0.5
-AXIS_RADIUS_RATIO = 25
+TIMELINE = "ts"
 
 # Mount frames as recorded on the tf stream.
 BASE_FRAME = "base_link"
@@ -137,11 +137,44 @@ 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 _tf_over(store: SqliteStore, window: Stream[Any]) -> Stream[TFMessage] | None:
+    """The recorded tf stream clipped to another stream's span.
+
+    Absolute bounds: relative ones anchor on each stream's own first observation.
+    """
+    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 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)
     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)
 
@@ -191,12 +224,8 @@ def _log_odometry(
     """Trace the sensor moving throughout the scene."""
     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(
@@ -205,7 +234,7 @@ def _log_odometry(
     if base is None:
         return
     rr.log(
-        "world/base_link",
+        "world/robot_body",
         rr.Transform3D(
             translation=[base.translation.x, base.translation.y, base.translation.z],
             quaternion=rr.Quaternion(
@@ -323,6 +352,11 @@ def _blueprint(crop: LocalCrop) -> rrb.Blueprint:
                 origin="world",
                 name="world",
                 contents=["+ $origin/**", "- $origin/local/**"],
+                # The graph buries the map it was built from.
+                overrides={
+                    "world/nodes": rrb.EntityBehavior(visible=False),
+                    "world/node_edges": rrb.EntityBehavior(visible=False),
+                },
             ),
             rrb.Vertical(
                 rrb.Spatial3DView(
@@ -553,6 +587,7 @@ def main(
         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)
 
         pose_tagged = lidar.align(odom, tolerance=align_tol).transform(
             FnTransformer(_attach_pose_from_odom)
@@ -570,6 +605,7 @@ def main(
                 support_min=support_min,
             )
         )
+        tf_sync = _TfSync(tf)
 
         configs = _parse_configs(config, wall_clearance, wall_buffer, wall_buffer_weight)
         ref_clearance = configs[0][0]
@@ -592,27 +628,9 @@ 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",
+                "world/robot_body/outline",
                 rr.Boxes3D(
                     half_sizes=[ROBOT_LENGTH / 2, ROBOT_WIDTH / 2, robot_height / 2],
                     colors=[(0, 255, 127)],
@@ -621,7 +639,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],
@@ -635,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(