From dba1e5a3827287d8660a887713cf7a2e96537887 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 23 Jun 2026 19:12:24 -0700 Subject: [PATCH 01/69] recorder: async callbacks + async pose_setter_for Convert the memory2 Recorder from thread/disposable rx subscriptions to manual async callbacks via process_observable, and let pose_setter_for methods be async (awaited in _resolve_pose). Update the fastlio and go2 recorders accordingly. --- .../sensors/lidar/fastlio2/recorder.py | 4 +-- dimos/memory2/module.py | 35 +++++++++++-------- .../go2/blueprints/smart/unitree_go2.py | 9 +++-- 3 files changed, 29 insertions(+), 19 deletions(-) diff --git a/dimos/hardware/sensors/lidar/fastlio2/recorder.py b/dimos/hardware/sensors/lidar/fastlio2/recorder.py index 08784cfe89..5f081d13e9 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/recorder.py +++ b/dimos/hardware/sensors/lidar/fastlio2/recorder.py @@ -45,13 +45,13 @@ class FastLio2Recorder(Recorder): _last_odom_pose: Pose | None = None @pose_setter_for("fastlio_odometry") - def _odom_pose(self, msg: Odometry) -> Pose | None: + async def _odom_pose(self, msg: Odometry) -> Pose | None: pose = getattr(msg, "pose", None) self._last_odom_pose = getattr(pose, "pose", None) if pose is not None else None return self._last_odom_pose @pose_setter_for("fastlio_lidar") - def _lidar_pose(self, msg: PointCloud2) -> Pose | None: + async def _lidar_pose(self, msg: PointCloud2) -> Pose | None: # Most-recent odometry pose, stamped directly (no tf). None before the # first odometry -> frame stored unposed, map-skipped. return self._last_odom_pose diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index 00cd46a1d0..f494a2070d 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -14,7 +14,7 @@ from __future__ import annotations -from collections.abc import Callable +from collections.abc import Awaitable, Callable import enum import inspect import os @@ -274,13 +274,14 @@ class RecorderConfig(MemoryModuleConfig): stream_remapping: dict[str, str] = Field(default_factory=dict) -PoseSetter = Callable[[Any], "Pose | None"] +PoseSetter = Callable[[Any], "Pose | None | Awaitable[Pose | None]"] def pose_setter_for(*stream_names: str) -> Callable[[Any], Any]: """Mark a method ``(self, msg) -> Pose | None`` as the pose setter for the - given recorded stream(s). Streams without a setter fall back to the tf-based - ``world <- frame_id`` lookup.""" + given recorded stream(s). The method may be sync or ``async def`` — the + recorder awaits it if it returns an awaitable. Streams without a setter fall + back to the tf-based ``world <- frame_id`` lookup.""" def decorate(fn: Any) -> Any: fn._pose_setter_for = tuple(stream_names) @@ -302,10 +303,11 @@ class MyRecorder(Recorder): Each stream's pose defaults to a ``world <- frame_id`` tf lookup; decorate a method with ``@pose_setter_for("stream")`` to source it elsewhere (e.g. from - an odometry stream):: + an odometry stream). Setters run on the module's event loop and may be + ``async def``:: @pose_setter_for("lidar") - def _lidar_pose(self, msg): + async def _lidar_pose(self, msg): return self._last_odom_pose """ @@ -368,12 +370,14 @@ def _port_to_stream(self, name: str, input_topic: In[Any], stream: Stream[Any]) already in world coords) fall back to ``config.default_frame_id`` — so every observation gets a robot-pose anchor when tf is publishing. - Registers the subscription as a disposable on this module. + Each port is recorded by an async callback dispatched on the module's + event loop via :meth:`process_observable`, which serialises invocations + and registers the subscription for cleanup on stop(). """ - def on_msg(msg: Any) -> None: + async def on_msg(msg: Any) -> None: ts = self._resolve_ts(name, msg) - pose = self._resolve_pose(name, msg, ts) + pose = await self._resolve_pose(name, msg, ts) if not pose: logger.warning( "[%s] No pose for time %s (msg ts: %s), storing without pose", @@ -383,7 +387,7 @@ def on_msg(msg: Any) -> None: ) stream.append(msg, ts=ts, pose=pose) - self.register_disposable(Disposable(input_topic.subscribe(on_msg))) + self.process_observable(input_topic.pure_observable(), on_msg) def _prepare_streams(self) -> None: """On APPEND, drop the streams this recorder is about to (re)write — the @@ -401,13 +405,16 @@ def _resolve_ts(self, name: str, msg: Any) -> float: """Timestamp to record *msg* at. Override to re-base onto another clock.""" return getattr(msg, "ts", None) or time.time() - def _resolve_pose(self, name: str, msg: Any, ts: float) -> Pose | None: + async def _resolve_pose(self, name: str, msg: Any, ts: float) -> Pose | None: """Pose to anchor *msg* with. Dispatches to the stream's - ``@pose_setter_for`` if one is defined, else falls back to a - ``world <- frame_id`` tf lookup.""" + ``@pose_setter_for`` if one is defined (awaited when async), else falls + back to a ``world <- frame_id`` tf lookup.""" setter = self._pose_setters.get(name) if setter is not None: - return cast("Pose | None", setter(msg)) + result = setter(msg) + if inspect.isawaitable(result): + result = await result + return cast("Pose | None", result) frame_id = getattr(msg, "frame_id", None) or self.config.default_frame_id transform = self.tf.get( self.config.root_frame, frame_id, time_point=ts, time_tolerance=self.config.tf_tolerance diff --git a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py index 9993e541c4..3351a3dc16 100644 --- a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py +++ b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py @@ -62,13 +62,16 @@ class Go2Memory(Recorder): _last_odom_pose: Pose | None = None @pose_setter_for("odom") - def _odom_pose(self, msg: PoseStamped) -> Pose | None: + async def _odom_pose(self, msg: PoseStamped) -> Pose | None: self._last_odom_pose = msg return self._last_odom_pose @pose_setter_for("lidar") - def _lidar_pose(self, msg: PointCloud2) -> Pose | None: - return self._last_odom_pose # should always exist (odom alwyas wins the race) + async def _lidar_pose(self, msg: PointCloud2) -> Pose | None: + # Yes, it doesn't make sense to register lidar at the odom pose because the + # go2 lidar is in the world frame, but map.py (for now) needs this. + # TODO: fix map.py to use a transform frame + return getattr(self, "_last_odom_pose", None) unitree_go2_markers = ( From af6a06be3071f37af2f9df7241f8b4e0c899b4f6 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 23 Jun 2026 23:49:58 -0700 Subject: [PATCH 02/69] recorder: require @pose_setter_for setters to be async Raise TypeError at decoration time if a non-async function is decorated, and always await the setter in _resolve_pose. --- dimos/memory2/module.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index f494a2070d..e16c1527e7 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -274,16 +274,20 @@ class RecorderConfig(MemoryModuleConfig): stream_remapping: dict[str, str] = Field(default_factory=dict) -PoseSetter = Callable[[Any], "Pose | None | Awaitable[Pose | None]"] +PoseSetter = Callable[[Any], "Awaitable[Pose | None]"] def pose_setter_for(*stream_names: str) -> Callable[[Any], Any]: - """Mark a method ``(self, msg) -> Pose | None`` as the pose setter for the - given recorded stream(s). The method may be sync or ``async def`` — the - recorder awaits it if it returns an awaitable. Streams without a setter fall - back to the tf-based ``world <- frame_id`` lookup.""" + """Mark an ``async def`` method ``(self, msg) -> Pose | None`` as the pose + setter for the given recorded stream(s). Streams without a setter fall back + to the tf-based ``world <- frame_id`` lookup.""" def decorate(fn: Any) -> Any: + if not inspect.iscoroutinefunction(fn): + raise TypeError( + f"@pose_setter_for must decorate an `async def` method; " + f"{getattr(fn, '__qualname__', fn)} is not async" + ) fn._pose_setter_for = tuple(stream_names) return fn @@ -406,15 +410,12 @@ def _resolve_ts(self, name: str, msg: Any) -> float: return getattr(msg, "ts", None) or time.time() async def _resolve_pose(self, name: str, msg: Any, ts: float) -> Pose | None: - """Pose to anchor *msg* with. Dispatches to the stream's - ``@pose_setter_for`` if one is defined (awaited when async), else falls - back to a ``world <- frame_id`` tf lookup.""" + """Pose to anchor *msg* with. Dispatches to the stream's (async) + ``@pose_setter_for`` if one is defined, else falls back to a + ``world <- frame_id`` tf lookup.""" setter = self._pose_setters.get(name) if setter is not None: - result = setter(msg) - if inspect.isawaitable(result): - result = await result - return cast("Pose | None", result) + return cast("Pose | None", await setter(msg)) frame_id = getattr(msg, "frame_id", None) or self.config.default_frame_id transform = self.tf.get( self.config.root_frame, frame_id, time_point=ts, time_tolerance=self.config.tf_tolerance From df4d802ebbf26daf0868536c50a5c0c9c3913444 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 15:10:35 +0800 Subject: [PATCH 03/69] - --- dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py index 3351a3dc16..e50cf8b8e3 100644 --- a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py +++ b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2.py @@ -68,8 +68,9 @@ async def _odom_pose(self, msg: PoseStamped) -> Pose | None: @pose_setter_for("lidar") async def _lidar_pose(self, msg: PointCloud2) -> Pose | None: - # Yes, it doesn't make sense to register lidar at the odom pose because the - # go2 lidar is in the world frame, but map.py (for now) needs this. + # go2 lidar (currently) is in world-frame + # so it doesn't make sense to register lidar at the odom pose + # but we do it anyways because map.py (for now) requires it # TODO: fix map.py to use a transform frame return getattr(self, "_last_odom_pose", None) From e8287a2daf448f3ec3546a1d3d329b8408cd9b26 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 14:51:57 +0800 Subject: [PATCH 04/69] add gsc_pgo loop-closure module pinned to public gsc_pgo repo --- .../components/loop_closure/gsc_pgo/module.py | 231 ++++++++++++++++++ 1 file changed, 231 insertions(+) create mode 100644 dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py new file mode 100644 index 0000000000..a0d0c22099 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py @@ -0,0 +1,231 @@ +# 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. + +"""Native C++ PGO module — faithful reimplementation of the original nav stack PGO. + +Uses GTSAM iSAM2 for pose graph optimization and PCL ICP for loop closure. +""" + +from __future__ import annotations + +from pathlib import Path + +from reactivex.disposable import Disposable + +from dimos.core.core import rpc +from dimos.core.native_module import NativeModule, NativeModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.navigation.jnav.msgs.Graph3D import Graph3D +from dimos.navigation.jnav.msgs.GraphDelta3D import GraphDelta3D +from dimos.navigation.jnav.msgs.Landmark import Landmark +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +class PGOConfig(NativeModuleConfig): + # C++ + nix flake live in the standalone repo github.com/jeff-hykin/gsc_pgo. + # Pinned to a rev for reproducibility; bump when the C++ changes. The build + # runs in this module dir and drops a `result` symlink here (gitignored). + cwd: str | None = str(Path(__file__).resolve().parent) + executable: str = "result/bin/pgo" + build_command: str | None = ( + 'nix build "github:jeff-hykin/gsc_pgo/494e7a1d657c3702ec805c9e3d251a2fe8bc9529#default"' + " --no-write-lock-file" + ) + + frame_id: str = "map" + child_frame_id: str = "odom" + body_frame: str = "base_link" + + # Keyframe detection + key_pose_delta_deg: float = 10.0 + key_pose_delta_trans: float = 0.5 + + # Loop closure + loop_search_radius: float = 3.0 + loop_time_thresh: float = 5.0 + loop_score_thresh: float = 0.15 + loop_submap_half_range: int = 5 + submap_resolution: float = 0.1 + min_loop_detect_duration: float = 2.0 + # Feature-poverty gate: skip loop search when the scan's descriptor + # vertical-structure std is below this (open grass can't place itself -> + # PGO no-op). 0 = off. Superseded by loop_min_occupancy/loop_min_degeneracy + # (structure overlaps too much between scenes to threshold cleanly). + min_descriptor_std: float = 0.0 + + # Structure-spread gate: require >= this many occupied Scan-Context cells. + # Open grass clusters returns near the sensor (few rings filled); built + # scenes spread out to range. Calibrated on go2 fastlio (1200-cell 20x60 + # descriptor): grassy ~70 vs gir_park ~88 vs downtown ~120 at equal point + # count -> measures spread, not density. 0 disables. + loop_min_occupancy: int = 80 + # Observability gate (Zhang 2016 / X-ICP degeneracy): reject a candidate + # whose source scan's smallest normalized normal-scatter eigenvalue is below + # this. Planar/degenerate (grass) -> ~0; ICP slides in-plane and reports low + # fitness for a bogus closure. Real scenes (incl. sparse gir_park) sit >0.15. + # 0 disables. + loop_min_degeneracy: float = 0.05 + + # Input mode: transform world-frame scans to body-frame using odom + unregister_input: bool = True + + # Debug global-map publishing — OFF by default. Emitted on the internal + # `_global_map` port (leading underscore) so it never autoconnects to a + # consumer's `global_map` In: the terrain_mapper is the planner's single + # authoritative global_map. Two producers on `global_map` made the costmap + # flicker. Set a rate > 0 only for viz/debug of the PGO's corrected cloud. + global_map_voxel_size: float = 0.1 + global_map_publish_rate: float = 0.0 + + # Scan Context place recognition (used by loop closure search) + use_scan_context: bool = True + scan_context_num_rings: int = 20 + scan_context_num_sectors: int = 60 + scan_context_max_range_m: float = 80.0 + scan_context_top_k: int = 10 + scan_context_match_threshold: float = 0.4 + scan_context_lidar_height_m: float = 2.0 + + # Skip ICP on candidates farther than this (m). 0 disables. + loop_candidate_max_distance_m: float = 30.0 + + # --- Tag (AprilTag/ArUco) loop closure -------------------------------- + use_tag_loop_closure: bool = False + # LCM channel of the static TF tree (dimos "pattern#msg_name" convention). + tf_static_channel: str = "/tf_static#tf2_msgs.TFMessage" + tag_loop_time_thresh: float = 5.0 + tag_assoc_max_dt: float = 0.1 + tag_buffer_window: float = 2.0 + # Anisotropic tag noise (variances, tag frame; normal = +z): in-plane/yaw + # tight, range/out-of-plane loose so tag range error can't distort z. + tag_var_inplane_trans_m2: float = 0.0025 + tag_var_range_trans_m2: float = 0.25 + tag_var_yaw_rot_rad2: float = 0.0025 + tag_var_outplane_rot_rad2: float = 0.04 + # 6-DOF Mahalanobis gate vs current estimate (chi^2 95% = 12.59). 0 = off. + tag_consistency_chi2: float = 0.0 + # Robust (Huber) kernel on all loop factors (lidar + tag). Off = original. + loop_robust_kernel: bool = False + loop_robust_huber_k: float = 1.345 + + # --- Landmark events (decoupled perceiver -> PGO factor-graph manager) ---- + # When set, the PGO ingests Landmark events on the `landmarks` In and + # attaches each as a graph landmark variable + a BetweenFactor(keyframe, + # landmark). Two sightings of the same landmark id share the variable, so a + # revisit closes the loop and GTSAM optimizes it jointly. A separate + # perceiver (utils/apriltag_perceiver.py) does the detection + noise/ + # confidence filtering and emits the Landmark events. Off by default. + use_landmarks: bool = False + # Anisotropic landmark observation noise (variances, landmark frame; normal = + # +z): in-plane/yaw tight, range/out-of-plane loose (mirrors the tag model). + landmark_var_inplane_trans_m2: float = 0.0025 + landmark_var_range_trans_m2: float = 0.25 + landmark_var_yaw_rot_rad2: float = 0.0025 + landmark_var_outplane_rot_rad2: float = 0.04 + landmark_assoc_max_dt: float = 0.2 + landmark_buffer_window: float = 3.0 + + # --- Gravity anchor ------------------------------------------------------ + # Pin keyframe 0 (whose orientation is gravity-aligned by the LIO front end) + # so landmark/loop closures cannot rotate the initial roll/pitch off gravity. + # The full pose is pinned (also the gauge reference); roll/pitch stiffness is + # the gravity component. Variances (smaller = stiffer). + gravity_anchor: bool = True + gravity_anchor_rp_var: float = 1e-12 + gravity_anchor_yaw_var: float = 1e-12 + gravity_anchor_trans_var: float = 1e-12 + # Per-keyframe gravity anchor (roll/pitch-only prior on EVERY keyframe). Anchoring + # only kf0 lets a big loop closure tilt inner keyframes' roll/pitch, which converts + # horizontal travel into vertical and corrupts z by tens of metres. Pinning every + # keyframe's roll/pitch to its gravity-aligned LIO orientation (yaw + translation + # left free) keeps the closure in-plane and preserves the z structure. + # Default OFF: the anisotropic odometry between-factor is the primary gravity- + # preservation mechanism (and still lets landmarks correct slow tilt drift). + # This absolute prior is a harder lock for when the front end's absolute tilt + # is trustworthy (e.g. ZUPT in the LIO estimator). + gravity_anchor_per_keyframe: bool = False + gravity_anchor_kf_rp_var: float = 1e-4 + + # Anisotropic odometry between-factor: the LIO relative roll/pitch is accurate + # (IMU sees gravity each step) but yaw drifts, so roll/pitch are stiff and yaw + # looser. This keeps a loop closure from sloshing its (mostly-yaw) correction + # into roll/pitch — a tilt that converts horizontal travel into vertical and + # corrupts z. Roll/pitch variance is small but nonzero, so landmarks can still + # correct slow tilt drift across the graph ("accurate but not perfect"). + odom_rot_rp_var: float = 1e-8 + odom_rot_yaw_var: float = 1e-5 + odom_trans_xy_var: float = 1e-4 + odom_trans_z_var: float = 1e-6 + + # Bounded FIFO depth: keep at most this many pending scans, dropping the + # oldest when full (<=0 = unbounded). Generous enough that an ack-gated eval + # replay never drops a scan, bounded enough to cap live latency/memory. + max_scan_queue: int = 100 + + debug: bool = False + + +class PGO(NativeModule): + """Pose graph optimization with loop closure using GTSAM iSAM2 + PCL ICP.""" + + config: PGOConfig + + # named "lidar" to match the LoopClosure spec; the binary pairs it with the + # latest odometry pose internally, so a raw sensor-frame scan is expected. + lidar: In[PointCloud2] + odometry: In[Odometry] + # Optional: tag detections (tag-in-optical pose + numeric id) for tag-based + # loop closure. Only consumed when config.use_tag_loop_closure is set. + tag_detections: In[Detection3DArray] + # Optional: decoupled Landmark events from a perceiver (e.g. the AprilTag + # perceiver). Only consumed when config.use_landmarks is set; each becomes a + # graph landmark variable + observation factor that GTSAM optimizes jointly. + landmarks: In[Landmark] + corrected_odometry: Out[Odometry] + correction: Out[Transform] + # Internal/debug only (off by default) — see global_map_publish_rate. Named + # with a leading underscore so autoconnect won't wire it to `global_map` Ins. + _global_map: Out[PointCloud2] + pose_graph: Out[Graph3D] + loop_closure_event: Out[GraphDelta3D] + + @rpc + def start(self) -> None: + super().start() + self.tf.publish( + Transform( + frame_id=self.config.frame_id, + child_frame_id=self.config.child_frame_id, + ) + ) + self.register_disposable( + Disposable( + self.correction.transport.subscribe(self._on_correction_for_tf, self.correction) + ) + ) + if self.config.debug: + logger.info("PGO native module started (C++ iSAM2 + PCL ICP)") + + def _on_correction_for_tf(self, msg: Transform) -> None: + self.tf.publish(msg) + + @rpc + def stop(self) -> None: + super().stop() From f9d3acb8f8e7f7a26c71b696366259a326d0ad29 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 14:52:11 +0800 Subject: [PATCH 05/69] add jnav PGO message types (Graph3D, GraphDelta3D, Landmark, Marker) --- dimos/navigation/jnav/msgs/Graph3D.hpp | 175 +++++++++++++++ dimos/navigation/jnav/msgs/Graph3D.py | 233 ++++++++++++++++++++ dimos/navigation/jnav/msgs/GraphDelta3D.hpp | 168 ++++++++++++++ dimos/navigation/jnav/msgs/GraphDelta3D.py | 198 +++++++++++++++++ dimos/navigation/jnav/msgs/Landmark.py | 200 +++++++++++++++++ dimos/navigation/jnav/msgs/Marker.py | 97 ++++++++ dimos/navigation/jnav/msgs/test_Landmark.py | 129 +++++++++++ dimos/navigation/jnav/msgs/test_Marker.py | 67 ++++++ 8 files changed, 1267 insertions(+) create mode 100644 dimos/navigation/jnav/msgs/Graph3D.hpp create mode 100644 dimos/navigation/jnav/msgs/Graph3D.py create mode 100644 dimos/navigation/jnav/msgs/GraphDelta3D.hpp create mode 100644 dimos/navigation/jnav/msgs/GraphDelta3D.py create mode 100644 dimos/navigation/jnav/msgs/Landmark.py create mode 100644 dimos/navigation/jnav/msgs/Marker.py create mode 100644 dimos/navigation/jnav/msgs/test_Landmark.py create mode 100644 dimos/navigation/jnav/msgs/test_Marker.py diff --git a/dimos/navigation/jnav/msgs/Graph3D.hpp b/dimos/navigation/jnav/msgs/Graph3D.hpp new file mode 100644 index 0000000000..a440127ae2 --- /dev/null +++ b/dimos/navigation/jnav/msgs/Graph3D.hpp @@ -0,0 +1,175 @@ +// Copyright 2026 Dimensional Inc. +// SPDX-License-Identifier: Apache-2.0 +// +// Typed C++ helper mirroring the Python `dimos.msgs.nav_msgs.Graph3D`. +// Canonical schema lives in `dimos/msgs/nav_msgs/Graph3D.ksy` — keep +// encode() in sync with that file (and with Graph3D.py.lcm_decode). +// +// Wire format (big-endian): +// +// uint64 edge_count +// uint64 node_count +// double timestamp // seconds since epoch +// per node (node_count): +// pose_stamped: +// double ts +// uint32 frame_id_len +// bytes frame_id (utf-8, no terminator) +// 7×double pos_x, pos_y, pos_z, quat_x, quat_y, quat_z, quat_w +// uint64 id +// uint64 metadata_id +// per edge (edge_count): +// uint64 start_id +// uint64 end_id +// double timestamp +// uint64 metadata_id +// +// Edges reference nodes by `id`, not by index. + +#pragma once + +#include +#include +#include +#include +#include + +#include + +namespace dimos { + +namespace graph3d_detail { + +// Host-order → big-endian byte writers. Avoid for portability +// (macOS uses different names) — write byte-by-byte from the top. + +inline void write_u32_be(std::vector& out, uint32_t v) { + out.push_back(static_cast((v >> 24) & 0xFF)); + out.push_back(static_cast((v >> 16) & 0xFF)); + out.push_back(static_cast((v >> 8) & 0xFF)); + out.push_back(static_cast( v & 0xFF)); +} + +inline void write_u64_be(std::vector& out, uint64_t v) { + for (int shift = 56; shift >= 0; shift -= 8) { + out.push_back(static_cast((v >> shift) & 0xFF)); + } +} + +inline void write_double_be(std::vector& out, double v) { + uint64_t bits; + std::memcpy(&bits, &v, sizeof(bits)); + write_u64_be(out, bits); +} + +inline void write_bytes(std::vector& out, const std::string& s) { + out.insert(out.end(), s.begin(), s.end()); +} + +} // namespace graph3d_detail + +class Graph3D { +public: + struct PoseStamped { + double ts = 0.0; + std::string frame_id; + double pos_x = 0.0, pos_y = 0.0, pos_z = 0.0; + double quat_x = 0.0, quat_y = 0.0, quat_z = 0.0, quat_w = 1.0; + }; + + struct Node3D { + PoseStamped pose; + uint64_t id = 0; + uint64_t metadata_id = 0; + }; + + struct Edge { + uint64_t start_id = 0; + uint64_t end_id = 0; + double timestamp = 0.0; + uint64_t metadata_id = 0; + }; + + Graph3D(std::string frame_id, double timestamp) + : frame_id_(std::move(frame_id)), timestamp_(timestamp) {} + + void reserve_nodes(size_t capacity) { nodes_.reserve(capacity); } + void reserve_edges(size_t capacity) { edges_.reserve(capacity); } + + // Add a node. The pose's frame_id defaults to the graph's frame_id — + // override per-node only if a node lives in a different frame. + void add_node(uint64_t id, uint64_t metadata_id, double pose_ts, + double pos_x, double pos_y, double pos_z, + double quat_x, double quat_y, double quat_z, double quat_w, + std::string node_frame_id = "") { + PoseStamped pose; + pose.ts = pose_ts; + pose.frame_id = node_frame_id.empty() ? frame_id_ : std::move(node_frame_id); + pose.pos_x = pos_x; pose.pos_y = pos_y; pose.pos_z = pos_z; + pose.quat_x = quat_x; pose.quat_y = quat_y; pose.quat_z = quat_z; pose.quat_w = quat_w; + nodes_.push_back({pose, id, metadata_id}); + } + + // Position-only convenience (orientation defaults to identity). + void add_node_xyz(uint64_t id, uint64_t metadata_id, double pose_ts, + double pos_x, double pos_y, double pos_z) { + add_node(id, metadata_id, pose_ts, pos_x, pos_y, pos_z, 0.0, 0.0, 0.0, 1.0); + } + + void add_edge(uint64_t start_id, uint64_t end_id, double edge_ts, + uint64_t metadata_id = 0) { + edges_.push_back({start_id, end_id, edge_ts, metadata_id}); + } + + size_t node_count() const { return nodes_.size(); } + size_t edge_count() const { return edges_.size(); } + const std::string& frame_id() const { return frame_id_; } + + std::vector encode() const { + using namespace graph3d_detail; + std::vector out; + // Conservative reservation: header + per-node fixed bytes + per-edge. + // frame_id strings add variable length on top — that just causes a + // realloc, not correctness issues. + out.reserve(24 + nodes_.size() * 84 + edges_.size() * 32); + write_u64_be(out, static_cast(edges_.size())); + write_u64_be(out, static_cast(nodes_.size())); + write_double_be(out, timestamp_); + for (const auto& n : nodes_) { + // pose_stamped first (per Graph3D.ksy) + write_double_be(out, n.pose.ts); + write_u32_be(out, static_cast(n.pose.frame_id.size())); + write_bytes(out, n.pose.frame_id); + write_double_be(out, n.pose.pos_x); + write_double_be(out, n.pose.pos_y); + write_double_be(out, n.pose.pos_z); + write_double_be(out, n.pose.quat_x); + write_double_be(out, n.pose.quat_y); + write_double_be(out, n.pose.quat_z); + write_double_be(out, n.pose.quat_w); + // then id, metadata_id + write_u64_be(out, n.id); + write_u64_be(out, n.metadata_id); + } + for (const auto& e : edges_) { + write_u64_be(out, e.start_id); + write_u64_be(out, e.end_id); + write_double_be(out, e.timestamp); + write_u64_be(out, e.metadata_id); + } + return out; + } + + int publish(lcm::LCM& lcm, const std::string& channel) const { + std::vector bytes = encode(); + return lcm.publish(channel, bytes.data(), static_cast(bytes.size())); + } + +private: + std::string frame_id_; + double timestamp_; + std::vector nodes_; + std::vector edges_; +}; + +} // namespace dimos diff --git a/dimos/navigation/jnav/msgs/Graph3D.py b/dimos/navigation/jnav/msgs/Graph3D.py new file mode 100644 index 0000000000..89d3b51e72 --- /dev/null +++ b/dimos/navigation/jnav/msgs/Graph3D.py @@ -0,0 +1,233 @@ +# 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. + +"""Graph3D: pose-graph / visibility-graph message with typed nodes and edges. + +Edges reference nodes by ``id`` (not list index), so producers are free +to reorder or re-emit nodes between snapshots. ``metadata_id`` is a +caller-defined enum — ex: for far_planner: 0=normal, 1=odom, 2=goal +""" + +from __future__ import annotations + +from dataclasses import dataclass +import struct +import time +from typing import TYPE_CHECKING, BinaryIO + +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.types.timestamped import Timestamped + +if TYPE_CHECKING: + from rerun._baseclasses import Archetype + + +_DEFAULT_NODE_COLORS: dict[int, tuple[int, int, int, int]] = { + 0: (180, 180, 180, 200), + 1: (0, 255, 0, 255), + 2: (255, 0, 0, 255), + 3: (255, 165, 0, 200), + 4: (0, 200, 255, 200), +} +_DEFAULT_NODE_COLOR = (200, 200, 200, 180) + +_DEFAULT_EDGE_COLORS: dict[int, tuple[int, int, int, int]] = { + 0: (0, 220, 100, 200), # odom / traversable — green + 1: (255, 180, 0, 220), # loop_closure / partial — yellow + 2: (255, 50, 50, 150), # blocked — red +} +_DEFAULT_EDGE_COLOR = (180, 180, 180, 180) + + +class Graph3D(Timestamped): + msg_name = "nav_msgs.Graph3D" + + @dataclass + class Node3D: + pose: PoseStamped + id: int = 0 + metadata_id: int = 0 + + @dataclass + class Edge: + start_id: int + end_id: int + timestamp: float = 0.0 + metadata_id: int = 0 + + ts: float + nodes: list[Node3D] + edges: list[Edge] + + def __init__( + self, + ts: float = 0.0, + nodes: list[Graph3D.Node3D] | None = None, + edges: list[Graph3D.Edge] | None = None, + ) -> None: + self.ts = ts if ts != 0 else time.time() + self.nodes = nodes if nodes is not None else [] + self.edges = edges if edges is not None else [] + + def lcm_encode(self) -> bytes: + # Field order matches Graph3D.ksy: edge_count, node_count, ts, + # nodes[] (pose, id, metadata_id), edges[]. + parts: list[bytes] = [] + parts.append(struct.pack(">QQd", len(self.edges), len(self.nodes), self.ts)) + for node in self.nodes: + frame_id_bytes = node.pose.frame_id.encode("utf-8") + parts.append(struct.pack(">d", node.pose.ts)) + parts.append(struct.pack(">I", len(frame_id_bytes))) + parts.append(frame_id_bytes) + parts.append( + struct.pack( + ">7d", + node.pose.position.x, + node.pose.position.y, + node.pose.position.z, + node.pose.orientation.x, + node.pose.orientation.y, + node.pose.orientation.z, + node.pose.orientation.w, + ) + ) + parts.append(struct.pack(">QQ", node.id, node.metadata_id)) + for edge in self.edges: + parts.append( + struct.pack(">QQdQ", edge.start_id, edge.end_id, edge.timestamp, edge.metadata_id) + ) + return b"".join(parts) + + @classmethod + def lcm_decode(cls, data: bytes | BinaryIO) -> Graph3D: + buf = data if isinstance(data, (bytes, bytearray)) else data.read() + offset = 0 + edge_count, node_count, graph_ts = struct.unpack_from(">QQd", buf, offset) + offset += 24 + + nodes: list[Graph3D.Node3D] = [] + for _ in range(node_count): + (pose_ts,) = struct.unpack_from(">d", buf, offset) + offset += 8 + (frame_id_len,) = struct.unpack_from(">I", buf, offset) + offset += 4 + frame_id = buf[offset : offset + frame_id_len].decode("utf-8") + offset += frame_id_len + px, py, pz, qx, qy, qz, qw = struct.unpack_from(">7d", buf, offset) + offset += 56 + node_id, metadata_id = struct.unpack_from(">QQ", buf, offset) + offset += 16 + pose = PoseStamped( + ts=pose_ts, + frame_id=frame_id, + position=Vector3(px, py, pz), + orientation=Quaternion(qx, qy, qz, qw), + ) + nodes.append(cls.Node3D(pose=pose, id=node_id, metadata_id=metadata_id)) + + edges: list[Graph3D.Edge] = [] + for _ in range(edge_count): + start_id, end_id, edge_ts, edge_metadata_id = struct.unpack_from(">QQdQ", buf, offset) + offset += 32 + edges.append( + cls.Edge( + start_id=start_id, + end_id=end_id, + timestamp=edge_ts, + metadata_id=edge_metadata_id, + ) + ) + + return cls(ts=graph_ts, nodes=nodes, edges=edges) + + def to_rerun( + self, + z_offset: float = 0.0, + radii: float = 0.12, + node_colors: dict[int, tuple[int, int, int, int]] | None = None, + ) -> Archetype: + """Default visualization: ``rr.Points3D`` of just the nodes. + + For nodes + edges in separate entity sub-paths, use + ``to_rerun_multi`` from a ``visual_override`` callback. + """ + import rerun as rr + + nc = node_colors if node_colors is not None else _DEFAULT_NODE_COLORS + positions = [ + [n.pose.position.x, n.pose.position.y, n.pose.position.z + z_offset] for n in self.nodes + ] + colors = [nc.get(n.metadata_id, _DEFAULT_NODE_COLOR) for n in self.nodes] + node_radii = [radii * 2.0 if n.metadata_id in (1, 2) else radii for n in self.nodes] + return rr.Points3D(positions, colors=colors, radii=node_radii) + + def to_rerun_multi( + self, + base_path: str, + z_offset: float = 0.0, + node_radius: float = 0.12, + edge_radius: float = 0.04, + node_colors: dict[int, tuple[int, int, int, int]] | None = None, + edge_colors: dict[int, tuple[int, int, int, int]] | None = None, + ) -> list[tuple[str, Archetype]]: + """Return ``[(base_path/nodes, Points3D), (base_path/edges, LineStrips3D)]``. + + Intended for use from ``visual_override`` callbacks where the + bridge supports the ``RerunMulti`` list-of-tuples form. + """ + import rerun as rr + + nc = node_colors if node_colors is not None else _DEFAULT_NODE_COLORS + ec = edge_colors if edge_colors is not None else _DEFAULT_EDGE_COLORS + + node_positions = [ + [n.pose.position.x, n.pose.position.y, n.pose.position.z + z_offset] for n in self.nodes + ] + node_colors_list = [nc.get(n.metadata_id, _DEFAULT_NODE_COLOR) for n in self.nodes] + node_radii = [ + node_radius * 2.0 if n.metadata_id in (1, 2) else node_radius for n in self.nodes + ] + nodes_archetype = rr.Points3D(node_positions, colors=node_colors_list, radii=node_radii) + + id_to_pose: dict[int, PoseStamped] = {n.id: n.pose for n in self.nodes} + strips: list[list[list[float]]] = [] + edge_colors_list: list[tuple[int, int, int, int]] = [] + for edge in self.edges: + start = id_to_pose.get(edge.start_id) + end = id_to_pose.get(edge.end_id) + if start is None or end is None: + continue + strips.append( + [ + [start.position.x, start.position.y, start.position.z + z_offset], + [end.position.x, end.position.y, end.position.z + z_offset], + ] + ) + edge_colors_list.append(ec.get(edge.metadata_id, _DEFAULT_EDGE_COLOR)) + edges_archetype = rr.LineStrips3D( + strips, colors=edge_colors_list, radii=[edge_radius] * len(strips) + ) + + return [ + (f"{base_path}/nodes", nodes_archetype), + (f"{base_path}/edges", edges_archetype), + ] + + def __len__(self) -> int: + return len(self.nodes) + + def __str__(self) -> str: + return f"Graph3D(nodes={len(self.nodes)}, edges={len(self.edges)})" diff --git a/dimos/navigation/jnav/msgs/GraphDelta3D.hpp b/dimos/navigation/jnav/msgs/GraphDelta3D.hpp new file mode 100644 index 0000000000..4db42eb71c --- /dev/null +++ b/dimos/navigation/jnav/msgs/GraphDelta3D.hpp @@ -0,0 +1,168 @@ +// Copyright 2026 Dimensional Inc. +// SPDX-License-Identifier: Apache-2.0 +// +// Typed C++ helper mirroring the Python `dimos.msgs.nav_msgs.GraphDelta3D`. +// +// Wire format (big-endian): +// +// uint64 node_count +// double timestamp // seconds since epoch +// per node (node_count): +// pose_stamped: // (same as Graph3D's node3d pose) +// double ts +// uint32 frame_id_len +// bytes frame_id (utf-8, no terminator) +// 7×double pos_x, pos_y, pos_z, quat_x, quat_y, quat_z, quat_w +// uint64 id +// uint64 metadata_id +// per transform (node_count): +// 7×double translation_x, translation_y, translation_z, +// rotation_x, rotation_y, rotation_z, rotation_w +// +// Two aligned arrays: ``transforms[i]`` is the SE(3) delta about to +// be applied to ``nodes[i]``. ``post_pose = transforms[i] * nodes[i].pose`` +// is the convention (left-multiply). +// +// `GraphDelta3D.py.lcm_decode` reads exactly this layout — keep in sync. + +#pragma once + +#include +#include +#include +#include +#include + +#include + +namespace dimos { + +namespace graph_delta3d_detail { + +inline void write_u32_be(std::vector& out, uint32_t v) { + out.push_back(static_cast((v >> 24) & 0xFF)); + out.push_back(static_cast((v >> 16) & 0xFF)); + out.push_back(static_cast((v >> 8) & 0xFF)); + out.push_back(static_cast( v & 0xFF)); +} + +inline void write_u64_be(std::vector& out, uint64_t v) { + for (int shift = 56; shift >= 0; shift -= 8) { + out.push_back(static_cast((v >> shift) & 0xFF)); + } +} + +inline void write_double_be(std::vector& out, double v) { + uint64_t bits; + std::memcpy(&bits, &v, sizeof(bits)); + write_u64_be(out, bits); +} + +inline void write_bytes(std::vector& out, const std::string& s) { + out.insert(out.end(), s.begin(), s.end()); +} + +} // namespace graph_delta3d_detail + +class GraphDelta3D { +public: + struct PoseStamped { + double ts = 0.0; + std::string frame_id; + double pos_x = 0.0, pos_y = 0.0, pos_z = 0.0; + double quat_x = 0.0, quat_y = 0.0, quat_z = 0.0, quat_w = 1.0; + }; + + struct Node3D { + PoseStamped pose; + uint64_t id = 0; + uint64_t metadata_id = 0; + }; + + struct Transform { + double translation_x = 0.0, translation_y = 0.0, translation_z = 0.0; + double rotation_x = 0.0, rotation_y = 0.0, rotation_z = 0.0, rotation_w = 1.0; + }; + + GraphDelta3D(std::string frame_id, double timestamp) + : frame_id_(std::move(frame_id)), timestamp_(timestamp) {} + + void reserve(size_t capacity) { + nodes_.reserve(capacity); + transforms_.reserve(capacity); + } + + // Add a node + its SE(3) delta. Pass empty `node_frame_id` to inherit + // the graph's frame_id. + void add(uint64_t id, uint64_t metadata_id, double pose_ts, + double pos_x, double pos_y, double pos_z, + double quat_x, double quat_y, double quat_z, double quat_w, + double translation_x, double translation_y, double translation_z, + double rotation_x, double rotation_y, double rotation_z, double rotation_w, + std::string node_frame_id = "") { + Node3D node; + node.id = id; + node.metadata_id = metadata_id; + node.pose.ts = pose_ts; + node.pose.frame_id = node_frame_id.empty() ? frame_id_ : std::move(node_frame_id); + node.pose.pos_x = pos_x; node.pose.pos_y = pos_y; node.pose.pos_z = pos_z; + node.pose.quat_x = quat_x; node.pose.quat_y = quat_y; + node.pose.quat_z = quat_z; node.pose.quat_w = quat_w; + nodes_.push_back(node); + + Transform tf; + tf.translation_x = translation_x; tf.translation_y = translation_y; tf.translation_z = translation_z; + tf.rotation_x = rotation_x; tf.rotation_y = rotation_y; + tf.rotation_z = rotation_z; tf.rotation_w = rotation_w; + transforms_.push_back(tf); + } + + size_t size() const { return nodes_.size(); } + bool empty() const { return nodes_.empty(); } + const std::string& frame_id() const { return frame_id_; } + + std::vector encode() const { + using namespace graph_delta3d_detail; + std::vector out; + out.reserve(16 + nodes_.size() * (84 + 56)); + write_u64_be(out, static_cast(nodes_.size())); + write_double_be(out, timestamp_); + for (const auto& n : nodes_) { + write_double_be(out, n.pose.ts); + write_u32_be(out, static_cast(n.pose.frame_id.size())); + write_bytes(out, n.pose.frame_id); + write_double_be(out, n.pose.pos_x); + write_double_be(out, n.pose.pos_y); + write_double_be(out, n.pose.pos_z); + write_double_be(out, n.pose.quat_x); + write_double_be(out, n.pose.quat_y); + write_double_be(out, n.pose.quat_z); + write_double_be(out, n.pose.quat_w); + write_u64_be(out, n.id); + write_u64_be(out, n.metadata_id); + } + for (const auto& t : transforms_) { + write_double_be(out, t.translation_x); + write_double_be(out, t.translation_y); + write_double_be(out, t.translation_z); + write_double_be(out, t.rotation_x); + write_double_be(out, t.rotation_y); + write_double_be(out, t.rotation_z); + write_double_be(out, t.rotation_w); + } + return out; + } + + int publish(lcm::LCM& lcm, const std::string& channel) const { + std::vector bytes = encode(); + return lcm.publish(channel, bytes.data(), static_cast(bytes.size())); + } + +private: + std::string frame_id_; + double timestamp_; + std::vector nodes_; + std::vector transforms_; +}; + +} // namespace dimos diff --git a/dimos/navigation/jnav/msgs/GraphDelta3D.py b/dimos/navigation/jnav/msgs/GraphDelta3D.py new file mode 100644 index 0000000000..36fccca459 --- /dev/null +++ b/dimos/navigation/jnav/msgs/GraphDelta3D.py @@ -0,0 +1,198 @@ +# 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. + +"""GraphDelta3D: per-node SE(3) transforms about to be applied to a list of nodes. + +Two aligned arrays: ``nodes[i]`` is the node, ``transforms[i]`` is the +SE(3) delta about to be applied to it. ``post_pose = transforms[i] * +nodes[i].pose`` is the convention (left-multiply). + +Use case: PGO publishes this on ``loop_closure_event`` when iSAM2 +smooths the pose graph — ``nodes[i]`` is the keyframe pre-smooth, +``transforms[i]`` is the delta iSAM2 just applied to it. Consumers can +re-derive post-poses or filter to large deltas. + +Wire format mirrors ``Graph3D`` conventions: big-endian, ``Node3D`` +serialization shared, ``Transform`` is just 7 f8s (translation + +quaternion). Custom binary, dispatched by the ``#nav_msgs.GraphDelta3D`` +channel-name suffix. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import struct +import time +from typing import TYPE_CHECKING, BinaryIO + +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.navigation.jnav.msgs.Graph3D import Graph3D +from dimos.types.timestamped import Timestamped + +if TYPE_CHECKING: + from rerun._baseclasses import Archetype + + +class GraphDelta3D(Timestamped): + msg_name = "nav_msgs.GraphDelta3D" + + # Reuse Graph3D's nested Node3D for wire-format consistency. A + # GraphDelta3D[i].node is byte-identical to a Graph3D.nodes[i]. + Node3D = Graph3D.Node3D + + @dataclass + class Transform: + """SE(3) transform — translation + rotation quaternion (xyzw).""" + + translation: Vector3 + rotation: Quaternion + + ts: float + nodes: list[Graph3D.Node3D] + transforms: list[Transform] + + def __init__( + self, + ts: float = 0.0, + nodes: list[Graph3D.Node3D] | None = None, + transforms: list[Transform] | None = None, + ) -> None: + self.ts = ts if ts != 0 else time.time() + self.nodes = nodes if nodes is not None else [] + self.transforms = transforms if transforms is not None else [] + if len(self.nodes) != len(self.transforms): + raise ValueError( + f"nodes ({len(self.nodes)}) and transforms ({len(self.transforms)}) " + "must be the same length — they're aligned arrays" + ) + + def lcm_encode(self) -> bytes: + parts: list[bytes] = [] + parts.append(struct.pack(">Qd", len(self.nodes), self.ts)) + for node in self.nodes: + frame_id_bytes = node.pose.frame_id.encode("utf-8") + parts.append(struct.pack(">d", node.pose.ts)) + parts.append(struct.pack(">I", len(frame_id_bytes))) + parts.append(frame_id_bytes) + parts.append( + struct.pack( + ">7d", + node.pose.position.x, + node.pose.position.y, + node.pose.position.z, + node.pose.orientation.x, + node.pose.orientation.y, + node.pose.orientation.z, + node.pose.orientation.w, + ) + ) + parts.append(struct.pack(">QQ", node.id, node.metadata_id)) + for transform in self.transforms: + parts.append( + struct.pack( + ">7d", + transform.translation.x, + transform.translation.y, + transform.translation.z, + transform.rotation.x, + transform.rotation.y, + transform.rotation.z, + transform.rotation.w, + ) + ) + return b"".join(parts) + + @classmethod + def lcm_decode(cls, data: bytes | BinaryIO) -> GraphDelta3D: + buf = data if isinstance(data, (bytes, bytearray)) else data.read() + offset = 0 + node_count, graph_ts = struct.unpack_from(">Qd", buf, offset) + offset += 16 + + nodes: list[Graph3D.Node3D] = [] + for _ in range(node_count): + (pose_ts,) = struct.unpack_from(">d", buf, offset) + offset += 8 + (frame_id_len,) = struct.unpack_from(">I", buf, offset) + offset += 4 + frame_id = buf[offset : offset + frame_id_len].decode("utf-8") + offset += frame_id_len + px, py, pz, qx, qy, qz, qw = struct.unpack_from(">7d", buf, offset) + offset += 56 + node_id, metadata_id = struct.unpack_from(">QQ", buf, offset) + offset += 16 + pose = PoseStamped( + ts=pose_ts, + frame_id=frame_id, + position=Vector3(px, py, pz), + orientation=Quaternion(qx, qy, qz, qw), + ) + nodes.append(Graph3D.Node3D(pose=pose, id=node_id, metadata_id=metadata_id)) + + transforms: list[GraphDelta3D.Transform] = [] + for _ in range(node_count): + tx, ty, tz, qx, qy, qz, qw = struct.unpack_from(">7d", buf, offset) + offset += 56 + transforms.append( + cls.Transform( + translation=Vector3(tx, ty, tz), + rotation=Quaternion(qx, qy, qz, qw), + ) + ) + + return cls(ts=graph_ts, nodes=nodes, transforms=transforms) + + def to_rerun( + self, + z_offset: float = 0.0, + arrow_scale: float = 1.0, + ) -> Archetype: + """Render each (node, transform) pair as an arrow from node.pose to post_pose. + + The arrow origin is the node's current position; the vector is + the translation component of the transform (scaled by + ``arrow_scale``). Rotation deltas aren't visualized by default — + callers wanting to see those can subclass. + """ + import rerun as rr + + if not self.nodes: + return rr.Arrows3D(origins=[], vectors=[]) + + origins = [] + vectors = [] + for node, transform in zip(self.nodes, self.transforms, strict=True): + origins.append( + [ + node.pose.position.x, + node.pose.position.y, + node.pose.position.z + z_offset, + ] + ) + vectors.append( + [ + transform.translation.x * arrow_scale, + transform.translation.y * arrow_scale, + transform.translation.z * arrow_scale, + ] + ) + return rr.Arrows3D(origins=origins, vectors=vectors) + + def __len__(self) -> int: + return len(self.nodes) + + def __str__(self) -> str: + return f"GraphDelta3D(nodes={len(self.nodes)})" diff --git a/dimos/navigation/jnav/msgs/Landmark.py b/dimos/navigation/jnav/msgs/Landmark.py new file mode 100644 index 0000000000..d299fa3404 --- /dev/null +++ b/dimos/navigation/jnav/msgs/Landmark.py @@ -0,0 +1,200 @@ +# 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. + +"""Landmark: a cross-map reference shared between maps by ``id``. + +Two maps that both observe the same ``id`` are bridgeable: MultiMap composes the +relative transform between their frames through the landmark. The pose is stored +relative to the frame it was observed in (a camera frame for an apriltag, an odom +frame for a relocalization event, the nearest cloud's frame for a UI click), NOT +the map root — so the uncertainty stays honest for a future gtsam graph. MultiMap +resolves it to the map root via the map's recorded tf at ``ts``. + +``id`` is URL-like and encodes the exact source so identical tag numbers from +different families/sizes don't false-bridge unrelated maps, e.g. +``apriltag://36h11/40cm/5`` or ``reloc://map0/dim_city``. + +``replacement`` (milliseconds) lets a perceiver publish a *corrective* landmark: +a consumer (e.g. the PGO factor-graph manager) wipes out any earlier landmark of +the same ``id`` observed within ``replacement`` ms before this one's ``ts`` and +keeps this one instead. ``0`` (the default) means "additive" — keep all prior +landmarks, which is what loop closure needs (the first and the revisit sighting +must both survive). A small positive window only supersedes the immediately-prior +noisy estimate of the *same* sighting, never the historical closure observation. + +Confidence is **per-axis**, not a single scalar: ``confidence_x/_y/_z`` (the +observed translation axes) and ``confidence_roll/_pitch/_yaw`` (the rotation +axes), each in ``[0, 1]``. A consumer maps each per-axis confidence to that +axis's measurement weight (variance ~ 1 / confidence), so a landmark can be well +observed on some DOF and weak on others — e.g. a face-on AprilTag pins in-plane +position + yaw tightly but its range and out-of-plane tilt are loose. ``confidence`` +(the scalar) is retained as a coarse overall value; the per-axis fields default +to it when unspecified. +""" + +from __future__ import annotations + +import struct +import time +from typing import BinaryIO + +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.types.timestamped import Timestamped + + +class Landmark(Timestamped): + msg_name = "jnav.Landmark" + + ts: float + id: str # URL-like, shared across maps — the bridge key + map_id: str # which map's frame system this is in + frame_id: str # the observation frame the pose is relative to + kind: str # "apriltag" | "aruco" | "reloc" | "ui_click" | "shortcut" + confidence: float # coarse overall confidence; per-axis fields default to it + pose: Pose + replacement: float # ms: wipe same-id landmarks newer than (ts - replacement); 0 = additive + # Per-axis confidence in [0, 1]: how much this landmark constrains each DOF. + confidence_x: float + confidence_y: float + confidence_z: float + confidence_roll: float + confidence_pitch: float + confidence_yaw: float + + def __init__( + self, + id: str = "", + map_id: str = "", + pose: Pose | None = None, + frame_id: str = "", + confidence: float = 0.0, + kind: str = "", + ts: float = 0.0, + replacement: float = 0.0, + confidence_x: float | None = None, + confidence_y: float | None = None, + confidence_z: float | None = None, + confidence_roll: float | None = None, + confidence_pitch: float | None = None, + confidence_yaw: float | None = None, + ) -> None: + self.ts = ts if ts != 0 else time.time() + self.id = id + self.map_id = map_id + self.frame_id = frame_id + self.kind = kind + self.confidence = confidence + self.pose = pose if pose is not None else Pose() + self.replacement = replacement + # Per-axis confidence defaults to the coarse overall confidence. + self.confidence_x = confidence if confidence_x is None else confidence_x + self.confidence_y = confidence if confidence_y is None else confidence_y + self.confidence_z = confidence if confidence_z is None else confidence_z + self.confidence_roll = confidence if confidence_roll is None else confidence_roll + self.confidence_pitch = confidence if confidence_pitch is None else confidence_pitch + self.confidence_yaw = confidence if confidence_yaw is None else confidence_yaw + + def axis_confidence(self) -> tuple[float, float, float, float, float, float]: + """Per-axis confidence as ``(x, y, z, roll, pitch, yaw)``.""" + return ( + self.confidence_x, + self.confidence_y, + self.confidence_z, + self.confidence_roll, + self.confidence_pitch, + self.confidence_yaw, + ) + + def lcm_encode(self) -> bytes: + parts: list[bytes] = [struct.pack(">d", self.ts)] + for text in (self.id, self.map_id, self.frame_id, self.kind): + encoded = text.encode("utf-8") + parts.append(struct.pack(">I", len(encoded))) + parts.append(encoded) + parts.append(struct.pack(">d", self.confidence)) + p = self.pose + parts.append( + struct.pack( + ">7d", + p.position.x, + p.position.y, + p.position.z, + p.orientation.x, + p.orientation.y, + p.orientation.z, + p.orientation.w, + ) + ) + parts.append(struct.pack(">d", self.replacement)) + parts.append( + struct.pack( + ">6d", + self.confidence_x, + self.confidence_y, + self.confidence_z, + self.confidence_roll, + self.confidence_pitch, + self.confidence_yaw, + ) + ) + return b"".join(parts) + + @classmethod + def lcm_decode(cls, data: bytes | BinaryIO) -> Landmark: + buf = data if isinstance(data, (bytes, bytearray)) else data.read() + offset = 0 + (ts,) = struct.unpack_from(">d", buf, offset) + offset += 8 + texts: list[str] = [] + for _ in range(4): + (length,) = struct.unpack_from(">I", buf, offset) + offset += 4 + texts.append(buf[offset : offset + length].decode("utf-8")) + offset += length + id_, map_id, frame_id, kind = texts + (confidence,) = struct.unpack_from(">d", buf, offset) + offset += 8 + px, py, pz, qx, qy, qz, qw = struct.unpack_from(">7d", buf, offset) + offset += 56 + pose = Pose() + pose.position = Vector3(px, py, pz) + pose.orientation = Quaternion(qx, qy, qz, qw) + # Backward-compatible: older encodings omit the trailing replacement field. + replacement = 0.0 + if offset + 8 <= len(buf): + (replacement,) = struct.unpack_from(">d", buf, offset) + offset += 8 + # Backward-compatible: older encodings omit the per-axis confidences; + # they then default to the coarse overall confidence (None -> confidence). + cx = cy = cz = croll = cpitch = cyaw = None + if offset + 48 <= len(buf): + cx, cy, cz, croll, cpitch, cyaw = struct.unpack_from(">6d", buf, offset) + return cls( + id=id_, + map_id=map_id, + pose=pose, + frame_id=frame_id, + confidence=confidence, + kind=kind, + ts=ts, + replacement=replacement, + confidence_x=cx, + confidence_y=cy, + confidence_z=cz, + confidence_roll=croll, + confidence_pitch=cpitch, + confidence_yaw=cyaw, + ) diff --git a/dimos/navigation/jnav/msgs/Marker.py b/dimos/navigation/jnav/msgs/Marker.py new file mode 100644 index 0000000000..a7ed3fef98 --- /dev/null +++ b/dimos/navigation/jnav/msgs/Marker.py @@ -0,0 +1,97 @@ +# 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. + +"""Marker: a named navigation marker — a pose in a frame, a marker name, and an +optional map name. + +Drives the objective handler's ``goal_marker`` (navigate to a named marker) and +``save_marker`` (persist the current/given pose under a name) streams, so callers +can reference waypoints by name (e.g. dim_city's ``test_waypoint_*``) and scope +them to a map, instead of passing a bare Point/Pose with no identity. +""" + +from __future__ import annotations + +import struct +import time +from typing import BinaryIO + +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.types.timestamped import Timestamped + + +class Marker(Timestamped): + msg_name = "jnav.Marker" + + ts: float + frame_id: str + pose: Pose + marker: str # marker name (identity) + map: str # optional map name this marker belongs to ("" = unspecified) + + def __init__( + self, + pose: Pose | None = None, + marker: str = "", + map: str = "", + ts: float = 0.0, + frame_id: str = "map", + ) -> None: + self.ts = ts if ts != 0 else time.time() + self.frame_id = frame_id + self.pose = pose if pose is not None else Pose() + self.marker = marker + self.map = map + + def lcm_encode(self) -> bytes: + parts: list[bytes] = [struct.pack(">d", self.ts)] + for text in (self.frame_id, self.marker, self.map): + encoded = text.encode("utf-8") + parts.append(struct.pack(">I", len(encoded))) + parts.append(encoded) + p = self.pose + parts.append( + struct.pack( + ">7d", + p.position.x, + p.position.y, + p.position.z, + p.orientation.x, + p.orientation.y, + p.orientation.z, + p.orientation.w, + ) + ) + return b"".join(parts) + + @classmethod + def lcm_decode(cls, data: bytes | BinaryIO) -> Marker: + buf = data if isinstance(data, (bytes, bytearray)) else data.read() + offset = 0 + (ts,) = struct.unpack_from(">d", buf, offset) + offset += 8 + texts: list[str] = [] + for _ in range(3): + (length,) = struct.unpack_from(">I", buf, offset) + offset += 4 + texts.append(buf[offset : offset + length].decode("utf-8")) + offset += length + frame_id, marker, map_name = texts + px, py, pz, qx, qy, qz, qw = struct.unpack_from(">7d", buf, offset) + pose = Pose() + pose.position = Vector3(px, py, pz) + pose.orientation = Quaternion(qx, qy, qz, qw) + return cls(pose=pose, marker=marker, map=map_name, ts=ts, frame_id=frame_id) diff --git a/dimos/navigation/jnav/msgs/test_Landmark.py b/dimos/navigation/jnav/msgs/test_Landmark.py new file mode 100644 index 0000000000..1df849151a --- /dev/null +++ b/dimos/navigation/jnav/msgs/test_Landmark.py @@ -0,0 +1,129 @@ +# 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. + +"""Roundtrip + field tests for the jnav Landmark message.""" + +from __future__ import annotations + +from dimos.memory2.codecs.base import codec_for +from dimos.memory2.codecs.lcm import LcmCodec +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.navigation.jnav.msgs.Landmark import Landmark + + +def _pose(x: float, y: float, z: float) -> Pose: + pose = Pose() + pose.position.x, pose.position.y, pose.position.z = x, y, z + return pose + + +def test_landmark_roundtrip_preserves_all_fields() -> None: + landmark = Landmark( + id="apriltag://36h11/40cm/5", + map_id="dim_city", + pose=_pose(1.5, -2.0, 0.3), + frame_id="camera_optical", + confidence=0.82, + kind="apriltag", + ts=1781565207.5, + ) + decoded = Landmark.lcm_decode(landmark.lcm_encode()) + + assert decoded.id == "apriltag://36h11/40cm/5" + assert decoded.map_id == "dim_city" + assert decoded.frame_id == "camera_optical" + assert decoded.kind == "apriltag" + assert decoded.confidence == 0.82 + assert decoded.ts == landmark.ts + assert decoded.pose.position.x == 1.5 + assert decoded.pose.position.y == -2.0 + assert decoded.pose.position.z == 0.3 + + +def test_landmark_defaults() -> None: + landmark = Landmark() + assert landmark.id == "" + assert landmark.map_id == "" + assert landmark.frame_id == "" + assert landmark.kind == "" + assert landmark.confidence == 0.0 + assert landmark.replacement == 0.0 # additive by default + assert landmark.ts > 0 # auto-stamped + + decoded = Landmark.lcm_decode(landmark.lcm_encode()) + assert decoded.id == "" and decoded.map_id == "" and decoded.kind == "" + assert decoded.replacement == 0.0 + + +def test_landmark_replacement_roundtrips() -> None: + landmark = Landmark( + id="apriltag://36h11/40cm/13", + pose=_pose(0.1, 0.2, 0.5), + kind="apriltag", + ts=1781565207.5, + replacement=2500.0, # ms: corrective landmark supersedes recent same-id ones + ) + decoded = Landmark.lcm_decode(landmark.lcm_encode()) + assert decoded.replacement == 2500.0 + + +def test_landmark_decode_tolerates_legacy_bytes_without_replacement() -> None: + """Bytes written before the replacement + per-axis fields decode to defaults.""" + landmark = Landmark(id="x", pose=_pose(0.0, 0.0, 0.0), confidence=0.7, replacement=999.0) + # Strip the trailing replacement double + the 6 per-axis confidence doubles. + legacy = landmark.lcm_encode()[: -(8 + 48)] + decoded = Landmark.lcm_decode(legacy) + assert decoded.id == "x" + assert decoded.replacement == 0.0 + # Per-axis confidences fall back to the coarse overall confidence. + assert decoded.axis_confidence() == (0.7, 0.7, 0.7, 0.7, 0.7, 0.7) + + +def test_per_axis_confidence_defaults_to_overall() -> None: + lm = Landmark(id="t", confidence=0.6) + assert lm.axis_confidence() == (0.6, 0.6, 0.6, 0.6, 0.6, 0.6) + + +def test_per_axis_confidence_roundtrips() -> None: + lm = Landmark( + id="apriltag://36h11/40cm/13", + pose=_pose(0.1, 0.2, 0.5), + kind="apriltag", + confidence=0.5, + confidence_x=0.9, + confidence_y=0.8, + confidence_z=0.7, + confidence_roll=0.2, + confidence_pitch=0.1, + confidence_yaw=0.95, + ) + d = Landmark.lcm_decode(lm.lcm_encode()) + assert d.axis_confidence() == (0.9, 0.8, 0.7, 0.2, 0.1, 0.95) + assert d.confidence == 0.5 # coarse overall preserved independently + + +def test_per_axis_confidence_can_be_silent_on_some_dof() -> None: + """A landmark can be confident on some DOF and zero on others.""" + lm = Landmark(id="t", pose=_pose(0.0, 0.0, 0.0), confidence_z=0.9, confidence_x=0.0) + d = Landmark.lcm_decode(lm.lcm_encode()) + assert d.confidence_z == 0.9 and d.confidence_x == 0.0 + + +def test_landmark_uses_lcm_codec_in_memory2() -> None: + codec = codec_for(Landmark) + assert isinstance(codec, LcmCodec) + landmark = Landmark(id="reloc://map0/dim_city", map_id="map0", confidence=0.6, kind="reloc") + decoded = codec.decode(codec.encode(landmark)) + assert decoded.id == "reloc://map0/dim_city" + assert decoded.confidence == 0.6 diff --git a/dimos/navigation/jnav/msgs/test_Marker.py b/dimos/navigation/jnav/msgs/test_Marker.py new file mode 100644 index 0000000000..61394999d7 --- /dev/null +++ b/dimos/navigation/jnav/msgs/test_Marker.py @@ -0,0 +1,67 @@ +# 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. + +# 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. + +"""Roundtrip + field tests for the jnav Marker message.""" + +from __future__ import annotations + +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.navigation.jnav.msgs.Marker import Marker + + +def _pose(x: float, y: float, z: float) -> Pose: + pose = Pose() + pose.position.x, pose.position.y, pose.position.z = x, y, z + return pose + + +def test_marker_roundtrip_preserves_all_fields() -> None: + marker = Marker( + pose=_pose(-52.93, -55.95, 0.4), + marker="test_waypoint_2", + map="dim_city", + ts=1781565207.5, + frame_id="map", + ) + decoded = Marker.lcm_decode(marker.lcm_encode()) + + assert decoded.marker == "test_waypoint_2" + assert decoded.map == "dim_city" + assert decoded.frame_id == "map" + assert decoded.ts == marker.ts + assert decoded.pose.position.x == -52.93 + assert decoded.pose.position.y == -55.95 + assert decoded.pose.position.z == 0.4 + + +def test_marker_defaults() -> None: + marker = Marker() + assert marker.marker == "" + assert marker.map == "" + assert marker.frame_id == "map" + assert marker.ts > 0 # auto-stamped + # empty optional strings survive the roundtrip + decoded = Marker.lcm_decode(marker.lcm_encode()) + assert decoded.marker == "" and decoded.map == "" + + +def test_marker_unicode_name() -> None: + decoded = Marker.lcm_decode(Marker(marker="café_door", map="map_α").lcm_encode()) # noqa: RUF001 + assert decoded.marker == "café_door" + assert decoded.map == "map_α" # noqa: RUF001 From 617e5f7172b3bb21364695d50cfd155fbbb113f9 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 14:52:23 +0800 Subject: [PATCH 06/69] add jnav utils for PGO (apriltags, recording_db, trajectory_metrics, kitti, voxel_map, module_loading) --- .../jnav/utils/apriltag_agreement.py | 156 ++++ dimos/navigation/jnav/utils/apriltags.py | 665 ++++++++++++++++++ dimos/navigation/jnav/utils/kitti.py | 94 +++ dimos/navigation/jnav/utils/module_loading.py | 57 ++ dimos/navigation/jnav/utils/recording_db.py | 107 +++ .../jnav/utils/test_apriltag_agreement.py | 74 ++ dimos/navigation/jnav/utils/test_kitti.py | 60 ++ .../jnav/utils/trajectory_metrics.py | 268 +++++++ dimos/navigation/jnav/utils/voxel_map.py | 66 ++ 9 files changed, 1547 insertions(+) create mode 100644 dimos/navigation/jnav/utils/apriltag_agreement.py create mode 100644 dimos/navigation/jnav/utils/apriltags.py create mode 100644 dimos/navigation/jnav/utils/kitti.py create mode 100644 dimos/navigation/jnav/utils/module_loading.py create mode 100644 dimos/navigation/jnav/utils/recording_db.py create mode 100644 dimos/navigation/jnav/utils/test_apriltag_agreement.py create mode 100644 dimos/navigation/jnav/utils/test_kitti.py create mode 100644 dimos/navigation/jnav/utils/trajectory_metrics.py create mode 100644 dimos/navigation/jnav/utils/voxel_map.py diff --git a/dimos/navigation/jnav/utils/apriltag_agreement.py b/dimos/navigation/jnav/utils/apriltag_agreement.py new file mode 100644 index 0000000000..bab26e2305 --- /dev/null +++ b/dimos/navigation/jnav/utils/apriltag_agreement.py @@ -0,0 +1,156 @@ +# 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. + +"""April-tag agreement: a drift-quality metric for a corrected trajectory. + +A fixed April tag observed many times along a trajectory should map to ONE world +position. Odometry drift scatters those estimates (the same tag lands in a +different spot each lap); a good loop-closure / PGO correction pulls them back +together. So the spread of a tag's repeated world-position estimates is a +ground-truth-free measure of trajectory consistency — and the *drop* in spread +from raw odometry to PGO-corrected odometry is how much PGO helped. + +Pure functions over plain arrays; no PGO, sim, or I/O here (the benchmark harness +in gsc_pgo/pgo_apriltag_benchmark.py feeds these from real hk_village data). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from dimos.navigation.jnav.utils.trajectory_metrics import PoseLookup + +# sightings further apart than this are treated as separate visits +VISIT_GAP_S = 20.0 + + +@dataclass(frozen=True) +class TagAgreement: + """Per-tag spread of repeated world-position estimates (metres).""" + + tag_id: int + observations: int + spread: float # RMS distance of estimates to their centroid + + +@dataclass(frozen=True) +class AgreementReport: + """Agreement across all multiply-observed tags.""" + + per_tag: tuple[TagAgreement, ...] + mean_spread: float # mean per-tag spread (metres); lower = better agreement + total_observations: int + + @property + def tag_count(self) -> int: + return len(self.per_tag) + + +def tag_spread(positions: np.ndarray) -> float: + """RMS distance of a tag's world-position estimates to their centroid.""" + if len(positions) < 2: + return 0.0 + centroid = positions.mean(axis=0) + return float(np.sqrt(np.mean(np.sum((positions - centroid) ** 2, axis=1)))) + + +def agreement_report( + tag_positions: dict[int, np.ndarray], *, min_observations: int = 2 +) -> AgreementReport: + """Score per-tag agreement from ``tag_id -> (N, 3) world positions``. + + Tags seen fewer than ``min_observations`` times carry no agreement signal and + are excluded from ``mean_spread``. + """ + per_tag: list[TagAgreement] = [] + total = 0 + for tag_id in sorted(tag_positions): + positions = np.asarray(tag_positions[tag_id], dtype=np.float64).reshape(-1, 3) + total += len(positions) + if len(positions) < min_observations: + continue + per_tag.append(TagAgreement(tag_id, len(positions), tag_spread(positions))) + mean_spread = float(np.mean([t.spread for t in per_tag])) if per_tag else 0.0 + return AgreementReport(tuple(per_tag), mean_spread, total) + + +def agreement_improvement(raw: AgreementReport, corrected: AgreementReport) -> float: + """Fractional drop in mean spread from ``raw`` to ``corrected`` (1.0 = perfect). + + Positive means the correction tightened tag agreement; negative means it made + it worse. Returns 0.0 if there's no raw spread to improve on. + """ + if raw.mean_spread <= 0.0: + return 0.0 + return (raw.mean_spread - corrected.mean_spread) / raw.mean_spread + + +def tag_world_positions( + sightings: dict[int, list[float]], pose_lookup: PoseLookup +) -> dict[int, np.ndarray]: + """Map each tag's sighting times to robot world positions (the proxy estimate).""" + positions: dict[int, np.ndarray] = {} + for tag_id, times in sightings.items(): + located = [p for p in (pose_lookup(t) for t in times) if p is not None] + if located: + positions[tag_id] = np.vstack(located) + return positions + + +def split_visits(times: list[float], *, gap_s: float) -> list[list[float]]: + """Cluster sighting timestamps into visits separated by gaps > ``gap_s``.""" + visits: list[list[float]] = [] + for timestamp in sorted(times): + if visits and timestamp - visits[-1][-1] <= gap_s: + visits[-1].append(timestamp) + else: + visits.append([timestamp]) + return visits + + +def paired_tag_visit_positions( + sightings: dict[int, list[float]], + raw_lookup: PoseLookup, + corrected_lookup: PoseLookup, + *, + gap_s: float, +) -> tuple[dict[int, np.ndarray], dict[int, np.ndarray]]: + """One robot position per tag VISIT under both pose sources, visits paired. + + Outdoors a tag stays visible while the robot walks tens of metres, so + per-sighting spread is dominated by viewing distance, not drift. Visits are + clustered on timestamps only, and a visit is kept only when BOTH pose + sources can place it — so raw and corrected reports always score the exact + same visit set, and a visit outside the pose graph's coverage drops out + instead of skewing one side. + """ + raw_positions: dict[int, np.ndarray] = {} + corrected_positions: dict[int, np.ndarray] = {} + for tag_id, times in sightings.items(): + raw_medians: list[np.ndarray] = [] + corrected_medians: list[np.ndarray] = [] + for visit_times in split_visits(times, gap_s=gap_s): + raw_located = [p for p in (raw_lookup(t) for t in visit_times) if p is not None] + corrected_located = [ + p for p in (corrected_lookup(t) for t in visit_times) if p is not None + ] + if raw_located and corrected_located: + raw_medians.append(np.median(np.vstack(raw_located), axis=0)) + corrected_medians.append(np.median(np.vstack(corrected_located), axis=0)) + if raw_medians: + raw_positions[tag_id] = np.vstack(raw_medians) + corrected_positions[tag_id] = np.vstack(corrected_medians) + return raw_positions, corrected_positions diff --git a/dimos/navigation/jnav/utils/apriltags.py b/dimos/navigation/jnav/utils/apriltags.py new file mode 100644 index 0000000000..9bf22cdf82 --- /dev/null +++ b/dimos/navigation/jnav/utils/apriltags.py @@ -0,0 +1,665 @@ +# 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. + +"""AprilTag detection over a mem2.db color stream. + +Detects tags in the `color_image` stream and rejects bad glimpses through several +independent gates — motion blur (per-tag sharpness), PnP misfit (reprojection +error), too-small / too-far / too-oblique views, and fast camera motion — then +clusters same-id detections in time and reduces each cluster to one robust pose +via a Huber-weighted refinement seeded at the cluster medoid. (Re)writes the +`april_tags` PoseStamped stream of tag-in-camera relative poses (solvePnP, +marker_id in each observation's tags) and returns those representatives for the +GTSAM solve. +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Callable, Iterable +from copy import copy +from dataclasses import dataclass, field +from itertools import pairwise +import json +import math +from pathlib import Path +import time +from typing import Any + +import cv2 +import numpy as np +from scipy.spatial.transform import Rotation + +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.Image import Image + +DEFAULT_MAX_DISTANCE_M = 1.0 +DEFAULT_MAX_VIEW_ANGLE_DEG = 45.0 +DEFAULT_CLUSTER_GAP_SEC = 5.0 +DEFAULT_ROTATION_WEIGHT_M_PER_RAD = 0.5 +# A blurry tag still solves a pose; these reject the bad glimpses up front. +DEFAULT_MIN_SHARPNESS = 60.0 # Laplacian variance over the tag ROI +DEFAULT_MAX_REPROJ_PX = 2.0 # RMS solvePnP corner reprojection error +DEFAULT_MIN_TAG_PX = 24.0 # tag side length in pixels (sqrt of quad area) +DEFAULT_MAX_LINEAR_SPEED_MPS = 0.5 +DEFAULT_MAX_ANGULAR_SPEED_DPS = 50.0 +DEFAULT_MIN_OBSERVATIONS = 3 # clusters thinner than this are unreliable +DEFAULT_HUBER_DELTA_M = 0.05 # residual past which a sample is down-weighted +_HUBER_ITERATIONS = 5 + +# One tag observation: ts, marker_id, t_cam_marker (xyz + xyzw quat), and the +# per-glimpse quality fields (sharpness, reproj_px, tag_px, speed). +Detection = dict[str, Any] + + +def make_detector(dictionary_name: str) -> cv2.aruco.ArucoDetector: + d = cv2.aruco.getPredefinedDictionary(getattr(cv2.aruco, dictionary_name)) + return cv2.aruco.ArucoDetector(d, cv2.aruco.DetectorParameters()) + + +def _object_points(marker_length_m: float) -> np.ndarray: + h = marker_length_m / 2.0 + return np.array([[-h, h, 0.0], [h, h, 0.0], [h, -h, 0.0], [-h, -h, 0.0]], dtype=np.float32) + + +def estimate_marker_pose( + corners_pixels: np.ndarray, + marker_length_m: float, + intrinsics: np.ndarray, + distortion: np.ndarray, +) -> tuple[np.ndarray, np.ndarray] | None: + """solvePnP a single tag -> (rotation_vector, translation_vector) in the + camera_optical frame, or None if it failed.""" + image_corners = corners_pixels.reshape(4, 1, 2).astype(np.float32) + found, rotation_vector, translation_vector = cv2.solvePnP( + _object_points(marker_length_m), + image_corners, + intrinsics, + distortion, + flags=cv2.SOLVEPNP_IPPE_SQUARE, + ) + return (rotation_vector, translation_vector) if found else None + + +def view_quality(t_cam_marker: list[float]) -> tuple[float, float]: + """(distance_m, view_angle_deg) for a tag pose in the camera optical frame. + + distance is the camera->tag range; view_angle is the angle between the line + of sight and the tag's surface normal (0 = perfectly head-on).""" + translation = np.array(t_cam_marker[:3], dtype=np.float64) + distance = float(np.linalg.norm(translation)) + normal = Rotation.from_quat(t_cam_marker[3:7]).as_matrix()[:, 2] + line_of_sight = translation / (distance + 1e-9) + cos_angle = abs(float(np.dot(line_of_sight, normal))) + view_angle = math.degrees(math.acos(min(1.0, cos_angle))) + return distance, view_angle + + +def cluster_by_time(detections: list[Detection], gap_sec: float) -> list[list[Detection]]: + """Group same-marker detections into clusters. A new cluster begins whenever + the time gap to the previous same-marker detection exceeds gap_sec.""" + clusters: list[list[Detection]] = [] + by_marker: dict[int, list[Detection]] = defaultdict(list) + for detection in detections: + by_marker[detection["marker_id"]].append(detection) + for marker_detections in by_marker.values(): + marker_detections.sort(key=lambda detection: detection["ts"]) + current = [marker_detections[0]] + for detection in marker_detections[1:]: + if detection["ts"] - current[-1]["ts"] > gap_sec: + clusters.append(current) + current = [detection] + else: + current.append(detection) + clusters.append(current) + return clusters + + +def _pose_distance(a: list[float], b: list[float], rotation_weight_m_per_rad: float) -> float: + translation = float(np.linalg.norm(np.array(a[:3]) - np.array(b[:3]))) + rotation = 2.0 * math.acos(min(1.0, abs(float(np.dot(a[3:7], b[3:7]))))) + return translation + rotation_weight_m_per_rad * rotation + + +def cluster_medoid(cluster: list[Detection], rotation_weight_m_per_rad: float) -> Detection: + """The detection whose pose is most central (min total spatial+rotational + distance to the rest) — a robust representative of the cluster.""" + poses = [detection["t_cam_marker"] for detection in cluster] + best_index, best_cost = 0, float("inf") + for i in range(len(poses)): + cost = sum( + _pose_distance(poses[i], poses[j], rotation_weight_m_per_rad) + for j in range(len(poses)) + if j != i + ) + if cost < best_cost: + best_cost, best_index = cost, i + return cluster[best_index] + + +def reprojection_error_px( + corners_pixels: np.ndarray, + rotation_vector: np.ndarray, + translation_vector: np.ndarray, + marker_length_m: float, + intrinsics: np.ndarray, + distortion: np.ndarray, +) -> float: + """RMS pixel distance between detected corners and the solvePnP pose reprojected + back onto the image — a direct measure of how well the pose explains the tag.""" + projected, _ = cv2.projectPoints( + _object_points(marker_length_m), rotation_vector, translation_vector, intrinsics, distortion + ) + measured = corners_pixels.reshape(4, 2).astype(np.float64) + diff = projected.reshape(4, 2) - measured + return float(np.sqrt(np.mean(np.sum(diff * diff, axis=1)))) + + +def tag_pixel_size(corners_pixels: np.ndarray) -> float: + """Tag side length in pixels (sqrt of the quad's image area); small = unreliable.""" + quad = corners_pixels.reshape(4, 2).astype(np.float32) + return float(math.sqrt(abs(cv2.contourArea(quad)))) + + +def tag_sharpness(gray: np.ndarray, corners_pixels: np.ndarray) -> float: + """Laplacian variance over the tag's bounding box — low under motion blur.""" + quad = corners_pixels.reshape(4, 2) + x_min, y_min = np.floor(quad.min(0)).astype(int) + x_max, y_max = np.ceil(quad.max(0)).astype(int) + height, width = gray.shape[:2] + x_min, y_min = max(int(x_min), 0), max(int(y_min), 0) + x_max, y_max = min(int(x_max), width), min(int(y_max), height) + if x_max - x_min < 2 or y_max - y_min < 2: + return 0.0 + return float(cv2.Laplacian(gray[y_min:y_max, x_min:x_max], cv2.CV_64F).var()) + + +def _pose_xyz_quat(pose: Any) -> np.ndarray | None: + """Best-effort (x,y,z,qx,qy,qz,qw) from an observation pose (tuple/list or a msg + with .position/.orientation); None if it isn't a usable pose.""" + if pose is None: + return None + if hasattr(pose, "position") and hasattr(pose, "orientation"): + position, orientation = pose.position, pose.orientation + return np.array( + [ + position.x, + position.y, + position.z, + orientation.x, + orientation.y, + orientation.z, + orientation.w, + ] + ) + try: + values = [float(component) for component in pose] + except TypeError: + return None + return np.array(values[:7]) if len(values) >= 7 else None + + +def _camera_speeds(images: list[Any]) -> tuple[dict[float, tuple[float, float]], bool]: + """Per-image (linear m/s, angular deg/s) from consecutive posed frames. The second + return is False when too few frames carry poses to estimate motion (gate disabled).""" + posed = [ + (float(obs.ts), pose) + for obs in images + if (pose := _pose_xyz_quat(getattr(obs, "pose", None))) is not None + ] + posed.sort(key=lambda item: item[0]) + speeds: dict[float, tuple[float, float]] = {} + for (timestamp_a, pose_a), (timestamp_b, pose_b) in pairwise(posed): + dt = timestamp_b - timestamp_a + if dt <= 0: + continue + linear = float(np.linalg.norm(pose_b[:3] - pose_a[:3])) / dt + cos_half = min(1.0, abs(float(np.dot(pose_a[3:7], pose_b[3:7])))) + angular = math.degrees(2.0 * math.acos(cos_half)) / dt + speeds[timestamp_b] = (linear, angular) + return speeds, len(posed) >= 2 + + +def _huber_weights(residuals: np.ndarray, delta: float) -> np.ndarray: + """IRLS Huber weights: 1 inside `delta`, decaying as delta/r past it.""" + weights = np.ones_like(residuals) + outside = residuals > delta + weights[outside] = delta / residuals[outside] + return weights + + +def robust_cluster_pose( + cluster: list[Detection], rotation_weight_m_per_rad: float, huber_delta_m: float +) -> Detection: + """Cluster representative: the medoid, then refined by Huber-weighted IRLS — a + weighted-mean translation and weighted quaternion mean (Markley eigen method), + re-weighting each iteration so a lingering bad glimpse keeps losing influence.""" + medoid = cluster_medoid(cluster, rotation_weight_m_per_rad) + if len(cluster) < 2: + return medoid + poses = np.array([detection["t_cam_marker"] for detection in cluster], dtype=np.float64) + translations, quaternions = poses[:, :3], poses[:, 3:7] + reference = np.array(medoid["t_cam_marker"][3:7]) + signs = np.sign(quaternions @ reference) + signs[signs == 0] = 1.0 + quaternions = quaternions * signs[:, None] + estimate_translation = np.array(medoid["t_cam_marker"][:3]) + estimate_quaternion = reference.copy() + delta_rad = huber_delta_m / max(rotation_weight_m_per_rad, 1e-9) + for _ in range(_HUBER_ITERATIONS): + weights_t = _huber_weights( + np.linalg.norm(translations - estimate_translation, axis=1), huber_delta_m + ) + estimate_translation = (weights_t[:, None] * translations).sum(0) / weights_t.sum() + angular_residual = 2.0 * np.arccos( + np.clip(np.abs(quaternions @ estimate_quaternion), 0.0, 1.0) + ) + weights_r = _huber_weights(angular_residual, delta_rad) + scatter = ( + weights_r[:, None, None] * np.einsum("ni,nj->nij", quaternions, quaternions) + ).sum(0) + estimate_quaternion = np.linalg.eigh(scatter)[1][:, -1] + if estimate_quaternion @ reference < 0: + estimate_quaternion = -estimate_quaternion + return { + **medoid, + "t_cam_marker": [*estimate_translation.tolist(), *estimate_quaternion.tolist()], + } + + +def detect_apriltags( + store: Any, + intrinsics: np.ndarray, + distortion: np.ndarray, + image_stream: str = "color_image", + stream_name: str = "april_tags", + marker_length: float = 0.10, + dictionary: str = "DICT_APRILTAG_36h11", + *, + max_distance_m: float = DEFAULT_MAX_DISTANCE_M, + max_view_angle_deg: float = DEFAULT_MAX_VIEW_ANGLE_DEG, + cluster_gap_sec: float = DEFAULT_CLUSTER_GAP_SEC, + rotation_weight_m_per_rad: float = DEFAULT_ROTATION_WEIGHT_M_PER_RAD, + min_sharpness: float = DEFAULT_MIN_SHARPNESS, + max_reproj_px: float = DEFAULT_MAX_REPROJ_PX, + min_tag_px: float = DEFAULT_MIN_TAG_PX, + max_linear_speed_mps: float = DEFAULT_MAX_LINEAR_SPEED_MPS, + max_angular_speed_dps: float = DEFAULT_MAX_ANGULAR_SPEED_DPS, + min_observations: int = DEFAULT_MIN_OBSERVATIONS, + huber_delta_m: float = DEFAULT_HUBER_DELTA_M, +) -> list[Detection]: + """Detect tags in `image_stream`, reject bad glimpses (blur, PnP misfit, small/ + far/oblique views, fast motion), cluster same-id detections by time, drop thin + clusters, and (re)write the `april_tags` stream from one Huber-refined medoid + representative per cluster. Returns that list of representatives.""" + detector = make_detector(dictionary) + raw_detections: list[Detection] = [] + images = store.stream(image_stream, Image).to_list() + speed_by_ts, speed_available = _camera_speeds(images) + for image_obs in images: + image = image_obs.data + bgr = image.numpy() if hasattr(image, "numpy") else np.asarray(image.data) + gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) if bgr.ndim == 3 else bgr + all_corners, marker_ids, _ = detector.detectMarkers(bgr) + if marker_ids is None: + continue + for corners, marker_id in zip(all_corners, marker_ids.flatten(), strict=False): + pose = estimate_marker_pose(corners, marker_length, intrinsics, distortion) + if pose is None: + continue + rotation_vector, translation_vector = pose + quaternion = Rotation.from_rotvec(rotation_vector.reshape(3)).as_quat() # x,y,z,w + translation = translation_vector.reshape(3) + tag_in_camera = [ + float(translation[0]), + float(translation[1]), + float(translation[2]), + float(quaternion[0]), + float(quaternion[1]), + float(quaternion[2]), + float(quaternion[3]), + ] + raw_detections.append( + { + "ts": float(image_obs.ts), + "marker_id": int(marker_id), + "t_cam_marker": tag_in_camera, + "sharpness": tag_sharpness(gray, corners), + "reproj_px": reprojection_error_px( + corners, + rotation_vector, + translation_vector, + marker_length, + intrinsics, + distortion, + ), + "tag_px": tag_pixel_size(corners), + "speed": speed_by_ts.get(float(image_obs.ts)), + } + ) + + # Per-glimpse gates; count rejections per reason so thresholds are tunable. + rejected: dict[str, int] = defaultdict(int) + kept: list[Detection] = [] + for detection in raw_detections: + if detection["sharpness"] < min_sharpness: + rejected["blur"] += 1 + continue + if detection["reproj_px"] > max_reproj_px: + rejected["reproj"] += 1 + continue + if detection["tag_px"] < min_tag_px: + rejected["small"] += 1 + continue + distance, view_angle = view_quality(detection["t_cam_marker"]) + if distance > max_distance_m: + rejected["far"] += 1 + continue + if view_angle > max_view_angle_deg: + rejected["oblique"] += 1 + continue + speed = detection["speed"] + if speed is not None and ( + speed[0] > max_linear_speed_mps or speed[1] > max_angular_speed_dps + ): + rejected["motion"] += 1 + continue + kept.append(detection) + + # One Huber-refined representative per time-clustered group; drop thin clusters. + detections: list[Detection] = [] + thin_clusters = 0 + for cluster in cluster_by_time(kept, cluster_gap_sec): + if len(cluster) < min_observations: + thin_clusters += 1 + continue + detections.append( + { + **robust_cluster_pose(cluster, rotation_weight_m_per_rad, huber_delta_m), + "n_observations": len(cluster), + } + ) + detections.sort(key=lambda detection: detection["ts"]) + + if stream_name in store.list_streams(): + store.delete_stream(stream_name) + april_tag_stream = store.stream(stream_name, PoseStamped) + for detection in detections: + tag_in_camera = detection["t_cam_marker"] + pose_stamped = PoseStamped( + ts=detection["ts"], position=tag_in_camera[:3], orientation=tag_in_camera[3:] + ) + april_tag_stream.append( + pose_stamped, + ts=detection["ts"], + pose=tuple(tag_in_camera), + tags={"marker_id": detection["marker_id"]}, + ) + + found_ids = sorted({detection["marker_id"] for detection in detections}) + gate_summary = ", ".join(f"{reason}={count}" for reason, count in sorted(rejected.items())) + if not speed_available: + gate_summary += (", " if gate_summary else "") + "motion-gate-off(no poses)" + print( + f" april_tags: {len(raw_detections)} raw -> {len(kept)} in-spec " + f"-> {len(detections)} clusters (dropped {thin_clusters} thin), " + f"markers {found_ids} (over {len(images)} images)" + ) + if gate_summary: + print(f" april_tags rejected: {gate_summary}") + return detections + + +# --------------------------------------------------------------------------- +# Streaming / incremental API +# --------------------------------------------------------------------------- + +_DEFAULT_MARKER_LENGTH = 0.10 +_DEFAULT_DICTIONARY = "DICT_APRILTAG_36h11" + + +@dataclass +class TagInfo: + tag_number: float + confidence: float + pose: list[float] # [x, y, z, qx, qy, qz, qw] relative to camera + ts: float + camera_frame_id: str = "" + + +@dataclass +class GroupingOptions: + tag_ids_to_ignore: list[float] = field(default_factory=list) + time_window: float = DEFAULT_CLUSTER_GAP_SEC + max_distance: float = DEFAULT_MAX_DISTANCE_M + min_confidence: float = 0.0 + + +class AprilTagger: + def __init__( + self, + camera_intrinsics: Any, + grouping_options: Any, + tags_ids_to_ignore: list[float], + ) -> None: + self.camera_intrinsics = camera_intrinsics + self.grouping_options = grouping_options + self.tags_ids_to_ignore = tags_ids_to_ignore + + def get_tag(self, img: Any) -> list[TagInfo]: + return pure_get_tags(img, self.camera_intrinsics) + + def iter_april_tag_groups( + self, + next_observation: Any, + options: Any = None, + state: dict[str, Any] | None = None, + ) -> tuple[list[TagInfo], dict[float, TagInfo], dict[float, TagInfo]]: + if state is None: + state = {} + if options is None: + options = self.grouping_options + recent_tags, tag_estimates, finished_tags = pure_iter_april_tag_groups( + next_observation, + self.camera_intrinsics, + options, + state.get("recent_tags", []), + state.get("tag_estimates", {}), + ) + state["recent_tags"] = recent_tags + state["tag_estimates"] = tag_estimates + return recent_tags, tag_estimates, finished_tags + + +def pure_get_tags(img: Any, camera_intrinsics: Any) -> list[TagInfo]: + """Detect AprilTags in an image and return relative poses via solvePnP. + + img: ndarray or an Image msg (frame_id/ts are carried onto each TagInfo). + camera_intrinsics: dict with "intrinsics" (3x3 ndarray), and optionally + "distortion", "marker_length", "dictionary". + """ + intrinsics = camera_intrinsics["intrinsics"] + distortion = camera_intrinsics.get("distortion", np.zeros(5)) + marker_length = camera_intrinsics.get("marker_length", _DEFAULT_MARKER_LENGTH) + dictionary = camera_intrinsics.get("dictionary", _DEFAULT_DICTIONARY) + + camera_frame_id = getattr(img, "frame_id", "") + ts = getattr(img, "ts", None) + if ts is None: + ts = time.time() + + if isinstance(img, np.ndarray): + bgr = img + elif hasattr(img, "numpy"): + bgr = img.numpy() + elif hasattr(img, "data"): + bgr = np.asarray(img.data) + else: + bgr = np.asarray(img) + if bgr.ndim == 2: + bgr = cv2.cvtColor(bgr, cv2.COLOR_GRAY2BGR) + + detector = make_detector(dictionary) + all_corners, marker_ids, _ = detector.detectMarkers(bgr) + if marker_ids is None: + return [] + + tags: list[TagInfo] = [] + for corners, marker_id in zip(all_corners, marker_ids.flatten(), strict=False): + pose_result = estimate_marker_pose(corners, marker_length, intrinsics, distortion) + if pose_result is None: + continue + rotation_vector, translation_vector = pose_result + quaternion = Rotation.from_rotvec(rotation_vector.reshape(3)).as_quat() + translation = translation_vector.reshape(3) + pose = [ + float(translation[0]), + float(translation[1]), + float(translation[2]), + float(quaternion[0]), + float(quaternion[1]), + float(quaternion[2]), + float(quaternion[3]), + ] + reproj = reprojection_error_px( + corners, + rotation_vector, + translation_vector, + marker_length, + intrinsics, + distortion, + ) + confidence = max(0.0, 1.0 - reproj / DEFAULT_MAX_REPROJ_PX) + tags.append( + TagInfo( + tag_number=float(marker_id), + confidence=confidence, + pose=pose, + ts=float(ts), + camera_frame_id=camera_frame_id, + ) + ) + return tags + + +def pure_iter_april_tag_groups( + next_observation: Any, + camera_intrinsics: Any, + options: Any, + recent_tags: list[TagInfo], + tag_estimates: dict[float, TagInfo], +) -> tuple[list[TagInfo], dict[float, TagInfo], dict[float, TagInfo]]: + """Returns (recent_tags, tag_estimates, finished_tags). + + finished_tags holds estimates whose tag id aged out of the time window this + iteration — no fresher observation can revise them, so they are final.""" + tag_estimates = copy(tag_estimates) + recent_tags = list(recent_tags) + + tags = pure_get_tags(next_observation.color_image, camera_intrinsics) + recent_tags.extend(tags) + + tag_ids_to_ignore = set(getattr(options, "tag_ids_to_ignore", [])) + time_window = getattr(options, "time_window", DEFAULT_CLUSTER_GAP_SEC) + max_distance = getattr(options, "max_distance", DEFAULT_MAX_DISTANCE_M) + min_confidence = getattr(options, "min_confidence", 0.0) + + # Purge tags outside the time window + if recent_tags: + latest_ts = max(tag.ts for tag in recent_tags) + recent_tags = [tag for tag in recent_tags if latest_ts - tag.ts <= time_window] + + all_recent_tag_ids = set( + tag.tag_number for tag in recent_tags if tag.tag_number not in tag_ids_to_ignore + ) + + # Tag ids with an estimate but no surviving observations are finalized + finished_tags = { + tag_id: estimate + for tag_id, estimate in tag_estimates.items() + if tag_id not in all_recent_tag_ids + } + for tag_id in finished_tags: + del tag_estimates[tag_id] + + for each_tag_id in all_recent_tag_ids: + tags_for_id = [tag for tag in recent_tags if tag.tag_number == each_tag_id] + + # Filter by distance and confidence + filtered = [ + tag + for tag in tags_for_id + if tag.confidence >= min_confidence + and float(np.linalg.norm(tag.pose[:3])) <= max_distance + ] + if not filtered: + continue + + # Pick the medoid (most central observation by pose distance) + if len(filtered) == 1: + tag_estimates[each_tag_id] = filtered[0] + else: + best_index = 0 + best_cost = float("inf") + for i, tag_i in enumerate(filtered): + cost = sum( + _pose_distance(tag_i.pose, tag_j.pose, DEFAULT_ROTATION_WEIGHT_M_PER_RAD) + for j, tag_j in enumerate(filtered) + if j != i + ) + if cost < best_cost: + best_cost = cost + best_index = i + tag_estimates[each_tag_id] = filtered[best_index] + + return recent_tags, tag_estimates, finished_tags + + +def load_intrinsics_json(path: Path) -> dict[str, Any]: + raw = json.loads(path.read_text()) + config: dict[str, Any] = { + "intrinsics": np.array(raw["intrinsics"], dtype=np.float64).reshape(3, 3), + "distortion": np.array(raw.get("distortion", []), dtype=np.float64), + } + for key in ("marker_length", "dictionary"): + if key in raw: + config[key] = raw[key] + return config + + +def sightings_from_observations( + observations: Iterable[tuple[int, float]], +) -> dict[int, list[float]]: + """Group ``(marker_id, timestamp)`` tag observations into ``tag_id -> [timestamps]``.""" + sightings: dict[int, list[float]] = {} + for marker_id, timestamp in observations: + sightings.setdefault(int(marker_id), []).append(float(timestamp)) + return sightings + + +def load_or_detect_sightings( + stored: Iterable[tuple[int, float]], + detect: Callable[[], Iterable[tuple[int, float]]], +) -> tuple[dict[int, list[float]], str]: + """``tag_id -> [sighting timestamps]`` from already-stored tag observations, + falling back to ``detect()`` when none are stored. + + Only marker_id + ts are used — tag positions are taken relative to the odom + stream, so stored poses are never read. The caller supplies the stored + iterable and the detect thunk, so this stays db-free. Returns the sightings + plus a source label (``"stream"`` or ``"detected"``).""" + sightings = sightings_from_observations(stored) + if sightings: + return sightings, "stream" + return sightings_from_observations(detect()), "detected" diff --git a/dimos/navigation/jnav/utils/kitti.py b/dimos/navigation/jnav/utils/kitti.py new file mode 100644 index 0000000000..d7358a976b --- /dev/null +++ b/dimos/navigation/jnav/utils/kitti.py @@ -0,0 +1,94 @@ +# 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. + +"""Official KITTI odometry error metric (translational % + rotational deg/m). + +The KITTI leaderboard metric: for every sub-sequence length in {100,200,...,800}m, +take every start frame, find the end frame ~that path-length away, and measure the +relative pose error between estimated and ground-truth. Average translational +error (as a fraction of length) and rotational error (rad per metre) over all +(start, length) pairs. Reported as translational % and rotational deg/m. This is +the devkit's `evaluate_odometry` algorithm. + +Poses are 4x4 body->world transforms (frame 0 = identity), one per frame. +""" + +from __future__ import annotations + +import numpy as np + +LENGTHS = (100.0, 200.0, 300.0, 400.0, 500.0, 600.0, 700.0, 800.0) +STEP = 10 # evaluate every 10th frame as a start (devkit convention) + + +def trajectory_distances(poses: list[np.ndarray]) -> list[float]: + """Cumulative path length at each frame.""" + distances = [0.0] + for index in range(1, len(poses)): + delta = poses[index][:3, 3] - poses[index - 1][:3, 3] + distances.append(distances[-1] + float(np.linalg.norm(delta))) + return distances + + +def _last_frame_for_length(distances: list[float], first: int, length: float) -> int: + target = distances[first] + length + for index in range(first, len(distances)): + if distances[index] >= target: + return index + return -1 + + +def _rotation_error(pose_error: np.ndarray) -> float: + trace = pose_error[0, 0] + pose_error[1, 1] + pose_error[2, 2] + return float(np.arccos(np.clip((trace - 1.0) / 2.0, -1.0, 1.0))) + + +def _translation_error(pose_error: np.ndarray) -> float: + return float(np.linalg.norm(pose_error[:3, 3])) + + +def kitti_odometry_error( + estimated: list[np.ndarray], ground_truth: list[np.ndarray] +) -> dict[str, float]: + """Average translational (%) and rotational (deg/m) error, devkit-style.""" + count = min(len(estimated), len(ground_truth)) + estimated, ground_truth = estimated[:count], ground_truth[:count] + distances = trajectory_distances(ground_truth) + + translational_errors: list[float] = [] + rotational_errors: list[float] = [] + for first in range(0, count, STEP): + for length in LENGTHS: + last = _last_frame_for_length(distances, first, length) + if last < 0: + continue + # Relative pose error: how far the estimated motion from first->last + # diverges from ground truth's. + gt_delta = np.linalg.inv(ground_truth[first]) @ ground_truth[last] + estimated_delta = np.linalg.inv(estimated[first]) @ estimated[last] + pose_error = np.linalg.inv(gt_delta) @ estimated_delta + translational_errors.append(_translation_error(pose_error) / length) + rotational_errors.append(_rotation_error(pose_error) / length) + + if not translational_errors: + return { + "translational_percent": float("nan"), + "rotational_deg_per_m": float("nan"), + "pairs": 0, + } + return { + "translational_percent": float(np.mean(translational_errors)) * 100.0, + "rotational_deg_per_m": float(np.degrees(np.mean(rotational_errors))), + "pairs": len(translational_errors), + } diff --git a/dimos/navigation/jnav/utils/module_loading.py b/dimos/navigation/jnav/utils/module_loading.py new file mode 100644 index 0000000000..3295abdea1 --- /dev/null +++ b/dimos/navigation/jnav/utils/module_loading.py @@ -0,0 +1,57 @@ +# 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. + +"""Dynamically load a Module class by file path and filter its config. + +Lets the loop-closure eval drivers swap in any module-under-test from a +``--module-path``/``--module-name`` pair without importing it statically. +""" + +from __future__ import annotations + +import importlib +from pathlib import Path +from typing import Any, get_type_hints + + +def load_module_class(module_path: Path, module_name: str) -> type: + """Import the class `module_name` from a python file. + + Resolved through the normal package system (path -> dotted name rooted at + the last `dimos` directory) so the class pickles/deploys into workers.""" + parts = module_path.resolve().with_suffix("").parts + if "dimos" not in parts: + raise SystemExit(f"--module-path must live inside the dimos package tree: {module_path}") + package_root = len(parts) - 1 - parts[::-1].index("dimos") + dotted = ".".join(parts[package_root:]) + module = importlib.import_module(dotted) + if not hasattr(module, module_name): + raise SystemExit(f"no class {module_name!r} in {dotted} ({module_path})") + return getattr(module, module_name) # type: ignore[no-any-return] + + +def filter_config_for_module(module_class: type, config: dict[str, Any]) -> dict[str, Any]: + """Drop config keys the module's Config class doesn't declare. + + DEFAULT_PGO_CONFIG carries cmu-specific knobs (use_scan_context, ...); + other loop-closure modules shouldn't crash on them.""" + try: + config_class = get_type_hints(module_class)["config"] + known = set(config_class.model_fields) + except Exception: + return config + dropped = sorted(set(config) - known) + if dropped: + print(f"config keys not in {config_class.__name__} (dropped): {dropped}") + return {key: value for key, value in config.items() if key in known} diff --git a/dimos/navigation/jnav/utils/recording_db.py b/dimos/navigation/jnav/utils/recording_db.py new file mode 100644 index 0000000000..c71ea06a4d --- /dev/null +++ b/dimos/navigation/jnav/utils/recording_db.py @@ -0,0 +1,107 @@ +# 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. + +"""Read recording streams (mem2.db) for loop-closure evaluation. + +SqliteStore access keyed by path (one shared store per db), stream iteration, +and a nearest-pose lookup over an odometry stream. Stream names are always +passed in by the caller — nothing here is tied to a particular recording layout. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import numpy as np + +from dimos.memory2.store.sqlite import SqliteStore +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.navigation.jnav.utils.trajectory_metrics import PoseLookup, trajectory_lookup + +# max |Δt| pairing a query time to an odom sample +ODOM_MATCH_TOLERANCE_S = 0.2 + +# Fixed-rate replay pacing + sizing caps (used by the eval replay drivers). +REPLAY_PUBLISH_HZ = 50.0 +REPLAY_DRAIN_MARGIN_S = 30.0 +MAX_REPLAY_SCANS = 4000 +MAX_REPLAY_ODOM = 16000 + +# One shared store per db path (SqliteStore can't take absolute paths through the +# replay adapter, so streams are read directly). +_stores: dict[str, SqliteStore] = {} + + +def store(db_path: Path) -> SqliteStore: + key = str(db_path) + cached = _stores.get(key) + if cached is None: + cached = SqliteStore(path=key, must_exist=True) + cached.start() + _stores[key] = cached + return cached + + +def list_streams(db_path: Path) -> list[str]: + return store(db_path).list_streams() + + +def stream_count(db_path: Path, stream_name: str) -> int: + return int(store(db_path).stream(stream_name).count()) + + +def iterate_stream( + db_path: Path, stream_name: str, *, stride: int = 1 +) -> Iterator[tuple[float, Any]]: + """Yield ``(timestamp, decoded message)``, decoding only kept rows.""" + stream: Any = store(db_path).stream(stream_name) + for index, observation in enumerate(stream): + if stride > 1 and index % stride: + continue + yield (float(observation.ts), observation.data) + + +def payload_pose(payload: Any) -> Pose: + """Normalize a recorded odometry payload to a ``Pose``. + + Handles both shapes found in recordings: ``Odometry`` (``.pose``) and + ``PoseStamped`` (flat position + ``.orientation``).""" + if hasattr(payload, "pose"): # Odometry + return payload.pose # type: ignore[no-any-return] + return Pose( # PoseStamped + payload.x, + payload.y, + payload.z, + payload.orientation.x, + payload.orientation.y, + payload.orientation.z, + payload.orientation.w, + ) + + +def odometry_lookup(db_path: Path, stream_name: str) -> PoseLookup: + """Nearest-position (``[x, y, z]``) lookup over a recording's odometry stream.""" + times: list[float] = [] + positions: list[list[float]] = [] + for timestamp, payload in iterate_stream(db_path, stream_name): + pose = payload_pose(payload) + times.append(timestamp) + positions.append([pose.position.x, pose.position.y, pose.position.z]) + return trajectory_lookup( + np.asarray(times, dtype=np.float64), + np.asarray(positions, dtype=np.float64), + ODOM_MATCH_TOLERANCE_S, + ) diff --git a/dimos/navigation/jnav/utils/test_apriltag_agreement.py b/dimos/navigation/jnav/utils/test_apriltag_agreement.py new file mode 100644 index 0000000000..fc177b3d3d --- /dev/null +++ b/dimos/navigation/jnav/utils/test_apriltag_agreement.py @@ -0,0 +1,74 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the April-tag agreement metric.""" + +from __future__ import annotations + +import numpy as np + +from dimos.navigation.jnav.utils.apriltag_agreement import ( + agreement_improvement, + agreement_report, + tag_spread, +) + + +def test_tag_spread_zero_for_identical() -> None: + assert tag_spread(np.array([[1.0, 2.0, 3.0]] * 4)) == 0.0 + + +def test_tag_spread_grows_with_scatter() -> None: + tight = tag_spread(np.array([[0.0, 0, 0], [0.1, 0, 0]])) + loose = tag_spread(np.array([[0.0, 0, 0], [2.0, 0, 0]])) + assert loose > tight + + +def test_single_observation_excluded() -> None: + report = agreement_report({7: np.array([[0.0, 0, 0]])}) + assert report.tag_count == 0 # one sighting carries no agreement signal + assert report.total_observations == 1 + assert report.mean_spread == 0.0 + + +def test_report_mean_spread() -> None: + report = agreement_report( + { + 1: np.array([[0.0, 0, 0], [0.0, 0, 0]]), # spread 0 + 2: np.array([[0.0, 0, 0], [2.0, 0, 0]]), # spread 1.0 + } + ) + assert report.tag_count == 2 + assert report.total_observations == 4 + assert abs(report.mean_spread - 0.5) < 1e-9 + + +def test_improvement_positive_when_corrected_tighter() -> None: + # Drifted: same tag scattered 4 m apart across laps. Corrected: pulled together. + raw = agreement_report({1: np.array([[0.0, 0, 0], [4.0, 0, 0]])}) + corrected = agreement_report({1: np.array([[0.0, 0, 0], [0.4, 0, 0]])}) + improvement = agreement_improvement(raw, corrected) + assert 0.85 < improvement <= 1.0 + + +def test_improvement_negative_when_corrected_worse() -> None: + raw = agreement_report({1: np.array([[0.0, 0, 0], [0.4, 0, 0]])}) + corrected = agreement_report({1: np.array([[0.0, 0, 0], [4.0, 0, 0]])}) + assert agreement_improvement(raw, corrected) < 0.0 + + +def test_improvement_zero_when_no_raw_spread() -> None: + raw = agreement_report({1: np.array([[0.0, 0, 0], [0.0, 0, 0]])}) + corrected = agreement_report({1: np.array([[0.0, 0, 0], [1.0, 0, 0]])}) + assert agreement_improvement(raw, corrected) == 0.0 diff --git a/dimos/navigation/jnav/utils/test_kitti.py b/dimos/navigation/jnav/utils/test_kitti.py new file mode 100644 index 0000000000..756e403cb6 --- /dev/null +++ b/dimos/navigation/jnav/utils/test_kitti.py @@ -0,0 +1,60 @@ +# 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. + +"""Devkit-verified self-tests for the official KITTI odometry error metric.""" + +from __future__ import annotations + +import numpy as np +from scipy.spatial.transform import Rotation + +from dimos.navigation.jnav.utils.kitti import kitti_odometry_error + +ZERO_ERROR_TOLERANCE = 1e-6 + + +def _straight_trajectory(scale: float = 1.0) -> list[np.ndarray]: + """A 1000 m straight run along x, one frame per metre, scaled along-track.""" + poses = [] + for index in range(1001): + pose = np.eye(4) + pose[0, 3] = float(index) * scale + poses.append(pose) + return poses + + +def test_identical_trajectory_has_zero_error() -> None: + ground_truth = _straight_trajectory() + result = kitti_odometry_error([pose.copy() for pose in ground_truth], ground_truth) + assert result["translational_percent"] < ZERO_ERROR_TOLERANCE + assert result["rotational_deg_per_m"] < ZERO_ERROR_TOLERANCE + + +def test_one_percent_scale_drift_reads_one_percent() -> None: + ground_truth = _straight_trajectory() + scaled = _straight_trajectory(scale=1.01) + result = kitti_odometry_error(scaled, ground_truth) + assert 0.8 < result["translational_percent"] < 1.2 + + +def test_constant_yaw_rate_has_nonzero_rotational_error() -> None: + ground_truth = _straight_trajectory() + yawed = [] + for index in range(1001): + pose = np.eye(4) + pose[:3, :3] = Rotation.from_euler("z", index * 0.01, degrees=True).as_matrix() + pose[0, 3] = float(index) + yawed.append(pose) + result = kitti_odometry_error(yawed, ground_truth) + assert result["rotational_deg_per_m"] > 0.0 diff --git a/dimos/navigation/jnav/utils/trajectory_metrics.py b/dimos/navigation/jnav/utils/trajectory_metrics.py new file mode 100644 index 0000000000..8701eb4785 --- /dev/null +++ b/dimos/navigation/jnav/utils/trajectory_metrics.py @@ -0,0 +1,268 @@ +# 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. + +"""Pure pose/trajectory math for loop-closure evaluation. + +Drift injection + correction and ground-truth(-free) trajectory metrics, all +operating on plain arrays and lookup callables (no db reads) so they can be +shared across the eval drivers. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable +from typing import Any + +import numpy as np +from scipy.spatial.transform import Rotation + +from dimos.navigation.jnav.utils.voxel_map import VoxelMap + +# (ts, x, y, z, qx, qy, qz, qw) per keyframe +GraphPose = tuple[float, float, float, float, float, float, float, float] +# time -> nearest pose: 3-vec [x,y,z] (PoseLookup) or 7-vec +quat (PoseLookup7) +PoseLookup = Callable[[float], "np.ndarray | None"] +PoseLookup7 = Callable[[float], "np.ndarray | None"] + + +def trajectory_lookup(times: np.ndarray, positions: np.ndarray, tolerance: float) -> PoseLookup: + """A ``time -> [x, y, z]`` nearest-sample lookup over a (ts, position) trajectory.""" + if len(times) == 0: + return lambda _timestamp: None + + def lookup(timestamp: float) -> np.ndarray | None: + index = int(np.argmin(np.abs(times - timestamp))) + if abs(float(times[index]) - timestamp) > tolerance: + return None + return np.asarray(positions[index], dtype=np.float64) + + return lookup + + +def graph_lookup(graph: list[tuple[float, float, float, float]]) -> PoseLookup: + """Nearest-keyframe lookup over an optimized pose graph (no tolerance). + + Keyframes only spawn on motion, so a parked robot legitimately maps to an + old keyframe — that keyframe IS its position (within the keyframe delta). + Unbounded nearest-in-time is sound when the graph is never truncated.""" + times = np.asarray([node[0] for node in graph], dtype=np.float64) + positions = np.asarray([[node[1], node[2], node[3]] for node in graph], dtype=np.float64) + return trajectory_lookup(times, positions, float("inf")) + + +# Voxel-agreement sampling: cap per-scan points so the map fits in memory +# without holding the whole recording at once. +VOXEL_SIZE_M = 0.2 +VOXEL_MAX_POINTS_PER_SCAN = 4000 + + +def drift_offset( + timestamp: float, t0: float, drift_per_sec: list[float] | np.ndarray +) -> np.ndarray: + """World translation injected at ``timestamp`` (grows linearly from ``t0``).""" + return np.asarray(drift_per_sec, dtype=np.float64) * (timestamp - t0) + + +def has_drift(drift_per_sec: list[float] | np.ndarray) -> bool: + return bool(np.any(np.asarray(drift_per_sec, dtype=np.float64))) + + +def drifted_lookup( + base_lookup: Callable[[float], np.ndarray | None], + drift_per_sec: list[float], + drift_t0: float, +) -> Callable[[float], np.ndarray | None]: + """Wrap a pose lookup so its xyz gets the same drift the module was fed. + + Pass-through when drift is zero. Works for both the 3-vec (xyz) and 7-vec + (xyz+quat) lookups — only the first three components are shifted.""" + if not has_drift(drift_per_sec): + return base_lookup + drift = np.asarray(drift_per_sec, dtype=np.float64) + + def lookup(timestamp: float) -> np.ndarray | None: + pose = base_lookup(timestamp) + if pose is None: + return None + shifted = np.array(pose, dtype=np.float64) + shifted[:3] += drift * (timestamp - drift_t0) + return shifted + + return lookup + + +def pose7_lookup(times: np.ndarray, poses: np.ndarray, tolerance: float) -> PoseLookup7: + """time -> nearest [x,y,z,qx,qy,qz,qw], or None past the tolerance.""" + + def lookup(timestamp: float) -> np.ndarray | None: + if len(times) == 0: + return None + index = int(np.argmin(np.abs(times - timestamp))) + if abs(float(times[index]) - timestamp) > tolerance: + return None + return np.asarray(poses[index], dtype=np.float64) + + return lookup + + +def drift_delta_lookup( + graph: list[GraphPose], raw_lookup: PoseLookup7 +) -> Callable[[float], tuple[np.ndarray, np.ndarray] | None]: + """time -> nearest keyframe's drift correction (R_delta, t_delta). + + The delta is computed at the KEYFRAME's own timestamp — + T_corrected(kf) * T_raw(kf_ts)^-1 — so raw and corrected poses describe the + same instant. (Comparing a keyframe pose against the raw pose at a nearby + scan's timestamp snaps scans onto keyframes and fakes a tighter map: an + identity correction must yield an identity delta.)""" + times: list[float] = [] + rotations: list[np.ndarray] = [] + translations: list[np.ndarray] = [] + for node in graph: + raw_pose = raw_lookup(node[0]) + if raw_pose is None: + continue + rotation_raw = Rotation.from_quat(raw_pose[3:7]).as_matrix() + rotation_corrected = Rotation.from_quat(node[4:8]).as_matrix() + rotation_delta = rotation_corrected @ rotation_raw.T + translation_delta = np.asarray(node[1:4]) - rotation_delta @ raw_pose[:3] + times.append(node[0]) + rotations.append(rotation_delta) + translations.append(translation_delta) + times_array = np.asarray(times, dtype=np.float64) + + def lookup(timestamp: float) -> tuple[np.ndarray, np.ndarray] | None: + if len(times_array) == 0: + return None + index = int(np.argmin(np.abs(times_array - timestamp))) + return rotations[index], translations[index] + + return lookup + + +def rigid_align_rmse(source: np.ndarray, target: np.ndarray) -> float: + """Absolute trajectory error: RMSE of ``source`` to ``target`` after a + best-fit rigid (rotation+translation, no scale) alignment (Kabsch/Umeyama). + Both are (N, 3). The alignment removes the gauge freedom — two trajectories + of the same shape in different world frames score 0.""" + if len(source) < 3: + return 0.0 + source_centroid = source.mean(axis=0) + target_centroid = target.mean(axis=0) + source_centered = source - source_centroid + target_centered = target - target_centroid + covariance = source_centered.T @ target_centered + u_matrix, _, vt_matrix = np.linalg.svd(covariance) + # Reflection fix so the result is a proper rotation (det = +1). + reflection = np.sign(np.linalg.det(vt_matrix.T @ u_matrix.T)) + correction = np.diag([1.0, 1.0, reflection]) + rotation = vt_matrix.T @ correction @ u_matrix.T + aligned = (rotation @ source_centered.T).T + target_centroid + return float(np.sqrt(np.mean(np.sum((aligned - target) ** 2, axis=1)))) + + +def trajectory_recovery_error( + graph: list[GraphPose], + gt_lookup: Callable[[float], np.ndarray | None], + drift_per_sec: list[float], + drift_t0: float, +) -> dict[str, float] | None: + """Drift-recovery ATE: how close the module's corrected keyframe trajectory + gets to the *un-drifted* ground-truth, vs the drifted input it was given. + + Only meaningful with injected drift (then GT = the un-drifted odom the drift + was added to). This is the right metric where tag/voxel agreement is weak — + e.g. KITTI's long single-loop trajectory. Returns None when drift is off or + too few keyframes resolve. ``trajectory_improvement`` = fraction of the drift + ATE removed (1.0 = perfect recovery, 0 = no help, negative = worse).""" + if not has_drift(drift_per_sec): + return None + drift = np.asarray(drift_per_sec, dtype=np.float64) + corrected_points: list[list[float]] = [] + gt_points: list[np.ndarray] = [] + drifted_points: list[np.ndarray] = [] + for node in graph: + timestamp = node[0] + gt_pose = gt_lookup(timestamp) + if gt_pose is None: + continue + gt_xyz = np.asarray(gt_pose, dtype=np.float64)[:3] + corrected_points.append([node[1], node[2], node[3]]) + gt_points.append(gt_xyz) + drifted_points.append(gt_xyz + drift * (timestamp - drift_t0)) + if len(gt_points) < 3: + return None + gt_array = np.asarray(gt_points) + drifted_ate = rigid_align_rmse(np.asarray(drifted_points), gt_array) + corrected_ate = rigid_align_rmse(np.asarray(corrected_points), gt_array) + improvement = (drifted_ate - corrected_ate) / drifted_ate if drifted_ate > 1e-9 else 0.0 + return { + "drifted_ate_m": drifted_ate, + "corrected_ate_m": corrected_ate, + "trajectory_improvement": improvement, + } + + +def lidar_voxel_agreement( + scans: Iterable[tuple[float, np.ndarray]], + raw_lookup: PoseLookup7, + graph: list[GraphPose], + *, + voxel_size: float = VOXEL_SIZE_M, + max_points_per_scan: int = VOXEL_MAX_POINTS_PER_SCAN, + drift_per_sec: list[float] | None = None, + drift_t0: float = 0.0, +) -> dict[str, Any]: + """Occupied-voxel counts of the lidar map, raw vs corrected. + + ``scans`` yields ``(timestamp, points)`` where ``points`` is an (N, 3+) array + registered in the raw odom world frame. Each scan is re-anchored by its + nearest keyframe's drift correction (see `drift_delta_lookup`), so a good + correction collapses doubled walls and the corrected map occupies fewer + voxels. ``improvement`` is the fractional voxel drop (positive = tighter). + When ``drift_per_sec`` is set the raw scans are first shifted into the same + drifted world the module solved in.""" + drift = np.asarray(drift_per_sec or [0.0, 0.0, 0.0], dtype=np.float64) + apply_drift = has_drift(drift) + delta_lookup = drift_delta_lookup(graph, raw_lookup) + raw_clouds: list[np.ndarray] = [] + corrected_clouds: list[np.ndarray] = [] + used = 0 + for timestamp, scan_points in scans: + delta = delta_lookup(timestamp) + if delta is None or raw_lookup(timestamp) is None: + continue + rotation_delta, translation_delta = delta + points = np.asarray(scan_points, dtype=np.float64)[:, :3] + if len(points) > max_points_per_scan: + points = points[:: -(-len(points) // max_points_per_scan)] + if apply_drift: + points = points + drift * (timestamp - drift_t0) + raw_clouds.append(points) + corrected_clouds.append(points @ rotation_delta.T + translation_delta) + used += 1 + + if not raw_clouds: + return {"status": "skipped: no scans matched both pose sources"} + raw_voxels = VoxelMap.from_points(np.vstack(raw_clouds), voxel_size).count + corrected_voxels = VoxelMap.from_points(np.vstack(corrected_clouds), voxel_size).count + improvement = (raw_voxels - corrected_voxels) / raw_voxels if raw_voxels else 0.0 + return { + "status": "ok", + "raw_voxels": raw_voxels, + "corrected_voxels": corrected_voxels, + "improvement": improvement, + "voxel_size_m": voxel_size, + "scans_used": used, + } diff --git a/dimos/navigation/jnav/utils/voxel_map.py b/dimos/navigation/jnav/utils/voxel_map.py new file mode 100644 index 0000000000..c3055d8c87 --- /dev/null +++ b/dimos/navigation/jnav/utils/voxel_map.py @@ -0,0 +1,66 @@ +# 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. + +"""Force a point cloud onto a discrete voxel grid for occupancy comparisons. + +A ``VoxelMap`` is just the set of occupied integer voxel keys at a fixed voxel +size. ``diff`` (symmetric-difference count, lower == better aligned) and the raw +occupied ``count`` are the building blocks for map-quality scores: drift smears a +map into more, mis-aligned voxels; a good loop closure collapses it back. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + + +@dataclass(frozen=True) +class VoxelMap: + voxel_size: float + keys: frozenset[tuple[int, int, int]] + + @classmethod + def from_points(cls, points: np.ndarray, voxel_size: float) -> VoxelMap: + if points.shape[0] == 0: + return cls(voxel_size=voxel_size, keys=frozenset()) + indices = np.floor(np.asarray(points)[:, :3] / voxel_size).astype(np.int64) + unique = np.unique(indices, axis=0) + return cls(voxel_size=voxel_size, keys=frozenset(map(tuple, unique))) + + @classmethod + def from_pointcloud(cls, cloud: PointCloud2, voxel_size: float) -> VoxelMap: + return cls.from_points(cloud.points_f32(), voxel_size) + + @property + def count(self) -> int: + """Number of occupied voxels.""" + return len(self.keys) + + def __len__(self) -> int: + return len(self.keys) + + def diff(self, other: VoxelMap) -> int: + """Symmetric-difference count: voxels occupied in exactly one of the two.""" + return len(self.keys ^ other.keys) + + def intersection(self, other: VoxelMap) -> int: + return len(self.keys & other.keys) + + def iou(self, other: VoxelMap) -> float: + union = len(self.keys | other.keys) + return len(self.keys & other.keys) / union if union else 0.0 From ac693ad6bfb190a421b26e9406d922c2fd1e0e29 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 14:52:31 +0800 Subject: [PATCH 07/69] add ivan_pgo, ivan_pgo_transformer and unrefined_pgo loop-closure modules --- .../loop_closure/gsc_pgo/.gitignore | 2 + .../gsc_pgo/hyperparam_analysis.md | 134 ++++ .../loop_closure/ivan_pgo/module.py | 637 ++++++++++++++++++ .../ivan_pgo_transformer/module.py | 263 ++++++++ .../unrefined_pgo/cpp/CMakeLists.txt | 49 ++ .../unrefined_pgo/cpp/commons.cpp | 8 + .../loop_closure/unrefined_pgo/cpp/commons.h | 29 + .../unrefined_pgo/cpp/dimos_native_module.hpp | 97 +++ .../loop_closure/unrefined_pgo/cpp/flake.lock | 144 ++++ .../loop_closure/unrefined_pgo/cpp/flake.nix | 70 ++ .../loop_closure/unrefined_pgo/cpp/main.cpp | 299 ++++++++ .../unrefined_pgo/cpp/point_cloud_utils.hpp | 170 +++++ .../unrefined_pgo/cpp/simple_pgo.cpp | 211 ++++++ .../unrefined_pgo/cpp/simple_pgo.h | 78 +++ .../loop_closure/unrefined_pgo/module.py | 286 ++++++++ 15 files changed, 2477 insertions(+) create mode 100644 dimos/navigation/jnav/components/loop_closure/gsc_pgo/.gitignore create mode 100644 dimos/navigation/jnav/components/loop_closure/gsc_pgo/hyperparam_analysis.md create mode 100644 dimos/navigation/jnav/components/loop_closure/ivan_pgo/module.py create mode 100644 dimos/navigation/jnav/components/loop_closure/ivan_pgo_transformer/module.py create mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/CMakeLists.txt create mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.cpp create mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.h create mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/dimos_native_module.hpp create mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.lock create mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.nix create mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/main.cpp create mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/point_cloud_utils.hpp create mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.cpp create mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.h create mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/module.py diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/.gitignore b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/.gitignore new file mode 100644 index 0000000000..750baebf41 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/.gitignore @@ -0,0 +1,2 @@ +result +result-* diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/hyperparam_analysis.md b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/hyperparam_analysis.md new file mode 100644 index 0000000000..8e7f608201 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/hyperparam_analysis.md @@ -0,0 +1,134 @@ +# gsc_pgo hyperparameter analysis + online tuning — sketch + +## The problem (from the eval campaign) + +cmu wins on big real-drift loops but **over-fires** elsewhere: +- `grassy_field` (05-32): **-0.369** tag with 110 closures — should be a no-op (+0.00x). + Open grass = feature-poor; ICP matches are garbage but get accepted. +- `gir_park1_2`: +0.578 with 23 closures, but unrefined got **+0.749 with 2 closures**. + cmu's extra closures are net noise — fewer, better closures would win. + +Goal: (a) **detect feature-poor / high-disagreement scans online and back off** +(ideally make PGO a no-op in grassy_field), and (b) accept **fewer, more +consistent** closures. Both reduce to: *trust loops less when the environment +can't support them, and only commit mutually-consistent ones.* + +## Signals cmu already computes (cheap to expose) + +Every candidate already produces these inside `searchForLoopPairs` / +`scan_context.cpp` — instrument the binary to log them per candidate (accepted +and rejected), keyed by ts: + +| signal | where | meaning | +|---|---|---| +| scan-context min cosine distance | `best_distance()` | place-match quality; high = no good match (new place OR feature-poor) | +| ICP fitness (`getFitnessScore`) | `searchForLoopPairs` | mean-sq inlier dist; high = bad geometric alignment | +| ICP converged + inlier ratio | PCL ICP | low inlier ratio = lots of disagreement (crowd/dynamic/degenerate) | +| candidate global distance | `loop_candidate_max_distance_m` gate | drift magnitude at the candidate | +| ring-key occupancy/entropy | scan-context descriptor | **feature richness** — near-uniform/empty descriptor = feature-poor | + +Two signals are *not* computed yet but are the highest-value adds: +- **registration Hessian eigenvalues** (J^T J of the ICP cost) — the canonical + degeneracy detector (below). +- **per-point residual distribution** of the scan-to-submap match — heavy tails = + dynamic objects (crowd). + +## Literature: detecting feature-poor / degenerate environments + +1. **Degeneracy factor (Zhang, Kaess, Singh — ICRA 2016, "On Degeneracy of + Optimization-based State Estimation").** THE standard. Eigen-decompose the + scan-matching information matrix `H = J^T J`; its smallest eigenvalue (or + `λ_min / λ_max` condition number) measures how well-constrained the solve is. + A featureless corridor/field has a small eigenvalue along the unconstrained + axis. → online "is this scan match trustworthy" scalar. Their **solution + remapping** projects the update out of degenerate directions; LOAM/LeGO-LOAM + use it. +2. **X-ICP (Tuna et al., T-RO 2023) — "Localizability-Aware ICP."** Per-axis + localizability from the optimization constraints; selectively adds soft + constraints only in observable directions. More principled than a scalar gate. +3. **Eigenvalue geometric features (Demantké 2011 / Weinmann 2015):** local + neighborhood covariance → linearity/planarity/scattering. Open grass scores + high "scattering" (3D-spread, no structure) — a direct feature-poverty flag. +4. **Scan-context descriptor entropy/occupancy** (cheap, already have the + descriptor): a low-entropy / sparsely-occupied ring-key = feature-poor scene + where cosine matching is unreliable. Raise the match bar or refuse loops. + +## Literature: scan disagreement / dynamic crowds + +5. **Robust kernels (Huber / Cauchy / Geman-McClure)** on the ICP residual — + down-weight the moving-people points. cmu currently uses raw fitness; a robust + cost + reporting the *inlier ratio* exposes "crowd present." +6. **Removert (Kim & Kim, IROS 2020):** range-image visibility differencing + removes dynamic points before matching. Strong for crowds. +7. **Stationarity / scan-to-scan consistency:** points that move between + consecutive frames are dynamic; a high dynamic-fraction → distrust this scan + for loop closure. (Cheap online proxy: residual after rigid alignment of + consecutive scans.) +8. **Learned moving-object segmentation (4DMOS, LMNet, MotionSeg3D)** — heavier; + masks people directly. Overkill for now but the ceiling. + +## Literature: loop-closure outlier rejection (the over-firing fix) + +The single highest-leverage change — cmu accepts each loop independently: + +9. **PCM — Pairwise Consistency Maximization (Mangelson et al., ICRA 2018).** + Build a consistency graph over candidate loops; keep only the **maximal + mutually-consistent clique**, reject the rest. Directly attacks "23 closures + where 2 good ones win." Cheap, backend-agnostic, no tuning. +10. **GNC — Graduated Non-Convexity (Yang et al., RA-L 2020) / Cauchy-IRLS.** + Backend robustly down-weights outlier loops during optimization. iSAM2-friendly. +11. **Switchable Constraints (Sünderhauf & Protzel 2012) / Dynamic Covariance + Scaling (Agarwal 2013):** the optimizer learns a switch/weight per loop and + can turn bad ones off. Lighter than PCM, online-friendly. + +## Online tuning design (signal → knob) + +Frame cmu's static thresholds as **initial values adapted by a measured +"environment trust" `τ ∈ [0,1]`** computed per keyframe from the signals above: + +``` +τ = f( degeneracy_factor, ringkey_entropy, icp_inlier_ratio, dynamic_fraction ) + # low when feature-poor / crowded / degenerate +``` + +Then drive the gates from `τ` (running statistics, not magic numbers): + +| knob (currently static) | online rule | +|---|---| +| `loop_score_thresh` (ICP gate) | tighten as `τ` drops; track running median/MAD of accepted-loop fitness and gate at `median − k·MAD` | +| `scan_context_match_threshold` | raise when ring-key entropy is low (cosine less discriminative in feature-poor scenes) | +| **accept/suppress** the loop | hard-gate: if `degeneracy_factor < ε` (Zhang) → **suppress** → PGO no-op in directions it can't observe (the grassy_field fix) | +| loop **batch** acceptance | replace independent accept with **PCM** over the recent candidate set → only commit the mutually-consistent subset (the gir_park1_2 fix) | +| loop noise model | already `Σ = fitness·I`; inflate further by `1/τ` so low-trust loops barely move the graph | + +Online-tuning techniques that fit (cheapest first): +- **Adaptive thresholding on running stats** (median/MAD, EWMA) — robust, no + training, the right default. +- **Switchable constraints / DCS** — lets the *optimizer* do the tuning; minimal code. +- **PCM** — not "tuning" but the structural fix for over-firing; do this first. +- (Avoid online Bayesian-opt/bandits for the fast gates — too slow/unstable to + converge within one run; reserve for offline meta-params.) + +## Proposed analysis (to validate before coding the online controller) + +1. **Instrument the binary** to emit per-candidate diagnostics (the table above) + to a sidecar jsonl. One pass per recording → the dataset that grounds `f(·)`. +2. **Static sweep** (via the new eval cell-caching + `--drift-per-sec`): vary + `loop_score_thresh`, `scan_context_match_threshold`, + `loop_candidate_max_distance_m`, `min_loop_detect_duration` per recording; + plot tag/voxel improvement & closure count. Expect grassy_field to be + monotone-better as the gate tightens (→ confirms "no-op is optimal there"). +3. **Correlate** accepted-vs-rejected closures with the signals: does + degeneracy_factor / ring-key entropy separate grassy_field's bad closures + from gir_park1_2's good ones? If yes, `τ` is learnable as a simple threshold. +4. **Prototype** the cheapest controller (PCM + degeneracy-gate + adaptive + `loop_score_thresh`) and re-run the full eval; target: grassy_field → ~0, + gir_park1_2 → ≥ unrefined's +0.749, no regression on the pointlio wins. + +## First concrete step + +Add the per-candidate diagnostic logging to `simple_pgo.cpp` (scan-context +distance, ICP fitness, inlier ratio, candidate distance) + compute the Zhang +degeneracy factor from the ICP Hessian. That single instrumentation pass +produces the data to decide which signals actually separate good/bad closures +here — everything else (the `f(·)` and the gates) follows from it. diff --git a/dimos/navigation/jnav/components/loop_closure/ivan_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/ivan_pgo/module.py new file mode 100644 index 0000000000..7c86f6dfb5 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/ivan_pgo/module.py @@ -0,0 +1,637 @@ +# 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. + +"""Pure-Python PGO ported from Ivan's `ivan/feat/go2loopclosure` branch +(`dimos/mapping/pgo.py`), adapted to the jnav LoopClosure spec. + +GTSAM iSAM2 pose graph + Open3D point-to-plane ICP loop verification + +KD-tree loop candidate search. Differences from the original: + * input renamed `registered_scan` -> `lidar` (LoopClosure spec) + * publishes `pose_graph: Out[Graph3D]` (optimized keyframes, odometry + + loop edges) and `loop_closure_event: Out[GraphDelta3D]` so the eval + harness can capture the corrected trajectory and count closures + * offline experiment helpers (transformers, two-pass voxel pipelines) + were left behind — this is just the runtime module +""" + +from __future__ import annotations + +from dataclasses import dataclass +import threading +import time +from typing import Any + +import gtsam # type: ignore[import-untyped] +import numpy as np +import open3d as o3d # type: ignore[import-untyped] +import open3d.core as o3c # type: ignore[import-untyped] +from reactivex.disposable import Disposable +from scipy.spatial import KDTree +from scipy.spatial.transform import Rotation + +from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +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 +from dimos.navigation.jnav.components.loop_closure.spec import LoopClosure +from dimos.navigation.jnav.msgs.Graph3D import Graph3D +from dimos.navigation.jnav.msgs.GraphDelta3D import GraphDelta3D +from dimos.utils.logging_config import setup_logger + +FRAME_MAP = "map" +FRAME_ODOM = "odom" +FRAME_BODY = "base_link" + +logger = setup_logger() + + +class PGOConfig(ModuleConfig): + world_frame: str = FRAME_MAP + + # Keyframe detection + key_pose_delta_trans: float = 0.5 + key_pose_delta_deg: float = 10.0 + + # Loop closure + loop_search_radius: float = 2.0 + loop_time_thresh: float = 20.0 + loop_score_thresh: float = 0.3 + loop_submap_half_range: int = 10 + min_icp_inliers: int = 10 + min_keyframes_for_loop_search: int = 10 + loop_closure_extra_iterations: int = 4 + submap_resolution: float = 0.2 + min_loop_detect_duration: float = 5.0 + + # Input mode + unregister_input: bool = True # Transform world-frame scans to body-frame using odom + + # Global map + publish_global_map: bool = True + global_map_publish_rate: float = 0.5 + global_map_voxel_size: float = 0.15 + + # ICP + max_icp_iterations: int = 50 + max_icp_correspondence_dist: float = 1.0 + + +@dataclass +class _KeyPose: + r_local: np.ndarray # 3x3 rotation in local/odom frame + t_local: np.ndarray # 3-vec translation in local/odom frame + r_global: np.ndarray # 3x3 corrected rotation + t_global: np.ndarray # 3-vec corrected translation + timestamp: float + body_cloud: np.ndarray # Nx3 points in body frame + + +def _icp( + source: np.ndarray, + target: np.ndarray, + max_iter: int = 50, + max_dist: float = 1.0, + tol: float = 1e-6, + min_inliers: int = 10, + init: np.ndarray | None = None, +) -> tuple[np.ndarray, float]: + """Point-to-plane ICP using Open3D's tensor pipeline. + + Returns ``(T, fitness)`` where ``fitness`` is mean squared inlier + distance (m²).""" + if len(source) < min_inliers or len(target) < min_inliers: + return np.eye(4), float("inf") + + cpu = o3c.Device("CPU:0") + src_pcd = o3d.t.geometry.PointCloud(o3c.Tensor(source.astype(np.float32), device=cpu)) + tgt_pcd = o3d.t.geometry.PointCloud(o3c.Tensor(target.astype(np.float32), device=cpu)) + + # Normals on the target enable point-to-plane ICP, which converges + # tighter than point-to-point on indoor scenes (walls give unambiguous + # normals that resolve the slide-along-wall ambiguity). + tgt_pcd.estimate_normals(max_nn=30, radius=0.3) + + init_T = ( + o3c.Tensor(init.astype(np.float64), dtype=o3c.float64, device=cpu) + if init is not None + else o3c.Tensor.eye(4, dtype=o3c.float64, device=cpu) + ) + + # Silence Open3D's "0 correspondence" warning — we deliberately use a + # tight max_correspondence_distance and reject loops with poor fitness. + with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Error): + result = o3d.t.pipelines.registration.icp( + source=src_pcd, + target=tgt_pcd, + max_correspondence_distance=max_dist, + init_source_to_target=init_T, + estimation_method=o3d.t.pipelines.registration.TransformationEstimationPointToPlane(), + criteria=o3d.t.pipelines.registration.ICPConvergenceCriteria( + relative_fitness=tol, + relative_rmse=tol, + max_iteration=max_iter, + ), + ) + + fitness_inlier_frac = float(result.fitness) + if fitness_inlier_frac == 0.0: + return np.eye(4), float("inf") + + rmse = float(result.inlier_rmse) + T = result.transformation.numpy() + return T, rmse * rmse + + +def _voxel_downsample(pts: np.ndarray, voxel_size: float) -> np.ndarray: + if len(pts) == 0 or voxel_size <= 0: + return pts + keys = np.floor(pts / voxel_size).astype(np.int32) + _, idx = np.unique(keys, axis=0, return_index=True) + return pts[idx] + + +class _SimplePGO: + def __init__(self, config: PGOConfig) -> None: + self._cfg = config + self._key_poses: list[_KeyPose] = [] + self._history_pairs: list[tuple[int, int]] = [] + self._cache_pairs: list[dict[str, Any]] = [] + self._r_offset = np.eye(3) + self._t_offset = np.zeros(3) + + params = gtsam.ISAM2Params() + params.setRelinearizeThreshold(0.01) + params.relinearizeSkip = 1 + self._isam2 = gtsam.ISAM2(params) + self._graph = gtsam.NonlinearFactorGraph() + self._values = gtsam.Values() + + def is_key_pose(self, r: np.ndarray, t: np.ndarray) -> bool: + if not self._key_poses: + return True + last = self._key_poses[-1] + delta_trans = np.linalg.norm(t - last.t_local) + # Angular distance via quaternion dot product + q_cur = Rotation.from_matrix(r).as_quat() # [x,y,z,w] + q_last = Rotation.from_matrix(last.r_local).as_quat() + dot = abs(np.dot(q_cur, q_last)) + delta_deg = np.degrees(2.0 * np.arccos(min(dot, 1.0))) + return bool( + delta_trans > self._cfg.key_pose_delta_trans or delta_deg > self._cfg.key_pose_delta_deg + ) + + def add_key_pose( + self, r_local: np.ndarray, t_local: np.ndarray, timestamp: float, body_cloud: np.ndarray + ) -> bool: + if not self.is_key_pose(r_local, t_local): + return False + + idx = len(self._key_poses) + init_r = self._r_offset @ r_local + init_t = self._r_offset @ t_local + self._t_offset + + pose = gtsam.Pose3(gtsam.Rot3(init_r), gtsam.Point3(init_t)) + self._values.insert(idx, pose) + + if idx == 0: + noise = gtsam.noiseModel.Diagonal.Variances(np.full(6, 1e-12)) + self._graph.add(gtsam.PriorFactorPose3(idx, pose, noise)) + else: + last = self._key_poses[-1] + r_between = last.r_local.T @ r_local + t_between = last.r_local.T @ (t_local - last.t_local) + noise = gtsam.noiseModel.Diagonal.Variances( + np.array([1e-6, 1e-6, 1e-6, 1e-4, 1e-4, 1e-6]) + ) + self._graph.add( + gtsam.BetweenFactorPose3( + idx - 1, idx, gtsam.Pose3(gtsam.Rot3(r_between), gtsam.Point3(t_between)), noise + ) + ) + + kp = _KeyPose( + r_local=r_local.copy(), + t_local=t_local.copy(), + r_global=init_r.copy(), + t_global=init_t.copy(), + timestamp=timestamp, + body_cloud=_voxel_downsample(body_cloud, self._cfg.submap_resolution), + ) + self._key_poses.append(kp) + return True + + def _get_submap(self, idx: int, half_range: int) -> np.ndarray: + lo = max(0, idx - half_range) + hi = min(len(self._key_poses) - 1, idx + half_range) + parts = [] + for i in range(lo, hi + 1): + kp = self._key_poses[i] + world = (kp.r_global @ kp.body_cloud.T).T + kp.t_global + parts.append(world) + if not parts: + return np.empty((0, 3)) + cloud = np.vstack(parts) + return _voxel_downsample(cloud, self._cfg.submap_resolution) + + def search_for_loops(self) -> None: + if len(self._key_poses) < self._cfg.min_keyframes_for_loop_search: + return + + # Rate limit + if self._history_pairs: + cur_time = self._key_poses[-1].timestamp + last_time = self._key_poses[self._history_pairs[-1][1]].timestamp + if cur_time - last_time < self._cfg.min_loop_detect_duration: + return + + cur_idx = len(self._key_poses) - 1 + cur_kp = self._key_poses[-1] + + # Build KD-tree of previous keyframe positions + positions = np.array([kp.t_global for kp in self._key_poses[:-1]]) + tree = KDTree(positions) + + idxs = tree.query_ball_point(cur_kp.t_global, self._cfg.loop_search_radius) + if not idxs: + return + + # Pick the spatially closest keyframe that's also old enough in time. + # query_ball_point doesn't sort, so we sort by distance ourselves. + candidates = [ + (float(np.linalg.norm(self._key_poses[i].t_global - cur_kp.t_global)), i) + for i in idxs + if abs(cur_kp.timestamp - self._key_poses[i].timestamp) > self._cfg.loop_time_thresh + ] + if not candidates: + return + candidates.sort() + loop_idx = candidates[0][1] + + # ICP verification + target = self._get_submap(loop_idx, self._cfg.loop_submap_half_range) + source = self._get_submap(cur_idx, 0) + + transform, fitness = _icp( + source, + target, + max_iter=self._cfg.max_icp_iterations, + max_dist=self._cfg.max_icp_correspondence_dist, + min_inliers=self._cfg.min_icp_inliers, + ) + if fitness > self._cfg.loop_score_thresh: + return + + # Compute relative pose + R_icp = transform[:3, :3] + t_icp = transform[:3, 3] + r_refined = R_icp @ cur_kp.r_global + t_refined = R_icp @ cur_kp.t_global + t_icp + r_offset = self._key_poses[loop_idx].r_global.T @ r_refined + t_offset = self._key_poses[loop_idx].r_global.T @ ( + t_refined - self._key_poses[loop_idx].t_global + ) + + self._cache_pairs.append( + { + "source": cur_idx, + "target": loop_idx, + "r_offset": r_offset, + "t_offset": t_offset, + "score": fitness, + } + ) + self._history_pairs.append((loop_idx, cur_idx)) + logger.info( + "Loop closure detected", + source=cur_idx, + target=loop_idx, + score=round(fitness, 4), + ) + + def smooth_and_update(self) -> None: + has_loop = bool(self._cache_pairs) + + for pair in self._cache_pairs: + # Pose3 noise model is [rx, ry, rz, x, y, z]. Use ICP fitness as + # the translation variance and a generous fixed rotation variance — + # loops shouldn't be trusted to fix rotation tightly. + trans_var = max(0.01, float(pair["score"])) # >= sigma_trans = 10 cm + rot_var = 0.05 # sigma_rot ~= 13 deg + noise = gtsam.noiseModel.Diagonal.Variances( + np.array([rot_var, rot_var, rot_var, trans_var, trans_var, trans_var]) + ) + self._graph.add( + gtsam.BetweenFactorPose3( + pair["target"], + pair["source"], + gtsam.Pose3(gtsam.Rot3(pair["r_offset"]), gtsam.Point3(pair["t_offset"])), + noise, + ) + ) + self._cache_pairs.clear() + + self._isam2.update(self._graph, self._values) + self._isam2.update() + if has_loop: + for _ in range(self._cfg.loop_closure_extra_iterations): + self._isam2.update() + self._graph = gtsam.NonlinearFactorGraph() + self._values = gtsam.Values() + + estimates = self._isam2.calculateBestEstimate() + for i in range(len(self._key_poses)): + pose = estimates.atPose3(i) + self._key_poses[i].r_global = pose.rotation().matrix() + self._key_poses[i].t_global = pose.translation() + + last = self._key_poses[-1] + self._r_offset = last.r_global @ last.r_local.T + self._t_offset = last.t_global - self._r_offset @ last.t_local + + def get_corrected_pose( + self, r_local: np.ndarray, t_local: np.ndarray + ) -> tuple[np.ndarray, np.ndarray]: + return self._r_offset @ r_local, self._r_offset @ t_local + self._t_offset + + def build_global_map(self, voxel_size: float) -> np.ndarray: + if not self._key_poses: + return np.empty((0, 3), dtype=np.float32) + parts = [] + for kp in self._key_poses: + world = (kp.r_global @ kp.body_cloud.T).T + kp.t_global + parts.append(world) + cloud = np.vstack(parts).astype(np.float32) + return _voxel_downsample(cloud, voxel_size) + + @property + def num_key_poses(self) -> int: + return len(self._key_poses) + + +def process_scan( + pgo: _SimplePGO, + cloud: PointCloud2, + r_local: np.ndarray, + t_local: np.ndarray, + ts: float, + unregister_input: bool, +) -> tuple[Odometry, Transform, bool] | None: + """Add a keyframe, run loop closure, return (corrected odom, map->odom tf, + keyframe_added) — or None on empty cloud. + + Caller must hold ``pgo``'s lock during this call.""" + points, _ = cloud.as_numpy() + if len(points) == 0: + return None + + if unregister_input: + # registered_scan is world-frame; transform back to body-frame. + body_pts = (r_local.T @ (points[:, :3].T - t_local[:, None])).T + else: + body_pts = points[:, :3] + + added = pgo.add_key_pose(r_local, t_local, ts, body_pts) + if added: + pgo.search_for_loops() + pgo.smooth_and_update() + + r_corr, t_corr = pgo.get_corrected_pose(r_local, t_local) + return ( + build_corrected_odometry(r_corr, t_corr, ts), + build_map_odom_tf(pgo._r_offset.copy(), pgo._t_offset.copy(), ts), + added, + ) + + +def build_corrected_odometry( + r: np.ndarray, + t: np.ndarray, + ts: float, + world_frame: str = FRAME_MAP, +) -> Odometry: + q = Rotation.from_matrix(r).as_quat() # [x,y,z,w] + return Odometry( + ts=ts, + frame_id=world_frame, + child_frame_id=FRAME_BODY, + pose=Pose( + position=[float(t[0]), float(t[1]), float(t[2])], + orientation=[float(q[0]), float(q[1]), float(q[2]), float(q[3])], + ), + ) + + +def build_map_odom_tf( + r_offset: np.ndarray, + t_offset: np.ndarray, + ts: float, + world_frame: str = FRAME_MAP, + odom_frame: str = FRAME_ODOM, +) -> Transform: + q = Rotation.from_matrix(r_offset).as_quat() # [x,y,z,w] + return Transform( + frame_id=world_frame, + child_frame_id=odom_frame, + translation=Vector3(float(t_offset[0]), float(t_offset[1]), float(t_offset[2])), + rotation=Quaternion(float(q[0]), float(q[1]), float(q[2]), float(q[3])), + ts=ts, + ) + + +def _keyframe_node(index: int, key_pose: _KeyPose, world_frame: str) -> Graph3D.Node3D: + q = Rotation.from_matrix(key_pose.r_global).as_quat() # [x,y,z,w] + return Graph3D.Node3D( + pose=PoseStamped( + ts=key_pose.timestamp, + frame_id=world_frame, + position=[float(v) for v in key_pose.t_global], + orientation=[float(q[0]), float(q[1]), float(q[2]), float(q[3])], + ), + id=index, + ) + + +class PGO(Module, LoopClosure): + """Pose graph optimization with loop closure (pure Python). + + Detects keyframes, performs loop closure via ICP + KD-tree search, and + optimizes the pose graph with GTSAM iSAM2. Publishes corrected odometry, + the optimized pose graph, loop-closure events, and an accumulated + global map.""" + + config: PGOConfig + + lidar: In[PointCloud2] + odometry: In[Odometry] + corrected_odometry: Out[Odometry] + pose_graph: Out[Graph3D] + loop_closure_event: Out[GraphDelta3D] + global_map: Out[PointCloud2] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._running = False + self._thread: threading.Thread | None = None + self._pgo: _SimplePGO | None = None + self._latest_r = np.eye(3) + self._latest_t = np.zeros(3) + self._latest_time = 0.0 + self._has_odom = False + self._last_global_map_time = 0.0 + self._published_loops = 0 + self._lock = threading.Lock() + # Protects _pgo mutations (add_key_pose, search_for_loops, + # smooth_and_update, build_global_map) against concurrent access + # from _on_scan and _publish_loop threads. + self._pgo_lock = threading.Lock() + + @rpc + def start(self) -> None: + super().start() + self._pgo = _SimplePGO(self.config) + # Identity map -> odom so consumers querying map -> body get a result + # before any loop-closure correction exists. + self.tf.publish(build_map_odom_tf(np.eye(3), np.zeros(3), time.time())) + self.register_disposable(Disposable(self.odometry.subscribe(self._on_odom))) + self.register_disposable(Disposable(self.lidar.subscribe(self._on_scan))) + self._running = True + if self.config.publish_global_map: + self._thread = threading.Thread(target=self._publish_global_map_loop, daemon=True) + self._thread.start() + logger.info( + "PGO module started (gtsam iSAM2, pure python)", + publish_global_map=self.config.publish_global_map, + ) + + @rpc + def stop(self) -> None: + self._running = False + if self._thread: + self._thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + super().stop() + + def _on_odom(self, msg: Odometry) -> None: + q = [ + msg.pose.orientation.x, + msg.pose.orientation.y, + msg.pose.orientation.z, + msg.pose.orientation.w, + ] + r = Rotation.from_quat(q).as_matrix() + t = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z]) + with self._lock: + self._latest_r = r + self._latest_t = t + self._latest_time = msg.ts if msg.ts else time.time() + self._has_odom = True + + def _on_scan(self, cloud: PointCloud2) -> None: + with self._lock: + if not self._has_odom: + return + r_local = self._latest_r.copy() + t_local = self._latest_t.copy() + ts = self._latest_time + + pgo = self._pgo + assert pgo is not None + + with self._pgo_lock: + result = process_scan(pgo, cloud, r_local, t_local, ts, self.config.unregister_input) + if result is None: + return + corrected_odom, tf_msg, keyframe_added = result + if keyframe_added: + graph_msg, loop_events = self._snapshot_graph(pgo, ts) + else: + graph_msg, loop_events = None, [] + + self.corrected_odometry.publish(corrected_odom) + self.tf.publish(tf_msg) + if graph_msg is not None: + self.pose_graph.publish(graph_msg) + for event in loop_events: + self.loop_closure_event.publish(event) + + def _snapshot_graph(self, pgo: _SimplePGO, ts: float) -> tuple[Graph3D, list[GraphDelta3D]]: + """The optimized graph (odometry chain + loop edges) and one + GraphDelta3D per loop pair not yet published. + + Caller must hold ``_pgo_lock``.""" + world_frame = self.config.world_frame + nodes = [ + _keyframe_node(index, key_pose, world_frame) + for index, key_pose in enumerate(pgo._key_poses) + ] + edges = [ + Graph3D.Edge( + start_id=index - 1, end_id=index, timestamp=pgo._key_poses[index].timestamp + ) + for index in range(1, len(pgo._key_poses)) + ] + edges += [ + Graph3D.Edge(start_id=target, end_id=source, metadata_id=1) + for target, source in pgo._history_pairs + ] + graph_msg = Graph3D(ts=ts, nodes=nodes, edges=edges) + + loop_events: list[GraphDelta3D] = [] + identity = GraphDelta3D.Transform( + translation=Vector3(0.0, 0.0, 0.0), rotation=Quaternion(0.0, 0.0, 0.0, 1.0) + ) + for target, source in pgo._history_pairs[self._published_loops :]: + loop_events.append( + GraphDelta3D( + ts=ts, + nodes=[ + _keyframe_node(target, pgo._key_poses[target], world_frame), + _keyframe_node(source, pgo._key_poses[source], world_frame), + ], + transforms=[identity, identity], + ) + ) + self._published_loops = len(pgo._history_pairs) + return graph_msg, loop_events + + def _publish_global_map_loop(self) -> None: + pgo = self._pgo + assert pgo is not None + rate = self.config.global_map_publish_rate + interval = 1.0 / rate if rate > 0 else 2.0 + + while self._running: + t0 = time.monotonic() + + if t0 - self._last_global_map_time > interval and pgo.num_key_poses > 0: + with self._pgo_lock: + cloud_np = pgo.build_global_map(self.config.global_map_voxel_size) + if len(cloud_np) > 0: + now = time.time() + self.global_map.publish( + PointCloud2.from_numpy( + cloud_np, frame_id=self.config.world_frame, timestamp=now + ) + ) + self._last_global_map_time = t0 + + elapsed = time.monotonic() - t0 + sleep_time = max(DEFAULT_THREAD_JOIN_TIMEOUT, interval - elapsed) + time.sleep(sleep_time) diff --git a/dimos/navigation/jnav/components/loop_closure/ivan_pgo_transformer/module.py b/dimos/navigation/jnav/components/loop_closure/ivan_pgo_transformer/module.py new file mode 100644 index 0000000000..17e7c84511 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/ivan_pgo_transformer/module.py @@ -0,0 +1,263 @@ +# 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. + +"""Ivan's CURRENT PGO (dimos/mapping/loop_closure/pgo.py, main, June 2026) +forced into an online LoopClosure module. + +That code ships as an offline memory2 Transformer, but its `_PGOState` core is +already incremental — one `process(pose, ts, world_cloud)` call per frame. This +wrapper brute-forces the online shape: every arriving scan is paired with the +latest odometry pose and pushed straight into `_PGOState`, then the corrected +odometry / pose graph / loop events are read back out of its (private) state. +No changes to the mapping code itself; this module owns the coupling. + +`_PGOState.process` expects world-frame registered scans (it unregisters them +internally via the odom pose) — set `input_is_registered: false` for raw +body-frame scans and they'll be pre-registered here first.""" + +from __future__ import annotations + +import threading +import time +from typing import Any + +import numpy as np +from scipy.spatial.transform import Rotation + +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.mapping.loop_closure.pgo import ( + PGOConfig as MappingPGOConfig, + _PGOState, + _pose3_to_transform, +) +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +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 +from dimos.navigation.jnav.components.loop_closure.spec import LoopClosure +from dimos.navigation.jnav.msgs.Graph3D import Graph3D +from dimos.navigation.jnav.msgs.GraphDelta3D import GraphDelta3D +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +class PGOConfig(ModuleConfig): + world_frame: str = "map" + odom_frame: str = "odom" + body_frame: str = "base_link" + # World-frame registered input (fastlio); false = raw body-frame scans. + input_is_registered: bool = True + + # Mirrors mapping/loop_closure PGOConfig — forwarded verbatim. + key_pose_delta_trans: float = 0.5 + key_pose_delta_deg: float = 10.0 + loop_search_radius: float = 2.0 + loop_time_thresh: float = 20.0 + loop_score_thresh: float = 0.3 + loop_submap_half_range: int = 10 + min_icp_inliers: int = 10 + min_keyframes_for_loop_search: int = 10 + loop_closure_extra_iterations: int = 4 + submap_resolution: float = 0.2 + min_loop_detect_duration: float = 5.0 + max_icp_iterations: int = 50 + max_icp_correspondence_dist: float = 1.0 + odom_rot_var: float = 1e-6 + odom_trans_var_xy: float = 1e-4 + odom_trans_var_z: float = 1e-6 + loop_rot_var: float = 0.05 + + +def _mapping_config(config: PGOConfig) -> MappingPGOConfig: + fields = {name: getattr(config, name) for name in MappingPGOConfig.model_fields} + return MappingPGOConfig(**fields) + + +def _pose3_node(index: int, pose: Any, ts: float, world_frame: str) -> Graph3D.Node3D: + translation = np.asarray(pose.translation()) + quaternion = Rotation.from_matrix(pose.rotation().matrix()).as_quat() # [x,y,z,w] + return Graph3D.Node3D( + pose=PoseStamped( + ts=ts, + frame_id=world_frame, + position=[float(v) for v in translation], + orientation=[float(v) for v in quaternion], + ), + id=index, + ) + + +class PGO(Module, LoopClosure): + """Online wrapper over Ivan's current incremental PGO core (`_PGOState`).""" + + config: PGOConfig + + lidar: In[PointCloud2] + odometry: In[Odometry] + corrected_odometry: Out[Odometry] + pose_graph: Out[Graph3D] + loop_closure_event: Out[GraphDelta3D] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._state: _PGOState | None = None + self._latest_odom: Odometry | None = None + self._published_loops = 0 + self._odom_lock = threading.Lock() + self._state_lock = threading.Lock() + self._unsub_odom: Any = None + self._unsub_lidar: Any = None + + @rpc + def start(self) -> None: + super().start() + self._state = _PGOState(_mapping_config(self.config)) + # Identity map -> odom so consumers querying map -> body get a result + # before any correction exists. + self.tf.publish(self._correction_tf(np.eye(4), time.time())) + self._unsub_odom = self.odometry.subscribe(self._on_odom) + self._unsub_lidar = self.lidar.subscribe(self._on_scan) + logger.info("PGO (ivan transformer core) started") + + @rpc + def stop(self) -> None: + if self._unsub_odom is not None: + self._unsub_odom.dispose() + if self._unsub_lidar is not None: + self._unsub_lidar.dispose() + super().stop() + + def _on_odom(self, msg: Odometry) -> None: + with self._odom_lock: + self._latest_odom = msg + + def _on_scan(self, cloud: PointCloud2) -> None: + import gtsam # type: ignore[import-untyped] + + with self._odom_lock: + odom = self._latest_odom + if odom is None or len(cloud) == 0: + return + state = self._state + assert state is not None + + position = odom.pose.position + orientation = odom.pose.orientation + local_pose = gtsam.Pose3( + gtsam.Rot3.Quaternion(orientation.w, orientation.x, orientation.y, orientation.z), + gtsam.Point3(position.x, position.y, position.z), + ) + ts = odom.ts if odom.ts else time.time() + + if not self.config.input_is_registered: + cloud = cloud.transform( + _pose3_to_transform( + local_pose, + ts=ts, + frame_id=self.config.world_frame, + child_frame_id=self.config.body_frame, + ) + ) + + with self._state_lock: + state.process(local_pose, ts, cloud) + keyframe_count = len(state._key_poses) + correction = state._world_correction + corrected = correction.compose(local_pose) + graph_msg = self._snapshot_graph(state, ts) if keyframe_count else None + loop_events = self._new_loop_events(state, ts) + + self._publish_corrected_odometry(corrected, ts) + self.tf.publish(self._correction_tf(correction.matrix(), ts)) + if graph_msg is not None: + self.pose_graph.publish(graph_msg) + for event in loop_events: + self.loop_closure_event.publish(event) + + def _publish_corrected_odometry(self, pose: Any, ts: float) -> None: + translation = np.asarray(pose.translation()) + quaternion = Rotation.from_matrix(pose.rotation().matrix()).as_quat() + self.corrected_odometry.publish( + Odometry( + ts=ts, + frame_id=self.config.world_frame, + child_frame_id=self.config.body_frame, + pose=Pose( + position=[float(v) for v in translation], + orientation=[float(v) for v in quaternion], + ), + ) + ) + + def _correction_tf(self, matrix: np.ndarray, ts: float) -> Transform: + quaternion = Rotation.from_matrix(matrix[:3, :3]).as_quat() + return Transform( + frame_id=self.config.world_frame, + child_frame_id=self.config.odom_frame, + translation=Vector3(float(matrix[0, 3]), float(matrix[1, 3]), float(matrix[2, 3])), + rotation=Quaternion(*[float(v) for v in quaternion]), + ts=ts, + ) + + def _snapshot_graph(self, state: _PGOState, ts: float) -> Graph3D: + """Optimized keyframes + odometry-chain and loop edges. + + Caller must hold ``_state_lock``.""" + world_frame = self.config.world_frame + nodes = [ + _pose3_node(index, key_pose.optimized, key_pose.timestamp, world_frame) + for index, key_pose in enumerate(state._key_poses) + ] + edges = [ + Graph3D.Edge( + start_id=index - 1, end_id=index, timestamp=state._key_poses[index].timestamp + ) + for index in range(1, len(state._key_poses)) + ] + edges += [ + Graph3D.Edge(start_id=pair.target, end_id=pair.source, metadata_id=1) + for pair in state._accepted_loops + ] + return Graph3D(ts=ts, nodes=nodes, edges=edges) + + def _new_loop_events(self, state: _PGOState, ts: float) -> list[GraphDelta3D]: + """One GraphDelta3D per accepted loop not yet published. + + Caller must hold ``_state_lock``.""" + world_frame = self.config.world_frame + identity = GraphDelta3D.Transform( + translation=Vector3(0.0, 0.0, 0.0), rotation=Quaternion(0.0, 0.0, 0.0, 1.0) + ) + events: list[GraphDelta3D] = [] + for pair in state._accepted_loops[self._published_loops :]: + source = state._key_poses[pair.source] + target = state._key_poses[pair.target] + events.append( + GraphDelta3D( + ts=ts, + nodes=[ + _pose3_node(pair.target, target.optimized, target.timestamp, world_frame), + _pose3_node(pair.source, source.optimized, source.timestamp, world_frame), + ], + transforms=[identity, identity], + ) + ) + self._published_loops = len(state._accepted_loops) + return events diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/CMakeLists.txt b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/CMakeLists.txt new file mode 100644 index 0000000000..8c7b6d5b94 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/CMakeLists.txt @@ -0,0 +1,49 @@ +cmake_minimum_required(VERSION 3.14) +project(pgo CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3") + +include(FetchContent) +FetchContent_Declare(dimos_lcm + GIT_REPOSITORY https://github.com/dimensionalOS/dimos-lcm.git + GIT_TAG main + GIT_SHALLOW TRUE +) +FetchContent_MakeAvailable(dimos_lcm) + +find_package(PkgConfig REQUIRED) +pkg_check_modules(LCM REQUIRED lcm) +find_package(Eigen3 REQUIRED) +find_package(PCL 1.8 REQUIRED COMPONENTS common filters kdtree registration io) +find_package(GTSAM REQUIRED) + +add_definitions(-DUSE_PCL) + +add_executable(pgo + main.cpp + simple_pgo.cpp + commons.cpp +) + +target_include_directories(pgo PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ${dimos_lcm_SOURCE_DIR}/generated/cpp_lcm_msgs + ${LCM_INCLUDE_DIRS} + ${EIGEN3_INCLUDE_DIR} + ${PCL_INCLUDE_DIRS} + ${GTSAM_INCLUDE_DIR} +) + +target_link_libraries(pgo PRIVATE + ${LCM_LIBRARIES} + ${PCL_LIBRARIES} + gtsam +) + +target_link_directories(pgo PRIVATE + ${LCM_LIBRARY_DIRS} +) + +install(TARGETS pgo DESTINATION bin) diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.cpp b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.cpp new file mode 100644 index 0000000000..f309a97a85 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.cpp @@ -0,0 +1,8 @@ +#include "commons.h" + +void PoseWithTime::setTime(int32_t _sec, uint32_t _nsec) +{ + sec = _sec; + nsec = _nsec; + second = static_cast(sec) + static_cast(nsec) / 1e9; +} diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.h b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.h new file mode 100644 index 0000000000..2f9244ede6 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.h @@ -0,0 +1,29 @@ +#pragma once +#include +#include +#include + +using PointType = pcl::PointXYZI; +using CloudType = pcl::PointCloud; +using PointVec = std::vector>; + +using M3D = Eigen::Matrix3d; +using V3D = Eigen::Vector3d; +using M3F = Eigen::Matrix3f; +using V3F = Eigen::Vector3f; +using M4F = Eigen::Matrix4f; +using V4F = Eigen::Vector4f; + +struct PoseWithTime { + V3D t; + M3D r; + int32_t sec; + uint32_t nsec; + double second; + void setTime(int32_t sec, uint32_t nsec); +}; + +struct CloudWithPose { + CloudType::Ptr cloud; + PoseWithTime pose; +}; diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/dimos_native_module.hpp b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/dimos_native_module.hpp new file mode 100644 index 0000000000..c22cca4326 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/dimos_native_module.hpp @@ -0,0 +1,97 @@ +// SmartNav Native Module helpers. +// Re-exports dimos NativeModule patterns for CLI arg parsing and LCM helpers. +// Based on dimos/hardware/sensors/lidar/common/dimos_native_module.hpp + +#pragma once + +#include +#include +#include +#include + +#include "std_msgs/Header.hpp" +#include "std_msgs/Time.hpp" + +namespace dimos { + +class NativeModule { +public: + NativeModule(int argc, char** argv) { + for (int i = 1; i < argc; ++i) { + std::string arg(argv[i]); + if (arg.size() > 2 && arg[0] == '-' && arg[1] == '-' && i + 1 < argc) { + args_[arg.substr(2)] = argv[++i]; + } + } + } + + /// Get the full LCM channel string for a declared port. + const std::string& topic(const std::string& port) const { + auto it = args_.find(port); + if (it == args_.end()) { + throw std::runtime_error("NativeModule: no topic for port '" + port + "'"); + } + return it->second; + } + + /// Get a string arg value, or a default if not present. + std::string arg(const std::string& key, const std::string& default_val = "") const { + auto it = args_.find(key); + return it != args_.end() ? it->second : default_val; + } + + /// Get a float arg value, or a default if not present. + float arg_float(const std::string& key, float default_val = 0.0f) const { + auto it = args_.find(key); + return it != args_.end() ? std::stof(it->second) : default_val; + } + + /// Get an int arg value, or a default if not present. + int arg_int(const std::string& key, int default_val = 0) const { + auto it = args_.find(key); + return it != args_.end() ? std::stoi(it->second) : default_val; + } + + /// Get a bool arg value, or a default if not present. + /// Present-but-unparseable values throw, matching arg_int/arg_float's + /// std::stoi/std::stof behaviour — a typo'd value or empty string is a + /// misconfiguration we want to surface immediately, not silently coerce + /// to false. + bool arg_bool(const std::string& key, bool default_val = false) const { + auto it = args_.find(key); + if (it == args_.end()) return default_val; + if (it->second == "true" || it->second == "1") return true; + if (it->second == "false" || it->second == "0") return false; + throw std::runtime_error( + "NativeModule: arg '--" + key + "' has unparseable bool value '" + + it->second + "' (expected true/false or 1/0)"); + } + + /// Check if a port/arg was provided. + bool has(const std::string& key) const { + return args_.count(key) > 0; + } + +private: + std::map args_; +}; + +/// Convert seconds (double) to a ROS-style Time message. +inline std_msgs::Time time_from_seconds(double t) { + std_msgs::Time ts; + ts.sec = static_cast(t); + ts.nsec = static_cast((t - ts.sec) * 1e9); + return ts; +} + +/// Build a stamped Header with auto-incrementing sequence number. +inline std_msgs::Header make_header(const std::string& frame_id, double ts) { + static std::atomic seq{0}; + std_msgs::Header h; + h.seq = seq.fetch_add(1, std::memory_order_relaxed); + h.stamp = time_from_seconds(ts); + h.frame_id = frame_id; + return h; +} + +} // namespace dimos diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.lock b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.lock new file mode 100644 index 0000000000..16898a2c2d --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.lock @@ -0,0 +1,144 @@ +{ + "nodes": { + "dimos-lcm": { + "flake": false, + "locked": { + "lastModified": 1769774949, + "narHash": "sha256-icRK7jerqNlwK1WZBrnIP04I2WozzFqTD7qsmnPxQuo=", + "owner": "dimensionalOS", + "repo": "dimos-lcm", + "rev": "0aa72b7b1bd3a65f50f5c03485ee9b728df56afe", + "type": "github" + }, + "original": { + "owner": "dimensionalOS", + "ref": "main", + "repo": "dimos-lcm", + "type": "github" + } + }, + "flake-utils": { + "inputs": { + "systems": "systems" + }, + "locked": { + "lastModified": 1731533236, + "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "gtsam-extended": { + "inputs": { + "flake-utils": [ + "flake-utils" + ], + "nixpkgs": [ + "nixpkgs" + ], + "nixpkgs-stable": "nixpkgs-stable" + }, + "locked": { + "lastModified": 1775168452, + "narHash": "sha256-fsUWPQIw+lk4VEQUOniiPLwiPoldWh8ZdnIdos202+I=", + "owner": "jeff-hykin", + "repo": "gtsam-extended", + "rev": "f4572a80b6339181693aee6029ca28153e59a993", + "type": "github" + }, + "original": { + "owner": "jeff-hykin", + "repo": "gtsam-extended", + "type": "github" + } + }, + "lcm-extended": { + "inputs": { + "flake-utils": [ + "flake-utils" + ], + "nixpkgs": [ + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1774902379, + "narHash": "sha256-gRFvEkbXCEoG4jEmsT+i0bMZ5kDHOtAaPsrbStXjdu4=", + "owner": "jeff-hykin", + "repo": "lcm_extended", + "rev": "7d12ad8546d3daae30528a6c28f2c9ff5b10baf7", + "type": "github" + }, + "original": { + "owner": "jeff-hykin", + "repo": "lcm_extended", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1777954456, + "narHash": "sha256-hGdgeU2Nk87RAuZyYjyDjFL6LK7dAZN5RE9+hrDTkDU=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "549bd84d6279f9852cae6225e372cc67fb91a4c1", + "type": "github" + }, + "original": { + "owner": "NixOS", + "ref": "nixos-unstable", + "repo": "nixpkgs", + "type": "github" + } + }, + "nixpkgs-stable": { + "locked": { + "lastModified": 1751274312, + "narHash": "sha256-/bVBlRpECLVzjV19t5KMdMFWSwKLtb5RyXdjz3LJT+g=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "50ab793786d9de88ee30ec4e4c24fb4236fc2674", + "type": "github" + }, + "original": { + "owner": "NixOS", + "ref": "nixos-24.11", + "repo": "nixpkgs", + "type": "github" + } + }, + "root": { + "inputs": { + "dimos-lcm": "dimos-lcm", + "flake-utils": "flake-utils", + "gtsam-extended": "gtsam-extended", + "lcm-extended": "lcm-extended", + "nixpkgs": "nixpkgs" + } + }, + "systems": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.nix b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.nix new file mode 100644 index 0000000000..ea8eae1327 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.nix @@ -0,0 +1,70 @@ +{ + description = "SmartNav PGO native module (pose graph optimization with iSAM2 + PCL ICP)"; + + inputs = { + nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable"; + flake-utils.url = "github:numtide/flake-utils"; + lcm-extended = { + url = "github:jeff-hykin/lcm_extended"; + inputs.nixpkgs.follows = "nixpkgs"; + inputs.flake-utils.follows = "flake-utils"; + }; + dimos-lcm = { + url = "github:dimensionalOS/dimos-lcm/main"; + flake = false; + }; + gtsam-extended = { + url = "github:jeff-hykin/gtsam-extended"; + inputs.nixpkgs.follows = "nixpkgs"; + inputs.flake-utils.follows = "flake-utils"; + }; + }; + + outputs = { self, nixpkgs, flake-utils, lcm-extended, dimos-lcm, gtsam-extended, ... }: + flake-utils.lib.eachDefaultSystem (system: + let + pkgs = import nixpkgs { inherit system; }; + lcm = lcm-extended.packages.${system}.lcm; + + gtsam-base = gtsam-extended.packages.${system}.gtsam-cpp; + gtsam = gtsam-base.overrideAttrs (_old: { + src = pkgs.fetchFromGitHub { + owner = "borglab"; + repo = "gtsam"; + rev = "1a9792a7ede244850a413739557635b606f295c0"; + sha256 = "sha256-zxm5TGVPW1vipFVpw01zcvKRw4mkh+5ZBCR1n6G466o="; + }; + env.NIX_CFLAGS_COMPILE = "-Wno-error=array-bounds"; + }); + in { + packages.default = pkgs.stdenv.mkDerivation { + pname = "smartnav-pgo"; + version = "0.1.0"; + src = ./.; + + nativeBuildInputs = [ pkgs.cmake pkgs.pkg-config ]; + buildInputs = [ + lcm + pkgs.glib + pkgs.eigen + pkgs.boost + pkgs.pcl + gtsam + ]; + + env.NIX_CFLAGS_COMPILE = "-Wno-error=array-bounds"; + + cmakeFlags = [ + "-DCMAKE_POLICY_VERSION_MINIMUM=3.5" + "-DFETCHCONTENT_SOURCE_DIR_DIMOS_LCM=${dimos-lcm}" + ]; + + # On macOS, libgtsam.4.dylib is referenced via @rpath but the binary + # has no LC_RPATH entries, so it fails to load at runtime. Add one + # pointing at the gtsam lib dir. + postInstall = pkgs.lib.optionalString pkgs.stdenv.isDarwin '' + ${pkgs.darwin.cctools}/bin/install_name_tool -add_rpath ${gtsam}/lib $out/bin/pgo + ''; + }; + }); +} diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/main.cpp b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/main.cpp new file mode 100644 index 0000000000..38314d4390 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/main.cpp @@ -0,0 +1,299 @@ +// PGO NativeModule — faithful port of pgo_node.cpp from ROS2 to LCM. +// Subscribes to registered_scan + odometry, runs SimplePGO (iSAM2 + PCL ICP), +// publishes corrected_odometry, global_map, and TF correction offset. + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "commons.h" +#include "simple_pgo.h" +#include "dimos_native_module.hpp" +#include "point_cloud_utils.hpp" + +#include "nav_msgs/Odometry.hpp" +#include "sensor_msgs/PointCloud2.hpp" +#include "geometry_msgs/Pose.hpp" +#include "geometry_msgs/Quaternion.hpp" +#include "geometry_msgs/Point.hpp" + +static std::atomic g_running{true}; +static void signal_handler(int) { g_running.store(false); } + +// Shared state between LCM callbacks and main loop +static std::mutex g_buffer_mutex; +static std::queue g_cloud_buffer; +static double g_last_message_time = 0.0; + +// Latest odometry for non-keyframe TF broadcasting +static std::mutex g_odom_mutex; +static M3D g_latest_r = M3D::Identity(); +static V3D g_latest_t = V3D::Zero(); +static double g_latest_time = 0.0; +static bool g_has_odom = false; + +class Handlers { +public: + void on_odometry(const lcm::ReceiveBuffer*, const std::string&, + const nav_msgs::Odometry* msg) { + M3D r = Eigen::Quaterniond( + msg->pose.pose.orientation.w, + msg->pose.pose.orientation.x, + msg->pose.pose.orientation.y, + msg->pose.pose.orientation.z + ).toRotationMatrix(); + + V3D t(msg->pose.pose.position.x, + msg->pose.pose.position.y, + msg->pose.pose.position.z); + + double ts = msg->header.stamp.sec + msg->header.stamp.nsec / 1e9; + + std::lock_guard lock(g_odom_mutex); + g_latest_r = r; + g_latest_t = t; + g_latest_time = ts; + g_has_odom = true; + } + + void on_registered_scan(const lcm::ReceiveBuffer*, const std::string&, + const sensor_msgs::PointCloud2* msg) { + std::lock_guard odom_lock(g_odom_mutex); + if (!g_has_odom) + return; + + double ts = g_latest_time; + + // Reject out-of-order messages + if (ts < g_last_message_time) + return; + g_last_message_time = ts; + + CloudWithPose cp; + cp.pose.r = g_latest_r; + cp.pose.t = g_latest_t; + cp.pose.setTime(static_cast(ts), + static_cast((ts - static_cast(ts)) * 1e9)); + + // Parse PointCloud2 to PCL + cp.cloud = CloudType::Ptr(new CloudType); + smartnav::to_pcl(*msg, *cp.cloud); + + std::lock_guard buf_lock(g_buffer_mutex); + g_cloud_buffer.push(cp); + } +}; + +static nav_msgs::Odometry build_odometry(const M3D& r, const V3D& t, double ts, + const std::string& frame_id, + const std::string& child_frame_id) { + nav_msgs::Odometry odom; + odom.header = dimos::make_header(frame_id, ts); + odom.child_frame_id = child_frame_id; + + Eigen::Quaterniond q(r); + odom.pose.pose.position.x = t.x(); + odom.pose.pose.position.y = t.y(); + odom.pose.pose.position.z = t.z(); + odom.pose.pose.orientation.x = q.x(); + odom.pose.pose.orientation.y = q.y(); + odom.pose.pose.orientation.z = q.z(); + odom.pose.pose.orientation.w = q.w(); + + return odom; +} + +int main(int argc, char** argv) +{ + signal(SIGTERM, signal_handler); + signal(SIGINT, signal_handler); + + dimos::NativeModule mod(argc, argv); + + // Port topics + std::string scan_topic = mod.topic("registered_scan"); + std::string odom_topic = mod.topic("odometry"); + std::string corrected_odom_topic = mod.topic("corrected_odometry"); + std::string global_map_topic = mod.topic("global_map"); + std::string tf_topic = mod.topic("pgo_tf"); + + // Config parameters + Config config; + config.key_pose_delta_deg = mod.arg_float("key_pose_delta_deg", 10.0f); + config.key_pose_delta_trans = mod.arg_float("key_pose_delta_trans", 0.5f); + config.loop_search_radius = mod.arg_float("loop_search_radius", 1.0f); + config.loop_time_tresh = mod.arg_float("loop_time_thresh", 60.0f); + config.loop_score_tresh = mod.arg_float("loop_score_thresh", 0.15f); + config.loop_submap_half_range = mod.arg_int("loop_submap_half_range", 5); + config.submap_resolution = mod.arg_float("submap_resolution", 0.1f); + config.min_loop_detect_duration = mod.arg_float("min_loop_detect_duration", 5.0f); + + // Node-level config + std::string world_frame = mod.arg("world_frame", "map"); + std::string local_frame = mod.arg("local_frame", "odom"); + float global_map_voxel_size = mod.arg_float("global_map_voxel_size", 0.1f); + float global_map_publish_rate = mod.arg_float("global_map_publish_rate", 1.0f); + // rate <= 0 disables global_map entirely (it's a big cloud that can overflow + // LCM's single-message limit and congest the bus). Consumers that only need + // corrected_odometry / pose_graph should turn it off. + bool publish_global_map = global_map_publish_rate > 0; + double global_map_interval = publish_global_map ? 1.0 / global_map_publish_rate : 0.0; + + // Unregister mode: transform world-frame scans to body-frame + bool unregister_input = mod.arg_bool("unregister_input", true); + + bool debug = mod.arg_bool("debug", false); + + pcl::console::setVerbosityLevel( + debug ? pcl::console::L_INFO : pcl::console::L_ERROR); + + SimplePGO pgo(config); + + lcm::LCM lcm; + if (!lcm.good()) { + fprintf(stderr, "PGO: LCM init failed\n"); + return 1; + } + + Handlers handlers; + lcm.subscribe(odom_topic, &Handlers::on_odometry, &handlers); + lcm.subscribe(scan_topic, &Handlers::on_registered_scan, &handlers); + + if (debug) { + fprintf(stderr, "PGO native module started\n"); + fprintf(stderr, " registered_scan: %s\n", scan_topic.c_str()); + fprintf(stderr, " odometry: %s\n", odom_topic.c_str()); + fprintf(stderr, " corrected_odometry: %s\n", corrected_odom_topic.c_str()); + fprintf(stderr, " global_map: %s\n", global_map_topic.c_str()); + fprintf(stderr, " pgo_tf: %s\n", tf_topic.c_str()); + } + + double last_global_map_time = 0.0; + int timer_period_ms = 50; // 20 Hz, matching original + + while (g_running.load()) { + // Drain all pending LCM messages + while (lcm.handleTimeout(0) > 0) {} + + // Check buffer + CloudWithPose cp; + bool has_data = false; + { + std::lock_guard lock(g_buffer_mutex); + if (!g_cloud_buffer.empty()) { + cp = g_cloud_buffer.front(); + // Drain entire queue (matching original: process oldest, discard rest) + while (!g_cloud_buffer.empty()) { + g_cloud_buffer.pop(); + } + has_data = true; + } + } + + if (!has_data) { + std::this_thread::sleep_for(std::chrono::milliseconds(timer_period_ms)); + continue; + } + + // Optionally transform world-frame scan to body-frame + if (unregister_input && cp.cloud && cp.cloud->size() > 0) { + CloudType::Ptr body_cloud(new CloudType); + // body = R_odom^T * (world_pts - t_odom) + M3D r_inv = cp.pose.r.transpose(); + for (const auto& pt : *cp.cloud) { + V3D world_pt(pt.x, pt.y, pt.z); + V3D body_pt = r_inv * (world_pt - cp.pose.t); + PointType bp; + bp.x = static_cast(body_pt.x()); + bp.y = static_cast(body_pt.y()); + bp.z = static_cast(body_pt.z()); + bp.intensity = pt.intensity; + body_cloud->push_back(bp); + } + cp.cloud = body_cloud; + } + + double cur_time = cp.pose.second; + + if (!pgo.addKeyPose(cp)) { + // Not a keyframe — still broadcast TF and corrected odom + M3D corr_r = pgo.offsetR() * cp.pose.r; + V3D corr_t = pgo.offsetR() * cp.pose.t + pgo.offsetT(); + + nav_msgs::Odometry corrected = build_odometry( + corr_r, corr_t, cur_time, world_frame, "base_link"); + lcm.publish(corrected_odom_topic, &corrected); + + nav_msgs::Odometry tf_msg = build_odometry( + pgo.offsetR(), pgo.offsetT(), cur_time, world_frame, local_frame); + lcm.publish(tf_topic, &tf_msg); + + std::this_thread::sleep_for(std::chrono::milliseconds(timer_period_ms)); + continue; + } + + // Keyframe added + pgo.searchForLoopPairs(); + pgo.smoothAndUpdate(); + + if (debug) { + fprintf(stderr, "PGO: keyframe %zu at (%.1f, %.1f, %.1f)\n", + pgo.keyPoses().size(), + cp.pose.t.x(), cp.pose.t.y(), cp.pose.t.z()); + } + + // Publish corrected odometry + M3D corr_r = pgo.offsetR() * cp.pose.r; + V3D corr_t = pgo.offsetR() * cp.pose.t + pgo.offsetT(); + nav_msgs::Odometry corrected = build_odometry( + corr_r, corr_t, cur_time, world_frame, "base_link"); + lcm.publish(corrected_odom_topic, &corrected); + + // Publish TF correction (map -> odom offset) + nav_msgs::Odometry tf_msg = build_odometry( + pgo.offsetR(), pgo.offsetT(), cur_time, world_frame, local_frame); + lcm.publish(tf_topic, &tf_msg); + + // Publish global map (throttled) + double now = cur_time; + if (publish_global_map && now - last_global_map_time >= global_map_interval) { + last_global_map_time = now; + + if (!pgo.keyPoses().empty()) { + CloudType::Ptr global_cloud(new CloudType); + for (size_t i = 0; i < pgo.keyPoses().size(); i++) { + CloudType::Ptr world_cloud(new CloudType); + pcl::transformPointCloud( + *pgo.keyPoses()[i].body_cloud, + *world_cloud, + pgo.keyPoses()[i].t_global, + Eigen::Quaterniond(pgo.keyPoses()[i].r_global)); + *global_cloud += *world_cloud; + } + + // Voxel downsample + CloudType::Ptr filtered(new CloudType); + pcl::VoxelGrid voxel; + voxel.setInputCloud(global_cloud); + voxel.setLeafSize(global_map_voxel_size, global_map_voxel_size, global_map_voxel_size); + voxel.filter(*filtered); + + sensor_msgs::PointCloud2 map_msg = smartnav::from_pcl(*filtered, world_frame, now); + lcm.publish(global_map_topic, &map_msg); + } + } + + std::this_thread::sleep_for(std::chrono::milliseconds(timer_period_ms)); + } + + if (debug) fprintf(stderr, "PGO native module shutting down\n"); + return 0; +} diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/point_cloud_utils.hpp b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/point_cloud_utils.hpp new file mode 100644 index 0000000000..0970e1f8de --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/point_cloud_utils.hpp @@ -0,0 +1,170 @@ +// Point cloud utility functions for SmartNav native modules. +// Provides PointCloud2 building/parsing helpers that work with dimos-lcm types. +// When USE_PCL is defined, also provides PCL interop utilities. + +#pragma once + +#include +#include +#include + +#include "sensor_msgs/PointCloud2.hpp" +#include "sensor_msgs/PointField.hpp" +#include "std_msgs/Header.hpp" + +#include "dimos_native_module.hpp" + +#ifdef USE_PCL +#include +#include +#include +#endif + +namespace smartnav { + +// Simple XYZI point structure (no PCL dependency) +struct PointXYZI { + float x, y, z, intensity; +}; + +// Build PointCloud2 from vector of XYZI points +inline sensor_msgs::PointCloud2 build_pointcloud2( + const std::vector& points, + const std::string& frame_id, + double timestamp +) { + sensor_msgs::PointCloud2 pc; + pc.header = dimos::make_header(frame_id, timestamp); + pc.height = 1; + pc.width = static_cast(points.size()); + pc.is_bigendian = 0; + pc.is_dense = 1; + + // Fields: x, y, z, intensity (all float32) + pc.fields_length = 4; + pc.fields.resize(4); + auto make_field = [](const std::string& name, int32_t offset) { + sensor_msgs::PointField f; + f.name = name; + f.offset = offset; + f.datatype = sensor_msgs::PointField::FLOAT32; + f.count = 1; + return f; + }; + pc.fields[0] = make_field("x", 0); + pc.fields[1] = make_field("y", 4); + pc.fields[2] = make_field("z", 8); + pc.fields[3] = make_field("intensity", 12); + + pc.point_step = 16; + pc.row_step = pc.point_step * pc.width; + pc.data_length = pc.row_step; + pc.data.resize(pc.data_length); + + for (size_t i = 0; i < points.size(); ++i) { + float* dst = reinterpret_cast(pc.data.data() + i * 16); + dst[0] = points[i].x; + dst[1] = points[i].y; + dst[2] = points[i].z; + dst[3] = points[i].intensity; + } + + return pc; +} + +// Parse PointCloud2 into vector of XYZI points +inline std::vector parse_pointcloud2(const sensor_msgs::PointCloud2& pc) { + std::vector points; + if (pc.width == 0 || pc.height == 0) return points; + + int num_points = pc.width * pc.height; + points.reserve(num_points); + + // Find field offsets + int x_off = -1, y_off = -1, z_off = -1, i_off = -1; + for (const auto& f : pc.fields) { + if (f.name == "x") x_off = f.offset; + else if (f.name == "y") y_off = f.offset; + else if (f.name == "z") z_off = f.offset; + else if (f.name == "intensity") i_off = f.offset; + } + + if (x_off < 0 || y_off < 0 || z_off < 0) return points; + + for (int n = 0; n < num_points; ++n) { + if (static_cast((n + 1) * pc.point_step) > pc.data.size()) break; + const uint8_t* base = pc.data.data() + n * pc.point_step; + PointXYZI p; + std::memcpy(&p.x, base + x_off, sizeof(float)); + std::memcpy(&p.y, base + y_off, sizeof(float)); + std::memcpy(&p.z, base + z_off, sizeof(float)); + if (i_off >= 0) std::memcpy(&p.intensity, base + i_off, sizeof(float)); + else p.intensity = 0.0f; + points.push_back(p); + } + + return points; +} + +// Get timestamp from PointCloud2 header +inline double get_timestamp(const sensor_msgs::PointCloud2& pc) { + return pc.header.stamp.sec + pc.header.stamp.nsec / 1e9; +} + +#ifdef USE_PCL +// Convert dimos-lcm PointCloud2 to PCL point cloud +inline void to_pcl(const sensor_msgs::PointCloud2& pc, + pcl::PointCloud& cloud) { + auto points = parse_pointcloud2(pc); + cloud.clear(); + cloud.reserve(points.size()); + for (const auto& p : points) { + pcl::PointXYZI pt; + pt.x = p.x; + pt.y = p.y; + pt.z = p.z; + pt.intensity = p.intensity; + cloud.push_back(pt); + } + cloud.width = cloud.size(); + cloud.height = 1; + cloud.is_dense = true; +} + +// Convert PCL point cloud to dimos-lcm PointCloud2 +inline sensor_msgs::PointCloud2 from_pcl( + const pcl::PointCloud& cloud, + const std::string& frame_id, + double timestamp +) { + std::vector points; + points.reserve(cloud.size()); + for (const auto& pt : cloud) { + points.push_back({pt.x, pt.y, pt.z, pt.intensity}); + } + return build_pointcloud2(points, frame_id, timestamp); +} +#endif + +// Quaternion to RPY conversion +inline void quat_to_rpy(double qx, double qy, double qz, double qw, + double& roll, double& pitch, double& yaw) { + // Roll (x-axis rotation) + double sinr_cosp = 2.0 * (qw * qx + qy * qz); + double cosr_cosp = 1.0 - 2.0 * (qx * qx + qy * qy); + roll = std::atan2(sinr_cosp, cosr_cosp); + + // Pitch (y-axis rotation) + double sinp = 2.0 * (qw * qy - qz * qx); + if (std::abs(sinp) >= 1.0) + pitch = std::copysign(M_PI / 2, sinp); + else + pitch = std::asin(sinp); + + // Yaw (z-axis rotation) + double siny_cosp = 2.0 * (qw * qz + qx * qy); + double cosy_cosp = 1.0 - 2.0 * (qy * qy + qz * qz); + yaw = std::atan2(siny_cosp, cosy_cosp); +} + +} // namespace smartnav diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.cpp b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.cpp new file mode 100644 index 0000000000..5fc18bf0e7 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.cpp @@ -0,0 +1,211 @@ +#include "simple_pgo.h" + +SimplePGO::SimplePGO(const Config &config) : m_config(config) +{ + gtsam::ISAM2Params isam2_params; + isam2_params.relinearizeThreshold = 0.01; + isam2_params.relinearizeSkip = 1; + m_isam2 = std::make_shared(isam2_params); + m_initial_values.clear(); + m_graph.resize(0); + m_r_offset.setIdentity(); + m_t_offset.setZero(); + + m_icp.setMaximumIterations(50); + m_icp.setMaxCorrespondenceDistance(10); + m_icp.setTransformationEpsilon(1e-6); + m_icp.setEuclideanFitnessEpsilon(1e-6); + m_icp.setRANSACIterations(0); +} + +bool SimplePGO::isKeyPose(const PoseWithTime &pose) +{ + if (m_key_poses.size() == 0) + return true; + const KeyPoseWithCloud &last_item = m_key_poses.back(); + double delta_trans = (pose.t - last_item.t_local).norm(); + double delta_deg = Eigen::Quaterniond(pose.r).angularDistance(Eigen::Quaterniond(last_item.r_local)) * 57.324; + if (delta_trans > m_config.key_pose_delta_trans || delta_deg > m_config.key_pose_delta_deg) + return true; + return false; +} +bool SimplePGO::addKeyPose(const CloudWithPose &cloud_with_pose) +{ + bool is_key_pose = isKeyPose(cloud_with_pose.pose); + if (!is_key_pose) + return false; + size_t idx = m_key_poses.size(); + M3D init_r = m_r_offset * cloud_with_pose.pose.r; + V3D init_t = m_r_offset * cloud_with_pose.pose.t + m_t_offset; + // 添加初始值 + m_initial_values.insert(idx, gtsam::Pose3(gtsam::Rot3(init_r), gtsam::Point3(init_t))); + if (idx == 0) + { + // 添加先验约束 + gtsam::noiseModel::Diagonal::shared_ptr noise = gtsam::noiseModel::Diagonal::Variances(gtsam::Vector6::Ones() * 1e-12); + m_graph.add(gtsam::PriorFactor(idx, gtsam::Pose3(gtsam::Rot3(init_r), gtsam::Point3(init_t)), noise)); + } + else + { + // 添加里程计约束 + const KeyPoseWithCloud &last_item = m_key_poses.back(); + M3D r_between = last_item.r_local.transpose() * cloud_with_pose.pose.r; + V3D t_between = last_item.r_local.transpose() * (cloud_with_pose.pose.t - last_item.t_local); + gtsam::noiseModel::Diagonal::shared_ptr noise = gtsam::noiseModel::Diagonal::Variances((gtsam::Vector(6) << 1e-6, 1e-6, 1e-6, 1e-4, 1e-4, 1e-6).finished()); + m_graph.add(gtsam::BetweenFactor(idx - 1, idx, gtsam::Pose3(gtsam::Rot3(r_between), gtsam::Point3(t_between)), noise)); + } + KeyPoseWithCloud item; + item.time = cloud_with_pose.pose.second; + item.r_local = cloud_with_pose.pose.r; + item.t_local = cloud_with_pose.pose.t; + item.body_cloud = cloud_with_pose.cloud; + item.r_global = init_r; + item.t_global = init_t; + m_key_poses.push_back(item); + return true; +} + +CloudType::Ptr SimplePGO::getSubMap(int idx, int half_range, double resolution) +{ + assert(idx >= 0 && idx < static_cast(m_key_poses.size())); + int min_idx = std::max(0, idx - half_range); + int max_idx = std::min(static_cast(m_key_poses.size()) - 1, idx + half_range); + + CloudType::Ptr ret(new CloudType); + for (int i = min_idx; i <= max_idx; i++) + { + + CloudType::Ptr body_cloud = m_key_poses[i].body_cloud; + CloudType::Ptr global_cloud(new CloudType); + pcl::transformPointCloud(*body_cloud, *global_cloud, m_key_poses[i].t_global, Eigen::Quaterniond(m_key_poses[i].r_global)); + *ret += *global_cloud; + } + if (resolution > 0) + { + pcl::VoxelGrid voxel_grid; + voxel_grid.setLeafSize(resolution, resolution, resolution); + voxel_grid.setInputCloud(ret); + voxel_grid.filter(*ret); + } + return ret; +} + +void SimplePGO::searchForLoopPairs() +{ + if (m_key_poses.size() < 10) + return; + if (m_config.min_loop_detect_duration > 0.0) + { + if (m_history_pairs.size() > 0) + { + double current_time = m_key_poses.back().time; + double last_time = m_key_poses[m_history_pairs.back().second].time; + if (current_time - last_time < m_config.min_loop_detect_duration) + return; + } + } + + size_t cur_idx = m_key_poses.size() - 1; + const KeyPoseWithCloud &last_item = m_key_poses.back(); + pcl::PointXYZ last_pose_pt; + last_pose_pt.x = last_item.t_global(0); + last_pose_pt.y = last_item.t_global(1); + last_pose_pt.z = last_item.t_global(2); + + pcl::PointCloud::Ptr key_poses_cloud(new pcl::PointCloud); + for (size_t i = 0; i < m_key_poses.size() - 1; i++) + { + pcl::PointXYZ pt; + pt.x = m_key_poses[i].t_global(0); + pt.y = m_key_poses[i].t_global(1); + pt.z = m_key_poses[i].t_global(2); + key_poses_cloud->push_back(pt); + } + pcl::KdTreeFLANN kdtree; + kdtree.setInputCloud(key_poses_cloud); + std::vector ids; + std::vector sqdists; + int neighbors = kdtree.radiusSearch(last_pose_pt, m_config.loop_search_radius, ids, sqdists); + if (neighbors == 0) + return; + + int loop_idx = -1; + for (size_t i = 0; i < ids.size(); i++) + { + int idx = ids[i]; + if (std::abs(last_item.time - m_key_poses[idx].time) > m_config.loop_time_tresh) + { + loop_idx = idx; + break; + } + } + + if (loop_idx == -1) + return; + + CloudType::Ptr target_cloud = getSubMap(loop_idx, m_config.loop_submap_half_range, m_config.submap_resolution); + CloudType::Ptr source_cloud = getSubMap(m_key_poses.size() - 1, 0, m_config.submap_resolution); + CloudType::Ptr align_cloud(new CloudType); + + m_icp.setInputSource(source_cloud); + m_icp.setInputTarget(target_cloud); + m_icp.align(*align_cloud); + + if (!m_icp.hasConverged() || m_icp.getFitnessScore() > m_config.loop_score_tresh) + return; + + M4F loop_transform = m_icp.getFinalTransformation(); + + LoopPair one_pair; + one_pair.source_id = cur_idx; + one_pair.target_id = loop_idx; + one_pair.score = m_icp.getFitnessScore(); + M3D r_refined = loop_transform.block<3, 3>(0, 0).cast() * m_key_poses[cur_idx].r_global; + V3D t_refined = loop_transform.block<3, 3>(0, 0).cast() * m_key_poses[cur_idx].t_global + loop_transform.block<3, 1>(0, 3).cast(); + one_pair.r_offset = m_key_poses[loop_idx].r_global.transpose() * r_refined; + one_pair.t_offset = m_key_poses[loop_idx].r_global.transpose() * (t_refined - m_key_poses[loop_idx].t_global); + m_cache_pairs.push_back(one_pair); + m_history_pairs.emplace_back(one_pair.target_id, one_pair.source_id); +} + +void SimplePGO::smoothAndUpdate() +{ + bool has_loop = !m_cache_pairs.empty(); + // 添加回环因子 + if (has_loop) + { + for (LoopPair &pair : m_cache_pairs) + { + m_graph.add(gtsam::BetweenFactor(pair.target_id, pair.source_id, + gtsam::Pose3(gtsam::Rot3(pair.r_offset), + gtsam::Point3(pair.t_offset)), + gtsam::noiseModel::Diagonal::Variances(gtsam::Vector6::Ones() * pair.score))); + } + std::vector().swap(m_cache_pairs); + } + // smooth and mapping + m_isam2->update(m_graph, m_initial_values); + m_isam2->update(); + if (has_loop) + { + m_isam2->update(); + m_isam2->update(); + m_isam2->update(); + m_isam2->update(); + } + m_graph.resize(0); + m_initial_values.clear(); + + // update key poses + gtsam::Values estimate_values = m_isam2->calculateBestEstimate(); + for (size_t i = 0; i < m_key_poses.size(); i++) + { + gtsam::Pose3 pose = estimate_values.at(i); + m_key_poses[i].r_global = pose.rotation().matrix().cast(); + m_key_poses[i].t_global = pose.translation().matrix().cast(); + } + // update offset + const KeyPoseWithCloud &last_item = m_key_poses.back(); + m_r_offset = last_item.r_global * last_item.r_local.transpose(); + m_t_offset = last_item.t_global - m_r_offset * last_item.t_local; +} diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.h b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.h new file mode 100644 index 0000000000..7f80c5b09a --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.h @@ -0,0 +1,78 @@ +#pragma once +#include "commons.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct KeyPoseWithCloud +{ + M3D r_local; + V3D t_local; + M3D r_global; + V3D t_global; + double time; + CloudType::Ptr body_cloud; +}; +struct LoopPair +{ + size_t source_id; + size_t target_id; + M3D r_offset; + V3D t_offset; + double score; +}; + +struct Config +{ + double key_pose_delta_deg = 10; + double key_pose_delta_trans = 1.0; + double loop_search_radius = 1.0; + double loop_time_tresh = 60.0; + double loop_score_tresh = 0.15; + int loop_submap_half_range = 5; + double submap_resolution = 0.1; + double min_loop_detect_duration = 10.0; +}; + +class SimplePGO +{ +public: + SimplePGO(const Config &config); + + bool isKeyPose(const PoseWithTime &pose); + + bool addKeyPose(const CloudWithPose &cloud_with_pose); + + bool hasLoop(){return m_cache_pairs.size() > 0;} + + void searchForLoopPairs(); + + void smoothAndUpdate(); + + CloudType::Ptr getSubMap(int idx, int half_range, double resolution); + std::vector> &historyPairs() { return m_history_pairs; } + std::vector &keyPoses() { return m_key_poses; } + + M3D offsetR() { return m_r_offset; } + V3D offsetT() { return m_t_offset; } + +private: + Config m_config; + std::vector m_key_poses; + std::vector> m_history_pairs; + std::vector m_cache_pairs; + M3D m_r_offset; + V3D m_t_offset; + std::shared_ptr m_isam2; + gtsam::Values m_initial_values; + gtsam::NonlinearFactorGraph m_graph; + pcl::IterativeClosestPoint m_icp; +}; diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/module.py new file mode 100644 index 0000000000..73a5fc59e4 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/module.py @@ -0,0 +1,286 @@ +# 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. + +"""The unrefined PGO — a frozen standalone snapshot of main's cmu_nav PGO +(GTSAM iSAM2 + PCL ICP C++ binary, the original gsc_pgo was refined from), +adapted to the LoopClosure spec without touching the binary. + +The cpp/ directory here is a copy of `cmu_nav/modules/pgo/cpp` at the time +of the snapshot, so later cmu_nav changes don't silently move this baseline. + +Spec adaptations, all Python-side: + * `lidar` — the binary expects a `registered_scan` topic; `_collect_topics` + aliases the lidar port's topic onto that arg name. + * `pose_graph` — the binary doesn't expose its internal graph, only the + current map->odom offset (`pgo_tf`) and corrected odometry. The wrapper + keyframes the RAW odometry stream (same delta gates as the binary) and + re-applies the LATEST offset to every keyframe on each correction update. + A single global offset can't reproduce iSAM2's per-keyframe smoothing — + but that offset is exactly what this PGO exposes to consumers, so the + synthesized graph is an honest picture of its output. + * `loop_closure_event` — emitted when the offset jumps by more than the + `_LOOP_EVENT_*` thresholds (the offset only moves materially when a loop + closure lands).""" + +from __future__ import annotations + +from dataclasses import dataclass +import math +from pathlib import Path +import threading +import time + +import numpy as np +from reactivex.disposable import Disposable +from scipy.spatial.transform import Rotation + +from dimos.core.core import rpc +from dimos.core.native_module import NativeModule, NativeModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +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 +from dimos.navigation.jnav.components.loop_closure.spec import LoopClosure +from dimos.navigation.jnav.msgs.Graph3D import Graph3D +from dimos.navigation.jnav.msgs.GraphDelta3D import GraphDelta3D +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + +# Offset jumps below these are smoothing noise, not loop closures. +_LOOP_EVENT_MIN_TRANS_M = 0.05 +_LOOP_EVENT_MIN_ROT_DEG = 1.0 + + +class PGOConfig(NativeModuleConfig): + cwd: str | None = str(Path(__file__).resolve().parent / "cpp") + # Absolute so the exists() check works from any worker cwd (skips rebuild). + executable: str = str(Path(__file__).resolve().parent / "cpp/result/bin/pgo") + # path:$PWD makes nix see this (git-untracked) copied directory. + build_command: str | None = 'nix build "path:$PWD#default" --no-write-lock-file' + + # Frame names + world_frame: str = "map" + local_frame: str = "odom" + + # Keyframe detection + key_pose_delta_deg: float = 10.0 + key_pose_delta_trans: float = 0.5 + + # Loop closure + loop_search_radius: float = 1.0 + loop_time_thresh: float = 60.0 + loop_score_thresh: float = 0.15 + loop_submap_half_range: int = 5 + submap_resolution: float = 0.1 + min_loop_detect_duration: float = 5.0 + + # Input mode: transform world-frame scans to body-frame using odom + unregister_input: bool = True + + # Global map publishing + global_map_voxel_size: float = 0.1 + global_map_publish_rate: float = 1.0 + + debug: bool = False + + +@dataclass +class _RawKeyframe: + ts: float + translation: np.ndarray # (3,) + rotation: np.ndarray # 3x3 + + +class PGO(NativeModule, LoopClosure): + """Pose graph optimization with loop closure using GTSAM iSAM2 + PCL ICP.""" + + config: PGOConfig + + lidar: In[PointCloud2] + odometry: In[Odometry] + corrected_odometry: Out[Odometry] + pose_graph: Out[Graph3D] + loop_closure_event: Out[GraphDelta3D] + global_map: Out[PointCloud2] + pgo_tf: Out[Odometry] + + def __init__(self, **kwargs: object) -> None: + super().__init__(**kwargs) + self._keyframes: list[_RawKeyframe] = [] + self._offset_rotation = np.eye(3) + self._offset_translation = np.zeros(3) + self._graph_lock = threading.Lock() + + def _collect_topics(self) -> dict[str, str]: + topics = super()._collect_topics() + # The binary asks for --registered_scan; feed it the lidar topic. + if "lidar" in topics: + topics["registered_scan"] = topics["lidar"] + return topics + + @rpc + def start(self) -> None: + super().start() + self.register_disposable( + Disposable(self.pgo_tf.transport.subscribe(self._on_tf_correction, self.pgo_tf)) + ) + self.register_disposable( + Disposable(self.odometry.transport.subscribe(self._on_raw_odometry, self.odometry)) + ) + # Seed identity TF so consumers can query map->body immediately. + self._publish_tf( + translation=(0.0, 0.0, 0.0), + rotation=(0.0, 0.0, 0.0, 1.0), + ts=time.time(), + ) + if self.config.debug: + logger.info("unrefined PGO native module started (C++ iSAM2 + PCL ICP)") + + @rpc + def stop(self) -> None: + super().stop() + + # --- TF passthrough (same as the cmu_nav wrapper) ----------------------- + + def _publish_tf( + self, + translation: tuple[float, float, float], + rotation: tuple[float, float, float, float], + ts: float, + ) -> None: + self.tf.publish( + Transform( + frame_id=self.config.world_frame, + child_frame_id=self.config.local_frame, + translation=Vector3(*translation), + rotation=Quaternion(*rotation), + ts=ts, + ) + ) + + def _on_tf_correction(self, msg: Odometry) -> None: + self._publish_tf( + translation=(msg.pose.position.x, msg.pose.position.y, msg.pose.position.z), + rotation=( + msg.pose.orientation.x, + msg.pose.orientation.y, + msg.pose.orientation.z, + msg.pose.orientation.w, + ), + ts=msg.ts or time.time(), + ) + + new_rotation = Rotation.from_quat( + [ + msg.pose.orientation.x, + msg.pose.orientation.y, + msg.pose.orientation.z, + msg.pose.orientation.w, + ] + ).as_matrix() + new_translation = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z]) + + with self._graph_lock: + delta_trans = float(np.linalg.norm(new_translation - self._offset_translation)) + cos_theta = float( + np.clip((np.trace(self._offset_rotation.T @ new_rotation) - 1.0) / 2.0, -1.0, 1.0) + ) + delta_deg = math.degrees(math.acos(cos_theta)) + self._offset_rotation = new_rotation + self._offset_translation = new_translation + is_loop = ( + delta_trans > _LOOP_EVENT_MIN_TRANS_M or delta_deg > _LOOP_EVENT_MIN_ROT_DEG + ) and bool(self._keyframes) + graph_msg = self._build_graph(msg.ts) if self._keyframes else None + event = self._build_loop_event(msg.ts) if is_loop else None + + if graph_msg is not None: + self.pose_graph.publish(graph_msg) + if event is not None: + self.loop_closure_event.publish(event) + + # --- synthesized pose graph (the binary doesn't expose its own) ---------- + + def _on_raw_odometry(self, msg: Odometry) -> None: + rotation = Rotation.from_quat( + [ + msg.pose.orientation.x, + msg.pose.orientation.y, + msg.pose.orientation.z, + msg.pose.orientation.w, + ] + ).as_matrix() + translation = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z]) + + with self._graph_lock: + if not self._is_keyframe(rotation, translation): + return + self._keyframes.append( + _RawKeyframe(ts=msg.ts, translation=translation, rotation=rotation) + ) + graph_msg = self._build_graph(msg.ts) + self.pose_graph.publish(graph_msg) + + def _is_keyframe(self, rotation: np.ndarray, translation: np.ndarray) -> bool: + if not self._keyframes: + return True + last = self._keyframes[-1] + delta_trans = float(np.linalg.norm(translation - last.translation)) + cos_theta = float(np.clip((np.trace(last.rotation.T @ rotation) - 1.0) / 2.0, -1.0, 1.0)) + delta_deg = math.degrees(math.acos(cos_theta)) + return ( + delta_trans > self.config.key_pose_delta_trans + or delta_deg > self.config.key_pose_delta_deg + ) + + def _corrected_node(self, index: int, keyframe: _RawKeyframe) -> Graph3D.Node3D: + rotation = self._offset_rotation @ keyframe.rotation + translation = self._offset_rotation @ keyframe.translation + self._offset_translation + quaternion = Rotation.from_matrix(rotation).as_quat() + return Graph3D.Node3D( + pose=PoseStamped( + ts=keyframe.ts, + frame_id=self.config.world_frame, + position=[float(v) for v in translation], + orientation=[float(v) for v in quaternion], + ), + id=index, + ) + + def _build_graph(self, ts: float) -> Graph3D: + """Caller must hold ``_graph_lock``.""" + nodes = [ + self._corrected_node(index, keyframe) for index, keyframe in enumerate(self._keyframes) + ] + edges = [ + Graph3D.Edge(start_id=index - 1, end_id=index, timestamp=self._keyframes[index].ts) + for index in range(1, len(self._keyframes)) + ] + return Graph3D(ts=ts, nodes=nodes, edges=edges) + + def _build_loop_event(self, ts: float) -> GraphDelta3D: + """Caller must hold ``_graph_lock``.""" + latest_index = len(self._keyframes) - 1 + identity = GraphDelta3D.Transform( + translation=Vector3(0.0, 0.0, 0.0), rotation=Quaternion(0.0, 0.0, 0.0, 1.0) + ) + return GraphDelta3D( + ts=ts, + nodes=[self._corrected_node(latest_index, self._keyframes[latest_index])], + transforms=[identity], + ) From 13f97dae5c3e7e0bfd1a6e4a22980ef71734eb67 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 14:52:43 +0800 Subject: [PATCH 08/69] add loop-closure eval/benchmark scripts and spec --- .../loop_closure/benchmark_table.py | 343 +++++++ .../jnav/components/loop_closure/eval.py | 924 ++++++++++++++++++ .../jnav/components/loop_closure/eval_all.py | 366 +++++++ .../loop_closure/eval_ground_truth_tag.py | 309 ++++++ .../components/loop_closure/eval_kitti.py | 192 ++++ .../components/loop_closure/eval_kitti_all.py | 122 +++ .../jnav/components/loop_closure/spec.py | 33 + 7 files changed, 2289 insertions(+) create mode 100644 dimos/navigation/jnav/components/loop_closure/benchmark_table.py create mode 100644 dimos/navigation/jnav/components/loop_closure/eval.py create mode 100644 dimos/navigation/jnav/components/loop_closure/eval_all.py create mode 100644 dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py create mode 100644 dimos/navigation/jnav/components/loop_closure/eval_kitti.py create mode 100644 dimos/navigation/jnav/components/loop_closure/eval_kitti_all.py create mode 100644 dimos/navigation/jnav/components/loop_closure/spec.py diff --git a/dimos/navigation/jnav/components/loop_closure/benchmark_table.py b/dimos/navigation/jnav/components/loop_closure/benchmark_table.py new file mode 100644 index 0000000000..102ebfdc25 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/benchmark_table.py @@ -0,0 +1,343 @@ +# 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. + +"""Resumable PGO benchmark over (environment x implementation/config). + +Fills one big table: every environment (go2 recordings, hk_village; kitti once +converted) scored against every PGO column (gsc_pgo in several configs, the +basic unrefined_pgo, ivan_pgo, ivan_pgo_transformer). Each cell is one eval.py +subprocess (sequential — they share the isolated LCM bus). + +CHECKPOINTED: a cell is skipped when its summary.json already exists and its +fingerprint still matches the db (size+mtime) and EVAL_VERSION. Kill it any +time and rerun — only missing/stale cells recompute. `--force` recomputes all. + +The universal score is **voxel agreement** (re-anchoring scans onto the +corrected trajectory should collapse double walls — ground-truth-free and +needs no camera), so tagless environments (kitti, hk_village) still get a +real number. April-tag agreement is reported additionally wherever a camera + +intrinsics sidecar exists. + +Usage: + uv run python dimos/navigation/jnav/components/loop_closure/benchmark_table.py + ... [--go2-root ~/datasets/go2_recordings] [--with-hk-village] [--force] + ... [--only-env NAME] [--only-col NAME] [--table-only] +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass, field +import json +from pathlib import Path +import subprocess +import sys +import time +from typing import Any + +LOOP_CLOSURE_DIR = Path(__file__).resolve().parent +EVAL_PY = LOOP_CLOSURE_DIR / "eval.py" +RESULTS_DIR = LOOP_CLOSURE_DIR / "eval_results" +TABLE_PATH = RESULTS_DIR / "benchmark_table.md" + +DEFAULT_GO2_ROOT = Path("~/datasets/go2_recordings").expanduser() +# loop_closure/modules/jnav/navigation/dimos/ -> parents[4] is repo root. +LFS_DATA_DIR = LOOP_CLOSURE_DIR.parents[4] / "data" # repo_root/data (hk_village*.db) + +# Shared loop-closure thresholds for the gsc_pgo configs (mirrors recordings_eval). +_CMU_BASE: dict[str, Any] = { + "loop_search_radius": 3.0, + "loop_time_thresh": 5.0, + "min_loop_detect_duration": 2.0, + "key_pose_delta_trans": 0.5, +} + + +@dataclass(frozen=True) +class Column: + """One implementation+config = one table column.""" + + name: str # column label + results-suffix (disambiguates same-class modules) + module_dir: str # subdir under loop_closure/ holding module.py (class PGO) + overrides: dict[str, Any] = field(default_factory=dict) + + +COLUMNS: list[Column] = [ + Column("cmu_stock", "gsc_pgo", {}), + Column("cmu_scan_context", "gsc_pgo", {**_CMU_BASE, "use_scan_context": True}), + Column("cmu_radius", "gsc_pgo", {**_CMU_BASE, "use_scan_context": False}), + Column( + "cmu_scan_context_far", + "gsc_pgo", + { + **_CMU_BASE, + "use_scan_context": True, + "loop_candidate_max_distance_m": 0.0, + "loop_score_thresh": 10000.0, + }, + ), + Column("unrefined", "unrefined_pgo", {}), + Column("ivan", "ivan_pgo", {"publish_global_map": False}), + Column("ivan_transformer", "ivan_pgo_transformer", {}), +] + + +@dataclass(frozen=True) +class Environment: + """One dataset = one table row.""" + + name: str # results-dir recording key + db_path: Path + odom_stream: str + lidar_stream: str + camera_stream: str | None = None + intrinsics_json: Path | None = None + + +def discover_go2(root: Path) -> list[Environment]: + environments = [] + for db_path in sorted(root.glob("*/mem2.db")): + recording = db_path.parent + sidecar = recording / "camera_intrinsics.json" + environments.append( + Environment( + name=recording.name, + db_path=db_path, + odom_stream="fastlio_odometry", + lidar_stream="fastlio_lidar", + camera_stream="color_image", + intrinsics_json=sidecar if sidecar.exists() else None, + ) + ) + return environments + + +def discover_hk_village(data_dir: Path) -> list[Environment]: + # hk_village LFS dbs publish world-frame lidar + PoseStamped odom; no camera + # intrinsics sidecar, so they score on voxel agreement alone. + environments = [] + for db_path in sorted(data_dir.glob("hk_village*.db")): + environments.append( + Environment( + name=db_path.stem, + db_path=db_path, + odom_stream="odom", + lidar_stream="lidar", + camera_stream=None, + intrinsics_json=None, + ) + ) + return environments + + +def cell_dir(environment: Environment, column: Column) -> Path: + # Mirrors eval.py's out_dir formula: __.PGO[.]. + return RESULTS_DIR / f"{environment.name}__{column.module_dir}.PGO.{column.name}" + + +def cell_is_fresh(environment: Environment, column: Column) -> bool: + summary_path = cell_dir(environment, column) / "summary.json" + if not summary_path.exists(): + return False + try: + summary = json.loads(summary_path.read_text()) + except (json.JSONDecodeError, OSError): + return False + fingerprint = summary.get("fingerprint", {}) + stat = environment.db_path.stat() + return ( + fingerprint.get("db_bytes") == stat.st_size + and fingerprint.get("db_mtime") == int(stat.st_mtime) + and fingerprint.get("version") is not None + ) + + +def _kill_zombies() -> None: + """Clear leftover native processes / workers that can wedge the next cell.""" + subprocess.run( + "lsof -ti tcp:7766 2>/dev/null | xargs kill -9 2>/dev/null;" + ' pkill -9 -f "bin/pgo|scene_lidar" 2>/dev/null', + shell=True, + check=False, + ) + + +def run_cell(environment: Environment, column: Column) -> bool: + command = [ + sys.executable, + "-u", + str(EVAL_PY), + "--db-path", + str(environment.db_path), + "--odom-stream", + environment.odom_stream, + "--lidar-stream", + environment.lidar_stream, + "--module-path", + str(LOOP_CLOSURE_DIR / column.module_dir / "module.py"), + "--module-name", + "PGO", + "--recording-name", + environment.name, + "--results-suffix", + column.name, + "--with-rrd", + "false", + "--lockstep", + "true", + ] + if environment.camera_stream is not None: + command += ["--camera-stream", environment.camera_stream] + if environment.intrinsics_json is not None: + command += ["--camera-intrinsics-json-path", str(environment.intrinsics_json)] + if column.overrides: + command += ["--pgo-config-json", json.dumps(column.overrides)] + print(f"\n=== {environment.name} x {column.name} ===", flush=True) + result = subprocess.run(command, check=False) + print(f"=== {environment.name} x {column.name} exit: {result.returncode} ===", flush=True) + return result.returncode == 0 + + +def _fmt(value: float | None, places: int = 3, signed: bool = True) -> str: + if value is None: + return "—" + return f"{value:+.{places}f}" if signed else f"{value:.{places}f}" + + +def render_table(environments: list[Environment]) -> Path: + cells: dict[tuple[str, str], dict[str, Any]] = {} + for summary_path in RESULTS_DIR.glob("*/summary.json"): + recording, _, module_key = summary_path.parent.name.rpartition("__") + try: + summary = json.loads(summary_path.read_text()) + except (json.JSONDecodeError, OSError): + continue + cells[(recording, module_key)] = summary.get("scores", {}) + + column_keys = [f"{column.module_dir}.PGO.{column.name}" for column in COLUMNS] + header = "| environment | " + " | ".join(column.name for column in COLUMNS) + " |" + sep = "|" + "---|" * (len(COLUMNS) + 1) + lines = [ + "# PGO benchmark — environments x implementations", + "", + "Each cell: **voxel improvement** (fractional drop in occupied 0.2 m voxels " + "after re-anchoring scans onto the corrected trajectory; the universal, " + "ground-truth-free score) — then `tag:` where a camera " + "exists, and `cl`. Higher is better; `—` = not yet run / N/A.", + "", + header, + sep, + ] + for environment in environments: + row_cells = [] + for column_key in column_keys: + scores = cells.get((environment.name, column_key)) + if scores is None: + row_cells.append("—") + continue + voxel = _fmt(scores.get("voxel_improvement")) + tag = scores.get("tag_improvement") + closures = scores.get("closures") + text = voxel + if tag is not None: + text += f" tag:{tag:+.2f}" + if closures is not None: + text += f" cl{closures}" + row_cells.append(text) + lines.append(f"| {environment.name} | " + " | ".join(row_cells) + " |") + + # Per-column mean voxel improvement (over environments that have a number). + lines += ["", "## Mean voxel improvement per column", ""] + lines.append("| " + " | ".join(column.name for column in COLUMNS) + " |") + lines.append("|" + "---|" * len(COLUMNS)) + means = [] + for column_key in column_keys: + values = [ + cells[(environment.name, column_key)]["voxel_improvement"] + for environment in environments + if (environment.name, column_key) in cells + and cells[(environment.name, column_key)].get("voxel_improvement") is not None + ] + means.append(f"{sum(values) / len(values):+.3f}" if values else "—") + lines.append("| " + " | ".join(means) + " |") + lines.append("") + + RESULTS_DIR.mkdir(exist_ok=True) + TABLE_PATH.write_text("\n".join(lines) + "\n") + return TABLE_PATH + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--go2-root", type=Path, default=DEFAULT_GO2_ROOT) + parser.add_argument("--with-hk-village", action="store_true") + parser.add_argument("--data-dir", type=Path, default=LFS_DATA_DIR) + parser.add_argument("--only-env", help="comma-separated environment names") + parser.add_argument("--only-col", help="comma-separated column names") + parser.add_argument("--force", action="store_true", help="recompute fresh cells too") + parser.add_argument( + "--attempts", type=int, default=2, help="retries per cell on transient RPC timeouts" + ) + parser.add_argument("--table-only", action="store_true", help="render from cache, run nothing") + args = parser.parse_args() + + environments = discover_go2(args.go2_root.expanduser()) + if args.with_hk_village: + environments += discover_hk_village(args.data_dir.expanduser()) + if args.only_env: + wanted = {name.strip() for name in args.only_env.split(",")} + environments = [environment for environment in environments if environment.name in wanted] + + columns = COLUMNS + if args.only_col: + wanted = {name.strip() for name in args.only_col.split(",")} + columns = [column for column in COLUMNS if column.name in wanted] + + if args.table_only: + print(f"table -> {render_table(environments)}") + return + + total = len(environments) * len(columns) + print(f"benchmark: {len(environments)} environments x {len(columns)} columns = {total} cells") + done = skipped = failed = 0 + for environment in environments: + for column in columns: + if not args.force and cell_is_fresh(environment, column): + skipped += 1 + print(f"skip (fresh): {environment.name} x {column.name}", flush=True) + continue + # Retry transient macOS LCM startup-RPC timeouts; a fresh process + # almost always gets past them. Kill zombies between attempts. + ok = False + for attempt in range(1, args.attempts + 1): + ok = run_cell(environment, column) + if ok: + break + _kill_zombies() + if attempt < args.attempts: + print( + f"retry {attempt + 1}/{args.attempts}: {environment.name} x {column.name}" + ) + time.sleep(5) + done += 1 if ok else 0 + failed += 0 if ok else 1 + render_table(environments) # refresh after every cell — live + crash-safe + + table = render_table(environments) + print(f"\ncells: {done} ran, {skipped} cached, {failed} failed") + print(f"table -> {table}") + + +if __name__ == "__main__": + main() diff --git a/dimos/navigation/jnav/components/loop_closure/eval.py b/dimos/navigation/jnav/components/loop_closure/eval.py new file mode 100644 index 0000000000..02afb44749 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/eval.py @@ -0,0 +1,924 @@ +# 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. + +"""Evaluate a loop-closure module against a recording. + +Two ground-truth-free scores, before vs after correction: + * April-tag agreement — a fixed tag re-seen along the run should map to one + world position; the spread of its per-visit robot positions measures drift. + Tags are taken as relative to the chosen odom stream (sighting time -> + nearest odom pose), so no static transforms or stored db poses are needed. + * Lidar-voxel agreement — re-anchoring the registered scans onto the + corrected trajectory should collapse double walls, so the corrected map + should occupy FEWER voxels than the raw one. + +Pipeline: + 1. April tags: read the db's `april_tags` stream (ts + marker_id only), or + detect them with sane defaults (medoid, blur/reproj/size/distance gates). + 2. Raw agreement over the raw odometry. + 3. Replay lidar + odom through the module (loaded dynamically from + --module-path/--module-name), capture its optimized pose graph. + 4. Corrected agreement + voxel agreement, written to + eval_results/__/summary.json (and an eval.rrd with the + raw + corrected trajectories when --with-rrd true). + +Usage: + uv run python dimos/navigation/jnav/components/loop_closure/eval.py \\ + --db-path ~/datasets/go2_recordings/2026-06-04_12-56pm-PST/mem2.db \\ + --odom-stream fastlio_odometry \\ + --camera-stream color_image \\ + --camera-intrinsics-json-path \\ + ~/datasets/go2_recordings/2026-06-04_12-56pm-PST/camera_intrinsics.json \\ + --module-path dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py \\ + --module-name PGO \\ + --pgo-config-json '{"use_scan_context": true}' \\ + --with-rrd true +""" + +from __future__ import annotations + +import argparse +import asyncio +from collections.abc import AsyncGenerator, Iterable +import importlib +import json +from pathlib import Path +import tempfile +import time +from typing import Any + +import numpy as np + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.coordination.module_coordinator import ModuleCoordinator +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.navigation.jnav.msgs.Graph3D import Graph3D +from dimos.navigation.jnav.msgs.GraphDelta3D import GraphDelta3D +from dimos.navigation.jnav.utils.apriltag_agreement import ( + VISIT_GAP_S, + AgreementReport, + agreement_improvement, + agreement_report, + paired_tag_visit_positions, +) +from dimos.navigation.jnav.utils.apriltags import ( + detect_apriltags, + load_intrinsics_json, + load_or_detect_sightings, +) +from dimos.navigation.jnav.utils.module_loading import ( + filter_config_for_module, + load_module_class, +) +from dimos.navigation.jnav.utils.recording_db import ( + MAX_REPLAY_ODOM, + MAX_REPLAY_SCANS, + ODOM_MATCH_TOLERANCE_S, + REPLAY_DRAIN_MARGIN_S, + REPLAY_PUBLISH_HZ, + iterate_stream, + list_streams, + odometry_lookup, + store, + stream_count, +) +from dimos.navigation.jnav.utils.trajectory_metrics import ( + GraphPose, + PoseLookup7, + drifted_lookup, + graph_lookup, + has_drift, + lidar_voxel_agreement, + pose7_lookup, + trajectory_recovery_error, +) + +RESULTS_DIR = Path(__file__).resolve().parent / "eval_results" +APRIL_TAGS_STREAM = "april_tags" +_RRD_MAX_PATH_POINTS = 5000 + +# Cap replayed scans fed to voxel agreement so the map fits in memory. +VOXEL_MAX_SCANS = 300 + +# Bump to invalidate every cached cell (scoring/replay semantics changed). +EVAL_VERSION = 1 + + +def cell_fingerprint( + db_path: Path, + pgo_config: dict[str, Any], + lidar_stream: str, + odom_stream: str, + drift_per_sec: list[float] | None = None, +) -> dict[str, Any]: + """Identity of a completed cell — the driver re-runs only when this changes + (db edited, config changed, streams changed, drift changed, or version).""" + stat = db_path.stat() + return { + "db_bytes": stat.st_size, + "db_mtime": int(stat.st_mtime), + "pgo_config": pgo_config, + "lidar_stream": lidar_stream, + "odom_stream": odom_stream, + "drift_per_sec": list(drift_per_sec or [0.0, 0.0, 0.0]), + "version": EVAL_VERSION, + } + + +# A known-good PGO config for replayed recordings: revisit gates loose enough +# for walks where the same spot is re-seen tens of seconds later. These now +# match gsc_pgo's own defaults (so it's a no-op for gsc_pgo); kept here to apply +# the same gates to the other PGO modules, which have different defaults. +DEFAULT_PGO_CONFIG: dict[str, Any] = { + "loop_search_radius": 3.0, + "loop_time_thresh": 5.0, + "min_loop_detect_duration": 2.0, + "key_pose_delta_trans": 0.5, + "use_scan_context": True, +} + + +class GraphCaptureConfig(ModuleConfig): + output_path: str = "" + + +class GraphCapture(Module): + """Captures the module's optimized pose graph WITH orientations + closures. + + Results are handed back via a JSON file written on teardown (modules run in + separate worker processes).""" + + config: GraphCaptureConfig + + pose_graph: In[Graph3D] + loop_closure_event: In[GraphDelta3D] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._graph: list[GraphPose] = [] + self._closures = 0 + + async def handle_pose_graph(self, msg: Graph3D) -> None: + self._graph = [ + ( + node.pose.ts, + node.pose.position.x, + node.pose.position.y, + node.pose.position.z, + node.pose.orientation.x, + node.pose.orientation.y, + node.pose.orientation.z, + node.pose.orientation.w, + ) + for node in msg.nodes + ] + + async def handle_loop_closure_event(self, msg: GraphDelta3D) -> None: + self._closures += 1 + + async def main(self) -> AsyncGenerator[None, None]: + yield + Path(self.config.output_path).write_text( + json.dumps({"graph": self._graph, "closures": self._closures}) + ) + + +class LockstepReplayConfig(ModuleConfig): + db: str = "" + lidar_stream: str = "lidar" + odometry_stream: str = "odom" + lidar_stride: int = 1 + odometry_stride: int = 2 + odom_publish_hz: float = 500.0 + ack_timeout_s: float = 30.0 + done_path: str = "" + # Artificial odometry drift: a constant-velocity world offset added to both + # odom poses and lidar clouds at time t (offset = drift_per_sec * (t - t0)). + # Consistent per-instant, so the trajectory warps over time — exactly the + # accumulating error loop closure is supposed to fix. [0,0,0] = no drift. + drift_per_sec: list[float] = [0.0, 0.0, 0.0] + drift_t0: float = 0.0 + + +class LockstepReplay(Module): + """Closed-loop replay: after each scan, wait for the module's + corrected_odometry ack before sending the next. + + Every module under test sees 100% of the (strided) scans regardless of + machine speed — wall clock varies, the data the module processes doesn't. + Odometry messages are cheap latest-state updates and stay fire-and-forget + (lightly paced). Writes a done-marker JSON (ack timeout count) at the end + so the host knows when to tear down. + + odom and lidar are merged into one time-sorted stream, so playback runs in + bursts: all odoms whose timestamps fall before the next scan are emitted + fire-and-forget (paced by odom_publish_hz), then one scan is sent and the + loop blocks on its ack. The only guarantee is one ack-wait per scan; the + odom burst size per gap is data-dependent (~ odom_rate / lidar_rate).""" + + config: LockstepReplayConfig + + lidar: Out[PointCloud2] + odometry: Out[Odometry] + corrected_odometry: In[Odometry] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._ack_count = 0 + self._ack_event: asyncio.Event | None = None + + async def handle_corrected_odometry(self, msg: Odometry) -> None: + self._ack_count += 1 + if self._ack_event is not None: + self._ack_event.set() + + def _load(self) -> list[tuple[float, str, Any]]: + db_path = Path(self.config.db) + merged: list[tuple[float, str, Any]] = [] + for timestamp, pose in iterate_stream( + db_path, self.config.odometry_stream, stride=self.config.odometry_stride + ): + merged.append((timestamp, "odom", pose)) + for timestamp, cloud in iterate_stream( + db_path, self.config.lidar_stream, stride=self.config.lidar_stride + ): + merged.append((timestamp, "lidar", cloud)) + merged.sort(key=lambda item: item[0]) + return merged + + async def main(self) -> AsyncGenerator[None, None]: + messages = await asyncio.to_thread(self._load) + self._task = asyncio.create_task(self._replay(messages)) + yield + self._task.cancel() + + async def _replay(self, messages: list[tuple[float, str, Any]]) -> None: + odom_period = 1.0 / self.config.odom_publish_hz + timeouts = 0 + scans_sent = 0 + drift = np.asarray(self.config.drift_per_sec, dtype=np.float64) + t0 = self.config.drift_t0 + apply_drift = has_drift(drift) + # Timestamps of scans the module never acked — the frames it (likely) + # skipped. Recorded for reproducibility of partial runs. + skipped_scan_ts: list[float] = [] + for timestamp, kind, payload in messages: + if kind == "odom": + pose = RateReplay._payload_pose(payload) + if apply_drift: + offset = drift * (timestamp - t0) + pose = Pose( + position=[ + pose.position.x + offset[0], + pose.position.y + offset[1], + pose.position.z + offset[2], + ], + orientation=[ + pose.orientation.x, + pose.orientation.y, + pose.orientation.z, + pose.orientation.w, + ], + ) + self.odometry.publish( + Odometry( + ts=timestamp, + frame_id="map", + child_frame_id="base_link", + pose=pose, + ) + ) + await asyncio.sleep(odom_period) + continue + + acks_before = self._ack_count + self._ack_event = asyncio.Event() + points = payload.points_f32() + if apply_drift: + points = points + (drift * (timestamp - t0)).astype(np.float32) + self.lidar.publish(PointCloud2.from_numpy(points, frame_id="map", timestamp=timestamp)) + scans_sent += 1 + deadline = time.monotonic() + self.config.ack_timeout_s + while self._ack_count == acks_before: + remaining = deadline - time.monotonic() + if remaining <= 0: + timeouts += 1 + skipped_scan_ts.append(timestamp) + break + try: + await asyncio.wait_for(self._ack_event.wait(), timeout=remaining) + except TimeoutError: + continue + self._ack_event.clear() + if scans_sent % _PROGRESS_EVERY_N_SCANS == 0: + # Periodic progress so a capped run still reports coverage. + Path(self.config.done_path + ".progress").write_text( + json.dumps(self._stats(timeouts, scans_sent, skipped_scan_ts)) + ) + + Path(self.config.done_path).write_text( + json.dumps(self._stats(timeouts, scans_sent, skipped_scan_ts)) + ) + + @staticmethod + def _stats(timeouts: int, scans_sent: int, skipped_scan_ts: list[float]) -> dict[str, Any]: + return { + "timeouts": timeouts, + "scans_sent": scans_sent, + "skipped_scan_ts": skipped_scan_ts, + } + + +class RateReplayConfig(ModuleConfig): + db: str = "" + lidar_stream: str = "lidar" + odometry_stream: str = "odom" + lidar_stride: int = 1 + odometry_stride: int = 2 + publish_hz: float = 40.0 + + +class RateReplay(Module): + """Legacy fixed-rate replay: publishes world-frame lidar + odometry at a set + Hz with timestamps preserved (no ack pacing — wall-clock dependent). + + Works for both odometry payload shapes found in recordings: ``Odometry`` + (go2 ``fastlio_odometry``) and ``PoseStamped`` (hk_village ``odom``). + """ + + config: RateReplayConfig + + lidar: Out[PointCloud2] + odometry: Out[Odometry] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.done = False + + def _load(self) -> list[tuple[float, str, Any]]: + db_path = Path(self.config.db) + merged: list[tuple[float, str, Any]] = [] + for timestamp, pose in iterate_stream( + db_path, self.config.odometry_stream, stride=self.config.odometry_stride + ): + merged.append((timestamp, "odom", pose)) + for timestamp, cloud in iterate_stream( + db_path, self.config.lidar_stream, stride=self.config.lidar_stride + ): + merged.append((timestamp, "lidar", cloud)) + merged.sort(key=lambda item: item[0]) + return merged + + @staticmethod + def _payload_pose(payload: Any) -> Pose: + if hasattr(payload, "pose"): # Odometry + return payload.pose # type: ignore[no-any-return] + return Pose( # PoseStamped + payload.x, + payload.y, + payload.z, + payload.orientation.x, + payload.orientation.y, + payload.orientation.z, + payload.orientation.w, + ) + + async def main(self) -> AsyncGenerator[None, None]: + messages = await asyncio.to_thread(self._load) + self._task = asyncio.create_task(self._replay(messages)) + yield + self._task.cancel() + + async def _replay(self, messages: list[tuple[float, str, Any]]) -> None: + period = 1.0 / self.config.publish_hz + for timestamp, kind, payload in messages: + if kind == "odom": + self.odometry.publish( + Odometry( + ts=timestamp, + frame_id="map", + child_frame_id="base_link", + pose=self._payload_pose(payload), + ) + ) + else: + self.lidar.publish( + PointCloud2.from_numpy( + payload.points_f32(), frame_id="map", timestamp=timestamp + ) + ) + await asyncio.sleep(period) + self.done = True + + +# Run cap scales with the workload: a per-scan budget (well above any sane +# processing time, below the 30s ack timeout) plus fixed startup overhead. +LOCKSTEP_PER_SCAN_BUDGET_S = 2.0 +LOCKSTEP_BASE_OVERHEAD_S = 120.0 +LOCKSTEP_POLL_S = 5.0 +LOCKSTEP_DRAIN_S = 10.0 +_PROGRESS_EVERY_N_SCANS = 200 + + +def run_module_graph( + db_path: Path, + module_class: type, + config_overrides: dict[str, Any], + *, + lidar_stream: str, + odom_stream: str, + lockstep: bool = True, + drift_per_sec: list[float] | None = None, + drift_t0: float = 0.0, +) -> tuple[list[GraphPose], int, dict[str, Any]]: + """Replay the recording through the module; return its optimized pose graph + (with orientations), loop-closure count, and replay stats. + + lockstep=True (default) paces scans on the module's corrected_odometry + acks — machine-speed independent. lockstep=False is the legacy fixed-rate + wall-clock replay. drift_per_sec injects a constant-velocity world offset + into the replayed odom+lidar (see LockstepReplayConfig).""" + drift_per_sec = drift_per_sec or [0.0, 0.0, 0.0] + output_path = Path(tempfile.gettempdir()) / f"jnav_lc_eval_{db_path.parent.name}.json" + output_path.unlink(missing_ok=True) + done_path = Path(tempfile.gettempdir()) / f"jnav_lc_eval_done_{db_path.parent.name}.json" + done_path.unlink(missing_ok=True) + Path(str(done_path) + ".progress").unlink(missing_ok=True) + lidar_stride = max(1, -(-stream_count(db_path, lidar_stream) // MAX_REPLAY_SCANS)) + odometry_stride = max(1, -(-stream_count(db_path, odom_stream) // MAX_REPLAY_ODOM)) + n_messages = stream_count(db_path, odom_stream) // odometry_stride + n_messages += stream_count(db_path, lidar_stream) // lidar_stride + + if lockstep: + replay_blueprint = LockstepReplay.blueprint( + db=str(db_path), + lidar_stream=lidar_stream, + odometry_stream=odom_stream, + lidar_stride=lidar_stride, + odometry_stride=odometry_stride, + done_path=str(done_path), + drift_per_sec=drift_per_sec, + drift_t0=drift_t0, + ) + else: + replay_blueprint = RateReplay.blueprint( + db=str(db_path), + lidar_stream=lidar_stream, + odometry_stream=odom_stream, + lidar_stride=lidar_stride, + odometry_stride=odometry_stride, + publish_hz=REPLAY_PUBLISH_HZ, + ) + + blueprint = autoconnect( + replay_blueprint, + module_class.blueprint(**config_overrides), # type: ignore[attr-defined] + GraphCapture.blueprint(output_path=str(output_path)), + ) + coordinator = ModuleCoordinator.build(blueprint) + mode = "lockstep" if lockstep else f"fixed-rate {REPLAY_PUBLISH_HZ}Hz" + print( + f"replaying {n_messages} messages through {module_class.__name__}" + f" ({mode}, lidar stride {lidar_stride}, odom stride {odometry_stride})" + ) + replay_stats: dict[str, Any] = {"mode": mode} + try: + if lockstep: + # Per-frame budget: the cap scales with how many scans are fed. + n_scans = stream_count(db_path, lidar_stream) // lidar_stride + max_run_s = n_scans * LOCKSTEP_PER_SCAN_BUDGET_S + LOCKSTEP_BASE_OVERHEAD_S + started = time.monotonic() + while not done_path.exists(): + elapsed = time.monotonic() - started + if elapsed > max_run_s: + replay_stats["hit_max_run_s"] = max_run_s + print( + f"lockstep replay hit the per-frame cap" + f" ({n_scans} scans x {LOCKSTEP_PER_SCAN_BUDGET_S}s" + f" + {LOCKSTEP_BASE_OVERHEAD_S}s = {round(max_run_s)}s) — stopping early" + ) + break + if int(elapsed) % 60 < LOCKSTEP_POLL_S and elapsed > LOCKSTEP_POLL_S: + print(f" ... lockstep replay running ({round(elapsed)}s)") + time.sleep(LOCKSTEP_POLL_S) + progress_path = Path(str(done_path) + ".progress") + if done_path.exists(): + replay_stats.update(json.loads(done_path.read_text())) + elif progress_path.exists(): + # Capped run: last periodic progress still tells us coverage + # and which frames the module never acked. + replay_stats.update(json.loads(progress_path.read_text())) + replay_stats["partial"] = True + progress_path.unlink(missing_ok=True) + time.sleep(LOCKSTEP_DRAIN_S) + else: + time.sleep(n_messages / REPLAY_PUBLISH_HZ + REPLAY_DRAIN_MARGIN_S) + finally: + coordinator.stop() + + if not output_path.exists(): + raise SystemExit(f"{module_class.__name__} produced no pose graph output") + data = json.loads(output_path.read_text()) + graph = [tuple(row) for row in data["graph"]] + return graph, int(data["closures"]), replay_stats # type: ignore[return-value] + + +def odometry_pose7_lookup(db_path: Path, odom_stream: str) -> PoseLookup7: + times: list[float] = [] + poses: list[list[float]] = [] + for timestamp, payload in iterate_stream(db_path, odom_stream): + pose = RateReplay._payload_pose(payload) + times.append(timestamp) + poses.append( + [ + pose.position.x, + pose.position.y, + pose.position.z, + pose.orientation.x, + pose.orientation.y, + pose.orientation.z, + pose.orientation.w, + ] + ) + return pose7_lookup( + np.asarray(times, dtype=np.float64), + np.asarray(poses, dtype=np.float64), + ODOM_MATCH_TOLERANCE_S, + ) + + +def _subsampled_path(positions: np.ndarray) -> np.ndarray: + stride = max(1, len(positions) // _RRD_MAX_PATH_POINTS) + return positions[::stride] + + +def write_trajectory_rrd(rrd_path: Path, raw_positions: np.ndarray, graph: list[GraphPose]) -> None: + """Raw + corrected trajectories (no tag rendering — tags are scored as + odom-relative, not placed in 3D).""" + import rerun as rr + + rr.init("jnav_loop_closure_eval", spawn=False) + rr.save(str(rrd_path)) + rr.log( + "trajectory/raw_odom", + rr.LineStrips3D([_subsampled_path(raw_positions)], colors=[[200, 90, 90]]), + static=True, + ) + corrected = np.asarray([[node[1], node[2], node[3]] for node in graph], dtype=np.float64) + rr.log( + "trajectory/corrected", + rr.LineStrips3D([_subsampled_path(corrected)], colors=[[90, 200, 120]]), + static=True, + ) + + +def _report_dict(report: AgreementReport) -> dict[str, Any]: + return { + "mean_spread_m": report.mean_spread, + "total_observations": report.total_observations, + "per_tag": [ + {"tag_id": tag.tag_id, "observations": tag.observations, "spread_m": tag.spread} + for tag in report.per_tag + ], + } + + +def evaluate( + db_path: Path, + *, + odom_stream: str, + camera_stream: str | None, + intrinsics_json: Path | None, + module_path: Path, + module_name: str, + pgo_config: dict[str, Any], + with_rrd: bool, + lidar_stream: str, + lockstep: bool = True, + results_suffix: str = "", + recording_name: str | None = None, + drift_per_sec: list[float] | None = None, + ignore_tags: set[int] | None = None, +) -> dict[str, Any]: + streams = list_streams(db_path) + for required in (odom_stream, lidar_stream): + if required not in streams: + raise SystemExit(f"no stream {required!r} in {db_path} (have: {streams})") + + module_class = load_module_class(module_path, module_name) + pgo_config = filter_config_for_module(module_class, pgo_config) + + # Artificial drift: the module is fed odom+lidar with a constant-velocity + # world offset added at each time; the raw-baseline scoring must apply the + # SAME offset so it compares against what the module actually saw. + drift_per_sec = drift_per_sec or [0.0, 0.0, 0.0] + drift_t0 = next(iterate_stream(db_path, odom_stream))[0] if has_drift(drift_per_sec) else 0.0 + + # April-tag agreement needs a camera + intrinsics; voxel agreement does not. + # Datasets without either (kitti-360, bare lidar recordings) still score on + # voxel agreement alone, so the same harness fills every table cell. + sightings: dict[int, list[float]] = {} + tag_source = "none" + have_camera = camera_stream is not None and camera_stream in streams + if have_camera and intrinsics_json is not None and intrinsics_json.exists(): + assert camera_stream is not None # narrowed by have_camera + camera = camera_stream + intrinsics_config = load_intrinsics_json(intrinsics_json) + db_store = store(db_path) + stored_stream = ( + db_store.stream(APRIL_TAGS_STREAM) + if APRIL_TAGS_STREAM in db_store.list_streams() + else [] + ) + stored = ((int(obs.tags["marker_id"]), float(obs.ts)) for obs in stored_stream) + + def detect() -> Iterable[tuple[int, float]]: + detections = detect_apriltags( + db_store, + intrinsics_config["intrinsics"], + intrinsics_config["distortion"], + image_stream=camera, + stream_name=APRIL_TAGS_STREAM, + marker_length=intrinsics_config.get("marker_length", 0.10), + dictionary=intrinsics_config.get("dictionary", "DICT_APRILTAG_36h11"), + ) + return ((int(d["marker_id"]), float(d["ts"])) for d in detections) + + sightings, tag_source = load_or_detect_sightings(stored, detect) + # Drop dynamic/unreliable tags (e.g. a tag on a moving object) so their + # motion isn't mistaken for trajectory drift. huge_loop_realsense tag #17 is + # dynamic; all others are static. + if ignore_tags: + dropped = sorted(tag_id for tag_id in sightings if tag_id in ignore_tags) + for tag_id in dropped: + del sightings[tag_id] + if dropped: + print(f"ignoring tags {dropped} (declared dynamic/unreliable)") + n_sightings = sum(len(times) for times in sightings.values()) + if sightings: + print(f"april tags ({tag_source}): {n_sightings} sightings across ids {sorted(sightings)}") + else: + print("no April tags (camera/intrinsics absent or none detected) — voxel agreement only") + + started = time.monotonic() + graph, closures, replay_stats = run_module_graph( + db_path, + module_class, + pgo_config, + lidar_stream=lidar_stream, + odom_stream=odom_stream, + lockstep=lockstep, + drift_per_sec=drift_per_sec, + drift_t0=drift_t0, + ) + runtime_s = time.monotonic() - started + if not graph: + raise SystemExit(f"{module_name} produced an empty pose graph") + + # The module solved on drifted input, so its graph lives in the drifted + # world; the raw baselines must be drifted to match (see drift_per_sec). + raw_xyz_lookup = drifted_lookup(odometry_lookup(db_path, odom_stream), drift_per_sec, drift_t0) + raw_pose7_lookup = drifted_lookup( + odometry_pose7_lookup(db_path, odom_stream), drift_per_sec, drift_t0 + ) + + xyz_graph = [(node[0], node[1], node[2], node[3]) for node in graph] + if sightings: + raw_tag_positions, corrected_tag_positions = paired_tag_visit_positions( + sightings, + raw_xyz_lookup, + graph_lookup(xyz_graph), + gap_s=VISIT_GAP_S, + ) + raw_report = agreement_report(raw_tag_positions) + corrected_report = agreement_report(corrected_tag_positions) + improvement: float | None = agreement_improvement(raw_report, corrected_report) + else: + raw_report = agreement_report({}) + corrected_report = agreement_report({}) + improvement = None # no tags — tag agreement is N/A for this cell + + voxel_stride = max(1, -(-stream_count(db_path, lidar_stream) // VOXEL_MAX_SCANS)) + voxel = lidar_voxel_agreement( + ( + (timestamp, cloud.points_f32()) + for timestamp, cloud in iterate_stream(db_path, lidar_stream, stride=voxel_stride) + ), + raw_pose7_lookup, + graph, + drift_per_sec=drift_per_sec, + drift_t0=drift_t0, + ) + + # Drift-recovery ATE: corrected trajectory vs the UN-drifted ground truth + # (the odom before drift was injected). Only meaningful with --drift-per-sec; + # the right metric where tag/voxel agreement is weak (e.g. KITTI's long loop). + trajectory = trajectory_recovery_error( + graph, odometry_lookup(db_path, odom_stream), drift_per_sec, drift_t0 + ) + if trajectory is not None: + print( + f" drift recovery: {trajectory['drifted_ate_m']:.2f}" + f" -> {trajectory['corrected_ate_m']:.2f} m ATE" + f" ({trajectory['trajectory_improvement']:+.3f})" + ) + + # Key by package + class — several loop-closure modules are all named PGO. + # results_suffix (dot-joined, NOT "__" which delimits the recording name) + # separates runs that differ in inputs, e.g. fastlio vs pointlio odometry. + module_package = module_class.__module__.rsplit(".", 2)[-2] + module_key = f"{module_package}.{module_name}" + ( + f".{results_suffix}" if results_suffix else "" + ) + # db.parent.name is the recording dir for go2; LFS dbs (hk_village) sit + # directly in data/, so an explicit recording_name avoids cell collisions. + out_dir = RESULTS_DIR / f"{recording_name or db_path.parent.name}__{module_key}" + out_dir.mkdir(parents=True, exist_ok=True) + rrd_path = out_dir / "eval.rrd" + if with_rrd: + raw_positions = np.asarray( + [ + [pose.position.x, pose.position.y, pose.position.z] + for _, payload in iterate_stream(db_path, odom_stream) + for pose in [RateReplay._payload_pose(payload)] + ], + dtype=np.float64, + ) + write_trajectory_rrd(rrd_path, raw_positions, graph) + + summary = { + "db": str(db_path), + "odom_stream": odom_stream, + "camera_stream": camera_stream, + "lidar_stream": lidar_stream, + "module": {"path": str(module_path), "name": module_name}, + "pgo_config": pgo_config, + "drift_per_sec": list(drift_per_sec), + "fingerprint": cell_fingerprint( + db_path, pgo_config, lidar_stream, odom_stream, drift_per_sec + ), + "replay": replay_stats, + "april_tags": { + "source": tag_source, + "sightings": n_sightings, + "ids": sorted(sightings), + }, + "scores": { + "raw_spread_m": raw_report.mean_spread if sightings else None, + "corrected_spread_m": corrected_report.mean_spread if sightings else None, + "tag_improvement": improvement, + "voxel_improvement": voxel.get("improvement"), + "trajectory_improvement": trajectory["trajectory_improvement"] if trajectory else None, + "drifted_ate_m": trajectory["drifted_ate_m"] if trajectory else None, + "corrected_ate_m": trajectory["corrected_ate_m"] if trajectory else None, + "closures": closures, + "keyframes": len(graph), + "runtime_s": round(runtime_s, 1), + }, + "raw_agreement": _report_dict(raw_report), + "corrected_agreement": _report_dict(corrected_report), + "voxel_agreement": voxel, + "rrd": str(rrd_path) if with_rrd else None, + "evaluated_at": time.strftime("%Y-%m-%d %H:%M:%S"), + } + (out_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + + print(f"\nresults -> {out_dir / 'summary.json'}") + if sightings: + print( + f" tag spread: {raw_report.mean_spread:.3f}" + f" -> {corrected_report.mean_spread:.3f} m" + ) + print(f" tag improvement: {improvement:+.3f} (1.0 = perfect)") + else: + print(" tag improvement: n/a (no tags)") + if voxel.get("status") == "ok": + print( + f" voxel agreement: {voxel['raw_voxels']} -> {voxel['corrected_voxels']} voxels" + f" ({voxel['improvement']:+.3f}, {voxel['scans_used']} scans @ {voxel['voxel_size_m']}m)" + ) + else: + print(f" voxel agreement: {voxel.get('status')}") + print(f" closures: {closures}, keyframes: {len(graph)}") + if with_rrd: + print(f" rrd: {rrd_path}") + return summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--db-path", type=Path, required=True) + parser.add_argument("--odom-stream", required=True) + parser.add_argument( + "--camera-stream", default=None, help="omit for tagless datasets (voxel agreement only)" + ) + parser.add_argument( + "--camera-intrinsics-json-path", + type=Path, + default=None, + help="omit for tagless datasets (voxel agreement only)", + ) + parser.add_argument("--module-path", type=Path, required=True) + parser.add_argument("--module-name", required=True) + parser.add_argument( + "--pgo-config-json", + help="inline JSON of module config overrides (default: scan_context variant)", + ) + parser.add_argument("--with-rrd", default="false", choices=["true", "false"]) + parser.add_argument( + "--lidar-stream", + default="fastlio_lidar", + help="lidar stream replayed into the module alongside the odometry", + ) + parser.add_argument( + "--lockstep", + default="true", + choices=["true", "false"], + help="pace scans on corrected_odometry acks (machine-independent); false = fixed-rate", + ) + parser.add_argument( + "--results-suffix", + default="", + help="extra results-dir key for runs with different inputs (e.g. pointlio)", + ) + parser.add_argument( + "--recording-name", + default=None, + help="results-dir recording key (default: db parent dir name)", + ) + parser.add_argument( + "--drift-per-sec", + default=None, + help="inject odom drift as a constant world velocity 'x,y,z' in m/s " + "(offset = this * (t - t0), added to odom+lidar). e.g. '0.01,0,0'", + ) + parser.add_argument( + "--ignore-tags", + default=None, + help="comma-separated April-tag ids to drop from scoring (dynamic/unreliable " + "tags whose motion would look like drift). e.g. '17'", + ) + args = parser.parse_args() + + drift_per_sec = ( + [float(v) for v in args.drift_per_sec.split(",")] if args.drift_per_sec else None + ) + if drift_per_sec is not None and len(drift_per_sec) != 3: + raise SystemExit(f"--drift-per-sec must be 'x,y,z', got {args.drift_per_sec!r}") + + ignore_tags = ( + {int(tag_id) for tag_id in args.ignore_tags.split(",")} if args.ignore_tags else None + ) + + db_path = args.db_path.expanduser() + if not db_path.exists(): + raise SystemExit(f"no such db: {db_path}") + intrinsics_json = ( + args.camera_intrinsics_json_path.expanduser() + if args.camera_intrinsics_json_path is not None + else None + ) + if intrinsics_json is not None and not intrinsics_json.exists(): + raise SystemExit(f"no such intrinsics json: {intrinsics_json}") + + pgo_config = dict(DEFAULT_PGO_CONFIG) + if args.pgo_config_json: + pgo_config.update(json.loads(args.pgo_config_json)) + + evaluate( + db_path, + odom_stream=args.odom_stream, + camera_stream=args.camera_stream, + intrinsics_json=intrinsics_json, + module_path=args.module_path, + module_name=args.module_name, + pgo_config=pgo_config, + with_rrd=args.with_rrd == "true", + lidar_stream=args.lidar_stream, + lockstep=args.lockstep == "true", + results_suffix=args.results_suffix, + recording_name=args.recording_name, + drift_per_sec=drift_per_sec, + ignore_tags=ignore_tags, + ) + + +if __name__ == "__main__": + # Re-import under the canonical dotted name so module classes defined here + # (GraphCapture) deploy into workers with a picklable __module__. + importlib.import_module("dimos.navigation.jnav.components.loop_closure.eval").main() diff --git a/dimos/navigation/jnav/components/loop_closure/eval_all.py b/dimos/navigation/jnav/components/loop_closure/eval_all.py new file mode 100644 index 0000000000..47df33d081 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/eval_all.py @@ -0,0 +1,366 @@ +# 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. + +"""Run the loop-closure eval over every PGO module and render a comparison table. + +Each module is evaluated in its own subprocess (sequentially — they share LCM) +via eval.py with lockstep replay, then all eval_results summaries are rendered +to eval_results/comparison.md. A failed module doesn't stop the rest. + +Usage: + uv run python dimos/navigation/jnav/components/loop_closure/eval_all.py \\ + --db-path ~/datasets/go2_recordings//mem2.db \\ + [--camera-intrinsics-json-path ] # default: sidecar next to db + [--only gsc_pgo,ivan_pgo] # subset by directory name + [--with-rrd true] [--lockstep true] +""" + +from __future__ import annotations + +import argparse +from collections import defaultdict +import json +from pathlib import Path +import subprocess +import sys +from typing import Any + +LOOP_CLOSURE_DIR = Path(__file__).resolve().parent +RESULTS_DIR = LOOP_CLOSURE_DIR / "eval_results" +TABLE_PATH = RESULTS_DIR / "comparison.md" +SIDECAR_NAME = "camera_intrinsics.json" + +# (directory, per-module config overrides). All classes are named PGO. +# global_map is disabled everywhere — nothing in the eval consumes it, and the +# native modules' big global_map clouds overflow LCM and congest the +# corrected_odometry ack channel the lockstep replay waits on (rate 0 = off +# for the native binaries; publish_global_map=False for the python ivan). +MODULES: list[tuple[str, dict[str, Any]]] = [ + ("gsc_pgo", {"use_scan_context": True, "global_map_publish_rate": 0.0}), + ("ivan_pgo", {"publish_global_map": False}), + ("ivan_pgo_transformer", {}), + ("unrefined_pgo", {"global_map_publish_rate": 0.0}), +] + + +def _fmt_spread(s: dict[str, Any]) -> str: + raw, corrected = s.get("raw_spread_m"), s.get("corrected_spread_m") + if raw is None or corrected is None: # tagless dataset (KITTI, bare lidar) + return "—" + return f"{raw:.2f} -> {corrected:.2f}" + + +TABLE_COLUMNS = [ + ("tag spread (m)", _fmt_spread), + ( + "tag improvement", + lambda s: f"{s['tag_improvement']:+.3f}" if s.get("tag_improvement") is not None else "—", + ), + ( + "voxel improvement", + lambda s: f"{s['voxel_improvement']:+.3f}" + if s.get("voxel_improvement") is not None + else "—", + ), + ( + "drift recovery", # ATE improvement vs un-drifted GT (only with --drift-per-sec) + lambda s: f"{s['trajectory_improvement']:+.3f}" + if s.get("trajectory_improvement") is not None + else "—", + ), + ("closures", lambda s: str(s["closures"])), + ("keyframes", lambda s: str(s["keyframes"])), + ("runtime (s)", lambda s: str(s["runtime_s"])), +] + + +def run_one( + module_dir: str, + overrides: dict[str, Any], + *, + db_path: Path, + odom_stream: str, + camera_stream: str, + lidar_stream: str, + intrinsics_json: Path | None, + with_rrd: str, + lockstep: str, + results_suffix: str = "", + drift_per_sec: str = "", + ignore_tags: str = "", +) -> bool: + command = [ + sys.executable, + "-u", + str(LOOP_CLOSURE_DIR / "eval.py"), + "--db-path", + str(db_path), + "--odom-stream", + odom_stream, + "--camera-stream", + camera_stream, + "--lidar-stream", + lidar_stream, + "--module-path", + str(LOOP_CLOSURE_DIR / module_dir / "module.py"), + "--module-name", + "PGO", + "--with-rrd", + with_rrd, + "--lockstep", + lockstep, + ] + if intrinsics_json is not None: + command += ["--camera-intrinsics-json-path", str(intrinsics_json)] + if drift_per_sec: + command += ["--drift-per-sec", drift_per_sec] + if ignore_tags: + command += ["--ignore-tags", ignore_tags] + if results_suffix: + command += ["--results-suffix", results_suffix] + if overrides: + command += ["--pgo-config-json", json.dumps(overrides)] + print(f"\n=== {module_dir} ===", flush=True) + result = subprocess.run(command, check=False) + print(f"=== {module_dir} exit: {result.returncode} ===", flush=True) + return result.returncode == 0 + + +def _row_label(module_key: str, replay: dict[str, Any]) -> str: + label = module_key + scans = replay.get("scans_sent") + timeouts = replay.get("timeouts") + if scans is not None: + label += f" ({scans} scans, {timeouts} ack-timeouts)" if timeouts else f" ({scans} scans)" + if replay.get("hit_max_run_s"): + label += " [hit run cap]" + return label + + +def _section_for(recording: str, module_key: str) -> str: + """Which breakdown section a result belongs to.""" + if recording.startswith("hk_village"): + return "hk_village" + if "drift" in module_key: + return "drift" + return "real" + + +def _tag_improvement_key(labelled_row: tuple[str, dict[str, Any]]) -> float: + # Best first; fall back to drift-recovery for tagless rows, then sink to the + # bottom if neither is present. + scores = labelled_row[1] + value = scores.get("tag_improvement") + if value is None: + value = scores.get("trajectory_improvement") + return float(value) if value is not None else float("-inf") + + +def _render_section( + by_recording: dict[str, list[tuple[str, dict[str, Any]]]], sort_reverse: bool +) -> list[str]: + lines: list[str] = [] + for recording in sorted(by_recording): + rows = sorted(by_recording[recording], key=_tag_improvement_key, reverse=sort_reverse) + lines += [f"### {recording}", ""] + lines.append("| module | " + " | ".join(name for name, _ in TABLE_COLUMNS) + " | replay |") + lines.append("|" + "---|" * (len(TABLE_COLUMNS) + 2)) + for module_label, scores in rows: + cells = [render(scores) for _, render in TABLE_COLUMNS] + lines.append(f"| {module_label} | " + " | ".join(cells) + f" | {scores['_mode']} |") + lines.append("") + return lines + + +def _base_module(module_key: str) -> str: + """`gsc_pgo.PGO.drift0p05_0_0` -> `gsc_pgo`.""" + return module_key.split(".")[0] + + +def _conclusion(real: dict[str, list[tuple[str, dict[str, Any]]]]) -> list[str]: + """Data-driven summary: count per-recording tag-improvement wins per module.""" + wins: dict[str, int] = defaultdict(int) + tagged_recordings = 0 + for rows in real.values(): + tagged = [(label, s) for label, s in rows if s.get("tag_improvement") is not None] + if not tagged: + continue + tagged_recordings += 1 + best_label = max(tagged, key=lambda item: item[1]["tag_improvement"])[0] + wins[_base_module(best_label)] += 1 + if not tagged_recordings: + return [] + ranking = sorted(wins.items(), key=lambda item: item[1], reverse=True) + leader = ranking[0][0] if ranking else "n/a" + lines = ["## Conclusion", ""] + lines.append( + f"Across {tagged_recordings} tagged real recordings, the best per-recording " + f"April-tag improvement was won by:" + ) + lines.append("") + for module, count in ranking: + lines.append(f"- **{module}** — best on {count}/{tagged_recordings}") + lines += [ + "", + f"`{leader}` is the strongest loop-closure PGO overall on real recordings; " + "`unrefined_pgo` (the pass-through baseline) is the floor. On tagless data " + "(hk_village) modules are ranked by voxel improvement, and under artificial " + "drift by drift-recovery (fraction of injected ATE removed).", + "", + ] + return lines + + +def render_table() -> Path: + buckets: dict[str, dict[str, list[tuple[str, dict[str, Any]]]]] = { + "real": defaultdict(list), + "hk_village": defaultdict(list), + "drift": defaultdict(list), + } + for summary_path in sorted(RESULTS_DIR.glob("*/summary.json")): + summary = json.loads(summary_path.read_text()) + # Skip non-PGO summaries (e.g. the cross-recording ground-truth tag eval, + # which has a different shape and its own write-up). + if "scores" not in summary: + continue + recording, _, module_key = summary_path.parent.name.rpartition("__") + replay = summary.get("replay", {}) + row = {**summary["scores"], "_mode": replay.get("mode", "fixed-rate (pre-lockstep)")} + section = _section_for(recording, module_key) + buckets[section][recording].append((_row_label(module_key, replay), row)) + + lines = [ + "# Loop-closure PGO comparison", + "", + "Higher improvement = better. **Tag improvement**: fractional drop in per-visit", + "April-tag position spread (1.0 = perfect). **Voxel improvement**: fractional drop", + "in occupied 0.2 m voxels after re-anchoring scans onto the corrected trajectory.", + "**Drift recovery**: fraction of injected ATE removed vs the un-drifted ground truth.", + "", + ] + if buckets["real"]: + lines += ["## Real recordings (clean input)", ""] + lines += _render_section(buckets["real"], sort_reverse=True) + if buckets["hk_village"]: + lines += ["## hk_village (tagless — voxel agreement only)", ""] + lines += _render_section(buckets["hk_village"], sort_reverse=True) + if buckets["drift"]: + lines += ["## Artificial drift robustness", ""] + lines += _render_section(buckets["drift"], sort_reverse=True) + + kitti_table = RESULTS_DIR / "kitti_comparison.md" + if kitti_table.exists(): + lines += [ + "## KITTI (official odometry error)", + "", + "Scored with the official KITTI translational (%) / rotational (deg/m) error " + "(lower = better); see [kitti_comparison.md](kitti_comparison.md).", + "", + ] + + lines += _conclusion(buckets["real"]) + RESULTS_DIR.mkdir(exist_ok=True) + TABLE_PATH.write_text("\n".join(lines)) + return TABLE_PATH + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--db-path", type=Path, required=True) + parser.add_argument("--odom-stream", default="fastlio_odometry") + parser.add_argument("--camera-stream", default="color_image") + parser.add_argument("--lidar-stream", default="fastlio_lidar") + parser.add_argument( + "--camera-intrinsics-json-path", + type=Path, + help=f"default: {SIDECAR_NAME} next to the db", + ) + parser.add_argument("--only", help="comma-separated module directory names") + parser.add_argument("--with-rrd", default="true", choices=["true", "false"]) + parser.add_argument("--lockstep", default="true", choices=["true", "false"]) + parser.add_argument( + "--results-suffix", + default="", + help="extra results-dir key for runs with different inputs (e.g. pointlio)", + ) + parser.add_argument( + "--drift-per-sec", + default="", + help="inject constant-velocity world drift 'x,y,z' (m/s) into odom+lidar; " + "stress-tests how each PGO corrects known drift. Auto-tags the results dir.", + ) + parser.add_argument( + "--ignore-tags", + default="", + help="comma-separated April-tag ids to drop from scoring (e.g. '17' for the " + "dynamic tag in huge_loop_realsense)", + ) + args = parser.parse_args() + + # Keep drifted runs in their own results dirs so they don't clobber the + # clean-input comparison (e.g. drift 0.1,0,0 -> suffix '...drift0p1x'). + results_suffix = args.results_suffix + if args.drift_per_sec: + drift_tag = "drift" + args.drift_per_sec.replace(".", "p").replace(",", "_") + results_suffix = f"{results_suffix}.{drift_tag}" if results_suffix else drift_tag + + db_path = args.db_path.expanduser() + if not db_path.exists(): + raise SystemExit(f"no such db: {db_path}") + # Tagless datasets (KITTI, bare lidar): empty --camera-stream -> no intrinsics + # needed, voxel-agreement scoring only. + tagless = not args.camera_stream + intrinsics_json: Path | None = None + if not tagless: + sidecar: Path = args.camera_intrinsics_json_path or db_path.parent / SIDECAR_NAME + intrinsics_json = sidecar.expanduser() + if not intrinsics_json.exists(): + raise SystemExit(f"no intrinsics json: {intrinsics_json}") + + selected = MODULES + if args.only: + wanted = {name.strip() for name in args.only.split(",")} + selected = [(name, overrides) for name, overrides in MODULES if name in wanted] + missing = wanted - {name for name, _ in selected} + if missing: + raise SystemExit( + f"unknown modules: {sorted(missing)} (have: {[m for m, _ in MODULES]})" + ) + + outcomes = {} + for module_dir, overrides in selected: + outcomes[module_dir] = run_one( + module_dir, + overrides, + db_path=db_path, + odom_stream=args.odom_stream, + camera_stream=args.camera_stream, + lidar_stream=args.lidar_stream, + intrinsics_json=intrinsics_json, + with_rrd=args.with_rrd, + lockstep=args.lockstep, + results_suffix=results_suffix, + drift_per_sec=args.drift_per_sec, + ignore_tags=args.ignore_tags, + ) + + table = render_table() + print(f"\ntable -> {table}") + failed = [name for name, ok in outcomes.items() if not ok] + if failed: + raise SystemExit(f"failed modules: {failed}") + + +if __name__ == "__main__": + main() diff --git a/dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py b/dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py new file mode 100644 index 0000000000..5644fb8add --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py @@ -0,0 +1,309 @@ +# 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. + +"""Cross-recording ground-truth tag eval (todo item 4). + +Run the best PGO on huge_loop_realsense to get corrected April-tag world +positions (treated as ground truth, since cmu reaches ~0.14 m tag spread there), +then measure how close each PGO moves huge_loop_go2's tags toward those GT +locations — a CROSS-recording metric, unlike the within-recording agreement. + +Both recordings see the same physical tags. The pipeline per recording: + 1. run a PGO -> corrected keyframe trajectory (graph). + 2. read each tag sighting's pose-in-optical from the recorded april_tags + stream (PoseStamped: t_optical_tag). + 3. place each sighting in the world: + tag_world = T_world_base(corrected, t) · T_base_optical · t_optical_tag + where T_base_optical = the recording's `optical_in_base` extrinsic. + 4. average each tag's world positions -> one centroid per tag (the estimate). +Then Umeyama-align go2's tag constellation onto the realsense GT constellation +over shared tag ids; residual RMSE = how well the PGO recovered the true tag +geometry (lower = better; an uncorrected / wrong PGO leaves higher residual). + +Tag 17 is dynamic on huge_loop_realsense; it is dropped via --ignore-tags. + +Results are written to eval_results/__ground_truth_tag/summary.json +(per-module residual + closures, plus the GT tag ids), and the frame-composition +core is covered by --self-test. +""" + +from __future__ import annotations + +import argparse +from collections import defaultdict +from collections.abc import Callable +import json +from pathlib import Path +from typing import Any + +import numpy as np +from scipy.spatial.transform import Rotation + +from dimos.navigation.jnav.components.loop_closure.eval import ( + odometry_pose7_lookup, + run_module_graph, +) +from dimos.navigation.jnav.utils.module_loading import ( + filter_config_for_module, + load_module_class, +) +from dimos.navigation.jnav.utils.recording_db import store +from dimos.navigation.jnav.utils.trajectory_metrics import ( + drift_delta_lookup, + rigid_align_rmse, +) + + +def pose7_to_matrix(pose7: np.ndarray | list[float]) -> np.ndarray: + """[x,y,z, qx,qy,qz,qw] -> 4x4 homogeneous transform.""" + pose = np.asarray(pose7, dtype=np.float64) + transform = np.eye(4) + transform[:3, :3] = Rotation.from_quat(pose[3:7]).as_matrix() + transform[:3, 3] = pose[:3] + return transform + + +def tag_world_position( + robot_pose7: np.ndarray | list[float], + optical_in_base7: np.ndarray | list[float], + tag_in_optical7: np.ndarray | list[float], +) -> np.ndarray: + """World xyz of a tag from the corrected robot pose, the camera->base + extrinsic, and the tag-in-optical detection. Composes + T_world_base · T_base_optical · T_optical_tag and returns the translation.""" + transform = ( + pose7_to_matrix(robot_pose7) + @ pose7_to_matrix(optical_in_base7) + @ pose7_to_matrix(tag_in_optical7) + ) + return transform[:3, 3] + + +def read_optical_in_base(intrinsics_json: Path) -> list[float]: + raw = json.loads(Path(intrinsics_json).read_text()) + extrinsic = raw.get("optical_in_base") + if extrinsic is None: + raise SystemExit(f"no optical_in_base (camera->base extrinsic) in {intrinsics_json}") + return list(extrinsic) + + +def tag_in_camera_sightings(db_path: Path) -> list[tuple[float, int, list[float]]]: + """(ts, marker_id, tag-in-optical pose7) per April-tag observation, read + straight from the recorded april_tags PoseStamped stream (no re-detection).""" + sightings: list[tuple[float, int, list[float]]] = [] + for observation in store(db_path).stream("april_tags"): + pose_tuple = observation.pose_tuple + if pose_tuple is None: + continue + sightings.append( + (float(observation.ts), int(observation.tags["marker_id"]), list(pose_tuple)) + ) + return sightings + + +def corrected_pose7_at( + timestamp: float, + raw_pose7_lookup: Callable[[float], np.ndarray | None], + delta_lookup: Callable[[float], tuple[np.ndarray, np.ndarray] | None], +) -> np.ndarray | None: + """Corrected robot pose7 at a time: apply the nearest keyframe's drift + correction (R_delta, t_delta) to the raw odom pose.""" + raw = raw_pose7_lookup(timestamp) + delta = delta_lookup(timestamp) + if raw is None or delta is None: + return None + rotation_delta, translation_delta = delta + raw = np.asarray(raw, dtype=np.float64) + rotation_corrected = rotation_delta @ Rotation.from_quat(raw[3:7]).as_matrix() + translation_corrected = rotation_delta @ raw[:3] + translation_delta + return np.array([*translation_corrected, *Rotation.from_matrix(rotation_corrected).as_quat()]) + + +def tag_constellation( + db_path: Path, + module_path: Path, + module_name: str, + config: dict[str, Any], + *, + lidar_stream: str, + odom_stream: str, + optical_in_base: list[float], + ignore_tags: set[int], +) -> tuple[dict[int, np.ndarray], int]: + """Run the PGO, then place each tag sighting in the world via the corrected + trajectory + extrinsic, averaging per tag. Returns {marker_id: world centroid} + and the closure count.""" + module_class = load_module_class(module_path, module_name) + config = filter_config_for_module(module_class, config) + graph, closures, _ = run_module_graph( + db_path, + module_class, + config, + lidar_stream=lidar_stream, + odom_stream=odom_stream, + lockstep=True, + ) + raw_pose7_lookup = odometry_pose7_lookup(db_path, odom_stream) + delta_lookup = drift_delta_lookup(graph, raw_pose7_lookup) + by_tag: dict[int, list[np.ndarray]] = defaultdict(list) + for timestamp, marker_id, tag_pose7 in tag_in_camera_sightings(db_path): + if marker_id in ignore_tags: + continue + robot_pose7 = corrected_pose7_at(timestamp, raw_pose7_lookup, delta_lookup) + if robot_pose7 is None: + continue + by_tag[marker_id].append(tag_world_position(robot_pose7, optical_in_base, tag_pose7)) + constellation = {tag: np.mean(positions, axis=0) for tag, positions in by_tag.items()} + return constellation, closures + + +def constellation_residual( + gt: dict[int, np.ndarray], test: dict[int, np.ndarray] +) -> tuple[float, list[int]]: + """Umeyama-align the test tag constellation onto the GT; residual RMSE over + shared tags (lower = the PGO recovered the true tag geometry better).""" + shared = sorted(set(gt) & set(test)) + if len(shared) < 3: + return float("nan"), shared + gt_points = np.asarray([gt[tag] for tag in shared]) + test_points = np.asarray([test[tag] for tag in shared]) + return rigid_align_rmse(test_points, gt_points), shared + + +def _self_test() -> None: + identity = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] + # 1) everything identity, tag 2m ahead in optical -> world = 2m ahead. + world = tag_world_position(identity, identity, [0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 1.0]) + assert np.allclose(world, [0.0, 0.0, 2.0]), world + + # 2) robot translated by (10,5,0), no rotation: tag world shifts by that. + world = tag_world_position( + [10.0, 5.0, 0.0, 0.0, 0.0, 0.0, 1.0], identity, [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0] + ) + assert np.allclose(world, [11.0, 5.0, 0.0]), world + + # 3) robot yawed 90deg (+z): a tag 1m ahead in base (x) lands 1m to +y of robot. + yaw90 = [0.0, 0.0, 0.0, *Rotation.from_euler("z", 90, degrees=True).as_quat()] + world = tag_world_position(yaw90, identity, [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]) + assert np.allclose(world, [0.0, 1.0, 0.0], atol=1e-9), world + + # 4) optical->base rotation (-0.5,0.5,-0.5,0.5) maps optical +z (forward) to + # base +x (forward). Tag 3m ahead in optical -> 3m ahead in base/world. + optical_in_base = [0.0, 0.0, 0.0, -0.5, 0.5, -0.5, 0.5] + world = tag_world_position(identity, optical_in_base, [0.0, 0.0, 3.0, 0.0, 0.0, 0.0, 1.0]) + assert np.allclose(world, [3.0, 0.0, 0.0], atol=1e-9), world + + print("frame-composition self-test PASSED (4 cases)") + + +LOOP_CLOSURE_DIR = Path(__file__).resolve().parent +DEFAULT_GT_CONFIG = {"use_scan_context": True, "global_map_publish_rate": 0} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--gt-db", type=Path, required=True, help="recording for ground truth (realsense)" + ) + parser.add_argument("--gt-intrinsics", type=Path, required=True) + parser.add_argument("--gt-odom-stream", default="fastlio_odometry") + parser.add_argument("--gt-lidar-stream", default="fastlio_lidar") + parser.add_argument("--test-db", type=Path, required=True, help="recording to score (go2)") + parser.add_argument("--test-intrinsics", type=Path, required=True) + parser.add_argument("--test-odom-stream", default="fastlio_odometry") + parser.add_argument("--test-lidar-stream", default="fastlio_lidar") + parser.add_argument("--gt-module", default="gsc_pgo", help="PGO dir used to build GT") + parser.add_argument("--modules", default="gsc_pgo,ivan_pgo,ivan_pgo_transformer,unrefined_pgo") + parser.add_argument("--ignore-tags", default="17", help="dynamic tags to drop") + args = parser.parse_args() + + ignore_tags = {int(t) for t in args.ignore_tags.split(",")} if args.ignore_tags else set() + gt_extrinsic = read_optical_in_base(args.gt_intrinsics) + test_extrinsic = read_optical_in_base(args.test_intrinsics) + + def module_py(module_dir: str) -> Path: + return LOOP_CLOSURE_DIR / module_dir / "module.py" + + print(f"building GT constellation: {args.gt_module} on {args.gt_db.parent.name}") + gt_constellation, _ = tag_constellation( + args.gt_db, + module_py(args.gt_module), + "PGO", + dict(DEFAULT_GT_CONFIG), + lidar_stream=args.gt_lidar_stream, + odom_stream=args.gt_odom_stream, + optical_in_base=gt_extrinsic, + ignore_tags=ignore_tags, + ) + print(f" GT tags: {sorted(gt_constellation)}") + + print(f"\n{'module':24s} {'tag-to-GT residual (m)':>22s} {'shared tags':>12s} {'closures':>9s}") + print("-" * 70) + rows: list[dict[str, Any]] = [] + for module_dir in [m.strip() for m in args.modules.split(",")]: + try: + test_constellation, closures = tag_constellation( + args.test_db, + module_py(module_dir), + "PGO", + dict(DEFAULT_GT_CONFIG), + lidar_stream=args.test_lidar_stream, + odom_stream=args.test_odom_stream, + optical_in_base=test_extrinsic, + ignore_tags=ignore_tags, + ) + residual, shared = constellation_residual(gt_constellation, test_constellation) + rows.append( + { + "module": module_dir, + "residual_m": residual, + "shared_tags": len(shared), + "closures": closures, + } + ) + print(f"{module_dir:24s} {residual:>22.3f} {len(shared):>12d} {closures:>9d}") + except Exception as error: + print(f"{module_dir:24s} FAILED: {type(error).__name__}: {error}") + rows.append({"module": module_dir, "residual_m": None, "error": str(error)}) + + scored = [row for row in rows if row.get("residual_m") is not None] + best_module = min(scored, key=lambda row: row["residual_m"])["module"] if scored else None + if best_module is not None: + best = next(row for row in scored if row["module"] == best_module) + print( + f"\nbest (lowest residual to GT tag geometry): {best_module} ({best['residual_m']:.3f} m)" + ) + + out_dir = LOOP_CLOSURE_DIR / "eval_results" / f"{args.test_db.parent.name}__ground_truth_tag" + out_dir.mkdir(parents=True, exist_ok=True) + summary = { + "gt_db": str(args.gt_db), + "gt_module": args.gt_module, + "test_db": str(args.test_db), + "ignore_tags": sorted(ignore_tags), + "gt_tags": sorted(int(tag) for tag in gt_constellation), + "best_module": best_module, + "rows": rows, + } + (out_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + print(f"\nresults -> {out_dir / 'summary.json'}") + + +if __name__ == "__main__": + import sys + + if "--self-test" in sys.argv: + _self_test() + else: + main() diff --git a/dimos/navigation/jnav/components/loop_closure/eval_kitti.py b/dimos/navigation/jnav/components/loop_closure/eval_kitti.py new file mode 100644 index 0000000000..cf7e53fe32 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/eval_kitti.py @@ -0,0 +1,192 @@ +# 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. + +"""Evaluate one PGO on a KITTI db with the official KITTI odometry error. + +Parallels eval.py, but for the kitti_to_db.py recordings. Runs the module over +the db's drifty ICP odometry (fastlio_odometry) + registered scans (fastlio_lidar), +reads the ground truth from the db's gt_odometry stream, and reports translational +(%) and rotational (deg/m) error — the leaderboard metric — for the corrected +trajectory and the raw-odometry baseline. + +Usage: + uv run python dimos/navigation/jnav/components/loop_closure/eval_kitti.py \ + --db-path ~/datasets/kitti/kitti_seq07/mem2.db \ + --module-path dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py \ + [--module-name PGO] [--pgo-config-json '{...}'] +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np +from scipy.spatial.transform import Rotation + +from dimos.navigation.jnav.components.loop_closure.eval import run_module_graph +from dimos.navigation.jnav.utils.kitti import kitti_odometry_error +from dimos.navigation.jnav.utils.module_loading import ( + filter_config_for_module, + load_module_class, +) +from dimos.navigation.jnav.utils.recording_db import iterate_stream + +DEFAULT_CONFIG = {"use_scan_context": True, "global_map_publish_rate": 0} + +# Stream names produced by kitti_to_db (overridable on the CLI). +DEFAULT_LIDAR_STREAM = "fastlio_lidar" +DEFAULT_ODOM_STREAM = "fastlio_odometry" +DEFAULT_GT_STREAM = "gt_odometry" + + +def poses_from_stream(db_path: Path, stream: str) -> tuple[list[np.ndarray], list[float]]: + """Read an Odometry stream as (4x4 poses, timestamps).""" + poses: list[np.ndarray] = [] + times: list[float] = [] + for timestamp, message in iterate_stream(db_path, stream): + orientation = message.pose.orientation + position = message.pose.position + transform = np.eye(4) + transform[:3, :3] = Rotation.from_quat( + [orientation.x, orientation.y, orientation.z, orientation.w] + ).as_matrix() + transform[:3, 3] = [position.x, position.y, position.z] + poses.append(transform) + times.append(timestamp) + return poses, times + + +def _gt_at( + gt_poses: list[np.ndarray], gt_times: list[float], query: list[float] +) -> list[np.ndarray]: + times = np.asarray(gt_times) + return [gt_poses[int(np.argmin(np.abs(times - t)))] for t in query] + + +def corrected_trajectory( + db_path: Path, + module_path: Path, + module_name: str, + config: dict[str, Any], + *, + lidar_stream: str, + odom_stream: str, +): + module_class = load_module_class(module_path, module_name) + config = filter_config_for_module(module_class, config) + graph, closures, _ = run_module_graph( + db_path, + module_class, + config, + lidar_stream=lidar_stream, + odom_stream=odom_stream, + lockstep=True, + ) + poses, times = [], [] + for node in graph: + transform = np.eye(4) + transform[:3, :3] = Rotation.from_quat(node[4:8]).as_matrix() + transform[:3, 3] = node[1:4] + poses.append(transform) + times.append(node[0]) + return poses, times, closures + + +def evaluate( + db_path: Path, + module_path: Path, + module_name: str, + config: dict[str, Any], + results_suffix: str = "", + *, + lidar_stream: str = DEFAULT_LIDAR_STREAM, + odom_stream: str = DEFAULT_ODOM_STREAM, + gt_stream: str = DEFAULT_GT_STREAM, +) -> dict[str, Any]: + gt_poses, gt_times = poses_from_stream(db_path, gt_stream) + odom_poses, _ = poses_from_stream(db_path, odom_stream) + if not gt_poses: + raise SystemExit(f"no {gt_stream!r} stream in {db_path}") + baseline = kitti_odometry_error(odom_poses, gt_poses) + + corrected, corrected_times, closures = corrected_trajectory( + db_path, + module_path, + module_name, + config, + lidar_stream=lidar_stream, + odom_stream=odom_stream, + ) + error = kitti_odometry_error(corrected, _gt_at(gt_poses, gt_times, corrected_times)) + + print( + f"raw odometry: {baseline['translational_percent']:.2f}% transl / " + f"{baseline['rotational_deg_per_m']:.4f} deg/m" + ) + print( + f"after {module_path.parent.name}: {error['translational_percent']:.2f}% transl / " + f"{error['rotational_deg_per_m']:.4f} deg/m ({closures} closures)" + ) + + module_key = module_path.parent.name + (f".{results_suffix}" if results_suffix else "") + summary = { + "db": str(db_path), + "module": module_key, + "scores": { + "translational_percent": error["translational_percent"], + "rotational_deg_per_m": error["rotational_deg_per_m"], + "baseline_translational_percent": baseline["translational_percent"], + "baseline_rotational_deg_per_m": baseline["rotational_deg_per_m"], + "closures": closures, + "keyframes": len(corrected), + }, + } + recording = db_path.parent.name + out_dir = Path(__file__).resolve().parent / "eval_results" / f"{recording}__{module_key}" + out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / "kitti_summary.json").write_text(json.dumps(summary, indent=2)) + return summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--db-path", type=Path, required=True) + parser.add_argument("--module-path", type=Path, required=True) + parser.add_argument("--module-name", default="PGO") + parser.add_argument("--pgo-config-json", default="") + parser.add_argument("--results-suffix", default="") + parser.add_argument("--lidar-stream", default=DEFAULT_LIDAR_STREAM) + parser.add_argument("--odom-stream", default=DEFAULT_ODOM_STREAM) + parser.add_argument("--gt-stream", default=DEFAULT_GT_STREAM) + args = parser.parse_args() + config = dict(DEFAULT_CONFIG) + if args.pgo_config_json: + config.update(json.loads(args.pgo_config_json)) + evaluate( + args.db_path.expanduser(), + args.module_path, + args.module_name, + config, + args.results_suffix, + lidar_stream=args.lidar_stream, + odom_stream=args.odom_stream, + gt_stream=args.gt_stream, + ) + + +if __name__ == "__main__": + main() diff --git a/dimos/navigation/jnav/components/loop_closure/eval_kitti_all.py b/dimos/navigation/jnav/components/loop_closure/eval_kitti_all.py new file mode 100644 index 0000000000..3a16b799d1 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/eval_kitti_all.py @@ -0,0 +1,122 @@ +# 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. + +"""Run every PGO on KITTI db(s) and render an official-error comparison table. + +Parallels eval_all.py. For each KITTI db (one sequence) it evaluates every PGO +module in its own subprocess (they share LCM) via eval_kitti.py, then renders +translational (%) / rotational (deg/m) error per module to kitti_comparison.md, +sorted best-first (lower error = better), alongside the raw-odometry baseline. + +Usage: + uv run python dimos/navigation/jnav/components/loop_closure/eval_kitti_all.py \ + --recordings-dir ~/datasets/kitti/sequences # all kitti_seq*/mem2.db + [--db-path ~/datasets/kitti/sequences/kitti_seq07/mem2.db] # or one + [--only gsc_pgo,ivan_pgo] +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import subprocess +import sys +from typing import Any + +LOOP_CLOSURE_DIR = Path(__file__).resolve().parent +RESULTS_DIR = LOOP_CLOSURE_DIR / "eval_results" +TABLE_PATH = RESULTS_DIR / "kitti_comparison.md" +MODULES = ("gsc_pgo", "ivan_pgo", "ivan_pgo_transformer", "unrefined_pgo") + + +def run_one(db_path: Path, module_dir: str) -> bool: + command = [ + sys.executable, + "-u", + str(LOOP_CLOSURE_DIR / "eval_kitti.py"), + "--db-path", + str(db_path), + "--module-path", + str(LOOP_CLOSURE_DIR / module_dir / "module.py"), + "--module-name", + "PGO", + ] + print(f"\n=== {db_path.parent.name} / {module_dir} ===", flush=True) + return subprocess.run(command, check=False).returncode == 0 + + +def render_table() -> Path: + by_recording: dict[str, list[dict[str, Any]]] = {} + for summary_path in sorted(RESULTS_DIR.glob("*/kitti_summary.json")): + summary = json.loads(summary_path.read_text()) + recording = Path(summary["db"]).parent.name + by_recording.setdefault(recording, []).append( + {"module": summary["module"], **summary["scores"]} + ) + + lines = ["# KITTI official-error PGO comparison", ""] + lines += [ + "Translational error (%) and rotational error (deg/m), the KITTI", + "leaderboard metric (avg relative pose error over 100..800m sub-sequences).", + "**Lower is better.** raw odometry = the scan-to-scan ICP input the PGO refines.", + "", + ] + for recording in sorted(by_recording): + rows = sorted(by_recording[recording], key=lambda r: r["translational_percent"]) + lines += [f"## {recording}", ""] + lines.append("| module | transl % | rot deg/m | closures | keyframes |") + lines.append("|---|---|---|---|---|") + baseline = rows[0] + lines.append( + f"| _raw odometry_ | {baseline['baseline_translational_percent']:.2f} | " + f"{baseline['baseline_rotational_deg_per_m']:.4f} | — | — |" + ) + for row in rows: + lines.append( + f"| {row['module']} | {row['translational_percent']:.2f} | " + f"{row['rotational_deg_per_m']:.4f} | {row['closures']} | {row['keyframes']} |" + ) + lines.append("") + RESULTS_DIR.mkdir(exist_ok=True) + TABLE_PATH.write_text("\n".join(lines)) + return TABLE_PATH + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--recordings-dir", type=Path, default=Path("~/datasets/kitti/sequences")) + parser.add_argument( + "--db-path", type=Path, default=None, help="a single kitti db (overrides dir)" + ) + parser.add_argument("--only", help="comma-separated module dirs") + args = parser.parse_args() + + if args.db_path: + dbs = [args.db_path.expanduser()] + else: + dbs = sorted(args.recordings_dir.expanduser().glob("kitti_seq*/mem2.db")) + if not dbs: + raise SystemExit("no KITTI dbs found (run kitti_to_db.py first)") + modules = [m.strip() for m in args.only.split(",")] if args.only else list(MODULES) + + for db_path in dbs: + for module_dir in modules: + run_one(db_path, module_dir) + table = render_table() + print(f"\ntable -> {table}") + + +if __name__ == "__main__": + main() diff --git a/dimos/navigation/jnav/components/loop_closure/spec.py b/dimos/navigation/jnav/components/loop_closure/spec.py new file mode 100644 index 0000000000..d571192266 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/spec.py @@ -0,0 +1,33 @@ +# 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. + +from typing import Protocol + +from dimos.core.stream import In, Out +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.navigation.jnav.msgs.Graph3D import Graph3D +from dimos.navigation.jnav.msgs.GraphDelta3D import GraphDelta3D + + +class LoopClosure(Protocol): + # frame:sensor_link + lidar: In[PointCloud2] + odometry: In[Odometry] + + corrected_odometry: Out[Odometry] + # frame:map + pose_graph: Out[Graph3D] + # frame:map + loop_closure_event: Out[GraphDelta3D] From f17a86957e31e741d08e9159b584e1978373351d Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 14:52:51 +0800 Subject: [PATCH 09/69] add tf tree support to memory2 store (DbTf, write_tf_tree, lazy store.tf) --- dimos/memory2/db_tf.py | 180 ++++++++++++++++++++++++++++++++++++ dimos/memory2/store/base.py | 15 +++ 2 files changed, 195 insertions(+) create mode 100644 dimos/memory2/db_tf.py diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py new file mode 100644 index 0000000000..60b40242d4 --- /dev/null +++ b/dimos/memory2/db_tf.py @@ -0,0 +1,180 @@ +# 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. + +"""Transform lookups over the transforms recorded in a store. + +A store's ``tf`` member lazily reads every transform recorded under the ``tf`` +(and ``tf_static``) streams into a :class:`MultiTBuffer`, then answers +``store.tf.get(target_frame, source_frame, time)`` — composing multi-hop chains +(e.g. ``world -> map -> odom -> base_link -> mid360_link``) and interpolating to +the nearest recorded sample. This makes world-registration a real transform +lookup instead of assuming a single baked-in pose. + +``write_tf_tree`` populates those streams for a recording that lacks them. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +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.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.tf.tf import MultiTBuffer + +if TYPE_CHECKING: + from dimos.memory2.store.base import Store + +TF_STREAMS = ("tf", "tf_static") +# Larger than any single recording's span so the buffer never prunes loaded +# transforms (MultiTBuffer drops samples older than ts - buffer_size). +_NO_PRUNE = 1.0e15 + + +class DbTf: + """Transform lookups backed by the ``tf``/``tf_static`` streams of a store.""" + + def __init__(self, store: Store, stream_names: tuple[str, ...] = TF_STREAMS) -> None: + self._store = store + self._stream_names = stream_names + self._buffer: MultiTBuffer | None = None + + def _ensure_loaded(self) -> MultiTBuffer: + if self._buffer is not None: + return self._buffer + buffer = MultiTBuffer(buffer_size=_NO_PRUNE) + available = set(self._store.list_streams()) + for name in self._stream_names: + if name not in available: + continue + for observation in self._store.stream(name, TFMessage): + message = observation.data + transforms = getattr(message, "transforms", None) + if transforms is None: + transforms = [message] + buffer.receive_transform(*transforms) + self._buffer = buffer + return buffer + + def has_transforms(self) -> bool: + return bool(self._ensure_loaded().buffers) + + def get( + self, + target_frame: str, + source_frame: str, + time_point: float | None = None, + time_tolerance: float | None = None, + ) -> Transform | None: + """Transform that maps a point in ``source_frame`` into ``target_frame``. + + Returns ``None`` if no chain connects the two frames. Uses the buffer's + non-warning lookup so per-scan misses don't spam the log. + """ + buffer = self._ensure_loaded() + return buffer._get(target_frame, source_frame, time_point, time_tolerance) + + +def transform_matrix(transform: Transform) -> tuple[np.ndarray, np.ndarray]: + """Return ``(R, t)`` (3x3, 3) for ``transform`` so ``p_target = p_source @ R.T + t``.""" + rotation = transform.rotation + rotation_matrix = np.asarray(rotation.to_rotation_matrix(), float).reshape(3, 3) + translation = np.array( + [transform.translation.x, transform.translation.y, transform.translation.z], float + ) + return rotation_matrix, translation + + +def write_tf_tree( + store: Store, + *, + odom_stream: str, + odom_parent: str = "odom", + odom_child: str = "base_link", + root_links: tuple[tuple[str, str], ...] = (("world", "map"), ("map", "odom")), + sensor_child: str = "mid360_link", + sensor_translation: tuple[float, float, float] = (0.0, 0.0, 0.0), + sensor_rotation: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0), + static_period: float = 0.45, + stream_name: str = "tf", +) -> int: + """Populate ``store``'s tf stream from an odometry stream. + + - ``root_links`` and ``odom_child -> sensor_child`` are emitted as identity / + fixed transforms every ``static_period`` seconds across the recording span. + - ``odom_parent -> odom_child`` is emitted once per odometry sample, taken + from each observation's pose. + + Returns the number of tf observations written. + """ + import sqlite3 + + db_path = store.config.path + connection = sqlite3.connect(f"file:{db_path}?mode=ro&immutable=1", uri=True) + odom = np.array( + list( + connection.execute( + "select ts,pose_x,pose_y,pose_z,pose_qx,pose_qy,pose_qz,pose_qw " + f"from {odom_stream} order by ts" + ) + ), + float, + ) + connection.close() + if not len(odom): + raise ValueError(f"odom stream {odom_stream!r} is empty; cannot build tf tree") + + tf_stream = store.stream(stream_name, TFMessage) + written = 0 + + # dynamic: odom_parent -> odom_child, one per odometry sample + for row in odom: + ts = float(row[0]) + transform = Transform( + translation=Vector3(row[1], row[2], row[3]), + rotation=Quaternion(row[4], row[5], row[6], row[7]), + frame_id=odom_parent, + child_frame_id=odom_child, + ts=ts, + ) + tf_stream.append(TFMessage(transform), ts=ts) + written += 1 + + # static: root links + sensor mount, resampled every static_period + t0 = float(odom[0, 0]) + t1 = float(odom[-1, 0]) + + def statics_at(ts: float) -> list[Transform]: + links = [ + Transform(frame_id=parent, child_frame_id=child, ts=ts) for parent, child in root_links + ] + links.append( + Transform( + translation=Vector3(*sensor_translation), + rotation=Quaternion(*sensor_rotation), + frame_id=odom_child, + child_frame_id=sensor_child, + ts=ts, + ) + ) + return links + + for ts in np.arange(t0, t1 + static_period, static_period): + tf_stream.append(TFMessage(*statics_at(float(ts))), ts=float(ts)) + written += 1 + + return written diff --git a/dimos/memory2/store/base.py b/dimos/memory2/store/base.py index 7a7162a6d1..46be8b8cb5 100644 --- a/dimos/memory2/store/base.py +++ b/dimos/memory2/store/base.py @@ -29,6 +29,7 @@ from dimos.protocol.service.spec import BaseConfig, Configurable if TYPE_CHECKING: + from dimos.memory2.db_tf import DbTf from dimos.memory2.replay import Replay T = TypeVar("T") @@ -105,12 +106,26 @@ def __init__(self, **kwargs: Any) -> None: Configurable.__init__(self, **kwargs) CompositeResource.__init__(self) self._streams: dict[str, Stream[Any]] = {} + self._tf: DbTf | None = None @property def streams(self) -> StreamAccessor[Stream[Any]]: """Attribute-style access to streams: ``store.streams.name``.""" return StreamAccessor(self) + @property + def tf(self) -> DbTf: + """Transform lookups over the recording's ``tf``/``tf_static`` streams. + + ``store.tf.get(target_frame, source_frame, ts)`` composes multi-hop chains + (e.g. ``world -> ... -> mid360_link``) from the recorded transforms. + """ + if self._tf is None: + from dimos.memory2.db_tf import DbTf + + self._tf = DbTf(self) + return self._tf + def replay( self, *, From 1fb8536a97564cb3581aabf749ecd31b2d4519b2 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 14:52:58 +0800 Subject: [PATCH 10/69] add map post-processing pipeline (post_process, add_april, detect_tags, make_rrd) --- .../loop_closure/gsc_pgo/add_april.py | 216 +++++++ .../loop_closure/gsc_pgo/detect_tags.py | 175 +++++ .../loop_closure/gsc_pgo/make_rrd.py | 184 ++++++ .../loop_closure/gsc_pgo/post_process.py | 604 ++++++++++++++++++ 4 files changed, 1179 insertions(+) create mode 100644 dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py create mode 100644 dimos/navigation/jnav/components/loop_closure/gsc_pgo/detect_tags.py create mode 100644 dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py create mode 100644 dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py new file mode 100644 index 0000000000..010a02fc08 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py @@ -0,0 +1,216 @@ +# 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. + +"""Build the raw_april_tags + april_tags streams into a recording's mem2.db and +record the outcome in summary.json. --summary recomputes only the result section. + +Usage: + python dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py --rec PATH + [--summary] [--output PATH] [--camera color_image] [--intrinsics PATH] + [--tag-size 0.10] [--dict DICT_APRILTAG_36h11] [--dynamic 17] +""" + +import argparse +import json +from pathlib import Path +from typing import Any + +from dimos.navigation.jnav.utils import recording_db as rdb +from dimos.navigation.jnav.utils.apriltag_agreement import ( + VISIT_GAP_S, + split_visits, +) +from dimos.navigation.jnav.utils.apriltags import ( + ensure_april_streams, + gate_params, + load_intrinsics_json, +) + +RAW_STREAM = "raw_april_tags" +FILTERED_STREAM = "april_tags" +SUMMARY_NAME = "summary.json" +MIN_REVISITS = 2 # a tag seen on fewer visits than this carries no agreement signal +DEFAULT_MARKER_LENGTH_M = 0.10 +DEFAULT_DICTIONARY = "DICT_APRILTAG_36h11" + + +def _parse_dynamic(raw: str | None) -> list[int]: + if not raw: + return [] + return sorted({int(token) for token in raw.replace(",", " ").split()}) + + +def _times_by_tag(store: Any, stream_name: str) -> dict[int, list[float]]: + by_tag: dict[int, list[float]] = {} + for observation in store.stream(stream_name): + by_tag.setdefault(int(observation.tags["marker_id"]), []).append(float(observation.ts)) + return by_tag + + +def summarize(store: Any) -> dict[str, Any]: + """Per-tag raw detections + filtered visits from the existing streams, flagging + never-revisited tags. Prints the table and returns the result for summary.json.""" + streams = set(store.list_streams()) + raw_available = RAW_STREAM in streams + raw = _times_by_tag(store, RAW_STREAM) if raw_available else {} + filtered = _times_by_tag(store, FILTERED_STREAM) if FILTERED_STREAM in streams else {} + tag_ids = sorted(set(raw) | set(filtered)) + + print( + f" {'tag':>5} {'raw':>6} {'filtered':>9} {'revisits':>9} (visit gap {VISIT_GAP_S:.0f}s)" + ) + tags: list[dict[str, Any]] = [] + not_revisited: list[int] = [] + for tag_id in tag_ids: + raw_count = len(raw.get(tag_id, [])) + filtered_times = sorted(filtered.get(tag_id, [])) + visits = len(split_visits(filtered_times, gap_s=VISIT_GAP_S)) if filtered_times else 0 + revisited = visits >= MIN_REVISITS + if not revisited: + not_revisited.append(tag_id) + tags.append( + { + "tag_id": tag_id, + "raw": raw_count if raw_available else None, + "filtered": len(filtered_times), + "revisits": visits, + "revisited": revisited, + } + ) + raw_display = raw_count if raw_available else "-" + flag = " <-- NOT revisited" if not revisited else "" + print(f" {tag_id:>5} {raw_display!s:>6} {len(filtered_times):>9} {visits:>9}{flag}") + + print( + f" totals: {len(tag_ids)} tags | {len(tag_ids) - len(not_revisited)} revisited" + f" (>={MIN_REVISITS} visits) | {len(not_revisited)} not revisited" + ) + if not_revisited: + print(f" NOT revisited: {not_revisited}") + if not raw_available: + print(f" (no '{RAW_STREAM}' yet — raw counts N/A; run without --summary to build it)") + + return { + "visit_gap_s": VISIT_GAP_S, + "min_revisits": MIN_REVISITS, + "all_unfiltered_tag_ids": sorted(raw) if raw_available else sorted(filtered), + "total_tags": len(tag_ids), + "revisited": len(tag_ids) - len(not_revisited), + "not_revisited": not_revisited, + "raw_available": raw_available, + "tags": tags, + } + + +def _update_summary_json( + path: Path, + *, + filter_parameters: dict[str, Any] | None = None, + result: dict[str, Any] | None = None, +) -> Path: + """Merge the april_tags section into summary.json, preserving every other key.""" + data: dict[str, Any] = json.loads(path.read_text()) if path.exists() else {} + section: dict[str, Any] = data.get("april_tags", {}) + if filter_parameters is not None: + section["filter_parameters"] = filter_parameters + if result is not None: + section["result"] = result + data["april_tags"] = section + path.write_text(json.dumps(data, indent=2)) + return path + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--rec", type=Path, required=True, help="recording dir or mem2.db path") + parser.add_argument("--camera", default="color_image") + parser.add_argument("--intrinsics", type=Path, default=None) + parser.add_argument("--tag-size", type=float, default=None, help="marker length (m)") + parser.add_argument("--dict", dest="dictionary", default=None) + parser.add_argument( + "--summary", + action="store_true", + help="read-only on streams: recompute and write ONLY april_tags.result (a subset of" + " the full run, which also rebuilds streams + filter_parameters)", + ) + parser.add_argument( + "--output", + type=Path, + default=None, + help="summary.json path to write (default: /summary.json) — pick another" + " location to avoid overwriting an existing one", + ) + parser.add_argument( + "--dynamic", + default=None, + help="comma/space-separated tag ids on moving objects — kept in raw, dropped from filtered", + ) + args = parser.parse_args() + + recording = args.rec.expanduser() + db_path = recording if recording.name == "mem2.db" else recording / "mem2.db" + if not db_path.exists(): + parser.error(f"no mem2.db at {db_path}") + store = rdb.store(db_path) + summary_path = args.output.expanduser() if args.output else (db_path.parent / SUMMARY_NAME) + print(f"=== {db_path.parent.name} ===") + + if args.summary: + result = summarize(store) + _update_summary_json(summary_path, result=result) + print(f" updated {summary_path} april_tags.result") + return + + # Rebuild both streams and overwrite filter_parameters + result. + dynamic_tags = _parse_dynamic(args.dynamic) + intrinsics_path = (args.intrinsics or (db_path.parent / "camera_intrinsics.json")).expanduser() + config = load_intrinsics_json(intrinsics_path) + marker_length = ( + args.tag_size + if args.tag_size is not None + else config.get("marker_length", DEFAULT_MARKER_LENGTH_M) + ) + dictionary = args.dictionary or config.get("dictionary", DEFAULT_DICTIONARY) + ensure_april_streams( + store, + config["intrinsics"], + config["distortion"], + image_stream=args.camera, + marker_length=marker_length, + dictionary=dictionary, + raw_stream=RAW_STREAM, + filtered_stream=FILTERED_STREAM, + exclude_tags=dynamic_tags, + force=True, + ) + filter_parameters = { + "gates": gate_params(), + "marker_length_m": marker_length, + "dictionary": dictionary, + "camera_stream": args.camera, + "raw_stream": RAW_STREAM, + "filtered_stream": FILTERED_STREAM, + "dynamic_tags_excluded": dynamic_tags, + } + result = summarize(store) + _update_summary_json(summary_path, filter_parameters=filter_parameters, result=result) + print( + f" updated {summary_path} april_tags (filter_parameters + result); dynamic={dynamic_tags}" + ) + + +if __name__ == "__main__": + main() diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/detect_tags.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/detect_tags.py new file mode 100644 index 0000000000..71b11201cc --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/detect_tags.py @@ -0,0 +1,175 @@ +# 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. + +# Untyped analysis script: gtsam/open3d/cv2 lack type stubs. +# mypy: ignore-errors +"""Build the raw AprilTag stream: EVERY detection over the camera image, NO filtering whatsoever +(no blur/reproj/distance/angle/motion/size gate, no time-clustering). One row per per-frame +detection that yields a valid PnP pose. Each row carries its gate diagnostics in tags +(sharpness, reproj_px, tag_px, distance_m, view_angle_deg, lin_speed, ang_speed) so downstream +gate tuning in post_process.py needs no re-detection. + +Prints the raw per-marker histogram + visit structure (visit = sightings >30s apart). + +Usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/detect_tags.py --rec=PATH + [--camera=color_image] [--tag-size=0.10] + [--dict=DICT_APRILTAG_36h11] [--intrinsics=PATH] [--out=raw_april_tags] +""" + +import json +from pathlib import Path +import sys + +import cv2 +import numpy as np +from scipy.spatial.transform import Rotation + +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.sensor_msgs.Image import Image +from dimos.navigation.jnav.utils import recording_db as rdb +from dimos.navigation.jnav.utils.apriltags import ( + _camera_speeds, + estimate_marker_pose, + make_detector, + reprojection_error_px, + tag_pixel_size, + tag_sharpness, + view_quality, +) + + +def arg(flag, default=None): + return next((a.split("=", 1)[1] for a in sys.argv if a.startswith(flag + "=")), default) + + +REC_ARG = arg("--rec") +if not REC_ARG: + sys.exit( + "usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/detect_tags.py --rec=PATH [--camera=...] [--tag-size=...] " + "[--dict=...] [--intrinsics=PATH] [--out=...] (--rec is required)" + ) +REC = Path(REC_ARG).expanduser() +CAMERA = arg("--camera", "color_image") # camera image stream the tags are detected on +MARKER_LENGTH_M = float(arg("--tag-size", "0.10")) +DICTIONARY = arg("--dict", "DICT_APRILTAG_36h11") +STREAM = arg("--out", "raw_april_tags") +VISIT_GAP_S = 30.0 + +DB = REC / "mem2.db" +intr_path = Path(arg("--intrinsics", str(REC / "camera_intrinsics.json"))).expanduser() +intr = json.loads(intr_path.read_text()) +K = np.array(intr["intrinsics"], float).reshape(3, 3) +dist = np.array(intr.get("distortion", []), float) +st = rdb.store(DB) +if CAMERA not in st.list_streams(): + sys.exit(f"!! camera stream '{CAMERA}' not in db — available: {st.list_streams()}") + +detector = make_detector(DICTIONARY) +print(f"loading '{CAMERA}' frames...") +images = st.stream(CAMERA, Image).to_list() +speed_by_ts, speed_available = _camera_speeds(images) +print( + f"detecting over {len(images)} frames (unfiltered), tag_size={MARKER_LENGTH_M} m, dict={DICTIONARY}..." +) + +rows = [] +for image_obs in images: + image = image_obs.data + bgr = image.numpy() if hasattr(image, "numpy") else np.asarray(image.data) + gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) if bgr.ndim == 3 else bgr + all_corners, marker_ids, _ = detector.detectMarkers(bgr) + if marker_ids is None: + continue + for corners, marker_id in zip(all_corners, marker_ids.flatten(), strict=False): + pose = estimate_marker_pose(corners, MARKER_LENGTH_M, K, dist) + if pose is None: + continue + rvec, tvec = pose + quat = Rotation.from_rotvec(rvec.reshape(3)).as_quat() # x,y,z,w + t = tvec.reshape(3) + tcm = [ + float(t[0]), + float(t[1]), + float(t[2]), + float(quat[0]), + float(quat[1]), + float(quat[2]), + float(quat[3]), + ] + distance, view_angle = view_quality(tcm) + speed = speed_by_ts.get(float(image_obs.ts)) + rows.append( + { + "ts": float(image_obs.ts), + "marker_id": int(marker_id), + "tcm": tcm, + "sharpness": float(tag_sharpness(gray, corners)), + "reproj_px": float( + reprojection_error_px(corners, rvec, tvec, MARKER_LENGTH_M, K, dist) + ), + "tag_px": float(tag_pixel_size(corners)), + "distance_m": float(distance), + "view_angle_deg": float(view_angle), + "lin_speed": float(speed[0]) if speed else -1.0, + "ang_speed": float(speed[1]) if speed else -1.0, + } + ) +rows.sort(key=lambda r: r["ts"]) + +if STREAM in st.list_streams(): + st.delete_stream(STREAM) +out = st.stream(STREAM, PoseStamped) +for r in rows: + tcm = r["tcm"] + out.append( + PoseStamped(ts=r["ts"], position=tcm[:3], orientation=tcm[3:]), + ts=r["ts"], + pose=tuple(tcm), + tags={ + k: r[k] + for k in ( + "marker_id", + "sharpness", + "reproj_px", + "tag_px", + "distance_m", + "view_angle_deg", + "lin_speed", + "ang_speed", + ) + }, + ) +print(f"\nwrote {STREAM}: {len(rows)} unfiltered detections") + +by = {} +for r in rows: + by.setdefault(r["marker_id"], []).append(r) +print(f"\n=== RAW per-marker (visit = >{VISIT_GAP_S:.0f}s apart) ===") +print( + f"{'tag':>4} {'det':>4} {'visits':>6} {'dist_m':>12} {'sharp>=60%':>10} {'reproj<=2%':>10} {'span_s':>7}" +) +for mid in sorted(by): + rs = sorted(by[mid], key=lambda r: r["ts"]) + times = [r["ts"] for r in rs] + visits = [[times[0]]] + for tt in times[1:]: + (visits[-1].append(tt) if tt - visits[-1][-1] <= VISIT_GAP_S else visits.append([tt])) + d = [r["distance_m"] for r in rs] + sharp_ok = 100 * np.mean([r["sharpness"] >= 60 for r in rs]) + reproj_ok = 100 * np.mean([r["reproj_px"] <= 2 for r in rs]) + print( + f"{mid:>4} {len(rs):>4} {len(visits):>6} {min(d):5.2f}-{max(d):5.2f} " + f"{sharp_ok:9.0f}% {reproj_ok:9.0f}% {times[-1] - times[0]:7.0f}" + ) +print("\nmarkers present:", sorted(by)) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py new file mode 100644 index 0000000000..bfd2db540a --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py @@ -0,0 +1,184 @@ +# 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. + +# Untyped analysis script: gtsam/open3d/cv2 lack type stubs. +# mypy: ignore-errors +"""Combined comparison rrd: raw lidar cloud + EVERY gt_*_lidar version present in the db, each as +its own colored entity, plus AprilTag landmarks + trajectories. Re-run after adding a new GT method +and it picks the new stream up automatically. + +Importable: `build(...)` writes the rrd and returns its path (used by post_process.py). +Standalone: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py --rec=PATH [--lidar=...] [--odom=...] [--tags=...] [--out=...] +""" + +import json +from pathlib import Path +import sys + +from gtsam import Point3, Pose3, Rot3 +import numpy as np +import rerun as rr + +from dimos.navigation.jnav.utils import recording_db as rdb + +SCAN_STRIDE, VOXEL = 8, 0.10 +COLORS = {"raw": [220, 60, 60]} +PALETTE = [ + [60, 120, 230], + [60, 210, 90], + [230, 180, 50], + [200, 80, 220], + [80, 220, 220], + [240, 130, 60], +] +# same relaxed gates post_process uses, for placing landmark markers +GATE = dict(s=25.0, r=3.5, px=12.0, d=1.5, a=65.0, lv=1.5, av=150.0) + + +def build( + rec, + lidar_stream="pointlio_lidar", + odom_stream="pointlio_odometry", + tag_stream="raw_april_tags", + out_name="gt_compare.rrd", +): + rec = Path(rec).expanduser() + db = rec / "mem2.db" + out = rec / out_name + st = rdb.store(db) + intr = json.loads((rec / "camera_intrinsics.json").read_text()) + ext = np.array(intr["optical_in_base"], float) + Tbo = Pose3(Rot3.Quaternion(ext[6], ext[3], ext[4], ext[5]), Point3(ext[0], ext[1], ext[2])) + + def accumulate(name): + pts = [] + for i, o in enumerate(st.stream(name)): + if i % SCAN_STRIDE: + continue + xyz = np.asarray(o.data.points_f32()) + if len(xyz): + pts.append(xyz[::3]) + a = np.concatenate(pts, 0) + _, idx = np.unique(np.floor(a / VOXEL).astype(np.int64), axis=0, return_index=True) + return a[idx] + + def traj(name): + return np.array( + [ + [o.data.pose.position.x, o.data.pose.position.y, o.data.pose.position.z] + for o in st.stream(name) + ], + np.float32, + ) + + def landmarks(gt_odom): + gt = [ + ( + o.ts, + Pose3( + Rot3.Quaternion( + o.data.pose.orientation.w, + o.data.pose.orientation.x, + o.data.pose.orientation.y, + o.data.pose.orientation.z, + ), + Point3(o.data.pose.position.x, o.data.pose.position.y, o.data.pose.position.z), + ), + ) + for o in st.stream(gt_odom) + ] + gts = np.array([t for t, _ in gt]) + pos = {} + for obs in st.stream(tag_stream): + t = obs.tags + if not ( + float(t["sharpness"]) >= GATE["s"] + and float(t["reproj_px"]) <= GATE["r"] + and float(t["tag_px"]) >= GATE["px"] + and float(t["distance_m"]) <= GATE["d"] + and float(t["view_angle_deg"]) <= GATE["a"] + and (float(t["lin_speed"]) < 0 or float(t["lin_speed"]) <= GATE["lv"]) + and (float(t["ang_speed"]) < 0 or float(t["ang_speed"]) <= GATE["av"]) + ): + continue + ps = obs.data + cb = gt[int(np.argmin(np.abs(gts - float(obs.ts))))][1] + Tw = cb.compose(Tbo).compose( + Pose3( + Rot3.Quaternion( + ps.orientation.w, ps.orientation.x, ps.orientation.y, ps.orientation.z + ), + Point3(ps.x, ps.y, ps.z), + ) + ) + pos.setdefault(int(t["marker_id"]), []).append(np.asarray(Tw.translation())) + means = [np.mean(value, 0) for key, value in sorted(pos.items())] + lbl = [f"tag{key}" for key in sorted(pos)] + return np.array(means), lbl + + streams = st.list_streams() + gt_lidars = sorted(s for s in streams if s.startswith("gt_") and "_lidar" in s) + print("raw + GT lidar streams:", gt_lidars) + + rr.init("gt_compare") + rr.save(str(out)) + rr.log( + "raw/cloud", + rr.Points3D(accumulate(lidar_stream), colors=COLORS["raw"], radii=0.02), + static=True, + ) + rr.log( + "raw/trajectory", rr.LineStrips3D([traj(odom_stream)], colors=[255, 120, 120]), static=True + ) + for k, name in enumerate(gt_lidars): + color = PALETTE[k % len(PALETTE)] + cloud = accumulate(name) + rr.log(f"{name}/cloud", rr.Points3D(cloud, colors=color, radii=0.02), static=True) + print(f" logged {name}: {len(cloud):,} pts") + odom = name.replace("_lidar", "_odometry") + if odom in streams: + rr.log(f"{name}/trajectory", rr.LineStrips3D([traj(odom)], colors=color), static=True) + # landmarks placed against the first available gt odometry + gt_odoms = sorted(s for s in streams if s.startswith("gt_") and "_odometry" in s) + if gt_odoms: + lm, lbl = landmarks(gt_odoms[0]) + if len(lm): + rr.log( + "landmarks", + rr.Points3D(lm, colors=[255, 230, 0], radii=0.25, labels=lbl), + static=True, + ) + print(f" logged {len(lbl)} landmarks") + print("wrote", out) + return out + + +def _arg(flag, default=None): + return next((a.split("=", 1)[1] for a in sys.argv if a.startswith(flag + "=")), default) + + +if __name__ == "__main__": + rec_arg = _arg("--rec") + if not rec_arg: + sys.exit( + "usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py --rec=PATH [--lidar=...] [--odom=...] " + "[--tags=...] [--out=...] (--rec is required)" + ) + build( + rec_arg, + lidar_stream=_arg("--lidar", "pointlio_lidar"), + odom_stream=_arg("--odom", "pointlio_odometry"), + tag_stream=_arg("--tags", "raw_april_tags"), + out_name=_arg("--out", "gt_compare.rrd"), + ) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py new file mode 100644 index 0000000000..609b80f93b --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py @@ -0,0 +1,604 @@ +# 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. + +# Untyped analysis script: gtsam/open3d/cv2 lack type stubs. +# mypy: ignore-errors +"""AprilTag-loop-closed + ICP-refined ground-truth post-processing for a go2 recording. + +Source of tags = the UNFILTERED tag stream (build it first with detect_tags.py). Each raw +detection carries its gate diagnostics, so gates are applied here post-hoc (no re-detection) +and are easy to relax. Factors: one robust (best-reproj) observation per keyframe x marker -> +denser, balanced loop closure than one-medoid-per-visit. + +Two-stage solve: (1) GTSAM tag PGO -- anisotropic odometry between-factors (stiff roll/pitch + z +anchor gravity, loose yaw) + quality-weighted AprilTag landmark factors fix macro drift; +(2) ICP loop closures between spatially-close / temporally-distant lidar submaps anchor local +geometry. Writes _odometry / _lidar back into the recording db, optionally a .pc2.lcm +log of the corrected cloud, and opens a comparison rrd. + +Usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py [odom|lidar|both] --rec=PATH + [--lidar=pointlio_lidar] [--odom=pointlio_odometry] [--tags=raw_april_tags] + [--out=gt_pointlio] [--suffix=...] [--ignore-tags=17] [--no-icp] [--no-lcm] [--no-rrd] +""" + +import json +from pathlib import Path +import sqlite3 +import sys +import time + +from gtsam import ( + BetweenFactorPose3, + LevenbergMarquardtOptimizer, + LevenbergMarquardtParams, + NonlinearFactorGraph, + Point3, + Pose3, + PriorFactorPose3, + Rot3, + Symbol, + Values, + noiseModel, +) +import numpy as np + +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.navigation.jnav.utils import recording_db as rdb +from dimos.navigation.jnav.utils.apriltags import ( + DEFAULT_MAX_ANGULAR_SPEED_DPS, + DEFAULT_MAX_DISTANCE_M, + DEFAULT_MAX_LINEAR_SPEED_MPS, + DEFAULT_MAX_REPROJ_PX, + DEFAULT_MAX_VIEW_ANGLE_DEG, + DEFAULT_MIN_SHARPNESS, + DEFAULT_MIN_TAG_PX, +) + +VISIT_GAP_S = 30.0 +WHAT = sys.argv[1] if len(sys.argv) > 1 and not sys.argv[1].startswith("-") else "both" + + +def arg(flag, default=""): + return next((a.split("=", 1)[1] for a in sys.argv if a.startswith(flag + "=")), default) + + +REC_ARG = arg("--rec") +SUFFIX = arg("--suffix") +LIDAR_STREAM = arg("--lidar", "pointlio_lidar") # input lidar stream (world-registered scans) +ODOM_STREAM = arg("--odom", "pointlio_odometry") # input odometry stream (keyframe source) +RAW_STREAM = arg("--tags", "raw_april_tags") # input unfiltered AprilTag stream +IGNORE_TAGS = { + int(x) for x in arg("--ignore-tags").replace(",", " ").split() +} # dynamic/moving tags +OUT_PREFIX = arg("--out", "gt_pointlio") # output prefix -> _odometry / _lidar +WRITE_LCM = "--no-lcm" not in sys.argv # also emit _lidar.pc2.lcm of the corrected cloud +OPEN_RRD = "--no-rrd" not in sys.argv # build + open a comparison rrd at the end +LCM_VOXEL = float(arg("--lcm-voxel", "0.05")) # voxel size for the aggregated .pc2.lcm cloud +LCM_OUTLIER_NN = 20 # statistical outlier removal: neighbor count +LCM_OUTLIER_STD = 2.0 # ...and std-ratio threshold (lower = more aggressive) +LIDAR_FRAME = arg("--lidar-frame", "mid360_link") # frame the raw lidar scans live in +WORLD_FRAME = arg("--world-frame", "world") # frame to register scans into +USE_TF = "--no-tf" not in sys.argv # world-register via recording tf (fallback: obs.pose) +TF_TOL = float(arg("--tf-tol", "0.5")) # max seconds to the nearest recorded transform + +# Per-glimpse gates (speed == -1 means "unknown" and always passes). +GATE = dict( + min_sharpness=DEFAULT_MIN_SHARPNESS, + max_reproj_px=DEFAULT_MAX_REPROJ_PX, + min_tag_px=DEFAULT_MIN_TAG_PX, + max_distance_m=DEFAULT_MAX_DISTANCE_M, + max_view_angle_deg=DEFAULT_MAX_VIEW_ANGLE_DEG, + max_lin_speed=DEFAULT_MAX_LINEAR_SPEED_MPS, + max_ang_speed=DEFAULT_MAX_ANGULAR_SPEED_DPS, +) + +if not REC_ARG: + sys.exit( + "usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py [odom|lidar|both] --rec=PATH " + "[--lidar=...] [--odom=...] [--tags=...] [--out=...] [--suffix=...] " + "[--no-icp] [--no-lcm] [--no-rrd] (--rec is required: path to the recording dir)" + ) +REC = Path(REC_ARG).expanduser() +DB = REC / "mem2.db" +intr = json.loads((REC / "camera_intrinsics.json").read_text()) +ext = np.array(intr["optical_in_base"], float) +Tbo = Pose3(Rot3.Quaternion(ext[6], ext[3], ext[4], ext[5]), Point3(ext[0], ext[1], ext[2])) +st = rdb.store(DB) +if RAW_STREAM not in st.list_streams(): + sys.exit( + f"!! {RAW_STREAM} missing -- run detect_tags.py first to build the unfiltered tag stream." + ) +from dimos.memory2.db_tf import transform_matrix + +_TF_AVAILABLE = USE_TF and st.tf.has_transforms() + + +def world_points(obs): + """Nx3 world-registered points for a lidar observation. + + Primary: the recording's tf (``world <- LIDAR_FRAME``) at the scan time. + Fallback: the observation's stored pose as the transform. Otherwise the + scan is assumed already world-registered (legacy recordings). + """ + pts = np.asarray(obs.data.points_f32()) + if not len(pts): + return pts + if _TF_AVAILABLE: + tf = st.tf.get(WORLD_FRAME, LIDAR_FRAME, float(obs.ts), TF_TOL) + if tf is not None: + rot, trans = transform_matrix(tf) + return pts @ rot.T + trans + pose = getattr(obs, "pose", None) + if isinstance(pose, (tuple, list)) and len(pose) >= 7: + rot = Rot3.Quaternion(pose[6], pose[3], pose[4], pose[5]).matrix() + return pts @ rot.T + np.array(pose[:3], float) + return pts # already world-registered + + +print(f"recording: {REC}", flush=True) +if _TF_AVAILABLE: + print(f"world-registering {LIDAR_STREAM} via tf ({WORLD_FRAME} <- {LIDAR_FRAME})", flush=True) +print( + f"streams: tags={RAW_STREAM} odom={ODOM_STREAM} lidar={LIDAR_STREAM} -> out={OUT_PREFIX}{SUFFIX}", + flush=True, +) + + +def passes(t): + return ( + t["sharpness"] >= GATE["min_sharpness"] + and t["reproj_px"] <= GATE["max_reproj_px"] + and t["tag_px"] >= GATE["min_tag_px"] + and t["distance_m"] <= GATE["max_distance_m"] + and t["view_angle_deg"] <= GATE["max_view_angle_deg"] + and (t["lin_speed"] < 0 or t["lin_speed"] <= GATE["max_lin_speed"]) + and (t["ang_speed"] < 0 or t["ang_speed"] <= GATE["max_ang_speed"]) + ) + + +# read raw detections (pose + diagnostics) +print("reading tag detections...", flush=True) +raw = [] +for obs in st.stream(RAW_STREAM): + ps = obs.data + tg = obs.tags + raw.append( + dict( + ts=float(obs.ts), + marker_id=int(tg["marker_id"]), + T_cam_tag=Pose3( + Rot3.Quaternion( + ps.orientation.w, ps.orientation.x, ps.orientation.y, ps.orientation.z + ), + Point3(ps.x, ps.y, ps.z), + ), + reproj_px=float(tg["reproj_px"]), + **{ + k: float(tg[k]) + for k in ( + "sharpness", + "tag_px", + "distance_m", + "view_angle_deg", + "lin_speed", + "ang_speed", + ) + }, + ) + ) +gated = [t for t in raw if passes(t) and t["marker_id"] not in IGNORE_TAGS] + +# keyframes from raw odometry +con = sqlite3.connect(f"file:{DB}?mode=ro", uri=True) +fo = np.array( + list( + con.execute( + "select ts,pose_x,pose_y,pose_z,pose_qx,pose_qy,pose_qz,pose_qw " + f"from {ODOM_STREAM} order by ts" + ) + ), + float, +) +con.close() + + +def pr(r): + return Rot3.Quaternion(r[7], r[4], r[5], r[6]), np.array(r[1:4]) + + +ki = [0] +prev_rot, prev_pos = pr(fo[0]) +for i in range(1, len(fo)): + rot, pos = pr(fo[i]) + if ( + np.linalg.norm(pos - prev_pos) > 0.5 + or np.degrees(np.linalg.norm(Rot3.Logmap(prev_rot.inverse() * rot))) > 10 + ): + ki.append(i) + prev_rot, prev_pos = rot, pos +kfs = [pr(fo[i]) for i in ki] +kts = fo[ki, 0] +N = len(kfs) + +# one factor per keyframe x marker: keep the best-reproj detection in each bucket +bucket = {} +for t in gated: + kf = int(np.argmin(np.abs(kts - t["ts"]))) + key = (kf, t["marker_id"]) + if key not in bucket or t["reproj_px"] < bucket[key]["reproj_px"]: + bucket[key] = t + +# revisit report +raw_by, vis_by = {}, {} +for t in raw: + raw_by.setdefault(t["marker_id"], 0) + raw_by[t["marker_id"]] += 1 +for (_kf, mid), t in bucket.items(): + vis_by.setdefault(mid, []).append(t["ts"]) + + +def n_visits(times): + times = sorted(times) + visits = [[times[0]]] + for tt in times[1:]: + (visits[-1].append(tt) if tt - visits[-1][-1] <= VISIT_GAP_S else visits.append([tt])) + return len(visits) + + +print(f"gates: {GATE}") +print( + f"raw detections {len(raw)} -> {len(gated)} pass gates -> {len(bucket)} keyframe-tag factors\n" +) +print(f"{'tag':>4} | {'raw viewings':>12} | {'filtered revisits':>17}") +not_revisited = [] +for mid in sorted(raw_by): + nv = n_visits(vis_by[mid]) if mid in vis_by else 0 + flag = "" if nv >= 2 else " <-- NOT REVISITED" + print(f"{mid:>4} | {raw_by[mid]:>12} | {nv:>10} visit(s){flag}") + if nv < 2: + not_revisited.append(mid) +print( + f"\ntags NOT revisited (no loop-closure constraint): {not_revisited if not_revisited else 'none'}\n" +) + +# factor graph + solve +odom = noiseModel.Diagonal.Variances(np.array([1e-8, 1e-8, 1e-5, 1e-4, 1e-4, 1e-6])) +grav0 = noiseModel.Diagonal.Variances(np.array([1e-8, 1e-8, 1e-6, 1e-8, 1e-8, 1e-8])) + + +# quality weighting: planar-PnP pose error grows ~quadratically with range, and reproj_px is a +# direct misfit proxy. Inflate a glimpse's covariance by (dist/REF_D)^2 * (reproj/REF_R)^2 so a +# far/oblique/blurry tag pose contributes almost nothing while close, sharp ones dominate. +REF_D, REF_R = 0.4, 1.0 + + +def tn(Rbt, distance_m=REF_D, reproj_px=REF_R): + s = max((max(distance_m, 0.2) / REF_D) ** 2 * (max(reproj_px, 0.5) / REF_R) ** 2, 0.25) + Rm = Rbt.matrix() + c = np.zeros((6, 6)) + c[:3, :3] = Rm @ np.diag([0.04, 0.04, 0.0025]) @ Rm.T + c[3:, 3:] = Rm @ np.diag([0.0025, 0.0025, 0.25]) @ Rm.T + return noiseModel.Gaussian.Covariance(c * s) + + +print(f"building factor graph over {N} keyframes...", flush=True) +g = NonlinearFactorGraph() +v = Values() +for i in range(N): + rot, pos = kfs[i] + v.insert(i, Pose3(rot, Point3(pos))) + if i == 0: + g.add(PriorFactorPose3(0, Pose3(rot, Point3(pos)), grav0)) + else: + rot_prev, pos_prev = kfs[i - 1] + g.add( + BetweenFactorPose3( + i - 1, + i, + Pose3( + rot_prev.inverse() * rot, + Point3(rot_prev.inverse().rotate(Point3(pos - pos_prev))), + ), + odom, + ) + ) +seen = set() +for (kf, mid), t in sorted(bucket.items()): + rot, pos = kfs[kf] + kp = Pose3(rot, Point3(pos)) + T = Tbo.compose(t["T_cam_tag"]) + if mid not in seen: + seen.add(mid) + v.insert(Symbol("l", mid).key(), kp.compose(T)) + g.add( + BetweenFactorPose3( + kf, Symbol("l", mid).key(), T, tn(T.rotation(), t["distance_m"], t["reproj_px"]) + ) + ) + +print("solving stage 1 (tag PGO)...", flush=True) +lm_params = LevenbergMarquardtParams() +lm_params.setMaxIterations(200) +est = LevenbergMarquardtOptimizer(g, v, lm_params).optimize() +raw_kf = [Pose3(kfs[i][0], Point3(kfs[i][1])) for i in range(N)] + +# STAGE 2: ICP loop-closure refinement (tags=macro, lidar ICP=local anchor) +# Tag PGO has pulled revisits roughly together; now ICP the lidar submaps of spatially-close, +# temporally-distant keyframe pairs to add precise 6-DOF relative constraints, then re-solve. +ICP = "--no-icp" not in sys.argv +if ICP: + import open3d as o3d + from scipy.spatial import cKDTree + + ICP_RADIUS_M = 4.0 # tag-corrected positions must be within this to be a revisit candidate + ICP_MIN_DT_S = 25.0 # ...and at least this far apart in time (a real revisit, not adjacency) + ICP_MAX_CORR_M = 0.6 # ICP correspondence distance + ICP_VOXEL = 0.15 + ICP_FIT_MIN, ICP_RMSE_MAX = 0.45, 0.25 + SUBMAP_HALF_S = 1.0 # accumulate scans within +/- this of a keyframe time into its submap + + corr = [est.atPose3(i) for i in range(N)] + cpos = np.array([np.asarray(p.translation()) for p in corr]) + + # revisit candidate pairs + tree = cKDTree(cpos) + pairs = set() + for i, j in tree.query_pairs(ICP_RADIUS_M): + if abs(kts[i] - kts[j]) >= ICP_MIN_DT_S: + pairs.add((min(i, j), max(i, j))) + pairs = sorted(pairs, key=lambda p: np.linalg.norm(cpos[p[0]] - cpos[p[1]])) + involved = {k for p in pairs for k in p} + print( + f"ICP stage: {len(pairs)} revisit candidate pairs over {len(involved)} keyframes", + flush=True, + ) + + # build per-(involved)-keyframe body-frame submaps from the input lidar (odom-registered) + print("ICP stage: reading lidar submaps...", flush=True) + submap = {k: [] for k in involved} + scan_n = 0 + t_sub = time.time() + for obs in st.stream(LIDAR_STREAM): + scan_n += 1 + if scan_n % 20000 == 0: + print(f" read {scan_n} scans, {time.time() - t_sub:.0f}s", flush=True) + ts = float(obs.ts) + k = int(np.argmin(np.abs(kts - ts))) + if k not in submap or abs(kts[k] - ts) > SUBMAP_HALF_S: + continue + rot_k, pos_k = kfs[k] + world = world_points(obs) + submap[k].append((world - pos_k) @ rot_k.matrix()) # world -> kf-k body frame + pcd = {} + for k, chunks in submap.items(): + if not chunks: + continue + p = o3d.geometry.PointCloud() + p.points = o3d.utility.Vector3dVector(np.concatenate(chunks, 0).astype(np.float64)) + p = p.voxel_down_sample(ICP_VOXEL) + p.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=0.5, max_nn=30)) + pcd[k] = p + print(f"ICP stage: built {len(pcd)} submaps, registering {len(pairs)} pairs...", flush=True) + + icp_noise = noiseModel.Robust.Create( + noiseModel.mEstimator.Huber.Create(1.345), + noiseModel.Diagonal.Variances(np.array([4e-4, 4e-4, 4e-4, 2.5e-3, 2.5e-3, 2.5e-3])), + ) + added = 0 + t_icp = time.time() + for pair_n, (i, j) in enumerate(pairs): + if pair_n and pair_n % 5000 == 0: + print( + f" registered {pair_n}/{len(pairs)} pairs, {added} accepted, {time.time() - t_icp:.0f}s", + flush=True, + ) + if i not in pcd or j not in pcd: + continue + init = (corr[i].inverse() * corr[j]).matrix() # i<-j initial guess from tag correction + res = o3d.pipelines.registration.registration_icp( + pcd[j], + pcd[i], + ICP_MAX_CORR_M, + init, + o3d.pipelines.registration.TransformationEstimationPointToPlane(), + ) + if res.fitness >= ICP_FIT_MIN and res.inlier_rmse <= ICP_RMSE_MAX: + T = res.transformation + g.add(BetweenFactorPose3(i, j, Pose3(Rot3(T[:3, :3]), Point3(T[:3, 3])), icp_noise)) + added += 1 + print( + f"ICP stage: accepted {added}/{len(pairs)} loop closures " + f"(fit>={ICP_FIT_MIN}, rmse<={ICP_RMSE_MAX}m)", + flush=True, + ) + if added: + # est already holds every key (keyframe + landmark poses); warm-start from it + print("solving stage 2 (tag PGO + ICP closures)...", flush=True) + est = LevenbergMarquardtOptimizer(g, est, lm_params).optimize() + +C = [est.atPose3(i).compose(raw_kf[i].inverse()) for i in range(N)] +shift = max(float(np.linalg.norm(np.asarray(C[i].translation()))) for i in range(N)) +print( + f"PGO: N={N} keyframes, {len(bucket)} tag factors over {len(seen)} markers, " + f"max correction shift {shift:.1f} m", + flush=True, +) + + +def C_at(ts): + if ts <= kts[0]: + return C[0] + if ts >= kts[-1]: + return C[-1] + j = int(np.searchsorted(kts, ts)) + a, b = j - 1, j + al = (ts - kts[a]) / (kts[b] - kts[a]) + return C[a].compose(Pose3.Expmap(al * Pose3.Logmap(C[a].between(C[b])))) + + +def pose_tuple(P): + t = P.translation() + q = P.rotation().toQuaternion() + return (t[0], t[1], t[2], q.x(), q.y(), q.z(), q.w()) + + +if WHAT in ("odom", "both"): + name = f"{OUT_PREFIX}_odometry{SUFFIX}" + if name in st.list_streams(): + st.delete_stream(name) + out = st.stream(name, Odometry) + print(f"writing {name} ({len(fo)} poses)...", flush=True) + n = 0 + t0 = time.time() + for r in fo: + ts = float(r[0]) + P = C_at(ts).compose( + Pose3(Rot3.Quaternion(r[7], r[4], r[5], r[6]), Point3(r[1], r[2], r[3])) + ) + x, y, zz, qx, qy, qz, qw = pose_tuple(P) + out.append( + Odometry( + ts=ts, + frame_id="odom", + child_frame_id="base_link", + pose=Pose(x, y, zz, qx, qy, qz, qw), + ), + ts=ts, + pose=(x, y, zz, qx, qy, qz, qw), + ) + n += 1 + if n % 20000 == 0: + print(f" {n}/{len(fo)} poses, {time.time() - t0:.0f}s", flush=True) + print(f"wrote {name}: {len(fo)} poses in {time.time() - t0:.0f}s", flush=True) + +if WHAT in ("lidar", "both"): + name = f"{OUT_PREFIX}_lidar{SUFFIX}" + fo_ts = fo[:, 0] + + def base_pose(ts): + j = int(np.searchsorted(fo_ts, ts)) + j = min(max(j, 0), len(fo) - 1) + if j > 0 and abs(fo_ts[j - 1] - ts) < abs(fo_ts[j] - ts): + j -= 1 + r = fo[j] + return Pose3(Rot3.Quaternion(r[7], r[4], r[5], r[6]), Point3(r[1], r[2], r[3])) + + if name in st.list_streams(): + st.delete_stream(name) + out = st.stream(name, PointCloud2) + + # The db stream stays per-scan, but the .pc2.lcm is ONE aggregated cloud (voxel-downsampled + + # statistical-outlier-removed), not 5184 per-scan events. Intensity rides through open3d's + # voxel averaging via the color channel. Chunks are collapsed every CHUNK scans to bound memory. + if WRITE_LCM: + import open3d as o3d + + CHUNK = 1000 + agg_xyz, agg_i = [], [] # incrementally voxel-downsampled chunks + buf_xyz, buf_i = [], [] + have_inten = False + + def collapse(xyz_list, i_list, voxel): + pc = o3d.geometry.PointCloud() + pc.points = o3d.utility.Vector3dVector(np.concatenate(xyz_list).astype(np.float64)) + carry = bool(i_list) + if carry: + inten_col = np.concatenate(i_list).astype(np.float64)[:, None] + pc.colors = o3d.utility.Vector3dVector(np.repeat(inten_col, 3, axis=1)) + pc = pc.voxel_down_sample(voxel) + dx = np.asarray(pc.points, np.float32) + di = np.asarray(pc.colors, np.float32)[:, 0] if carry else None + return dx, di + + print(f"writing {name} (corrected lidar)...", flush=True) + n = 0 + t0 = time.time() + for obs in st.stream(LIDAR_STREAM): + ts = float(obs.ts) + Cts = C_at(ts) + Rm = Cts.rotation().matrix() + tv = np.asarray(Cts.translation()) + xyz = world_points(obs) + inten = obs.data.intensities_f32() + xyz2 = (xyz @ Rm.T + tv).astype(np.float32) + m = PointCloud2.from_numpy( + xyz2, frame_id="odom", intensities=(np.asarray(inten) if inten is not None else None) + ) + m.ts = ts # stamp the cloud (lcm_encode needs a non-None ts) + out.append(m, ts=ts, pose=pose_tuple(Cts.compose(base_pose(ts)))) + if WRITE_LCM: + buf_xyz.append(xyz2) + if inten is not None: + have_inten = True + buf_i.append(np.asarray(inten, np.float32)) + if len(buf_xyz) >= CHUNK: + dx, di = collapse(buf_xyz, buf_i if have_inten else [], LCM_VOXEL) + agg_xyz.append(dx) + if di is not None: + agg_i.append(di) + buf_xyz, buf_i = [], [] + n += 1 + if n % 2000 == 0: + print(f" {n} scans, {time.time() - t0:.0f}s", flush=True) + print(f"wrote {name}: {n} scans in {time.time() - t0:.0f}s", flush=True) + + if WRITE_LCM: + if buf_xyz: # flush remainder + dx, di = collapse(buf_xyz, buf_i if have_inten else [], LCM_VOXEL) + agg_xyz.append(dx) + if di is not None: + agg_i.append(di) + # final unified voxel pass over the per-chunk results, then statistical outlier removal + merged_i = agg_i if have_inten else [] + dx, di = collapse(agg_xyz, merged_i, LCM_VOXEL) + print( + f"aggregating .pc2.lcm: {len(dx):,} pts after voxel, removing outliers...", flush=True + ) + pc = o3d.geometry.PointCloud() + pc.points = o3d.utility.Vector3dVector(dx.astype(np.float64)) + if di is not None: + pc.colors = o3d.utility.Vector3dVector( + np.repeat(di.astype(np.float64)[:, None], 3, axis=1) + ) + pc, _keep = pc.remove_statistical_outlier(LCM_OUTLIER_NN, LCM_OUTLIER_STD) + merged_xyz = np.asarray(pc.points, np.float32) + merged_inten = np.asarray(pc.colors, np.float32)[:, 0] if di is not None else None + merged = PointCloud2.from_numpy(merged_xyz, frame_id="odom", intensities=merged_inten) + merged.ts = float(fo_ts[0]) + lcm_path = REC / f"{name}.pc2.lcm" + lcm_path.write_bytes(merged.lcm_encode()) + print( + f"wrote {lcm_path}: 1 aggregated cloud, {len(merged_xyz):,} pts " + f"(voxel {LCM_VOXEL} m, outlier nn={LCM_OUTLIER_NN}/std={LCM_OUTLIER_STD})", + flush=True, + ) + +# build + open the comparison rrd +if OPEN_RRD and WHAT in ("lidar", "both"): + import subprocess + + from dimos.navigation.jnav.components.loop_closure.gsc_pgo import make_rrd + + print("building comparison rrd...", flush=True) + rrd_path = make_rrd.build( + REC, lidar_stream=LIDAR_STREAM, odom_stream=ODOM_STREAM, tag_stream=RAW_STREAM + ) + rerun_bin = Path(sys.executable).parent / "rerun" + if rerun_bin.exists(): + subprocess.Popen([str(rerun_bin), str(rrd_path)]) + print(f"opened {rrd_path}", flush=True) + else: + print(f"rerun binary not found at {rerun_bin}; open manually: rerun {rrd_path}", flush=True) From 898fdef86d13b3a361cd65db4573e7386e426367 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 14:53:05 +0800 Subject: [PATCH 11/69] add jnav map post-processing and PGO migration docs --- experimental/docs/jnav/map_postprocessing.md | 75 ++++++++++ experimental/docs/jnav/pgo_migration_plan.md | 144 +++++++++++++++++++ 2 files changed, 219 insertions(+) create mode 100644 experimental/docs/jnav/map_postprocessing.md create mode 100644 experimental/docs/jnav/pgo_migration_plan.md diff --git a/experimental/docs/jnav/map_postprocessing.md b/experimental/docs/jnav/map_postprocessing.md new file mode 100644 index 0000000000..9311066e6c --- /dev/null +++ b/experimental/docs/jnav/map_postprocessing.md @@ -0,0 +1,75 @@ +# Map Postprocessing + +You recorded a run. The lidar map drifted. You want a clean one to compare against — ground truth, basically, without a motion-capture rig. + +This is the offline fix. Point it at a recording, it bends the trajectory back into shape using AprilTags it saw along the way, then snaps the local geometry together with ICP. Out comes a corrected map written back into the same recording. + +It runs on a `.db` after the fact. It is not part of the live nav stack and never touches the robot. + +## What you need in the recording + +A recording dir with `mem2.db` plus `camera_intrinsics.json`, and these streams inside the db: + +- a camera stream (`color_image`) +- odometry (`pointlio_odometry`) +- world-registered lidar scans (`pointlio_lidar`) +- AprilTags physically in the scene, sized and known + +The tags are the whole trick. Drift accumulates, but a tag you saw at minute 1 and again at minute 9 is the *same tag* — so the two sightings have to land in the same spot. That constraint is what pulls the map straight. + +## The three steps + +The scripts live in `dimos/navigation/jnav/components/loop_closure/gsc_pgo/`. + +**1. Detect the tags.** Run the camera frames through detection and write both the raw and the gated tag streams in one step: + +``` +python dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py --rec=PATH +``` + +`add_april.py` writes `raw_april_tags` (every detection, unfiltered, with its quality numbers attached — this is what postprocessing reads), the gated/clustered `april_tags` stream, and an `april_tags` section in the recording's `summary.json` (see below). Leaving `raw_april_tags` unfiltered matters: you tune the quality gates later without re-running detection, which is the slow part. (`detect_tags.py` writes just the raw stream if that's all you want; `--dynamic 17` keeps a moving tag in raw but drops it from the gated stream.) + +Inspect what was found without rebuilding anything — per tag, raw count and revisit count, flagging any tag never revisited (your fast check for whether a recording even has loop constraints): + +``` +python dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py --rec=PATH --summary +``` + +**2. Solve.** Two stages, one command: + +``` +python dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py both --rec=PATH +``` + +- **Tag PGO (GTSAM).** Odometry between-poses are stiff on roll/pitch and z (gravity isn't drifting) and loose on yaw, where the real error lives. The tag sightings are landmark factors, weighted by how good each detection was. This fixes the big-picture drift. +- **ICP refinement.** Lidar submaps that are close in space but far apart in time get aligned to each other. This cleans up the local geometry the tags don't directly constrain. + +It writes `gt_pointlio_odometry` and `gt_pointlio_lidar` back into the db, optionally a `.pc2.lcm` of the corrected cloud, and opens a comparison view. Run `odom`, `lidar`, or `both`. + +**3. Look at it.** `post_process.py` opens the rrd for you, but you can rebuild it anytime: + +``` +python dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py --rec=PATH +``` + +Raw cloud in red, every `gt_*` version in its own color, tag landmarks marked. Add another correction method and it shows up automatically — good for comparing approaches side by side. + +## What `add_april.py` records in `summary.json` + +It merges an `april_tags` section into the recording's `summary.json` (other keys preserved; it never touches `camera_intrinsics.json`): + +- `filter_parameters` — the exact gate thresholds used, marker size, dictionary, and any dynamic tags excluded. +- `result` — `all_unfiltered_tag_ids`, and per tag the raw detection count, filtered visits, and revisit count, plus the list of tags never revisited. + +So the gates and the outcome are auditable after the fact. `--summary` recomputes just the `result` from the existing streams (read-only). + +## Knobs worth knowing + +- `--no-icp` — tag PGO only, skip the ICP stage. +- `--no-lcm` / `--no-rrd` — skip the cloud export / the viewer. +- `--out=NAME` — output prefix, if you want to keep several corrections in one db. +- Tag quality gates (sharpness, reprojection error, distance, view angle, motion blur) are single-sourced in `dimos/navigation/jnav/utils/apriltags.py` (the `DEFAULT_*` constants); `post_process.py` and the eval both import them. They're relaxed by default to keep more sightings. Tighten them there if a bad tag pose is yanking the map around. + +## When it won't help + +No tags in the scene, or tags seen only once, means no loop constraints — you get ICP cleanup and not much else. Same story if the camera never got a clean look at a tag. Garbage detections in, garbage map out; that's what the gates are for. diff --git a/experimental/docs/jnav/pgo_migration_plan.md b/experimental/docs/jnav/pgo_migration_plan.md new file mode 100644 index 0000000000..a35efdc31f --- /dev/null +++ b/experimental/docs/jnav/pgo_migration_plan.md @@ -0,0 +1,144 @@ +# PGO → jnav Loop-Closure Migration Plan + +Goal: land the pose-graph-optimization (PGO) / loop-closure work on +`jeff/feat/jnav_pgo` in the new `jnav` layout, extract the C++ + nix flake into a +standalone `github:jeff-hykin/gsc_pgo` repo, and merge in the offline AprilTag +map-postprocessing tooling. + +## Source material + +Two branches hold the pieces; neither alone is the target. + +| Source | Has | Layout | +|---|---|---| +| `jeff/feat/jnav` | `jnav/{msgs,utils}/`, `loop_closure/{gsc_pgo,ivan_pgo,ivan_pgo_transformer,unrefined_pgo}`, `eval.py`, `eval_all.py`, Scan-Context + Landmark C++ | `dimos/navigation/jnav/modules/...` (inline C++) | +| `jeff/feat/better_pgo` | postprocessing scripts (`add_april.py`, `detect_tags.py`, `post_process.py`, `make_rrd.py`), `map_postprocessing.md` doc | `dimos/navigation/nav_stack/modules/pgo/scripts/...` | +| `jeff/feat/jnav_pgo` (current) | nothing yet — only `nav_stack→cmu_nav` rename on top of `main` | — | + +`jeff/feat/jnav` is the more evolved navigation reorg; `better_pgo` is the source +for the AprilTag postprocessing scripts/doc. The migration is the union of both. + +## Target layout (on `jeff/feat/jnav_pgo`) + +``` +dimos/navigation/jnav/ + msgs/ # Graph3D/GraphDelta3D/Landmark/Marker (.py + .hpp + tests) + utils/ # ALL pgo utils (apriltags, trajectory_metrics, recording_db, ...) + components/ # renamed from jnav's "modules/" + loop_closure/ + eval.py + eval_all.py + spec.py # LoopClosure Protocol + gsc_pgo/ + module.py # NativeModule -> builds external gsc_pgo flake + post_process.py # ported from better_pgo scripts/post_process.py + (add_april.py, detect_tags.py, make_rrd.py) # the other postprocess scripts + ivan_pgo/ ivan_pgo_transformer/ unrefined_pgo/ # comparison baselines (eval_all needs them) +experimental/docs/jnav/ + map_postprocessing.md # adapted from docs/capabilities/navigation/map_postprocessing.md +``` + +External repo: +``` +github:jeff-hykin/gsc_pgo # all C++ + flake.nix/flake.lock + flake.nix, flake.lock, CMakeLists.txt + main.cpp, simple_pgo.{cpp,h}, scan_context.{cpp,h}, + commons.{cpp,h}, point_cloud_utils.hpp, dimos_native_module.hpp, + pgo_landmark_test.cpp + msgs/ (Graph3D.hpp, GraphDelta3D.hpp, Landmark.hpp) # C++ wire helpers +``` + +## Decisions (confirmed by Jeff, 2026-06-22) + +1. **Base branch.** ✅ Base `jnav_pgo` on `jeff/feat/jnav`. BUT `post_process.py` + must come from `better_pgo` (not jnav, which lacks it). +2. **`components/` scope.** ⏳ pending (whole `jnav/modules/`→`components/` vs only + `loop_closure`). +3. **`gsc_pgo` repo.** ✅ **public**, created at + `github.com/jeff-hykin/gsc_pgo`, initial rev `494e7a1d657c3702ec805c9e3d251a2fe8bc9529`. + Flake input will pin to that rev: `github:jeff-hykin/gsc_pgo/494e7a1...#default`. + +## Work breakdown + +### Phase 0 — branch setup +- Confirm `jnav_pgo` base (decision 1). If basing on jnav: merge/cherry-pick the + `dimos/navigation/jnav/` tree onto current `jnav_pgo` (which already has the + `cmu_nav` rename). Resolve any overlap with cmu_nav's own `pgo` module. + +### Phase 1 — external `gsc_pgo` repo +- `gh repo create jeff-hykin/gsc_pgo` (visibility per decision 3). *(outward-facing — confirm first)* +- Move `gsc_pgo/cpp/*` (C++, CMakeLists, flake.nix, flake.lock, `msgs/*.hpp`, + `pgo_landmark_test.cpp`) into the new repo. Keep the flake's + `lcm-extended` / `dimos-lcm` / `gtsam-extended` inputs as-is. +- `flake.nix`: `src = ./.;` already self-contained — switching to a tracked git + repo *fixes* the untracked-files gotcha that forced `path:$PWD#default` + (see memory: nix untracked-files gotcha). Verify `nix build .#default` from a + clean checkout. +- Push, tag/record the rev for the flake pin. + +### Phase 2 — `gsc_pgo` module.py +- Update `PGOConfig`: + - `build_command = 'nix build "github:jeff-hykin/gsc_pgo/#default" --no-write-lock-file'` + - drop the `path:$PWD` workaround comment; `cwd` no longer needs a local `cpp/`. + - `executable` resolves from the nix `result/bin/pgo` (confirm NativeModule + out-of-tree build path handling). +- Keep `In/Out` stream wiring and all loop-closure hyperparameters unchanged. +- Confirm C++ `msgs/*.hpp` wire format still matches `jnav/msgs/*.py` decoders + (they're hand-synced; the `.hpp` headers travel with the C++ repo). + +### Phase 3 — msgs/ + utils/ +- `jnav/msgs/`: already present on jnav (Graph3D, GraphDelta3D, Landmark, Marker + + `.hpp` + tests). Bring over as-is. The `.hpp` files are duplicated into the + external repo (C++ side); decide whether the canonical `.hpp` lives in + `gsc_pgo` and `jnav/msgs/*.hpp` is a mirror, or vice-versa. Recommend: canonical + in `gsc_pgo`, keep a note in `jnav/msgs` pointing at it. +- `jnav/utils/`: bring all utils. `better_pgo`'s `eval_utils/` (apriltags, + apriltag_agreement, trajectory_metrics, recording_db, voxel_map, module_loading) + are already absorbed into jnav `utils/` — verify no newer logic in + `better_pgo` was lost (diff the 6 overlapping files). + +### Phase 4 — eval.py / eval_all.py +- Copy from jnav `loop_closure/`. Fix imports to `dimos.navigation.jnav.utils.*` + and `...jnav.msgs.*` (already correct on jnav). +- `eval_all.py` enumerates the comparison modules (`gsc_pgo`, `ivan_pgo`, + `ivan_pgo_transformer`, `unrefined_pgo`) → port all four component dirs so eval + doesn't break. (If baselines are unwanted, prune eval_all's list instead.) + +### Phase 5 — postprocessing scripts (from better_pgo) +- Port `scripts/post_process.py` → `components/loop_closure/gsc_pgo/post_process.py`. +- Port `add_april.py`, `detect_tags.py`, `make_rrd.py` alongside it (the doc's + 3-step flow references all of them). +- Rewrite their internal imports: `eval_utils.apriltags` → `dimos.navigation.jnav.utils.apriltags`, etc. +- Tag quality gates are single-sourced in `utils/apriltags.py` (`DEFAULT_*`); keep + post_process importing from there (don't duplicate constants). + +### Phase 6 — docs +- Adapt `docs/capabilities/navigation/map_postprocessing.md` → + `experimental/docs/jnav/map_postprocessing.md`. +- Rewrite every script path in the doc: + `dimos/navigation/nav_stack/modules/pgo/scripts/X.py` + → `dimos/navigation/jnav/components/loop_closure/gsc_pgo/X.py`. +- Update the `eval_utils/apriltags.py` reference → `jnav/utils/apriltags.py`. +- Add a back-link from the navigation readme if appropriate. + +### Phase 7 — wiring + tests +- `spec.py` `LoopClosure` Protocol: bring over, update any module paths. +- Blueprints: if any blueprint references the old `nav_stack`/`pgo` module path, + repoint to `jnav.components.loop_closure.gsc_pgo`. Regenerate + `all_blueprints.py` via `pytest dimos/robot/test_all_blueprints_generation.py`. +- Port tests (`test_pgo_synthetic_drift.py`, msgs tests, utils tests). Run + `uv run pytest` for the loop_closure + msgs + utils dirs. +- `ruff check --fix && ruff format`. + +## Risks / cut-corners to flag +- **C++ wire-format drift**: `.hpp` (C++) and `.py` (decode) are hand-synced. Once + the `.hpp` lives in an external repo, a change there can silently desync the + Python decoder. Mitigate: keep a `test_Graph3D` round-trip test in `jnav/msgs` + that builds against the pinned flake rev, or at least a schema-comment check. +- **Private flake in CI**: if `gsc_pgo` is private, dimos CI nix builds will fail + without a token — do not add a token to CI without asking (secrets rule). +- **eval_results/**: jnav carries committed `eval_results/*/summary.json` snapshots. + Decide whether to bring those (history/benchmarks) or regenerate. +- **Baseline PGO impls**: `ivan_pgo*` / `unrefined_pgo` carry their own inline C++ + (`unrefined_pgo/cpp`). If only `gsc_pgo` goes external, the layout is asymmetric + — acceptable (they're experimental baselines) but worth noting. From ab31c4b72a3005082f122bd5b0f5939418165556 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 00:27:07 -0700 Subject: [PATCH 12/69] recorder: count + warn on frames dropped by LATEST coalescing process_observable gains an optional on_drop callback fired once per message dropped by the dispatcher's single-slot LATEST mailbox. The Recorder uses it to count dropped frames per stream and log a throttled warning, so a slow sink no longer loses data silently. --- dimos/core/module.py | 18 ++++++++++++++---- dimos/memory2/module.py | 22 +++++++++++++++++++++- 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index 26a2b6f893..fdf6d23061 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -477,14 +477,16 @@ def process_observable( self, observable: "Observable[Any]", async_cb: Callable[[Any], Any], + on_drop: Callable[[], None] | None = None, ) -> "DisposableBase": """Subscribe `async_cb` (an async function) to `observable`, dispatching each emitted value onto self._loop. Invocations are serialized through a - per-subscription dispatcher task with LATEST coalescing. The subscription + per-subscription dispatcher task with LATEST coalescing. `on_drop`, if + given, fires once per message dropped by that coalescing. The subscription is registered for cleanup on stop().""" if not inspect.iscoroutinefunction(async_cb): raise TypeError("process_observable requires an `async def` callback") - on_msg, dispatcher_disp = self._make_async_dispatch(async_cb) + on_msg, dispatcher_disp = self._make_async_dispatch(async_cb, on_drop) sub = observable.subscribe(on_msg) return self.register_disposable(CompositeDisposable(sub, dispatcher_disp)) @@ -635,7 +637,9 @@ def _auto_bind_handlers(self) -> None: self.process_observable(in_stream.pure_observable(), handler) def _make_async_dispatch( - self, async_handler: Callable[[Any], Any] + self, + async_handler: Callable[[Any], Any], + on_drop: Callable[[], None] | None = None, ) -> tuple[Callable[[Any], None], "DisposableBase"]: """Build a sync callback that delivers `msg` into a single-slot LATEST mailbox drained by a dedicated dispatcher task on `self._loop`. @@ -645,7 +649,9 @@ def _make_async_dispatch( awaits). - If messages arrive faster than the handler can process them, intermediate messages are dropped and only the most recent unprocessed - message is kept (LATEST policy). + message is kept (LATEST policy). `on_drop`, if given, is called once + per dropped message (on the loop thread) so callers that need every + message can surface the loss. - The returned Disposable cancels the dispatcher task. """ loop = self._loop @@ -685,6 +691,10 @@ def on_msg(msg: Any) -> None: return def _set() -> None: + # A slot that still holds an unconsumed value is about to be + # overwritten — that queued message is being dropped (LATEST). + if slot["has_value"] and on_drop is not None: + on_drop() slot["value"] = msg slot["has_value"] = True event.set() diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index e16c1527e7..cc3aebc67d 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -318,6 +318,9 @@ async def _lidar_pose(self, msg): config: RecorderConfig _pose_setters: dict[str, Any] = {} + # Per-stream count of frames lost to the dispatcher's LATEST coalescing + # (sink slower than input). Populated lazily as drops happen. + _dropped_frames: dict[str, int] = {} @rpc def start(self) -> None: @@ -330,6 +333,7 @@ def start(self) -> None: return self._pose_setters = self._collect_pose_setters() + self._dropped_frames = {} # TODO: store reset API/logic is not implemented yet. This module # shouldn't need to know about files (SqliteStore specific), and @@ -391,7 +395,23 @@ async def on_msg(msg: Any) -> None: ) stream.append(msg, ts=ts, pose=pose) - self.process_observable(input_topic.pure_observable(), on_msg) + self.process_observable( + input_topic.pure_observable(), on_msg, on_drop=lambda: self._on_frame_dropped(name) + ) + + def _on_frame_dropped(self, name: str) -> None: + """A frame for *name* was dropped because the sink couldn't keep up with + the input rate (dispatcher LATEST coalescing). Count it and warn — once, + then on each power-of-ten — so silent data loss is visible without + flooding the log.""" + count = self._dropped_frames.get(name, 0) + 1 + self._dropped_frames[name] = count + if count == 1 or count % 1000 == 0: + logger.warning( + "[%s] Recorder dropped %d frame(s) — sink slower than input; recording is lossy", + name, + count, + ) def _prepare_streams(self) -> None: """On APPEND, drop the streams this recorder is about to (re)write — the From d6c0c664bc8a77def35a5b1d02c8d06be476eff7 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 15:32:38 +0800 Subject: [PATCH 13/69] remove comparision modules --- .../loop_closure/ivan_pgo/module.py | 637 ------------------ .../ivan_pgo_transformer/module.py | 263 -------- .../unrefined_pgo/cpp/CMakeLists.txt | 49 -- .../unrefined_pgo/cpp/commons.cpp | 8 - .../loop_closure/unrefined_pgo/cpp/commons.h | 29 - .../unrefined_pgo/cpp/dimos_native_module.hpp | 97 --- .../loop_closure/unrefined_pgo/cpp/flake.lock | 144 ---- .../loop_closure/unrefined_pgo/cpp/flake.nix | 70 -- .../loop_closure/unrefined_pgo/cpp/main.cpp | 299 -------- .../unrefined_pgo/cpp/point_cloud_utils.hpp | 170 ----- .../unrefined_pgo/cpp/simple_pgo.cpp | 211 ------ .../unrefined_pgo/cpp/simple_pgo.h | 78 --- .../loop_closure/unrefined_pgo/module.py | 286 -------- 13 files changed, 2341 deletions(-) delete mode 100644 dimos/navigation/jnav/components/loop_closure/ivan_pgo/module.py delete mode 100644 dimos/navigation/jnav/components/loop_closure/ivan_pgo_transformer/module.py delete mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/CMakeLists.txt delete mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.cpp delete mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.h delete mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/dimos_native_module.hpp delete mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.lock delete mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.nix delete mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/main.cpp delete mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/point_cloud_utils.hpp delete mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.cpp delete mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.h delete mode 100644 dimos/navigation/jnav/components/loop_closure/unrefined_pgo/module.py diff --git a/dimos/navigation/jnav/components/loop_closure/ivan_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/ivan_pgo/module.py deleted file mode 100644 index 7c86f6dfb5..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/ivan_pgo/module.py +++ /dev/null @@ -1,637 +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. - -"""Pure-Python PGO ported from Ivan's `ivan/feat/go2loopclosure` branch -(`dimos/mapping/pgo.py`), adapted to the jnav LoopClosure spec. - -GTSAM iSAM2 pose graph + Open3D point-to-plane ICP loop verification + -KD-tree loop candidate search. Differences from the original: - * input renamed `registered_scan` -> `lidar` (LoopClosure spec) - * publishes `pose_graph: Out[Graph3D]` (optimized keyframes, odometry + - loop edges) and `loop_closure_event: Out[GraphDelta3D]` so the eval - harness can capture the corrected trajectory and count closures - * offline experiment helpers (transformers, two-pass voxel pipelines) - were left behind — this is just the runtime module -""" - -from __future__ import annotations - -from dataclasses import dataclass -import threading -import time -from typing import Any - -import gtsam # type: ignore[import-untyped] -import numpy as np -import open3d as o3d # type: ignore[import-untyped] -import open3d.core as o3c # type: ignore[import-untyped] -from reactivex.disposable import Disposable -from scipy.spatial import KDTree -from scipy.spatial.transform import Rotation - -from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT -from dimos.core.core import rpc -from dimos.core.module import Module, ModuleConfig -from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -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 -from dimos.navigation.jnav.components.loop_closure.spec import LoopClosure -from dimos.navigation.jnav.msgs.Graph3D import Graph3D -from dimos.navigation.jnav.msgs.GraphDelta3D import GraphDelta3D -from dimos.utils.logging_config import setup_logger - -FRAME_MAP = "map" -FRAME_ODOM = "odom" -FRAME_BODY = "base_link" - -logger = setup_logger() - - -class PGOConfig(ModuleConfig): - world_frame: str = FRAME_MAP - - # Keyframe detection - key_pose_delta_trans: float = 0.5 - key_pose_delta_deg: float = 10.0 - - # Loop closure - loop_search_radius: float = 2.0 - loop_time_thresh: float = 20.0 - loop_score_thresh: float = 0.3 - loop_submap_half_range: int = 10 - min_icp_inliers: int = 10 - min_keyframes_for_loop_search: int = 10 - loop_closure_extra_iterations: int = 4 - submap_resolution: float = 0.2 - min_loop_detect_duration: float = 5.0 - - # Input mode - unregister_input: bool = True # Transform world-frame scans to body-frame using odom - - # Global map - publish_global_map: bool = True - global_map_publish_rate: float = 0.5 - global_map_voxel_size: float = 0.15 - - # ICP - max_icp_iterations: int = 50 - max_icp_correspondence_dist: float = 1.0 - - -@dataclass -class _KeyPose: - r_local: np.ndarray # 3x3 rotation in local/odom frame - t_local: np.ndarray # 3-vec translation in local/odom frame - r_global: np.ndarray # 3x3 corrected rotation - t_global: np.ndarray # 3-vec corrected translation - timestamp: float - body_cloud: np.ndarray # Nx3 points in body frame - - -def _icp( - source: np.ndarray, - target: np.ndarray, - max_iter: int = 50, - max_dist: float = 1.0, - tol: float = 1e-6, - min_inliers: int = 10, - init: np.ndarray | None = None, -) -> tuple[np.ndarray, float]: - """Point-to-plane ICP using Open3D's tensor pipeline. - - Returns ``(T, fitness)`` where ``fitness`` is mean squared inlier - distance (m²).""" - if len(source) < min_inliers or len(target) < min_inliers: - return np.eye(4), float("inf") - - cpu = o3c.Device("CPU:0") - src_pcd = o3d.t.geometry.PointCloud(o3c.Tensor(source.astype(np.float32), device=cpu)) - tgt_pcd = o3d.t.geometry.PointCloud(o3c.Tensor(target.astype(np.float32), device=cpu)) - - # Normals on the target enable point-to-plane ICP, which converges - # tighter than point-to-point on indoor scenes (walls give unambiguous - # normals that resolve the slide-along-wall ambiguity). - tgt_pcd.estimate_normals(max_nn=30, radius=0.3) - - init_T = ( - o3c.Tensor(init.astype(np.float64), dtype=o3c.float64, device=cpu) - if init is not None - else o3c.Tensor.eye(4, dtype=o3c.float64, device=cpu) - ) - - # Silence Open3D's "0 correspondence" warning — we deliberately use a - # tight max_correspondence_distance and reject loops with poor fitness. - with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Error): - result = o3d.t.pipelines.registration.icp( - source=src_pcd, - target=tgt_pcd, - max_correspondence_distance=max_dist, - init_source_to_target=init_T, - estimation_method=o3d.t.pipelines.registration.TransformationEstimationPointToPlane(), - criteria=o3d.t.pipelines.registration.ICPConvergenceCriteria( - relative_fitness=tol, - relative_rmse=tol, - max_iteration=max_iter, - ), - ) - - fitness_inlier_frac = float(result.fitness) - if fitness_inlier_frac == 0.0: - return np.eye(4), float("inf") - - rmse = float(result.inlier_rmse) - T = result.transformation.numpy() - return T, rmse * rmse - - -def _voxel_downsample(pts: np.ndarray, voxel_size: float) -> np.ndarray: - if len(pts) == 0 or voxel_size <= 0: - return pts - keys = np.floor(pts / voxel_size).astype(np.int32) - _, idx = np.unique(keys, axis=0, return_index=True) - return pts[idx] - - -class _SimplePGO: - def __init__(self, config: PGOConfig) -> None: - self._cfg = config - self._key_poses: list[_KeyPose] = [] - self._history_pairs: list[tuple[int, int]] = [] - self._cache_pairs: list[dict[str, Any]] = [] - self._r_offset = np.eye(3) - self._t_offset = np.zeros(3) - - params = gtsam.ISAM2Params() - params.setRelinearizeThreshold(0.01) - params.relinearizeSkip = 1 - self._isam2 = gtsam.ISAM2(params) - self._graph = gtsam.NonlinearFactorGraph() - self._values = gtsam.Values() - - def is_key_pose(self, r: np.ndarray, t: np.ndarray) -> bool: - if not self._key_poses: - return True - last = self._key_poses[-1] - delta_trans = np.linalg.norm(t - last.t_local) - # Angular distance via quaternion dot product - q_cur = Rotation.from_matrix(r).as_quat() # [x,y,z,w] - q_last = Rotation.from_matrix(last.r_local).as_quat() - dot = abs(np.dot(q_cur, q_last)) - delta_deg = np.degrees(2.0 * np.arccos(min(dot, 1.0))) - return bool( - delta_trans > self._cfg.key_pose_delta_trans or delta_deg > self._cfg.key_pose_delta_deg - ) - - def add_key_pose( - self, r_local: np.ndarray, t_local: np.ndarray, timestamp: float, body_cloud: np.ndarray - ) -> bool: - if not self.is_key_pose(r_local, t_local): - return False - - idx = len(self._key_poses) - init_r = self._r_offset @ r_local - init_t = self._r_offset @ t_local + self._t_offset - - pose = gtsam.Pose3(gtsam.Rot3(init_r), gtsam.Point3(init_t)) - self._values.insert(idx, pose) - - if idx == 0: - noise = gtsam.noiseModel.Diagonal.Variances(np.full(6, 1e-12)) - self._graph.add(gtsam.PriorFactorPose3(idx, pose, noise)) - else: - last = self._key_poses[-1] - r_between = last.r_local.T @ r_local - t_between = last.r_local.T @ (t_local - last.t_local) - noise = gtsam.noiseModel.Diagonal.Variances( - np.array([1e-6, 1e-6, 1e-6, 1e-4, 1e-4, 1e-6]) - ) - self._graph.add( - gtsam.BetweenFactorPose3( - idx - 1, idx, gtsam.Pose3(gtsam.Rot3(r_between), gtsam.Point3(t_between)), noise - ) - ) - - kp = _KeyPose( - r_local=r_local.copy(), - t_local=t_local.copy(), - r_global=init_r.copy(), - t_global=init_t.copy(), - timestamp=timestamp, - body_cloud=_voxel_downsample(body_cloud, self._cfg.submap_resolution), - ) - self._key_poses.append(kp) - return True - - def _get_submap(self, idx: int, half_range: int) -> np.ndarray: - lo = max(0, idx - half_range) - hi = min(len(self._key_poses) - 1, idx + half_range) - parts = [] - for i in range(lo, hi + 1): - kp = self._key_poses[i] - world = (kp.r_global @ kp.body_cloud.T).T + kp.t_global - parts.append(world) - if not parts: - return np.empty((0, 3)) - cloud = np.vstack(parts) - return _voxel_downsample(cloud, self._cfg.submap_resolution) - - def search_for_loops(self) -> None: - if len(self._key_poses) < self._cfg.min_keyframes_for_loop_search: - return - - # Rate limit - if self._history_pairs: - cur_time = self._key_poses[-1].timestamp - last_time = self._key_poses[self._history_pairs[-1][1]].timestamp - if cur_time - last_time < self._cfg.min_loop_detect_duration: - return - - cur_idx = len(self._key_poses) - 1 - cur_kp = self._key_poses[-1] - - # Build KD-tree of previous keyframe positions - positions = np.array([kp.t_global for kp in self._key_poses[:-1]]) - tree = KDTree(positions) - - idxs = tree.query_ball_point(cur_kp.t_global, self._cfg.loop_search_radius) - if not idxs: - return - - # Pick the spatially closest keyframe that's also old enough in time. - # query_ball_point doesn't sort, so we sort by distance ourselves. - candidates = [ - (float(np.linalg.norm(self._key_poses[i].t_global - cur_kp.t_global)), i) - for i in idxs - if abs(cur_kp.timestamp - self._key_poses[i].timestamp) > self._cfg.loop_time_thresh - ] - if not candidates: - return - candidates.sort() - loop_idx = candidates[0][1] - - # ICP verification - target = self._get_submap(loop_idx, self._cfg.loop_submap_half_range) - source = self._get_submap(cur_idx, 0) - - transform, fitness = _icp( - source, - target, - max_iter=self._cfg.max_icp_iterations, - max_dist=self._cfg.max_icp_correspondence_dist, - min_inliers=self._cfg.min_icp_inliers, - ) - if fitness > self._cfg.loop_score_thresh: - return - - # Compute relative pose - R_icp = transform[:3, :3] - t_icp = transform[:3, 3] - r_refined = R_icp @ cur_kp.r_global - t_refined = R_icp @ cur_kp.t_global + t_icp - r_offset = self._key_poses[loop_idx].r_global.T @ r_refined - t_offset = self._key_poses[loop_idx].r_global.T @ ( - t_refined - self._key_poses[loop_idx].t_global - ) - - self._cache_pairs.append( - { - "source": cur_idx, - "target": loop_idx, - "r_offset": r_offset, - "t_offset": t_offset, - "score": fitness, - } - ) - self._history_pairs.append((loop_idx, cur_idx)) - logger.info( - "Loop closure detected", - source=cur_idx, - target=loop_idx, - score=round(fitness, 4), - ) - - def smooth_and_update(self) -> None: - has_loop = bool(self._cache_pairs) - - for pair in self._cache_pairs: - # Pose3 noise model is [rx, ry, rz, x, y, z]. Use ICP fitness as - # the translation variance and a generous fixed rotation variance — - # loops shouldn't be trusted to fix rotation tightly. - trans_var = max(0.01, float(pair["score"])) # >= sigma_trans = 10 cm - rot_var = 0.05 # sigma_rot ~= 13 deg - noise = gtsam.noiseModel.Diagonal.Variances( - np.array([rot_var, rot_var, rot_var, trans_var, trans_var, trans_var]) - ) - self._graph.add( - gtsam.BetweenFactorPose3( - pair["target"], - pair["source"], - gtsam.Pose3(gtsam.Rot3(pair["r_offset"]), gtsam.Point3(pair["t_offset"])), - noise, - ) - ) - self._cache_pairs.clear() - - self._isam2.update(self._graph, self._values) - self._isam2.update() - if has_loop: - for _ in range(self._cfg.loop_closure_extra_iterations): - self._isam2.update() - self._graph = gtsam.NonlinearFactorGraph() - self._values = gtsam.Values() - - estimates = self._isam2.calculateBestEstimate() - for i in range(len(self._key_poses)): - pose = estimates.atPose3(i) - self._key_poses[i].r_global = pose.rotation().matrix() - self._key_poses[i].t_global = pose.translation() - - last = self._key_poses[-1] - self._r_offset = last.r_global @ last.r_local.T - self._t_offset = last.t_global - self._r_offset @ last.t_local - - def get_corrected_pose( - self, r_local: np.ndarray, t_local: np.ndarray - ) -> tuple[np.ndarray, np.ndarray]: - return self._r_offset @ r_local, self._r_offset @ t_local + self._t_offset - - def build_global_map(self, voxel_size: float) -> np.ndarray: - if not self._key_poses: - return np.empty((0, 3), dtype=np.float32) - parts = [] - for kp in self._key_poses: - world = (kp.r_global @ kp.body_cloud.T).T + kp.t_global - parts.append(world) - cloud = np.vstack(parts).astype(np.float32) - return _voxel_downsample(cloud, voxel_size) - - @property - def num_key_poses(self) -> int: - return len(self._key_poses) - - -def process_scan( - pgo: _SimplePGO, - cloud: PointCloud2, - r_local: np.ndarray, - t_local: np.ndarray, - ts: float, - unregister_input: bool, -) -> tuple[Odometry, Transform, bool] | None: - """Add a keyframe, run loop closure, return (corrected odom, map->odom tf, - keyframe_added) — or None on empty cloud. - - Caller must hold ``pgo``'s lock during this call.""" - points, _ = cloud.as_numpy() - if len(points) == 0: - return None - - if unregister_input: - # registered_scan is world-frame; transform back to body-frame. - body_pts = (r_local.T @ (points[:, :3].T - t_local[:, None])).T - else: - body_pts = points[:, :3] - - added = pgo.add_key_pose(r_local, t_local, ts, body_pts) - if added: - pgo.search_for_loops() - pgo.smooth_and_update() - - r_corr, t_corr = pgo.get_corrected_pose(r_local, t_local) - return ( - build_corrected_odometry(r_corr, t_corr, ts), - build_map_odom_tf(pgo._r_offset.copy(), pgo._t_offset.copy(), ts), - added, - ) - - -def build_corrected_odometry( - r: np.ndarray, - t: np.ndarray, - ts: float, - world_frame: str = FRAME_MAP, -) -> Odometry: - q = Rotation.from_matrix(r).as_quat() # [x,y,z,w] - return Odometry( - ts=ts, - frame_id=world_frame, - child_frame_id=FRAME_BODY, - pose=Pose( - position=[float(t[0]), float(t[1]), float(t[2])], - orientation=[float(q[0]), float(q[1]), float(q[2]), float(q[3])], - ), - ) - - -def build_map_odom_tf( - r_offset: np.ndarray, - t_offset: np.ndarray, - ts: float, - world_frame: str = FRAME_MAP, - odom_frame: str = FRAME_ODOM, -) -> Transform: - q = Rotation.from_matrix(r_offset).as_quat() # [x,y,z,w] - return Transform( - frame_id=world_frame, - child_frame_id=odom_frame, - translation=Vector3(float(t_offset[0]), float(t_offset[1]), float(t_offset[2])), - rotation=Quaternion(float(q[0]), float(q[1]), float(q[2]), float(q[3])), - ts=ts, - ) - - -def _keyframe_node(index: int, key_pose: _KeyPose, world_frame: str) -> Graph3D.Node3D: - q = Rotation.from_matrix(key_pose.r_global).as_quat() # [x,y,z,w] - return Graph3D.Node3D( - pose=PoseStamped( - ts=key_pose.timestamp, - frame_id=world_frame, - position=[float(v) for v in key_pose.t_global], - orientation=[float(q[0]), float(q[1]), float(q[2]), float(q[3])], - ), - id=index, - ) - - -class PGO(Module, LoopClosure): - """Pose graph optimization with loop closure (pure Python). - - Detects keyframes, performs loop closure via ICP + KD-tree search, and - optimizes the pose graph with GTSAM iSAM2. Publishes corrected odometry, - the optimized pose graph, loop-closure events, and an accumulated - global map.""" - - config: PGOConfig - - lidar: In[PointCloud2] - odometry: In[Odometry] - corrected_odometry: Out[Odometry] - pose_graph: Out[Graph3D] - loop_closure_event: Out[GraphDelta3D] - global_map: Out[PointCloud2] - - def __init__(self, **kwargs: Any) -> None: - super().__init__(**kwargs) - self._running = False - self._thread: threading.Thread | None = None - self._pgo: _SimplePGO | None = None - self._latest_r = np.eye(3) - self._latest_t = np.zeros(3) - self._latest_time = 0.0 - self._has_odom = False - self._last_global_map_time = 0.0 - self._published_loops = 0 - self._lock = threading.Lock() - # Protects _pgo mutations (add_key_pose, search_for_loops, - # smooth_and_update, build_global_map) against concurrent access - # from _on_scan and _publish_loop threads. - self._pgo_lock = threading.Lock() - - @rpc - def start(self) -> None: - super().start() - self._pgo = _SimplePGO(self.config) - # Identity map -> odom so consumers querying map -> body get a result - # before any loop-closure correction exists. - self.tf.publish(build_map_odom_tf(np.eye(3), np.zeros(3), time.time())) - self.register_disposable(Disposable(self.odometry.subscribe(self._on_odom))) - self.register_disposable(Disposable(self.lidar.subscribe(self._on_scan))) - self._running = True - if self.config.publish_global_map: - self._thread = threading.Thread(target=self._publish_global_map_loop, daemon=True) - self._thread.start() - logger.info( - "PGO module started (gtsam iSAM2, pure python)", - publish_global_map=self.config.publish_global_map, - ) - - @rpc - def stop(self) -> None: - self._running = False - if self._thread: - self._thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) - super().stop() - - def _on_odom(self, msg: Odometry) -> None: - q = [ - msg.pose.orientation.x, - msg.pose.orientation.y, - msg.pose.orientation.z, - msg.pose.orientation.w, - ] - r = Rotation.from_quat(q).as_matrix() - t = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z]) - with self._lock: - self._latest_r = r - self._latest_t = t - self._latest_time = msg.ts if msg.ts else time.time() - self._has_odom = True - - def _on_scan(self, cloud: PointCloud2) -> None: - with self._lock: - if not self._has_odom: - return - r_local = self._latest_r.copy() - t_local = self._latest_t.copy() - ts = self._latest_time - - pgo = self._pgo - assert pgo is not None - - with self._pgo_lock: - result = process_scan(pgo, cloud, r_local, t_local, ts, self.config.unregister_input) - if result is None: - return - corrected_odom, tf_msg, keyframe_added = result - if keyframe_added: - graph_msg, loop_events = self._snapshot_graph(pgo, ts) - else: - graph_msg, loop_events = None, [] - - self.corrected_odometry.publish(corrected_odom) - self.tf.publish(tf_msg) - if graph_msg is not None: - self.pose_graph.publish(graph_msg) - for event in loop_events: - self.loop_closure_event.publish(event) - - def _snapshot_graph(self, pgo: _SimplePGO, ts: float) -> tuple[Graph3D, list[GraphDelta3D]]: - """The optimized graph (odometry chain + loop edges) and one - GraphDelta3D per loop pair not yet published. - - Caller must hold ``_pgo_lock``.""" - world_frame = self.config.world_frame - nodes = [ - _keyframe_node(index, key_pose, world_frame) - for index, key_pose in enumerate(pgo._key_poses) - ] - edges = [ - Graph3D.Edge( - start_id=index - 1, end_id=index, timestamp=pgo._key_poses[index].timestamp - ) - for index in range(1, len(pgo._key_poses)) - ] - edges += [ - Graph3D.Edge(start_id=target, end_id=source, metadata_id=1) - for target, source in pgo._history_pairs - ] - graph_msg = Graph3D(ts=ts, nodes=nodes, edges=edges) - - loop_events: list[GraphDelta3D] = [] - identity = GraphDelta3D.Transform( - translation=Vector3(0.0, 0.0, 0.0), rotation=Quaternion(0.0, 0.0, 0.0, 1.0) - ) - for target, source in pgo._history_pairs[self._published_loops :]: - loop_events.append( - GraphDelta3D( - ts=ts, - nodes=[ - _keyframe_node(target, pgo._key_poses[target], world_frame), - _keyframe_node(source, pgo._key_poses[source], world_frame), - ], - transforms=[identity, identity], - ) - ) - self._published_loops = len(pgo._history_pairs) - return graph_msg, loop_events - - def _publish_global_map_loop(self) -> None: - pgo = self._pgo - assert pgo is not None - rate = self.config.global_map_publish_rate - interval = 1.0 / rate if rate > 0 else 2.0 - - while self._running: - t0 = time.monotonic() - - if t0 - self._last_global_map_time > interval and pgo.num_key_poses > 0: - with self._pgo_lock: - cloud_np = pgo.build_global_map(self.config.global_map_voxel_size) - if len(cloud_np) > 0: - now = time.time() - self.global_map.publish( - PointCloud2.from_numpy( - cloud_np, frame_id=self.config.world_frame, timestamp=now - ) - ) - self._last_global_map_time = t0 - - elapsed = time.monotonic() - t0 - sleep_time = max(DEFAULT_THREAD_JOIN_TIMEOUT, interval - elapsed) - time.sleep(sleep_time) diff --git a/dimos/navigation/jnav/components/loop_closure/ivan_pgo_transformer/module.py b/dimos/navigation/jnav/components/loop_closure/ivan_pgo_transformer/module.py deleted file mode 100644 index 17e7c84511..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/ivan_pgo_transformer/module.py +++ /dev/null @@ -1,263 +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. - -"""Ivan's CURRENT PGO (dimos/mapping/loop_closure/pgo.py, main, June 2026) -forced into an online LoopClosure module. - -That code ships as an offline memory2 Transformer, but its `_PGOState` core is -already incremental — one `process(pose, ts, world_cloud)` call per frame. This -wrapper brute-forces the online shape: every arriving scan is paired with the -latest odometry pose and pushed straight into `_PGOState`, then the corrected -odometry / pose graph / loop events are read back out of its (private) state. -No changes to the mapping code itself; this module owns the coupling. - -`_PGOState.process` expects world-frame registered scans (it unregisters them -internally via the odom pose) — set `input_is_registered: false` for raw -body-frame scans and they'll be pre-registered here first.""" - -from __future__ import annotations - -import threading -import time -from typing import Any - -import numpy as np -from scipy.spatial.transform import Rotation - -from dimos.core.core import rpc -from dimos.core.module import Module, ModuleConfig -from dimos.core.stream import In, Out -from dimos.mapping.loop_closure.pgo import ( - PGOConfig as MappingPGOConfig, - _PGOState, - _pose3_to_transform, -) -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -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 -from dimos.navigation.jnav.components.loop_closure.spec import LoopClosure -from dimos.navigation.jnav.msgs.Graph3D import Graph3D -from dimos.navigation.jnav.msgs.GraphDelta3D import GraphDelta3D -from dimos.utils.logging_config import setup_logger - -logger = setup_logger() - - -class PGOConfig(ModuleConfig): - world_frame: str = "map" - odom_frame: str = "odom" - body_frame: str = "base_link" - # World-frame registered input (fastlio); false = raw body-frame scans. - input_is_registered: bool = True - - # Mirrors mapping/loop_closure PGOConfig — forwarded verbatim. - key_pose_delta_trans: float = 0.5 - key_pose_delta_deg: float = 10.0 - loop_search_radius: float = 2.0 - loop_time_thresh: float = 20.0 - loop_score_thresh: float = 0.3 - loop_submap_half_range: int = 10 - min_icp_inliers: int = 10 - min_keyframes_for_loop_search: int = 10 - loop_closure_extra_iterations: int = 4 - submap_resolution: float = 0.2 - min_loop_detect_duration: float = 5.0 - max_icp_iterations: int = 50 - max_icp_correspondence_dist: float = 1.0 - odom_rot_var: float = 1e-6 - odom_trans_var_xy: float = 1e-4 - odom_trans_var_z: float = 1e-6 - loop_rot_var: float = 0.05 - - -def _mapping_config(config: PGOConfig) -> MappingPGOConfig: - fields = {name: getattr(config, name) for name in MappingPGOConfig.model_fields} - return MappingPGOConfig(**fields) - - -def _pose3_node(index: int, pose: Any, ts: float, world_frame: str) -> Graph3D.Node3D: - translation = np.asarray(pose.translation()) - quaternion = Rotation.from_matrix(pose.rotation().matrix()).as_quat() # [x,y,z,w] - return Graph3D.Node3D( - pose=PoseStamped( - ts=ts, - frame_id=world_frame, - position=[float(v) for v in translation], - orientation=[float(v) for v in quaternion], - ), - id=index, - ) - - -class PGO(Module, LoopClosure): - """Online wrapper over Ivan's current incremental PGO core (`_PGOState`).""" - - config: PGOConfig - - lidar: In[PointCloud2] - odometry: In[Odometry] - corrected_odometry: Out[Odometry] - pose_graph: Out[Graph3D] - loop_closure_event: Out[GraphDelta3D] - - def __init__(self, **kwargs: Any) -> None: - super().__init__(**kwargs) - self._state: _PGOState | None = None - self._latest_odom: Odometry | None = None - self._published_loops = 0 - self._odom_lock = threading.Lock() - self._state_lock = threading.Lock() - self._unsub_odom: Any = None - self._unsub_lidar: Any = None - - @rpc - def start(self) -> None: - super().start() - self._state = _PGOState(_mapping_config(self.config)) - # Identity map -> odom so consumers querying map -> body get a result - # before any correction exists. - self.tf.publish(self._correction_tf(np.eye(4), time.time())) - self._unsub_odom = self.odometry.subscribe(self._on_odom) - self._unsub_lidar = self.lidar.subscribe(self._on_scan) - logger.info("PGO (ivan transformer core) started") - - @rpc - def stop(self) -> None: - if self._unsub_odom is not None: - self._unsub_odom.dispose() - if self._unsub_lidar is not None: - self._unsub_lidar.dispose() - super().stop() - - def _on_odom(self, msg: Odometry) -> None: - with self._odom_lock: - self._latest_odom = msg - - def _on_scan(self, cloud: PointCloud2) -> None: - import gtsam # type: ignore[import-untyped] - - with self._odom_lock: - odom = self._latest_odom - if odom is None or len(cloud) == 0: - return - state = self._state - assert state is not None - - position = odom.pose.position - orientation = odom.pose.orientation - local_pose = gtsam.Pose3( - gtsam.Rot3.Quaternion(orientation.w, orientation.x, orientation.y, orientation.z), - gtsam.Point3(position.x, position.y, position.z), - ) - ts = odom.ts if odom.ts else time.time() - - if not self.config.input_is_registered: - cloud = cloud.transform( - _pose3_to_transform( - local_pose, - ts=ts, - frame_id=self.config.world_frame, - child_frame_id=self.config.body_frame, - ) - ) - - with self._state_lock: - state.process(local_pose, ts, cloud) - keyframe_count = len(state._key_poses) - correction = state._world_correction - corrected = correction.compose(local_pose) - graph_msg = self._snapshot_graph(state, ts) if keyframe_count else None - loop_events = self._new_loop_events(state, ts) - - self._publish_corrected_odometry(corrected, ts) - self.tf.publish(self._correction_tf(correction.matrix(), ts)) - if graph_msg is not None: - self.pose_graph.publish(graph_msg) - for event in loop_events: - self.loop_closure_event.publish(event) - - def _publish_corrected_odometry(self, pose: Any, ts: float) -> None: - translation = np.asarray(pose.translation()) - quaternion = Rotation.from_matrix(pose.rotation().matrix()).as_quat() - self.corrected_odometry.publish( - Odometry( - ts=ts, - frame_id=self.config.world_frame, - child_frame_id=self.config.body_frame, - pose=Pose( - position=[float(v) for v in translation], - orientation=[float(v) for v in quaternion], - ), - ) - ) - - def _correction_tf(self, matrix: np.ndarray, ts: float) -> Transform: - quaternion = Rotation.from_matrix(matrix[:3, :3]).as_quat() - return Transform( - frame_id=self.config.world_frame, - child_frame_id=self.config.odom_frame, - translation=Vector3(float(matrix[0, 3]), float(matrix[1, 3]), float(matrix[2, 3])), - rotation=Quaternion(*[float(v) for v in quaternion]), - ts=ts, - ) - - def _snapshot_graph(self, state: _PGOState, ts: float) -> Graph3D: - """Optimized keyframes + odometry-chain and loop edges. - - Caller must hold ``_state_lock``.""" - world_frame = self.config.world_frame - nodes = [ - _pose3_node(index, key_pose.optimized, key_pose.timestamp, world_frame) - for index, key_pose in enumerate(state._key_poses) - ] - edges = [ - Graph3D.Edge( - start_id=index - 1, end_id=index, timestamp=state._key_poses[index].timestamp - ) - for index in range(1, len(state._key_poses)) - ] - edges += [ - Graph3D.Edge(start_id=pair.target, end_id=pair.source, metadata_id=1) - for pair in state._accepted_loops - ] - return Graph3D(ts=ts, nodes=nodes, edges=edges) - - def _new_loop_events(self, state: _PGOState, ts: float) -> list[GraphDelta3D]: - """One GraphDelta3D per accepted loop not yet published. - - Caller must hold ``_state_lock``.""" - world_frame = self.config.world_frame - identity = GraphDelta3D.Transform( - translation=Vector3(0.0, 0.0, 0.0), rotation=Quaternion(0.0, 0.0, 0.0, 1.0) - ) - events: list[GraphDelta3D] = [] - for pair in state._accepted_loops[self._published_loops :]: - source = state._key_poses[pair.source] - target = state._key_poses[pair.target] - events.append( - GraphDelta3D( - ts=ts, - nodes=[ - _pose3_node(pair.target, target.optimized, target.timestamp, world_frame), - _pose3_node(pair.source, source.optimized, source.timestamp, world_frame), - ], - transforms=[identity, identity], - ) - ) - self._published_loops = len(state._accepted_loops) - return events diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/CMakeLists.txt b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/CMakeLists.txt deleted file mode 100644 index 8c7b6d5b94..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/CMakeLists.txt +++ /dev/null @@ -1,49 +0,0 @@ -cmake_minimum_required(VERSION 3.14) -project(pgo CXX) - -set(CMAKE_CXX_STANDARD 17) -set(CMAKE_CXX_STANDARD_REQUIRED ON) -set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3") - -include(FetchContent) -FetchContent_Declare(dimos_lcm - GIT_REPOSITORY https://github.com/dimensionalOS/dimos-lcm.git - GIT_TAG main - GIT_SHALLOW TRUE -) -FetchContent_MakeAvailable(dimos_lcm) - -find_package(PkgConfig REQUIRED) -pkg_check_modules(LCM REQUIRED lcm) -find_package(Eigen3 REQUIRED) -find_package(PCL 1.8 REQUIRED COMPONENTS common filters kdtree registration io) -find_package(GTSAM REQUIRED) - -add_definitions(-DUSE_PCL) - -add_executable(pgo - main.cpp - simple_pgo.cpp - commons.cpp -) - -target_include_directories(pgo PRIVATE - ${CMAKE_CURRENT_SOURCE_DIR} - ${dimos_lcm_SOURCE_DIR}/generated/cpp_lcm_msgs - ${LCM_INCLUDE_DIRS} - ${EIGEN3_INCLUDE_DIR} - ${PCL_INCLUDE_DIRS} - ${GTSAM_INCLUDE_DIR} -) - -target_link_libraries(pgo PRIVATE - ${LCM_LIBRARIES} - ${PCL_LIBRARIES} - gtsam -) - -target_link_directories(pgo PRIVATE - ${LCM_LIBRARY_DIRS} -) - -install(TARGETS pgo DESTINATION bin) diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.cpp b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.cpp deleted file mode 100644 index f309a97a85..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.cpp +++ /dev/null @@ -1,8 +0,0 @@ -#include "commons.h" - -void PoseWithTime::setTime(int32_t _sec, uint32_t _nsec) -{ - sec = _sec; - nsec = _nsec; - second = static_cast(sec) + static_cast(nsec) / 1e9; -} diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.h b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.h deleted file mode 100644 index 2f9244ede6..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/commons.h +++ /dev/null @@ -1,29 +0,0 @@ -#pragma once -#include -#include -#include - -using PointType = pcl::PointXYZI; -using CloudType = pcl::PointCloud; -using PointVec = std::vector>; - -using M3D = Eigen::Matrix3d; -using V3D = Eigen::Vector3d; -using M3F = Eigen::Matrix3f; -using V3F = Eigen::Vector3f; -using M4F = Eigen::Matrix4f; -using V4F = Eigen::Vector4f; - -struct PoseWithTime { - V3D t; - M3D r; - int32_t sec; - uint32_t nsec; - double second; - void setTime(int32_t sec, uint32_t nsec); -}; - -struct CloudWithPose { - CloudType::Ptr cloud; - PoseWithTime pose; -}; diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/dimos_native_module.hpp b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/dimos_native_module.hpp deleted file mode 100644 index c22cca4326..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/dimos_native_module.hpp +++ /dev/null @@ -1,97 +0,0 @@ -// SmartNav Native Module helpers. -// Re-exports dimos NativeModule patterns for CLI arg parsing and LCM helpers. -// Based on dimos/hardware/sensors/lidar/common/dimos_native_module.hpp - -#pragma once - -#include -#include -#include -#include - -#include "std_msgs/Header.hpp" -#include "std_msgs/Time.hpp" - -namespace dimos { - -class NativeModule { -public: - NativeModule(int argc, char** argv) { - for (int i = 1; i < argc; ++i) { - std::string arg(argv[i]); - if (arg.size() > 2 && arg[0] == '-' && arg[1] == '-' && i + 1 < argc) { - args_[arg.substr(2)] = argv[++i]; - } - } - } - - /// Get the full LCM channel string for a declared port. - const std::string& topic(const std::string& port) const { - auto it = args_.find(port); - if (it == args_.end()) { - throw std::runtime_error("NativeModule: no topic for port '" + port + "'"); - } - return it->second; - } - - /// Get a string arg value, or a default if not present. - std::string arg(const std::string& key, const std::string& default_val = "") const { - auto it = args_.find(key); - return it != args_.end() ? it->second : default_val; - } - - /// Get a float arg value, or a default if not present. - float arg_float(const std::string& key, float default_val = 0.0f) const { - auto it = args_.find(key); - return it != args_.end() ? std::stof(it->second) : default_val; - } - - /// Get an int arg value, or a default if not present. - int arg_int(const std::string& key, int default_val = 0) const { - auto it = args_.find(key); - return it != args_.end() ? std::stoi(it->second) : default_val; - } - - /// Get a bool arg value, or a default if not present. - /// Present-but-unparseable values throw, matching arg_int/arg_float's - /// std::stoi/std::stof behaviour — a typo'd value or empty string is a - /// misconfiguration we want to surface immediately, not silently coerce - /// to false. - bool arg_bool(const std::string& key, bool default_val = false) const { - auto it = args_.find(key); - if (it == args_.end()) return default_val; - if (it->second == "true" || it->second == "1") return true; - if (it->second == "false" || it->second == "0") return false; - throw std::runtime_error( - "NativeModule: arg '--" + key + "' has unparseable bool value '" - + it->second + "' (expected true/false or 1/0)"); - } - - /// Check if a port/arg was provided. - bool has(const std::string& key) const { - return args_.count(key) > 0; - } - -private: - std::map args_; -}; - -/// Convert seconds (double) to a ROS-style Time message. -inline std_msgs::Time time_from_seconds(double t) { - std_msgs::Time ts; - ts.sec = static_cast(t); - ts.nsec = static_cast((t - ts.sec) * 1e9); - return ts; -} - -/// Build a stamped Header with auto-incrementing sequence number. -inline std_msgs::Header make_header(const std::string& frame_id, double ts) { - static std::atomic seq{0}; - std_msgs::Header h; - h.seq = seq.fetch_add(1, std::memory_order_relaxed); - h.stamp = time_from_seconds(ts); - h.frame_id = frame_id; - return h; -} - -} // namespace dimos diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.lock b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.lock deleted file mode 100644 index 16898a2c2d..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.lock +++ /dev/null @@ -1,144 +0,0 @@ -{ - "nodes": { - "dimos-lcm": { - "flake": false, - "locked": { - "lastModified": 1769774949, - "narHash": "sha256-icRK7jerqNlwK1WZBrnIP04I2WozzFqTD7qsmnPxQuo=", - "owner": "dimensionalOS", - "repo": "dimos-lcm", - "rev": "0aa72b7b1bd3a65f50f5c03485ee9b728df56afe", - "type": "github" - }, - "original": { - "owner": "dimensionalOS", - "ref": "main", - "repo": "dimos-lcm", - "type": "github" - } - }, - "flake-utils": { - "inputs": { - "systems": "systems" - }, - "locked": { - "lastModified": 1731533236, - "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", - "owner": "numtide", - "repo": "flake-utils", - "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", - "type": "github" - }, - "original": { - "owner": "numtide", - "repo": "flake-utils", - "type": "github" - } - }, - "gtsam-extended": { - "inputs": { - "flake-utils": [ - "flake-utils" - ], - "nixpkgs": [ - "nixpkgs" - ], - "nixpkgs-stable": "nixpkgs-stable" - }, - "locked": { - "lastModified": 1775168452, - "narHash": "sha256-fsUWPQIw+lk4VEQUOniiPLwiPoldWh8ZdnIdos202+I=", - "owner": "jeff-hykin", - "repo": "gtsam-extended", - "rev": "f4572a80b6339181693aee6029ca28153e59a993", - "type": "github" - }, - "original": { - "owner": "jeff-hykin", - "repo": "gtsam-extended", - "type": "github" - } - }, - "lcm-extended": { - "inputs": { - "flake-utils": [ - "flake-utils" - ], - "nixpkgs": [ - "nixpkgs" - ] - }, - "locked": { - "lastModified": 1774902379, - "narHash": "sha256-gRFvEkbXCEoG4jEmsT+i0bMZ5kDHOtAaPsrbStXjdu4=", - "owner": "jeff-hykin", - "repo": "lcm_extended", - "rev": "7d12ad8546d3daae30528a6c28f2c9ff5b10baf7", - "type": "github" - }, - "original": { - "owner": "jeff-hykin", - "repo": "lcm_extended", - "type": "github" - } - }, - "nixpkgs": { - "locked": { - "lastModified": 1777954456, - "narHash": "sha256-hGdgeU2Nk87RAuZyYjyDjFL6LK7dAZN5RE9+hrDTkDU=", - "owner": "NixOS", - "repo": "nixpkgs", - "rev": "549bd84d6279f9852cae6225e372cc67fb91a4c1", - "type": "github" - }, - "original": { - "owner": "NixOS", - "ref": "nixos-unstable", - "repo": "nixpkgs", - "type": "github" - } - }, - "nixpkgs-stable": { - "locked": { - "lastModified": 1751274312, - "narHash": "sha256-/bVBlRpECLVzjV19t5KMdMFWSwKLtb5RyXdjz3LJT+g=", - "owner": "NixOS", - "repo": "nixpkgs", - "rev": "50ab793786d9de88ee30ec4e4c24fb4236fc2674", - "type": "github" - }, - "original": { - "owner": "NixOS", - "ref": "nixos-24.11", - "repo": "nixpkgs", - "type": "github" - } - }, - "root": { - "inputs": { - "dimos-lcm": "dimos-lcm", - "flake-utils": "flake-utils", - "gtsam-extended": "gtsam-extended", - "lcm-extended": "lcm-extended", - "nixpkgs": "nixpkgs" - } - }, - "systems": { - "locked": { - "lastModified": 1681028828, - "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", - "owner": "nix-systems", - "repo": "default", - "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", - "type": "github" - }, - "original": { - "owner": "nix-systems", - "repo": "default", - "type": "github" - } - } - }, - "root": "root", - "version": 7 -} diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.nix b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.nix deleted file mode 100644 index ea8eae1327..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/flake.nix +++ /dev/null @@ -1,70 +0,0 @@ -{ - description = "SmartNav PGO native module (pose graph optimization with iSAM2 + PCL ICP)"; - - inputs = { - nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable"; - flake-utils.url = "github:numtide/flake-utils"; - lcm-extended = { - url = "github:jeff-hykin/lcm_extended"; - inputs.nixpkgs.follows = "nixpkgs"; - inputs.flake-utils.follows = "flake-utils"; - }; - dimos-lcm = { - url = "github:dimensionalOS/dimos-lcm/main"; - flake = false; - }; - gtsam-extended = { - url = "github:jeff-hykin/gtsam-extended"; - inputs.nixpkgs.follows = "nixpkgs"; - inputs.flake-utils.follows = "flake-utils"; - }; - }; - - outputs = { self, nixpkgs, flake-utils, lcm-extended, dimos-lcm, gtsam-extended, ... }: - flake-utils.lib.eachDefaultSystem (system: - let - pkgs = import nixpkgs { inherit system; }; - lcm = lcm-extended.packages.${system}.lcm; - - gtsam-base = gtsam-extended.packages.${system}.gtsam-cpp; - gtsam = gtsam-base.overrideAttrs (_old: { - src = pkgs.fetchFromGitHub { - owner = "borglab"; - repo = "gtsam"; - rev = "1a9792a7ede244850a413739557635b606f295c0"; - sha256 = "sha256-zxm5TGVPW1vipFVpw01zcvKRw4mkh+5ZBCR1n6G466o="; - }; - env.NIX_CFLAGS_COMPILE = "-Wno-error=array-bounds"; - }); - in { - packages.default = pkgs.stdenv.mkDerivation { - pname = "smartnav-pgo"; - version = "0.1.0"; - src = ./.; - - nativeBuildInputs = [ pkgs.cmake pkgs.pkg-config ]; - buildInputs = [ - lcm - pkgs.glib - pkgs.eigen - pkgs.boost - pkgs.pcl - gtsam - ]; - - env.NIX_CFLAGS_COMPILE = "-Wno-error=array-bounds"; - - cmakeFlags = [ - "-DCMAKE_POLICY_VERSION_MINIMUM=3.5" - "-DFETCHCONTENT_SOURCE_DIR_DIMOS_LCM=${dimos-lcm}" - ]; - - # On macOS, libgtsam.4.dylib is referenced via @rpath but the binary - # has no LC_RPATH entries, so it fails to load at runtime. Add one - # pointing at the gtsam lib dir. - postInstall = pkgs.lib.optionalString pkgs.stdenv.isDarwin '' - ${pkgs.darwin.cctools}/bin/install_name_tool -add_rpath ${gtsam}/lib $out/bin/pgo - ''; - }; - }); -} diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/main.cpp b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/main.cpp deleted file mode 100644 index 38314d4390..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/main.cpp +++ /dev/null @@ -1,299 +0,0 @@ -// PGO NativeModule — faithful port of pgo_node.cpp from ROS2 to LCM. -// Subscribes to registered_scan + odometry, runs SimplePGO (iSAM2 + PCL ICP), -// publishes corrected_odometry, global_map, and TF correction offset. - -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include - -#include "commons.h" -#include "simple_pgo.h" -#include "dimos_native_module.hpp" -#include "point_cloud_utils.hpp" - -#include "nav_msgs/Odometry.hpp" -#include "sensor_msgs/PointCloud2.hpp" -#include "geometry_msgs/Pose.hpp" -#include "geometry_msgs/Quaternion.hpp" -#include "geometry_msgs/Point.hpp" - -static std::atomic g_running{true}; -static void signal_handler(int) { g_running.store(false); } - -// Shared state between LCM callbacks and main loop -static std::mutex g_buffer_mutex; -static std::queue g_cloud_buffer; -static double g_last_message_time = 0.0; - -// Latest odometry for non-keyframe TF broadcasting -static std::mutex g_odom_mutex; -static M3D g_latest_r = M3D::Identity(); -static V3D g_latest_t = V3D::Zero(); -static double g_latest_time = 0.0; -static bool g_has_odom = false; - -class Handlers { -public: - void on_odometry(const lcm::ReceiveBuffer*, const std::string&, - const nav_msgs::Odometry* msg) { - M3D r = Eigen::Quaterniond( - msg->pose.pose.orientation.w, - msg->pose.pose.orientation.x, - msg->pose.pose.orientation.y, - msg->pose.pose.orientation.z - ).toRotationMatrix(); - - V3D t(msg->pose.pose.position.x, - msg->pose.pose.position.y, - msg->pose.pose.position.z); - - double ts = msg->header.stamp.sec + msg->header.stamp.nsec / 1e9; - - std::lock_guard lock(g_odom_mutex); - g_latest_r = r; - g_latest_t = t; - g_latest_time = ts; - g_has_odom = true; - } - - void on_registered_scan(const lcm::ReceiveBuffer*, const std::string&, - const sensor_msgs::PointCloud2* msg) { - std::lock_guard odom_lock(g_odom_mutex); - if (!g_has_odom) - return; - - double ts = g_latest_time; - - // Reject out-of-order messages - if (ts < g_last_message_time) - return; - g_last_message_time = ts; - - CloudWithPose cp; - cp.pose.r = g_latest_r; - cp.pose.t = g_latest_t; - cp.pose.setTime(static_cast(ts), - static_cast((ts - static_cast(ts)) * 1e9)); - - // Parse PointCloud2 to PCL - cp.cloud = CloudType::Ptr(new CloudType); - smartnav::to_pcl(*msg, *cp.cloud); - - std::lock_guard buf_lock(g_buffer_mutex); - g_cloud_buffer.push(cp); - } -}; - -static nav_msgs::Odometry build_odometry(const M3D& r, const V3D& t, double ts, - const std::string& frame_id, - const std::string& child_frame_id) { - nav_msgs::Odometry odom; - odom.header = dimos::make_header(frame_id, ts); - odom.child_frame_id = child_frame_id; - - Eigen::Quaterniond q(r); - odom.pose.pose.position.x = t.x(); - odom.pose.pose.position.y = t.y(); - odom.pose.pose.position.z = t.z(); - odom.pose.pose.orientation.x = q.x(); - odom.pose.pose.orientation.y = q.y(); - odom.pose.pose.orientation.z = q.z(); - odom.pose.pose.orientation.w = q.w(); - - return odom; -} - -int main(int argc, char** argv) -{ - signal(SIGTERM, signal_handler); - signal(SIGINT, signal_handler); - - dimos::NativeModule mod(argc, argv); - - // Port topics - std::string scan_topic = mod.topic("registered_scan"); - std::string odom_topic = mod.topic("odometry"); - std::string corrected_odom_topic = mod.topic("corrected_odometry"); - std::string global_map_topic = mod.topic("global_map"); - std::string tf_topic = mod.topic("pgo_tf"); - - // Config parameters - Config config; - config.key_pose_delta_deg = mod.arg_float("key_pose_delta_deg", 10.0f); - config.key_pose_delta_trans = mod.arg_float("key_pose_delta_trans", 0.5f); - config.loop_search_radius = mod.arg_float("loop_search_radius", 1.0f); - config.loop_time_tresh = mod.arg_float("loop_time_thresh", 60.0f); - config.loop_score_tresh = mod.arg_float("loop_score_thresh", 0.15f); - config.loop_submap_half_range = mod.arg_int("loop_submap_half_range", 5); - config.submap_resolution = mod.arg_float("submap_resolution", 0.1f); - config.min_loop_detect_duration = mod.arg_float("min_loop_detect_duration", 5.0f); - - // Node-level config - std::string world_frame = mod.arg("world_frame", "map"); - std::string local_frame = mod.arg("local_frame", "odom"); - float global_map_voxel_size = mod.arg_float("global_map_voxel_size", 0.1f); - float global_map_publish_rate = mod.arg_float("global_map_publish_rate", 1.0f); - // rate <= 0 disables global_map entirely (it's a big cloud that can overflow - // LCM's single-message limit and congest the bus). Consumers that only need - // corrected_odometry / pose_graph should turn it off. - bool publish_global_map = global_map_publish_rate > 0; - double global_map_interval = publish_global_map ? 1.0 / global_map_publish_rate : 0.0; - - // Unregister mode: transform world-frame scans to body-frame - bool unregister_input = mod.arg_bool("unregister_input", true); - - bool debug = mod.arg_bool("debug", false); - - pcl::console::setVerbosityLevel( - debug ? pcl::console::L_INFO : pcl::console::L_ERROR); - - SimplePGO pgo(config); - - lcm::LCM lcm; - if (!lcm.good()) { - fprintf(stderr, "PGO: LCM init failed\n"); - return 1; - } - - Handlers handlers; - lcm.subscribe(odom_topic, &Handlers::on_odometry, &handlers); - lcm.subscribe(scan_topic, &Handlers::on_registered_scan, &handlers); - - if (debug) { - fprintf(stderr, "PGO native module started\n"); - fprintf(stderr, " registered_scan: %s\n", scan_topic.c_str()); - fprintf(stderr, " odometry: %s\n", odom_topic.c_str()); - fprintf(stderr, " corrected_odometry: %s\n", corrected_odom_topic.c_str()); - fprintf(stderr, " global_map: %s\n", global_map_topic.c_str()); - fprintf(stderr, " pgo_tf: %s\n", tf_topic.c_str()); - } - - double last_global_map_time = 0.0; - int timer_period_ms = 50; // 20 Hz, matching original - - while (g_running.load()) { - // Drain all pending LCM messages - while (lcm.handleTimeout(0) > 0) {} - - // Check buffer - CloudWithPose cp; - bool has_data = false; - { - std::lock_guard lock(g_buffer_mutex); - if (!g_cloud_buffer.empty()) { - cp = g_cloud_buffer.front(); - // Drain entire queue (matching original: process oldest, discard rest) - while (!g_cloud_buffer.empty()) { - g_cloud_buffer.pop(); - } - has_data = true; - } - } - - if (!has_data) { - std::this_thread::sleep_for(std::chrono::milliseconds(timer_period_ms)); - continue; - } - - // Optionally transform world-frame scan to body-frame - if (unregister_input && cp.cloud && cp.cloud->size() > 0) { - CloudType::Ptr body_cloud(new CloudType); - // body = R_odom^T * (world_pts - t_odom) - M3D r_inv = cp.pose.r.transpose(); - for (const auto& pt : *cp.cloud) { - V3D world_pt(pt.x, pt.y, pt.z); - V3D body_pt = r_inv * (world_pt - cp.pose.t); - PointType bp; - bp.x = static_cast(body_pt.x()); - bp.y = static_cast(body_pt.y()); - bp.z = static_cast(body_pt.z()); - bp.intensity = pt.intensity; - body_cloud->push_back(bp); - } - cp.cloud = body_cloud; - } - - double cur_time = cp.pose.second; - - if (!pgo.addKeyPose(cp)) { - // Not a keyframe — still broadcast TF and corrected odom - M3D corr_r = pgo.offsetR() * cp.pose.r; - V3D corr_t = pgo.offsetR() * cp.pose.t + pgo.offsetT(); - - nav_msgs::Odometry corrected = build_odometry( - corr_r, corr_t, cur_time, world_frame, "base_link"); - lcm.publish(corrected_odom_topic, &corrected); - - nav_msgs::Odometry tf_msg = build_odometry( - pgo.offsetR(), pgo.offsetT(), cur_time, world_frame, local_frame); - lcm.publish(tf_topic, &tf_msg); - - std::this_thread::sleep_for(std::chrono::milliseconds(timer_period_ms)); - continue; - } - - // Keyframe added - pgo.searchForLoopPairs(); - pgo.smoothAndUpdate(); - - if (debug) { - fprintf(stderr, "PGO: keyframe %zu at (%.1f, %.1f, %.1f)\n", - pgo.keyPoses().size(), - cp.pose.t.x(), cp.pose.t.y(), cp.pose.t.z()); - } - - // Publish corrected odometry - M3D corr_r = pgo.offsetR() * cp.pose.r; - V3D corr_t = pgo.offsetR() * cp.pose.t + pgo.offsetT(); - nav_msgs::Odometry corrected = build_odometry( - corr_r, corr_t, cur_time, world_frame, "base_link"); - lcm.publish(corrected_odom_topic, &corrected); - - // Publish TF correction (map -> odom offset) - nav_msgs::Odometry tf_msg = build_odometry( - pgo.offsetR(), pgo.offsetT(), cur_time, world_frame, local_frame); - lcm.publish(tf_topic, &tf_msg); - - // Publish global map (throttled) - double now = cur_time; - if (publish_global_map && now - last_global_map_time >= global_map_interval) { - last_global_map_time = now; - - if (!pgo.keyPoses().empty()) { - CloudType::Ptr global_cloud(new CloudType); - for (size_t i = 0; i < pgo.keyPoses().size(); i++) { - CloudType::Ptr world_cloud(new CloudType); - pcl::transformPointCloud( - *pgo.keyPoses()[i].body_cloud, - *world_cloud, - pgo.keyPoses()[i].t_global, - Eigen::Quaterniond(pgo.keyPoses()[i].r_global)); - *global_cloud += *world_cloud; - } - - // Voxel downsample - CloudType::Ptr filtered(new CloudType); - pcl::VoxelGrid voxel; - voxel.setInputCloud(global_cloud); - voxel.setLeafSize(global_map_voxel_size, global_map_voxel_size, global_map_voxel_size); - voxel.filter(*filtered); - - sensor_msgs::PointCloud2 map_msg = smartnav::from_pcl(*filtered, world_frame, now); - lcm.publish(global_map_topic, &map_msg); - } - } - - std::this_thread::sleep_for(std::chrono::milliseconds(timer_period_ms)); - } - - if (debug) fprintf(stderr, "PGO native module shutting down\n"); - return 0; -} diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/point_cloud_utils.hpp b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/point_cloud_utils.hpp deleted file mode 100644 index 0970e1f8de..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/point_cloud_utils.hpp +++ /dev/null @@ -1,170 +0,0 @@ -// Point cloud utility functions for SmartNav native modules. -// Provides PointCloud2 building/parsing helpers that work with dimos-lcm types. -// When USE_PCL is defined, also provides PCL interop utilities. - -#pragma once - -#include -#include -#include - -#include "sensor_msgs/PointCloud2.hpp" -#include "sensor_msgs/PointField.hpp" -#include "std_msgs/Header.hpp" - -#include "dimos_native_module.hpp" - -#ifdef USE_PCL -#include -#include -#include -#endif - -namespace smartnav { - -// Simple XYZI point structure (no PCL dependency) -struct PointXYZI { - float x, y, z, intensity; -}; - -// Build PointCloud2 from vector of XYZI points -inline sensor_msgs::PointCloud2 build_pointcloud2( - const std::vector& points, - const std::string& frame_id, - double timestamp -) { - sensor_msgs::PointCloud2 pc; - pc.header = dimos::make_header(frame_id, timestamp); - pc.height = 1; - pc.width = static_cast(points.size()); - pc.is_bigendian = 0; - pc.is_dense = 1; - - // Fields: x, y, z, intensity (all float32) - pc.fields_length = 4; - pc.fields.resize(4); - auto make_field = [](const std::string& name, int32_t offset) { - sensor_msgs::PointField f; - f.name = name; - f.offset = offset; - f.datatype = sensor_msgs::PointField::FLOAT32; - f.count = 1; - return f; - }; - pc.fields[0] = make_field("x", 0); - pc.fields[1] = make_field("y", 4); - pc.fields[2] = make_field("z", 8); - pc.fields[3] = make_field("intensity", 12); - - pc.point_step = 16; - pc.row_step = pc.point_step * pc.width; - pc.data_length = pc.row_step; - pc.data.resize(pc.data_length); - - for (size_t i = 0; i < points.size(); ++i) { - float* dst = reinterpret_cast(pc.data.data() + i * 16); - dst[0] = points[i].x; - dst[1] = points[i].y; - dst[2] = points[i].z; - dst[3] = points[i].intensity; - } - - return pc; -} - -// Parse PointCloud2 into vector of XYZI points -inline std::vector parse_pointcloud2(const sensor_msgs::PointCloud2& pc) { - std::vector points; - if (pc.width == 0 || pc.height == 0) return points; - - int num_points = pc.width * pc.height; - points.reserve(num_points); - - // Find field offsets - int x_off = -1, y_off = -1, z_off = -1, i_off = -1; - for (const auto& f : pc.fields) { - if (f.name == "x") x_off = f.offset; - else if (f.name == "y") y_off = f.offset; - else if (f.name == "z") z_off = f.offset; - else if (f.name == "intensity") i_off = f.offset; - } - - if (x_off < 0 || y_off < 0 || z_off < 0) return points; - - for (int n = 0; n < num_points; ++n) { - if (static_cast((n + 1) * pc.point_step) > pc.data.size()) break; - const uint8_t* base = pc.data.data() + n * pc.point_step; - PointXYZI p; - std::memcpy(&p.x, base + x_off, sizeof(float)); - std::memcpy(&p.y, base + y_off, sizeof(float)); - std::memcpy(&p.z, base + z_off, sizeof(float)); - if (i_off >= 0) std::memcpy(&p.intensity, base + i_off, sizeof(float)); - else p.intensity = 0.0f; - points.push_back(p); - } - - return points; -} - -// Get timestamp from PointCloud2 header -inline double get_timestamp(const sensor_msgs::PointCloud2& pc) { - return pc.header.stamp.sec + pc.header.stamp.nsec / 1e9; -} - -#ifdef USE_PCL -// Convert dimos-lcm PointCloud2 to PCL point cloud -inline void to_pcl(const sensor_msgs::PointCloud2& pc, - pcl::PointCloud& cloud) { - auto points = parse_pointcloud2(pc); - cloud.clear(); - cloud.reserve(points.size()); - for (const auto& p : points) { - pcl::PointXYZI pt; - pt.x = p.x; - pt.y = p.y; - pt.z = p.z; - pt.intensity = p.intensity; - cloud.push_back(pt); - } - cloud.width = cloud.size(); - cloud.height = 1; - cloud.is_dense = true; -} - -// Convert PCL point cloud to dimos-lcm PointCloud2 -inline sensor_msgs::PointCloud2 from_pcl( - const pcl::PointCloud& cloud, - const std::string& frame_id, - double timestamp -) { - std::vector points; - points.reserve(cloud.size()); - for (const auto& pt : cloud) { - points.push_back({pt.x, pt.y, pt.z, pt.intensity}); - } - return build_pointcloud2(points, frame_id, timestamp); -} -#endif - -// Quaternion to RPY conversion -inline void quat_to_rpy(double qx, double qy, double qz, double qw, - double& roll, double& pitch, double& yaw) { - // Roll (x-axis rotation) - double sinr_cosp = 2.0 * (qw * qx + qy * qz); - double cosr_cosp = 1.0 - 2.0 * (qx * qx + qy * qy); - roll = std::atan2(sinr_cosp, cosr_cosp); - - // Pitch (y-axis rotation) - double sinp = 2.0 * (qw * qy - qz * qx); - if (std::abs(sinp) >= 1.0) - pitch = std::copysign(M_PI / 2, sinp); - else - pitch = std::asin(sinp); - - // Yaw (z-axis rotation) - double siny_cosp = 2.0 * (qw * qz + qx * qy); - double cosy_cosp = 1.0 - 2.0 * (qy * qy + qz * qz); - yaw = std::atan2(siny_cosp, cosy_cosp); -} - -} // namespace smartnav diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.cpp b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.cpp deleted file mode 100644 index 5fc18bf0e7..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.cpp +++ /dev/null @@ -1,211 +0,0 @@ -#include "simple_pgo.h" - -SimplePGO::SimplePGO(const Config &config) : m_config(config) -{ - gtsam::ISAM2Params isam2_params; - isam2_params.relinearizeThreshold = 0.01; - isam2_params.relinearizeSkip = 1; - m_isam2 = std::make_shared(isam2_params); - m_initial_values.clear(); - m_graph.resize(0); - m_r_offset.setIdentity(); - m_t_offset.setZero(); - - m_icp.setMaximumIterations(50); - m_icp.setMaxCorrespondenceDistance(10); - m_icp.setTransformationEpsilon(1e-6); - m_icp.setEuclideanFitnessEpsilon(1e-6); - m_icp.setRANSACIterations(0); -} - -bool SimplePGO::isKeyPose(const PoseWithTime &pose) -{ - if (m_key_poses.size() == 0) - return true; - const KeyPoseWithCloud &last_item = m_key_poses.back(); - double delta_trans = (pose.t - last_item.t_local).norm(); - double delta_deg = Eigen::Quaterniond(pose.r).angularDistance(Eigen::Quaterniond(last_item.r_local)) * 57.324; - if (delta_trans > m_config.key_pose_delta_trans || delta_deg > m_config.key_pose_delta_deg) - return true; - return false; -} -bool SimplePGO::addKeyPose(const CloudWithPose &cloud_with_pose) -{ - bool is_key_pose = isKeyPose(cloud_with_pose.pose); - if (!is_key_pose) - return false; - size_t idx = m_key_poses.size(); - M3D init_r = m_r_offset * cloud_with_pose.pose.r; - V3D init_t = m_r_offset * cloud_with_pose.pose.t + m_t_offset; - // 添加初始值 - m_initial_values.insert(idx, gtsam::Pose3(gtsam::Rot3(init_r), gtsam::Point3(init_t))); - if (idx == 0) - { - // 添加先验约束 - gtsam::noiseModel::Diagonal::shared_ptr noise = gtsam::noiseModel::Diagonal::Variances(gtsam::Vector6::Ones() * 1e-12); - m_graph.add(gtsam::PriorFactor(idx, gtsam::Pose3(gtsam::Rot3(init_r), gtsam::Point3(init_t)), noise)); - } - else - { - // 添加里程计约束 - const KeyPoseWithCloud &last_item = m_key_poses.back(); - M3D r_between = last_item.r_local.transpose() * cloud_with_pose.pose.r; - V3D t_between = last_item.r_local.transpose() * (cloud_with_pose.pose.t - last_item.t_local); - gtsam::noiseModel::Diagonal::shared_ptr noise = gtsam::noiseModel::Diagonal::Variances((gtsam::Vector(6) << 1e-6, 1e-6, 1e-6, 1e-4, 1e-4, 1e-6).finished()); - m_graph.add(gtsam::BetweenFactor(idx - 1, idx, gtsam::Pose3(gtsam::Rot3(r_between), gtsam::Point3(t_between)), noise)); - } - KeyPoseWithCloud item; - item.time = cloud_with_pose.pose.second; - item.r_local = cloud_with_pose.pose.r; - item.t_local = cloud_with_pose.pose.t; - item.body_cloud = cloud_with_pose.cloud; - item.r_global = init_r; - item.t_global = init_t; - m_key_poses.push_back(item); - return true; -} - -CloudType::Ptr SimplePGO::getSubMap(int idx, int half_range, double resolution) -{ - assert(idx >= 0 && idx < static_cast(m_key_poses.size())); - int min_idx = std::max(0, idx - half_range); - int max_idx = std::min(static_cast(m_key_poses.size()) - 1, idx + half_range); - - CloudType::Ptr ret(new CloudType); - for (int i = min_idx; i <= max_idx; i++) - { - - CloudType::Ptr body_cloud = m_key_poses[i].body_cloud; - CloudType::Ptr global_cloud(new CloudType); - pcl::transformPointCloud(*body_cloud, *global_cloud, m_key_poses[i].t_global, Eigen::Quaterniond(m_key_poses[i].r_global)); - *ret += *global_cloud; - } - if (resolution > 0) - { - pcl::VoxelGrid voxel_grid; - voxel_grid.setLeafSize(resolution, resolution, resolution); - voxel_grid.setInputCloud(ret); - voxel_grid.filter(*ret); - } - return ret; -} - -void SimplePGO::searchForLoopPairs() -{ - if (m_key_poses.size() < 10) - return; - if (m_config.min_loop_detect_duration > 0.0) - { - if (m_history_pairs.size() > 0) - { - double current_time = m_key_poses.back().time; - double last_time = m_key_poses[m_history_pairs.back().second].time; - if (current_time - last_time < m_config.min_loop_detect_duration) - return; - } - } - - size_t cur_idx = m_key_poses.size() - 1; - const KeyPoseWithCloud &last_item = m_key_poses.back(); - pcl::PointXYZ last_pose_pt; - last_pose_pt.x = last_item.t_global(0); - last_pose_pt.y = last_item.t_global(1); - last_pose_pt.z = last_item.t_global(2); - - pcl::PointCloud::Ptr key_poses_cloud(new pcl::PointCloud); - for (size_t i = 0; i < m_key_poses.size() - 1; i++) - { - pcl::PointXYZ pt; - pt.x = m_key_poses[i].t_global(0); - pt.y = m_key_poses[i].t_global(1); - pt.z = m_key_poses[i].t_global(2); - key_poses_cloud->push_back(pt); - } - pcl::KdTreeFLANN kdtree; - kdtree.setInputCloud(key_poses_cloud); - std::vector ids; - std::vector sqdists; - int neighbors = kdtree.radiusSearch(last_pose_pt, m_config.loop_search_radius, ids, sqdists); - if (neighbors == 0) - return; - - int loop_idx = -1; - for (size_t i = 0; i < ids.size(); i++) - { - int idx = ids[i]; - if (std::abs(last_item.time - m_key_poses[idx].time) > m_config.loop_time_tresh) - { - loop_idx = idx; - break; - } - } - - if (loop_idx == -1) - return; - - CloudType::Ptr target_cloud = getSubMap(loop_idx, m_config.loop_submap_half_range, m_config.submap_resolution); - CloudType::Ptr source_cloud = getSubMap(m_key_poses.size() - 1, 0, m_config.submap_resolution); - CloudType::Ptr align_cloud(new CloudType); - - m_icp.setInputSource(source_cloud); - m_icp.setInputTarget(target_cloud); - m_icp.align(*align_cloud); - - if (!m_icp.hasConverged() || m_icp.getFitnessScore() > m_config.loop_score_tresh) - return; - - M4F loop_transform = m_icp.getFinalTransformation(); - - LoopPair one_pair; - one_pair.source_id = cur_idx; - one_pair.target_id = loop_idx; - one_pair.score = m_icp.getFitnessScore(); - M3D r_refined = loop_transform.block<3, 3>(0, 0).cast() * m_key_poses[cur_idx].r_global; - V3D t_refined = loop_transform.block<3, 3>(0, 0).cast() * m_key_poses[cur_idx].t_global + loop_transform.block<3, 1>(0, 3).cast(); - one_pair.r_offset = m_key_poses[loop_idx].r_global.transpose() * r_refined; - one_pair.t_offset = m_key_poses[loop_idx].r_global.transpose() * (t_refined - m_key_poses[loop_idx].t_global); - m_cache_pairs.push_back(one_pair); - m_history_pairs.emplace_back(one_pair.target_id, one_pair.source_id); -} - -void SimplePGO::smoothAndUpdate() -{ - bool has_loop = !m_cache_pairs.empty(); - // 添加回环因子 - if (has_loop) - { - for (LoopPair &pair : m_cache_pairs) - { - m_graph.add(gtsam::BetweenFactor(pair.target_id, pair.source_id, - gtsam::Pose3(gtsam::Rot3(pair.r_offset), - gtsam::Point3(pair.t_offset)), - gtsam::noiseModel::Diagonal::Variances(gtsam::Vector6::Ones() * pair.score))); - } - std::vector().swap(m_cache_pairs); - } - // smooth and mapping - m_isam2->update(m_graph, m_initial_values); - m_isam2->update(); - if (has_loop) - { - m_isam2->update(); - m_isam2->update(); - m_isam2->update(); - m_isam2->update(); - } - m_graph.resize(0); - m_initial_values.clear(); - - // update key poses - gtsam::Values estimate_values = m_isam2->calculateBestEstimate(); - for (size_t i = 0; i < m_key_poses.size(); i++) - { - gtsam::Pose3 pose = estimate_values.at(i); - m_key_poses[i].r_global = pose.rotation().matrix().cast(); - m_key_poses[i].t_global = pose.translation().matrix().cast(); - } - // update offset - const KeyPoseWithCloud &last_item = m_key_poses.back(); - m_r_offset = last_item.r_global * last_item.r_local.transpose(); - m_t_offset = last_item.t_global - m_r_offset * last_item.t_local; -} diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.h b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.h deleted file mode 100644 index 7f80c5b09a..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/cpp/simple_pgo.h +++ /dev/null @@ -1,78 +0,0 @@ -#pragma once -#include "commons.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -struct KeyPoseWithCloud -{ - M3D r_local; - V3D t_local; - M3D r_global; - V3D t_global; - double time; - CloudType::Ptr body_cloud; -}; -struct LoopPair -{ - size_t source_id; - size_t target_id; - M3D r_offset; - V3D t_offset; - double score; -}; - -struct Config -{ - double key_pose_delta_deg = 10; - double key_pose_delta_trans = 1.0; - double loop_search_radius = 1.0; - double loop_time_tresh = 60.0; - double loop_score_tresh = 0.15; - int loop_submap_half_range = 5; - double submap_resolution = 0.1; - double min_loop_detect_duration = 10.0; -}; - -class SimplePGO -{ -public: - SimplePGO(const Config &config); - - bool isKeyPose(const PoseWithTime &pose); - - bool addKeyPose(const CloudWithPose &cloud_with_pose); - - bool hasLoop(){return m_cache_pairs.size() > 0;} - - void searchForLoopPairs(); - - void smoothAndUpdate(); - - CloudType::Ptr getSubMap(int idx, int half_range, double resolution); - std::vector> &historyPairs() { return m_history_pairs; } - std::vector &keyPoses() { return m_key_poses; } - - M3D offsetR() { return m_r_offset; } - V3D offsetT() { return m_t_offset; } - -private: - Config m_config; - std::vector m_key_poses; - std::vector> m_history_pairs; - std::vector m_cache_pairs; - M3D m_r_offset; - V3D m_t_offset; - std::shared_ptr m_isam2; - gtsam::Values m_initial_values; - gtsam::NonlinearFactorGraph m_graph; - pcl::IterativeClosestPoint m_icp; -}; diff --git a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/module.py deleted file mode 100644 index 73a5fc59e4..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/unrefined_pgo/module.py +++ /dev/null @@ -1,286 +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. - -"""The unrefined PGO — a frozen standalone snapshot of main's cmu_nav PGO -(GTSAM iSAM2 + PCL ICP C++ binary, the original gsc_pgo was refined from), -adapted to the LoopClosure spec without touching the binary. - -The cpp/ directory here is a copy of `cmu_nav/modules/pgo/cpp` at the time -of the snapshot, so later cmu_nav changes don't silently move this baseline. - -Spec adaptations, all Python-side: - * `lidar` — the binary expects a `registered_scan` topic; `_collect_topics` - aliases the lidar port's topic onto that arg name. - * `pose_graph` — the binary doesn't expose its internal graph, only the - current map->odom offset (`pgo_tf`) and corrected odometry. The wrapper - keyframes the RAW odometry stream (same delta gates as the binary) and - re-applies the LATEST offset to every keyframe on each correction update. - A single global offset can't reproduce iSAM2's per-keyframe smoothing — - but that offset is exactly what this PGO exposes to consumers, so the - synthesized graph is an honest picture of its output. - * `loop_closure_event` — emitted when the offset jumps by more than the - `_LOOP_EVENT_*` thresholds (the offset only moves materially when a loop - closure lands).""" - -from __future__ import annotations - -from dataclasses import dataclass -import math -from pathlib import Path -import threading -import time - -import numpy as np -from reactivex.disposable import Disposable -from scipy.spatial.transform import Rotation - -from dimos.core.core import rpc -from dimos.core.native_module import NativeModule, NativeModuleConfig -from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -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 -from dimos.navigation.jnav.components.loop_closure.spec import LoopClosure -from dimos.navigation.jnav.msgs.Graph3D import Graph3D -from dimos.navigation.jnav.msgs.GraphDelta3D import GraphDelta3D -from dimos.utils.logging_config import setup_logger - -logger = setup_logger() - -# Offset jumps below these are smoothing noise, not loop closures. -_LOOP_EVENT_MIN_TRANS_M = 0.05 -_LOOP_EVENT_MIN_ROT_DEG = 1.0 - - -class PGOConfig(NativeModuleConfig): - cwd: str | None = str(Path(__file__).resolve().parent / "cpp") - # Absolute so the exists() check works from any worker cwd (skips rebuild). - executable: str = str(Path(__file__).resolve().parent / "cpp/result/bin/pgo") - # path:$PWD makes nix see this (git-untracked) copied directory. - build_command: str | None = 'nix build "path:$PWD#default" --no-write-lock-file' - - # Frame names - world_frame: str = "map" - local_frame: str = "odom" - - # Keyframe detection - key_pose_delta_deg: float = 10.0 - key_pose_delta_trans: float = 0.5 - - # Loop closure - loop_search_radius: float = 1.0 - loop_time_thresh: float = 60.0 - loop_score_thresh: float = 0.15 - loop_submap_half_range: int = 5 - submap_resolution: float = 0.1 - min_loop_detect_duration: float = 5.0 - - # Input mode: transform world-frame scans to body-frame using odom - unregister_input: bool = True - - # Global map publishing - global_map_voxel_size: float = 0.1 - global_map_publish_rate: float = 1.0 - - debug: bool = False - - -@dataclass -class _RawKeyframe: - ts: float - translation: np.ndarray # (3,) - rotation: np.ndarray # 3x3 - - -class PGO(NativeModule, LoopClosure): - """Pose graph optimization with loop closure using GTSAM iSAM2 + PCL ICP.""" - - config: PGOConfig - - lidar: In[PointCloud2] - odometry: In[Odometry] - corrected_odometry: Out[Odometry] - pose_graph: Out[Graph3D] - loop_closure_event: Out[GraphDelta3D] - global_map: Out[PointCloud2] - pgo_tf: Out[Odometry] - - def __init__(self, **kwargs: object) -> None: - super().__init__(**kwargs) - self._keyframes: list[_RawKeyframe] = [] - self._offset_rotation = np.eye(3) - self._offset_translation = np.zeros(3) - self._graph_lock = threading.Lock() - - def _collect_topics(self) -> dict[str, str]: - topics = super()._collect_topics() - # The binary asks for --registered_scan; feed it the lidar topic. - if "lidar" in topics: - topics["registered_scan"] = topics["lidar"] - return topics - - @rpc - def start(self) -> None: - super().start() - self.register_disposable( - Disposable(self.pgo_tf.transport.subscribe(self._on_tf_correction, self.pgo_tf)) - ) - self.register_disposable( - Disposable(self.odometry.transport.subscribe(self._on_raw_odometry, self.odometry)) - ) - # Seed identity TF so consumers can query map->body immediately. - self._publish_tf( - translation=(0.0, 0.0, 0.0), - rotation=(0.0, 0.0, 0.0, 1.0), - ts=time.time(), - ) - if self.config.debug: - logger.info("unrefined PGO native module started (C++ iSAM2 + PCL ICP)") - - @rpc - def stop(self) -> None: - super().stop() - - # --- TF passthrough (same as the cmu_nav wrapper) ----------------------- - - def _publish_tf( - self, - translation: tuple[float, float, float], - rotation: tuple[float, float, float, float], - ts: float, - ) -> None: - self.tf.publish( - Transform( - frame_id=self.config.world_frame, - child_frame_id=self.config.local_frame, - translation=Vector3(*translation), - rotation=Quaternion(*rotation), - ts=ts, - ) - ) - - def _on_tf_correction(self, msg: Odometry) -> None: - self._publish_tf( - translation=(msg.pose.position.x, msg.pose.position.y, msg.pose.position.z), - rotation=( - msg.pose.orientation.x, - msg.pose.orientation.y, - msg.pose.orientation.z, - msg.pose.orientation.w, - ), - ts=msg.ts or time.time(), - ) - - new_rotation = Rotation.from_quat( - [ - msg.pose.orientation.x, - msg.pose.orientation.y, - msg.pose.orientation.z, - msg.pose.orientation.w, - ] - ).as_matrix() - new_translation = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z]) - - with self._graph_lock: - delta_trans = float(np.linalg.norm(new_translation - self._offset_translation)) - cos_theta = float( - np.clip((np.trace(self._offset_rotation.T @ new_rotation) - 1.0) / 2.0, -1.0, 1.0) - ) - delta_deg = math.degrees(math.acos(cos_theta)) - self._offset_rotation = new_rotation - self._offset_translation = new_translation - is_loop = ( - delta_trans > _LOOP_EVENT_MIN_TRANS_M or delta_deg > _LOOP_EVENT_MIN_ROT_DEG - ) and bool(self._keyframes) - graph_msg = self._build_graph(msg.ts) if self._keyframes else None - event = self._build_loop_event(msg.ts) if is_loop else None - - if graph_msg is not None: - self.pose_graph.publish(graph_msg) - if event is not None: - self.loop_closure_event.publish(event) - - # --- synthesized pose graph (the binary doesn't expose its own) ---------- - - def _on_raw_odometry(self, msg: Odometry) -> None: - rotation = Rotation.from_quat( - [ - msg.pose.orientation.x, - msg.pose.orientation.y, - msg.pose.orientation.z, - msg.pose.orientation.w, - ] - ).as_matrix() - translation = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z]) - - with self._graph_lock: - if not self._is_keyframe(rotation, translation): - return - self._keyframes.append( - _RawKeyframe(ts=msg.ts, translation=translation, rotation=rotation) - ) - graph_msg = self._build_graph(msg.ts) - self.pose_graph.publish(graph_msg) - - def _is_keyframe(self, rotation: np.ndarray, translation: np.ndarray) -> bool: - if not self._keyframes: - return True - last = self._keyframes[-1] - delta_trans = float(np.linalg.norm(translation - last.translation)) - cos_theta = float(np.clip((np.trace(last.rotation.T @ rotation) - 1.0) / 2.0, -1.0, 1.0)) - delta_deg = math.degrees(math.acos(cos_theta)) - return ( - delta_trans > self.config.key_pose_delta_trans - or delta_deg > self.config.key_pose_delta_deg - ) - - def _corrected_node(self, index: int, keyframe: _RawKeyframe) -> Graph3D.Node3D: - rotation = self._offset_rotation @ keyframe.rotation - translation = self._offset_rotation @ keyframe.translation + self._offset_translation - quaternion = Rotation.from_matrix(rotation).as_quat() - return Graph3D.Node3D( - pose=PoseStamped( - ts=keyframe.ts, - frame_id=self.config.world_frame, - position=[float(v) for v in translation], - orientation=[float(v) for v in quaternion], - ), - id=index, - ) - - def _build_graph(self, ts: float) -> Graph3D: - """Caller must hold ``_graph_lock``.""" - nodes = [ - self._corrected_node(index, keyframe) for index, keyframe in enumerate(self._keyframes) - ] - edges = [ - Graph3D.Edge(start_id=index - 1, end_id=index, timestamp=self._keyframes[index].ts) - for index in range(1, len(self._keyframes)) - ] - return Graph3D(ts=ts, nodes=nodes, edges=edges) - - def _build_loop_event(self, ts: float) -> GraphDelta3D: - """Caller must hold ``_graph_lock``.""" - latest_index = len(self._keyframes) - 1 - identity = GraphDelta3D.Transform( - translation=Vector3(0.0, 0.0, 0.0), rotation=Quaternion(0.0, 0.0, 0.0, 1.0) - ) - return GraphDelta3D( - ts=ts, - nodes=[self._corrected_node(latest_index, self._keyframes[latest_index])], - transforms=[identity], - ) From 844f8f82e0a3b001f4e22eec3fae36ee47011079 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 15:38:40 +0800 Subject: [PATCH 14/69] address greptile review: guard SQL table names, lock tf lazy-load, close cached stores --- dimos/memory2/db_tf.py | 48 ++++++++++++------- .../loop_closure/gsc_pgo/post_process.py | 5 ++ dimos/navigation/jnav/utils/recording_db.py | 11 +++++ 3 files changed, 48 insertions(+), 16 deletions(-) diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py index 60b40242d4..7d498d1abe 100644 --- a/dimos/memory2/db_tf.py +++ b/dimos/memory2/db_tf.py @@ -26,6 +26,9 @@ from __future__ import annotations +import re +import sqlite3 +import threading from typing import TYPE_CHECKING import numpy as np @@ -43,6 +46,15 @@ # Larger than any single recording's span so the buffer never prunes loaded # transforms (MultiTBuffer drops samples older than ts - buffer_size). _NO_PRUNE = 1.0e15 +# SQLite can't parameterize table names, so caller-supplied stream names are +# interpolated; allow only safe identifiers to keep that injection-free. +_SAFE_TABLE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _safe_table(name: str) -> str: + if not _SAFE_TABLE.match(name): + raise ValueError(f"unsafe stream/table name: {name!r}") + return name class DbTf: @@ -52,23 +64,27 @@ def __init__(self, store: Store, stream_names: tuple[str, ...] = TF_STREAMS) -> self._store = store self._stream_names = stream_names self._buffer: MultiTBuffer | None = None + self._load_lock = threading.Lock() def _ensure_loaded(self) -> MultiTBuffer: if self._buffer is not None: return self._buffer - buffer = MultiTBuffer(buffer_size=_NO_PRUNE) - available = set(self._store.list_streams()) - for name in self._stream_names: - if name not in available: - continue - for observation in self._store.stream(name, TFMessage): - message = observation.data - transforms = getattr(message, "transforms", None) - if transforms is None: - transforms = [message] - buffer.receive_transform(*transforms) - self._buffer = buffer - return buffer + with self._load_lock: + if self._buffer is not None: # another thread loaded while we waited + return self._buffer + buffer = MultiTBuffer(buffer_size=_NO_PRUNE) + available = set(self._store.list_streams()) + for name in self._stream_names: + if name not in available: + continue + for observation in self._store.stream(name, TFMessage): + message = observation.data + transforms = getattr(message, "transforms", None) + if transforms is None: + transforms = [message] + buffer.receive_transform(*transforms) + self._buffer = buffer + return buffer def has_transforms(self) -> bool: return bool(self._ensure_loaded().buffers) @@ -86,6 +102,8 @@ def get( non-warning lookup so per-scan misses don't spam the log. """ buffer = self._ensure_loaded() + # _get is the non-warning lookup; public get() logs on every miss, which + # spams the log for per-scan registration where misses are expected. return buffer._get(target_frame, source_frame, time_point, time_tolerance) @@ -121,15 +139,13 @@ def write_tf_tree( Returns the number of tf observations written. """ - import sqlite3 - db_path = store.config.path connection = sqlite3.connect(f"file:{db_path}?mode=ro&immutable=1", uri=True) odom = np.array( list( connection.execute( "select ts,pose_x,pose_y,pose_z,pose_qx,pose_qy,pose_qz,pose_qw " - f"from {odom_stream} order by ts" + f"from {_safe_table(odom_stream)} order by ts" ) ), float, diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py index 609b80f93b..9f52ae8dfe 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py @@ -34,6 +34,7 @@ import json from pathlib import Path +import re import sqlite3 import sys import time @@ -79,6 +80,10 @@ def arg(flag, default=""): SUFFIX = arg("--suffix") LIDAR_STREAM = arg("--lidar", "pointlio_lidar") # input lidar stream (world-registered scans) ODOM_STREAM = arg("--odom", "pointlio_odometry") # input odometry stream (keyframe source) +# ODOM_STREAM is interpolated as a table name (SQLite can't parameterize those); +# reject anything that isn't a plain identifier to keep that injection-free. +if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", ODOM_STREAM): + raise ValueError(f"unsafe --odom stream name: {ODOM_STREAM!r}") RAW_STREAM = arg("--tags", "raw_april_tags") # input unfiltered AprilTag stream IGNORE_TAGS = { int(x) for x in arg("--ignore-tags").replace(",", " ").split() diff --git a/dimos/navigation/jnav/utils/recording_db.py b/dimos/navigation/jnav/utils/recording_db.py index c71ea06a4d..88f1aab647 100644 --- a/dimos/navigation/jnav/utils/recording_db.py +++ b/dimos/navigation/jnav/utils/recording_db.py @@ -21,6 +21,7 @@ from __future__ import annotations +import atexit from collections.abc import Iterator from pathlib import Path from typing import Any @@ -55,6 +56,16 @@ def store(db_path: Path) -> SqliteStore: return cached +def close_all() -> None: + """Stop every cached store, releasing their file handles.""" + for cached in _stores.values(): + cached.stop() + _stores.clear() + + +atexit.register(close_all) + + def list_streams(db_path: Path) -> list[str]: return store(db_path).list_streams() From 12d2b6a7d263fff15feef41bff952779155f12d5 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 16:07:33 +0800 Subject: [PATCH 15/69] port missing apriltag gate functions (ensure_april_streams, gate_params) so add_april imports resolve --- dimos/navigation/jnav/utils/apriltags.py | 253 ++++++++++++++++++++++- 1 file changed, 251 insertions(+), 2 deletions(-) diff --git a/dimos/navigation/jnav/utils/apriltags.py b/dimos/navigation/jnav/utils/apriltags.py index 9bf22cdf82..347065af38 100644 --- a/dimos/navigation/jnav/utils/apriltags.py +++ b/dimos/navigation/jnav/utils/apriltags.py @@ -421,9 +421,7 @@ def detect_apriltags( return detections -# --------------------------------------------------------------------------- # Streaming / incremental API -# --------------------------------------------------------------------------- _DEFAULT_MARKER_LENGTH = 0.10 _DEFAULT_DICTIONARY = "DICT_APRILTAG_36h11" @@ -663,3 +661,254 @@ def load_or_detect_sightings( if sightings: return sightings, "stream" return sightings_from_observations(detect()), "detected" + + +def detect_raw_detections( + store: Any, + intrinsics: np.ndarray, + distortion: np.ndarray, + *, + image_stream: str = "color_image", + marker_length: float = 0.10, + dictionary: str = "DICT_APRILTAG_36h11", +) -> tuple[list[Detection], bool, int]: + """Every valid-PnP tag detection over `image_stream`, unfiltered, with gate diagnostics.""" + detector = make_detector(dictionary) + raw_detections: list[Detection] = [] + images = store.stream(image_stream, Image).to_list() + speed_by_ts, speed_available = _camera_speeds(images) + for image_obs in images: + image = image_obs.data + bgr = image.numpy() if hasattr(image, "numpy") else np.asarray(image.data) + gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) if bgr.ndim == 3 else bgr + all_corners, marker_ids, _ = detector.detectMarkers(bgr) + if marker_ids is None: + continue + for corners, marker_id in zip(all_corners, marker_ids.flatten(), strict=False): + pose = estimate_marker_pose(corners, marker_length, intrinsics, distortion) + if pose is None: + continue + rotation_vector, translation_vector = pose + quaternion = Rotation.from_rotvec(rotation_vector.reshape(3)).as_quat() # x,y,z,w + translation = translation_vector.reshape(3) + tag_in_camera = [ + float(translation[0]), + float(translation[1]), + float(translation[2]), + float(quaternion[0]), + float(quaternion[1]), + float(quaternion[2]), + float(quaternion[3]), + ] + raw_detections.append( + { + "ts": float(image_obs.ts), + "marker_id": int(marker_id), + "t_cam_marker": tag_in_camera, + "sharpness": tag_sharpness(gray, corners), + "reproj_px": reprojection_error_px( + corners, + rotation_vector, + translation_vector, + marker_length, + intrinsics, + distortion, + ), + "tag_px": tag_pixel_size(corners), + "speed": speed_by_ts.get(float(image_obs.ts)), + } + ) + return raw_detections, speed_available, len(images) + + +def gate_detections( + raw_detections: list[Detection], + *, + max_distance_m: float = DEFAULT_MAX_DISTANCE_M, + max_view_angle_deg: float = DEFAULT_MAX_VIEW_ANGLE_DEG, + cluster_gap_sec: float = DEFAULT_CLUSTER_GAP_SEC, + rotation_weight_m_per_rad: float = DEFAULT_ROTATION_WEIGHT_M_PER_RAD, + min_sharpness: float = DEFAULT_MIN_SHARPNESS, + max_reproj_px: float = DEFAULT_MAX_REPROJ_PX, + min_tag_px: float = DEFAULT_MIN_TAG_PX, + max_linear_speed_mps: float = DEFAULT_MAX_LINEAR_SPEED_MPS, + max_angular_speed_dps: float = DEFAULT_MAX_ANGULAR_SPEED_DPS, + min_observations: int = DEFAULT_MIN_OBSERVATIONS, + huber_delta_m: float = DEFAULT_HUBER_DELTA_M, +) -> tuple[list[Detection], dict[str, int], int]: + """Gate + time-cluster raw detections into one Huber-refined medoid per cluster.""" + rejected: dict[str, int] = defaultdict(int) + kept: list[Detection] = [] + for detection in raw_detections: + if detection["sharpness"] < min_sharpness: + rejected["blur"] += 1 + continue + if detection["reproj_px"] > max_reproj_px: + rejected["reproj"] += 1 + continue + if detection["tag_px"] < min_tag_px: + rejected["small"] += 1 + continue + distance, view_angle = view_quality(detection["t_cam_marker"]) + if distance > max_distance_m: + rejected["far"] += 1 + continue + if view_angle > max_view_angle_deg: + rejected["oblique"] += 1 + continue + speed = detection["speed"] + if speed is not None and ( + speed[0] > max_linear_speed_mps or speed[1] > max_angular_speed_dps + ): + rejected["motion"] += 1 + continue + kept.append(detection) + + detections: list[Detection] = [] + thin_clusters = 0 + for cluster in cluster_by_time(kept, cluster_gap_sec): + if len(cluster) < min_observations: + thin_clusters += 1 + continue + detections.append( + { + **robust_cluster_pose(cluster, rotation_weight_m_per_rad, huber_delta_m), + "n_observations": len(cluster), + } + ) + detections.sort(key=lambda detection: detection["ts"]) + return detections, dict(rejected), thin_clusters + + +def _write_tag_stream( + store: Any, stream_name: str, detections: list[Detection], *, diagnostics: bool +) -> None: + """(Re)write a PoseStamped tag stream, with gate diagnostics when `diagnostics`.""" + if stream_name in store.list_streams(): + store.delete_stream(stream_name) + stream = store.stream(stream_name, PoseStamped) + for detection in detections: + pose = detection["t_cam_marker"] + tags: dict[str, Any] = {"marker_id": detection["marker_id"]} + if diagnostics: + speed = detection.get("speed") + distance, view_angle = view_quality(pose) + tags.update( + { + "sharpness": float(detection["sharpness"]), + "reproj_px": float(detection["reproj_px"]), + "tag_px": float(detection["tag_px"]), + "distance_m": float(distance), + "view_angle_deg": float(view_angle), + "lin_speed": float(speed[0]) if speed else -1.0, + "ang_speed": float(speed[1]) if speed else -1.0, + } + ) + stream.append( + PoseStamped(ts=detection["ts"], position=pose[:3], orientation=pose[3:]), + ts=detection["ts"], + pose=tuple(pose), + tags=tags, + ) + + +def write_april_streams( + store: Any, + raw_detections: list[Detection], + filtered_detections: list[Detection], + *, + raw_stream: str = "raw_april_tags", + filtered_stream: str = "april_tags", +) -> None: + """Persist the unfiltered `raw_april_tags` and the gated `april_tags` streams.""" + _write_tag_stream(store, raw_stream, raw_detections, diagnostics=True) + _write_tag_stream(store, filtered_stream, filtered_detections, diagnostics=False) + + +def _print_tag_summary( + stream_name: str, + n_raw: int, + detections: list[Detection], + rejected: dict[str, int], + thin_clusters: int, + speed_available: bool, + n_images: int, +) -> None: + in_spec = n_raw - sum(rejected.values()) + found_ids = sorted({detection["marker_id"] for detection in detections}) + gate_summary = ", ".join(f"{reason}={count}" for reason, count in sorted(rejected.items())) + if not speed_available: + gate_summary += (", " if gate_summary else "") + "motion-gate-off(no poses)" + print( + f" {stream_name}: {n_raw} raw -> {in_spec} in-spec " + f"-> {len(detections)} clusters (dropped {thin_clusters} thin), " + f"markers {found_ids} (over {n_images} images)" + ) + if gate_summary: + print(f" {stream_name} rejected: {gate_summary}") + + +def gate_params() -> dict[str, Any]: + """The current AprilTag gate thresholds, for recording in a sidecar.""" + return { + "max_distance_m": DEFAULT_MAX_DISTANCE_M, + "max_view_angle_deg": DEFAULT_MAX_VIEW_ANGLE_DEG, + "min_sharpness": DEFAULT_MIN_SHARPNESS, + "max_reproj_px": DEFAULT_MAX_REPROJ_PX, + "min_tag_px": DEFAULT_MIN_TAG_PX, + "max_linear_speed_mps": DEFAULT_MAX_LINEAR_SPEED_MPS, + "max_angular_speed_dps": DEFAULT_MAX_ANGULAR_SPEED_DPS, + "min_observations": DEFAULT_MIN_OBSERVATIONS, + "cluster_gap_sec": DEFAULT_CLUSTER_GAP_SEC, + "huber_delta_m": DEFAULT_HUBER_DELTA_M, + "rotation_weight_m_per_rad": DEFAULT_ROTATION_WEIGHT_M_PER_RAD, + } + + +def ensure_april_streams( + store: Any, + intrinsics: np.ndarray, + distortion: np.ndarray, + *, + image_stream: str = "color_image", + marker_length: float = 0.10, + dictionary: str = "DICT_APRILTAG_36h11", + raw_stream: str = "raw_april_tags", + filtered_stream: str = "april_tags", + exclude_tags: Iterable[int] = (), + force: bool = False, +) -> list[Detection]: + """Ensure both tag streams exist (`raw_april_tags`, `april_tags`); force=True rewrites. + + `exclude_tags` are kept in raw but dropped from filtered. Returns the filtered reps.""" + existing = set(store.list_streams()) + if not force and raw_stream in existing and filtered_stream in existing: + return [] + raw_detections, speed_available, n_images = detect_raw_detections( + store, + intrinsics, + distortion, + image_stream=image_stream, + marker_length=marker_length, + dictionary=dictionary, + ) + detections, rejected, thin_clusters = gate_detections(raw_detections) + excluded = set(exclude_tags) + if excluded: + detections = [d for d in detections if d["marker_id"] not in excluded] + write_april_streams( + store, raw_detections, detections, raw_stream=raw_stream, filtered_stream=filtered_stream + ) + print(f" {raw_stream}: {len(raw_detections)} detections (unfiltered)") + if excluded: + print(f" excluded dynamic tags from {filtered_stream}: {sorted(excluded)}") + _print_tag_summary( + filtered_stream, + len(raw_detections), + detections, + rejected, + thin_clusters, + speed_available, + n_images, + ) + return detections From 49bd14267eaa1a59d9e0a6e25d58b8667b084927 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 16:07:34 +0800 Subject: [PATCH 16/69] fix mypy strict errors in jnav (type annotations, SqliteStoreConfig.path, loop var) --- dimos/memory2/db_tf.py | 10 +++++++--- dimos/navigation/jnav/components/loop_closure/eval.py | 2 +- .../components/loop_closure/eval_ground_truth_tag.py | 5 +++-- .../jnav/components/loop_closure/eval_kitti.py | 2 +- 4 files changed, 12 insertions(+), 7 deletions(-) diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py index 7d498d1abe..ae4247aad0 100644 --- a/dimos/memory2/db_tf.py +++ b/dimos/memory2/db_tf.py @@ -33,6 +33,7 @@ import numpy as np +from dimos.memory2.store.sqlite import SqliteStoreConfig from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.geometry_msgs.Vector3 import Vector3 @@ -139,7 +140,10 @@ def write_tf_tree( Returns the number of tf observations written. """ - db_path = store.config.path + config = store.config + if not isinstance(config, SqliteStoreConfig): + raise TypeError("write_tf_tree reads the db directly and needs a SqliteStore") + db_path = config.path connection = sqlite3.connect(f"file:{db_path}?mode=ro&immutable=1", uri=True) odom = np.array( list( @@ -189,8 +193,8 @@ def statics_at(ts: float) -> list[Transform]: ) return links - for ts in np.arange(t0, t1 + static_period, static_period): - tf_stream.append(TFMessage(*statics_at(float(ts))), ts=float(ts)) + for static_ts in np.arange(t0, t1 + static_period, static_period): + tf_stream.append(TFMessage(*statics_at(float(static_ts))), ts=float(static_ts)) written += 1 return written diff --git a/dimos/navigation/jnav/components/loop_closure/eval.py b/dimos/navigation/jnav/components/loop_closure/eval.py index 02afb44749..27902331b0 100644 --- a/dimos/navigation/jnav/components/loop_closure/eval.py +++ b/dimos/navigation/jnav/components/loop_closure/eval.py @@ -640,7 +640,7 @@ def evaluate( camera = camera_stream intrinsics_config = load_intrinsics_json(intrinsics_json) db_store = store(db_path) - stored_stream = ( + stored_stream: Any = ( db_store.stream(APRIL_TAGS_STREAM) if APRIL_TAGS_STREAM in db_store.list_streams() else [] diff --git a/dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py b/dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py index 5644fb8add..3cb34b9018 100644 --- a/dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py +++ b/dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py @@ -87,7 +87,7 @@ def tag_world_position( @ pose7_to_matrix(optical_in_base7) @ pose7_to_matrix(tag_in_optical7) ) - return transform[:3, 3] + return np.asarray(transform[:3, 3]) def read_optical_in_base(intrinsics_json: Path) -> list[float]: @@ -102,7 +102,8 @@ def tag_in_camera_sightings(db_path: Path) -> list[tuple[float, int, list[float] """(ts, marker_id, tag-in-optical pose7) per April-tag observation, read straight from the recorded april_tags PoseStamped stream (no re-detection).""" sightings: list[tuple[float, int, list[float]]] = [] - for observation in store(db_path).stream("april_tags"): + tag_stream: Any = store(db_path).stream("april_tags") + for observation in tag_stream: pose_tuple = observation.pose_tuple if pose_tuple is None: continue diff --git a/dimos/navigation/jnav/components/loop_closure/eval_kitti.py b/dimos/navigation/jnav/components/loop_closure/eval_kitti.py index cf7e53fe32..b7dfac55a1 100644 --- a/dimos/navigation/jnav/components/loop_closure/eval_kitti.py +++ b/dimos/navigation/jnav/components/loop_closure/eval_kitti.py @@ -85,7 +85,7 @@ def corrected_trajectory( *, lidar_stream: str, odom_stream: str, -): +) -> tuple[list[Any], list[Any], Any]: module_class = load_module_class(module_path, module_name) config = filter_config_for_module(module_class, config) graph, closures, _ = run_module_graph( From bf56b945f48772ade9b68d05f8657350c959627a Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 16:07:36 +0800 Subject: [PATCH 17/69] drop section-marker comments to satisfy codebase check --- .../jnav/components/loop_closure/gsc_pgo/module.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py index a0d0c22099..f2d412ecb3 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py @@ -106,7 +106,7 @@ class PGOConfig(NativeModuleConfig): # Skip ICP on candidates farther than this (m). 0 disables. loop_candidate_max_distance_m: float = 30.0 - # --- Tag (AprilTag/ArUco) loop closure -------------------------------- + # Tag (AprilTag/ArUco) loop closure use_tag_loop_closure: bool = False # LCM channel of the static TF tree (dimos "pattern#msg_name" convention). tf_static_channel: str = "/tf_static#tf2_msgs.TFMessage" @@ -125,7 +125,7 @@ class PGOConfig(NativeModuleConfig): loop_robust_kernel: bool = False loop_robust_huber_k: float = 1.345 - # --- Landmark events (decoupled perceiver -> PGO factor-graph manager) ---- + # Landmark events (decoupled perceiver -> PGO factor-graph manager) # When set, the PGO ingests Landmark events on the `landmarks` In and # attaches each as a graph landmark variable + a BetweenFactor(keyframe, # landmark). Two sightings of the same landmark id share the variable, so a @@ -142,7 +142,7 @@ class PGOConfig(NativeModuleConfig): landmark_assoc_max_dt: float = 0.2 landmark_buffer_window: float = 3.0 - # --- Gravity anchor ------------------------------------------------------ + # Gravity anchor # Pin keyframe 0 (whose orientation is gravity-aligned by the LIO front end) # so landmark/loop closures cannot rotate the initial roll/pitch off gravity. # The full pose is pinned (also the gauge reference); roll/pitch stiffness is From 9d5e332a3a090bbc5c879f4f7b70f621a36bccd8 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 16:21:49 +0800 Subject: [PATCH 18/69] regenerate all_blueprints.py for jnav loop-closure modules --- dimos/robot/all_blueprints.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 22bf7d27e1..c79a3b75f4 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -172,6 +172,7 @@ "go2-teleop-module": "dimos.teleop.quest.quest_extensions.Go2TeleopModule", "google-maps-skill-container": "dimos.agents.skills.google_maps_skill_container.GoogleMapsSkillContainer", "gps-nav-skill-container": "dimos.agents.skills.gps_nav_skill.GpsNavSkillContainer", + "graph-capture": "dimos.navigation.jnav.components.loop_closure.eval.GraphCapture", "grasping-module": "dimos.manipulation.grasping.grasping.GraspingModule", "gstreamer-camera-module": "dimos.hardware.sensors.camera.gstreamer.gstreamer_camera.GstreamerCameraModule", "hosted-arm-teleop-module": "dimos.teleop.quest_hosted.hosted_extensions.HostedArmTeleopModule", @@ -183,6 +184,7 @@ "keyboard-teleop": "dimos.robot.unitree.keyboard_teleop.KeyboardTeleop", "keyboard-teleop-module": "dimos.teleop.keyboard.keyboard_teleop_module.KeyboardTeleopModule", "local-planner": "dimos.navigation.cmu_nav.modules.local_planner.local_planner.LocalPlanner", + "lockstep-replay": "dimos.navigation.jnav.components.loop_closure.eval.LockstepReplay", "manipulation-module": "dimos.manipulation.manipulation_module.ManipulationModule", "map": "dimos.robot.unitree.type.map.Map", "marker-detection-stream-module": "dimos.perception.fiducial.marker_detection_stream_module.MarkerDetectionStreamModule", @@ -210,12 +212,13 @@ "perceive-loop-skill": "dimos.perception.perceive_loop_skill.PerceiveLoopSkill", "person-follow-skill-container": "dimos.agents.skills.person_follow.PersonFollowSkillContainer", "person-tracker": "dimos.perception.detection.person_tracker.PersonTracker", - "pgo": "dimos.navigation.cmu_nav.modules.pgo.pgo.PGO", + "pgo": "dimos.navigation.jnav.components.loop_closure.gsc_pgo.module.PGO", "phone-teleop-module": "dimos.teleop.phone.phone_teleop_module.PhoneTeleopModule", "pick-and-place-module": "dimos.manipulation.pick_and_place_module.PickAndPlaceModule", "point-lio": "dimos.hardware.sensors.lidar.pointlio.module.PointLio", "pointlio-recorder": "dimos.hardware.sensors.lidar.pointlio.recorder.PointlioRecorder", "quest-teleop-module": "dimos.teleop.quest.quest_teleop_module.QuestTeleopModule", + "rate-replay": "dimos.navigation.jnav.components.loop_closure.eval.RateReplay", "ray-tracing-voxel-map": "dimos.mapping.ray_tracing.module.RayTracingVoxelMap", "real-sense-camera": "dimos.hardware.sensors.camera.realsense.camera.RealSenseCamera", "receiver-module": "dimos.utils.demo_image_encoding.ReceiverModule", From 60d114cbf0f4809d9331ab348457199c8bdb81e4 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 16:39:08 +0800 Subject: [PATCH 19/69] move jnav .gitignore up to jnav/ root --- .../jnav/{components/loop_closure/gsc_pgo => }/.gitignore | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename dimos/navigation/jnav/{components/loop_closure/gsc_pgo => }/.gitignore (100%) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/.gitignore b/dimos/navigation/jnav/.gitignore similarity index 100% rename from dimos/navigation/jnav/components/loop_closure/gsc_pgo/.gitignore rename to dimos/navigation/jnav/.gitignore From aae43194611760d1fa1186b78a8fbb2633d1b03f Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 16:42:51 +0800 Subject: [PATCH 20/69] add public MultiTBuffer.lookup (non-warning) and use it in DbTf instead of private _get --- dimos/memory2/db_tf.py | 4 +--- dimos/protocol/tf/tf.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py index ae4247aad0..7f52de854b 100644 --- a/dimos/memory2/db_tf.py +++ b/dimos/memory2/db_tf.py @@ -103,9 +103,7 @@ def get( non-warning lookup so per-scan misses don't spam the log. """ buffer = self._ensure_loaded() - # _get is the non-warning lookup; public get() logs on every miss, which - # spams the log for per-scan registration where misses are expected. - return buffer._get(target_frame, source_frame, time_point, time_tolerance) + return buffer.lookup(target_frame, source_frame, time_point, time_tolerance) def transform_matrix(transform: Transform) -> tuple[np.ndarray, np.ndarray]: diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index eb3d72b470..4dd345a530 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -214,6 +214,19 @@ def _wait_get( return None self._cv.wait(timeout=remaining) + def lookup( + self, + parent_frame: str, + child_frame: str, + time_point: float | None = None, + time_tolerance: float | None = None, + ) -> Transform | None: + """Composed transform lookup that does NOT log on a miss (unlike `get`). + + For high-frequency / best-effort lookups where misses are expected and a + per-call warning would spam the log (e.g. per-scan world-registration).""" + return self._get(parent_frame, child_frame, time_point, time_tolerance) + def get( self, parent_frame: str, From 5f3adfc58fa62e853ca02dca964b4c0cac4f827e Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 16:42:52 +0800 Subject: [PATCH 21/69] rename DEFAULT_MARKER_LENGTH_M -> DEFAULT_MARKER_LENGTH_METERS --- .../jnav/components/loop_closure/gsc_pgo/add_april.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py index 010a02fc08..818a71e8f3 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py @@ -41,7 +41,7 @@ FILTERED_STREAM = "april_tags" SUMMARY_NAME = "summary.json" MIN_REVISITS = 2 # a tag seen on fewer visits than this carries no agreement signal -DEFAULT_MARKER_LENGTH_M = 0.10 +DEFAULT_MARKER_LENGTH_METERS = 0.10 DEFAULT_DICTIONARY = "DICT_APRILTAG_36h11" @@ -181,7 +181,7 @@ def main() -> None: marker_length = ( args.tag_size if args.tag_size is not None - else config.get("marker_length", DEFAULT_MARKER_LENGTH_M) + else config.get("marker_length", DEFAULT_MARKER_LENGTH_METERS) ) dictionary = args.dictionary or config.get("dictionary", DEFAULT_DICTIONARY) ensure_april_streams( From 7ae443ca4d81ef27fcd6efbcdb1a23195f2c28e6 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 17:26:06 +0800 Subject: [PATCH 22/69] - --- .../gsc_pgo/hyperparam_analysis.md | 134 ------------------ 1 file changed, 134 deletions(-) delete mode 100644 dimos/navigation/jnav/components/loop_closure/gsc_pgo/hyperparam_analysis.md diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/hyperparam_analysis.md b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/hyperparam_analysis.md deleted file mode 100644 index 8e7f608201..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/hyperparam_analysis.md +++ /dev/null @@ -1,134 +0,0 @@ -# gsc_pgo hyperparameter analysis + online tuning — sketch - -## The problem (from the eval campaign) - -cmu wins on big real-drift loops but **over-fires** elsewhere: -- `grassy_field` (05-32): **-0.369** tag with 110 closures — should be a no-op (+0.00x). - Open grass = feature-poor; ICP matches are garbage but get accepted. -- `gir_park1_2`: +0.578 with 23 closures, but unrefined got **+0.749 with 2 closures**. - cmu's extra closures are net noise — fewer, better closures would win. - -Goal: (a) **detect feature-poor / high-disagreement scans online and back off** -(ideally make PGO a no-op in grassy_field), and (b) accept **fewer, more -consistent** closures. Both reduce to: *trust loops less when the environment -can't support them, and only commit mutually-consistent ones.* - -## Signals cmu already computes (cheap to expose) - -Every candidate already produces these inside `searchForLoopPairs` / -`scan_context.cpp` — instrument the binary to log them per candidate (accepted -and rejected), keyed by ts: - -| signal | where | meaning | -|---|---|---| -| scan-context min cosine distance | `best_distance()` | place-match quality; high = no good match (new place OR feature-poor) | -| ICP fitness (`getFitnessScore`) | `searchForLoopPairs` | mean-sq inlier dist; high = bad geometric alignment | -| ICP converged + inlier ratio | PCL ICP | low inlier ratio = lots of disagreement (crowd/dynamic/degenerate) | -| candidate global distance | `loop_candidate_max_distance_m` gate | drift magnitude at the candidate | -| ring-key occupancy/entropy | scan-context descriptor | **feature richness** — near-uniform/empty descriptor = feature-poor | - -Two signals are *not* computed yet but are the highest-value adds: -- **registration Hessian eigenvalues** (J^T J of the ICP cost) — the canonical - degeneracy detector (below). -- **per-point residual distribution** of the scan-to-submap match — heavy tails = - dynamic objects (crowd). - -## Literature: detecting feature-poor / degenerate environments - -1. **Degeneracy factor (Zhang, Kaess, Singh — ICRA 2016, "On Degeneracy of - Optimization-based State Estimation").** THE standard. Eigen-decompose the - scan-matching information matrix `H = J^T J`; its smallest eigenvalue (or - `λ_min / λ_max` condition number) measures how well-constrained the solve is. - A featureless corridor/field has a small eigenvalue along the unconstrained - axis. → online "is this scan match trustworthy" scalar. Their **solution - remapping** projects the update out of degenerate directions; LOAM/LeGO-LOAM - use it. -2. **X-ICP (Tuna et al., T-RO 2023) — "Localizability-Aware ICP."** Per-axis - localizability from the optimization constraints; selectively adds soft - constraints only in observable directions. More principled than a scalar gate. -3. **Eigenvalue geometric features (Demantké 2011 / Weinmann 2015):** local - neighborhood covariance → linearity/planarity/scattering. Open grass scores - high "scattering" (3D-spread, no structure) — a direct feature-poverty flag. -4. **Scan-context descriptor entropy/occupancy** (cheap, already have the - descriptor): a low-entropy / sparsely-occupied ring-key = feature-poor scene - where cosine matching is unreliable. Raise the match bar or refuse loops. - -## Literature: scan disagreement / dynamic crowds - -5. **Robust kernels (Huber / Cauchy / Geman-McClure)** on the ICP residual — - down-weight the moving-people points. cmu currently uses raw fitness; a robust - cost + reporting the *inlier ratio* exposes "crowd present." -6. **Removert (Kim & Kim, IROS 2020):** range-image visibility differencing - removes dynamic points before matching. Strong for crowds. -7. **Stationarity / scan-to-scan consistency:** points that move between - consecutive frames are dynamic; a high dynamic-fraction → distrust this scan - for loop closure. (Cheap online proxy: residual after rigid alignment of - consecutive scans.) -8. **Learned moving-object segmentation (4DMOS, LMNet, MotionSeg3D)** — heavier; - masks people directly. Overkill for now but the ceiling. - -## Literature: loop-closure outlier rejection (the over-firing fix) - -The single highest-leverage change — cmu accepts each loop independently: - -9. **PCM — Pairwise Consistency Maximization (Mangelson et al., ICRA 2018).** - Build a consistency graph over candidate loops; keep only the **maximal - mutually-consistent clique**, reject the rest. Directly attacks "23 closures - where 2 good ones win." Cheap, backend-agnostic, no tuning. -10. **GNC — Graduated Non-Convexity (Yang et al., RA-L 2020) / Cauchy-IRLS.** - Backend robustly down-weights outlier loops during optimization. iSAM2-friendly. -11. **Switchable Constraints (Sünderhauf & Protzel 2012) / Dynamic Covariance - Scaling (Agarwal 2013):** the optimizer learns a switch/weight per loop and - can turn bad ones off. Lighter than PCM, online-friendly. - -## Online tuning design (signal → knob) - -Frame cmu's static thresholds as **initial values adapted by a measured -"environment trust" `τ ∈ [0,1]`** computed per keyframe from the signals above: - -``` -τ = f( degeneracy_factor, ringkey_entropy, icp_inlier_ratio, dynamic_fraction ) - # low when feature-poor / crowded / degenerate -``` - -Then drive the gates from `τ` (running statistics, not magic numbers): - -| knob (currently static) | online rule | -|---|---| -| `loop_score_thresh` (ICP gate) | tighten as `τ` drops; track running median/MAD of accepted-loop fitness and gate at `median − k·MAD` | -| `scan_context_match_threshold` | raise when ring-key entropy is low (cosine less discriminative in feature-poor scenes) | -| **accept/suppress** the loop | hard-gate: if `degeneracy_factor < ε` (Zhang) → **suppress** → PGO no-op in directions it can't observe (the grassy_field fix) | -| loop **batch** acceptance | replace independent accept with **PCM** over the recent candidate set → only commit the mutually-consistent subset (the gir_park1_2 fix) | -| loop noise model | already `Σ = fitness·I`; inflate further by `1/τ` so low-trust loops barely move the graph | - -Online-tuning techniques that fit (cheapest first): -- **Adaptive thresholding on running stats** (median/MAD, EWMA) — robust, no - training, the right default. -- **Switchable constraints / DCS** — lets the *optimizer* do the tuning; minimal code. -- **PCM** — not "tuning" but the structural fix for over-firing; do this first. -- (Avoid online Bayesian-opt/bandits for the fast gates — too slow/unstable to - converge within one run; reserve for offline meta-params.) - -## Proposed analysis (to validate before coding the online controller) - -1. **Instrument the binary** to emit per-candidate diagnostics (the table above) - to a sidecar jsonl. One pass per recording → the dataset that grounds `f(·)`. -2. **Static sweep** (via the new eval cell-caching + `--drift-per-sec`): vary - `loop_score_thresh`, `scan_context_match_threshold`, - `loop_candidate_max_distance_m`, `min_loop_detect_duration` per recording; - plot tag/voxel improvement & closure count. Expect grassy_field to be - monotone-better as the gate tightens (→ confirms "no-op is optimal there"). -3. **Correlate** accepted-vs-rejected closures with the signals: does - degeneracy_factor / ring-key entropy separate grassy_field's bad closures - from gir_park1_2's good ones? If yes, `τ` is learnable as a simple threshold. -4. **Prototype** the cheapest controller (PCM + degeneracy-gate + adaptive - `loop_score_thresh`) and re-run the full eval; target: grassy_field → ~0, - gir_park1_2 → ≥ unrefined's +0.749, no regression on the pointlio wins. - -## First concrete step - -Add the per-candidate diagnostic logging to `simple_pgo.cpp` (scan-context -distance, ICP fitness, inlier ratio, candidate distance) + compute the Zhang -degeneracy factor from the ICP Hessian. That single instrumentation pass -produces the data to decide which signals actually separate good/bad closures -here — everything else (the `f(·)` and the gates) follows from it. From bc956e3754976c1d65ae8a8b57b69b11d2aad6f1 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 17:32:30 +0800 Subject: [PATCH 23/69] - --- .../navigation/jnav/components/loop_closure/gsc_pgo/module.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py index f2d412ecb3..c5abe1eafb 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py @@ -200,11 +200,11 @@ class PGO(NativeModule): landmarks: In[Landmark] corrected_odometry: Out[Odometry] correction: Out[Transform] + pose_graph: Out[Graph3D] + loop_closure_event: Out[GraphDelta3D] # Internal/debug only (off by default) — see global_map_publish_rate. Named # with a leading underscore so autoconnect won't wire it to `global_map` Ins. _global_map: Out[PointCloud2] - pose_graph: Out[Graph3D] - loop_closure_event: Out[GraphDelta3D] @rpc def start(self) -> None: From 257dd722a5598872ac3a0261739e955969a93af0 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 17:34:23 +0800 Subject: [PATCH 24/69] move gsc_pgo CLI scripts into gsc_pgo/scripts/ --- .../loop_closure/gsc_pgo/{ => scripts}/add_april.py | 2 +- .../loop_closure/gsc_pgo/{ => scripts}/detect_tags.py | 4 ++-- .../loop_closure/gsc_pgo/{ => scripts}/make_rrd.py | 4 ++-- .../loop_closure/gsc_pgo/{ => scripts}/post_process.py | 4 ++-- experimental/docs/jnav/map_postprocessing.md | 8 ++++---- 5 files changed, 11 insertions(+), 11 deletions(-) rename dimos/navigation/jnav/components/loop_closure/gsc_pgo/{ => scripts}/add_april.py (98%) rename dimos/navigation/jnav/components/loop_closure/gsc_pgo/{ => scripts}/detect_tags.py (98%) rename dimos/navigation/jnav/components/loop_closure/gsc_pgo/{ => scripts}/make_rrd.py (97%) rename dimos/navigation/jnav/components/loop_closure/gsc_pgo/{ => scripts}/post_process.py (99%) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py similarity index 98% rename from dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py rename to dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py index 818a71e8f3..f7309010e6 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py @@ -16,7 +16,7 @@ record the outcome in summary.json. --summary recomputes only the result section. Usage: - python dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py --rec PATH + python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py --rec PATH [--summary] [--output PATH] [--camera color_image] [--intrinsics PATH] [--tag-size 0.10] [--dict DICT_APRILTAG_36h11] [--dynamic 17] """ diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/detect_tags.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py similarity index 98% rename from dimos/navigation/jnav/components/loop_closure/gsc_pgo/detect_tags.py rename to dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py index 71b11201cc..52048bb01e 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/detect_tags.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py @@ -22,7 +22,7 @@ Prints the raw per-marker histogram + visit structure (visit = sightings >30s apart). -Usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/detect_tags.py --rec=PATH +Usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py --rec=PATH [--camera=color_image] [--tag-size=0.10] [--dict=DICT_APRILTAG_36h11] [--intrinsics=PATH] [--out=raw_april_tags] """ @@ -56,7 +56,7 @@ def arg(flag, default=None): REC_ARG = arg("--rec") if not REC_ARG: sys.exit( - "usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/detect_tags.py --rec=PATH [--camera=...] [--tag-size=...] " + "usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py --rec=PATH [--camera=...] [--tag-size=...] " "[--dict=...] [--intrinsics=PATH] [--out=...] (--rec is required)" ) REC = Path(REC_ARG).expanduser() diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py similarity index 97% rename from dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py rename to dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py index bfd2db540a..304d543299 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py @@ -19,7 +19,7 @@ and it picks the new stream up automatically. Importable: `build(...)` writes the rrd and returns its path (used by post_process.py). -Standalone: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py --rec=PATH [--lidar=...] [--odom=...] [--tags=...] [--out=...] +Standalone: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py --rec=PATH [--lidar=...] [--odom=...] [--tags=...] [--out=...] """ import json @@ -172,7 +172,7 @@ def _arg(flag, default=None): rec_arg = _arg("--rec") if not rec_arg: sys.exit( - "usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py --rec=PATH [--lidar=...] [--odom=...] " + "usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py --rec=PATH [--lidar=...] [--odom=...] " "[--tags=...] [--out=...] (--rec is required)" ) build( diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py similarity index 99% rename from dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py rename to dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py index 9f52ae8dfe..4b9a1bd7ab 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py @@ -27,7 +27,7 @@ geometry. Writes _odometry / _lidar back into the recording db, optionally a .pc2.lcm log of the corrected cloud, and opens a comparison rrd. -Usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py [odom|lidar|both] --rec=PATH +Usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py [odom|lidar|both] --rec=PATH [--lidar=pointlio_lidar] [--odom=pointlio_odometry] [--tags=raw_april_tags] [--out=gt_pointlio] [--suffix=...] [--ignore-tags=17] [--no-icp] [--no-lcm] [--no-rrd] """ @@ -112,7 +112,7 @@ def arg(flag, default=""): if not REC_ARG: sys.exit( - "usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py [odom|lidar|both] --rec=PATH " + "usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py [odom|lidar|both] --rec=PATH " "[--lidar=...] [--odom=...] [--tags=...] [--out=...] [--suffix=...] " "[--no-icp] [--no-lcm] [--no-rrd] (--rec is required: path to the recording dir)" ) diff --git a/experimental/docs/jnav/map_postprocessing.md b/experimental/docs/jnav/map_postprocessing.md index 9311066e6c..4f60e042e5 100644 --- a/experimental/docs/jnav/map_postprocessing.md +++ b/experimental/docs/jnav/map_postprocessing.md @@ -24,7 +24,7 @@ The scripts live in `dimos/navigation/jnav/components/loop_closure/gsc_pgo/`. **1. Detect the tags.** Run the camera frames through detection and write both the raw and the gated tag streams in one step: ``` -python dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py --rec=PATH +python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py --rec=PATH ``` `add_april.py` writes `raw_april_tags` (every detection, unfiltered, with its quality numbers attached — this is what postprocessing reads), the gated/clustered `april_tags` stream, and an `april_tags` section in the recording's `summary.json` (see below). Leaving `raw_april_tags` unfiltered matters: you tune the quality gates later without re-running detection, which is the slow part. (`detect_tags.py` writes just the raw stream if that's all you want; `--dynamic 17` keeps a moving tag in raw but drops it from the gated stream.) @@ -32,13 +32,13 @@ python dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py --rec= Inspect what was found without rebuilding anything — per tag, raw count and revisit count, flagging any tag never revisited (your fast check for whether a recording even has loop constraints): ``` -python dimos/navigation/jnav/components/loop_closure/gsc_pgo/add_april.py --rec=PATH --summary +python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py --rec=PATH --summary ``` **2. Solve.** Two stages, one command: ``` -python dimos/navigation/jnav/components/loop_closure/gsc_pgo/post_process.py both --rec=PATH +python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py both --rec=PATH ``` - **Tag PGO (GTSAM).** Odometry between-poses are stiff on roll/pitch and z (gravity isn't drifting) and loose on yaw, where the real error lives. The tag sightings are landmark factors, weighted by how good each detection was. This fixes the big-picture drift. @@ -49,7 +49,7 @@ It writes `gt_pointlio_odometry` and `gt_pointlio_lidar` back into the db, optio **3. Look at it.** `post_process.py` opens the rrd for you, but you can rebuild it anytime: ``` -python dimos/navigation/jnav/components/loop_closure/gsc_pgo/make_rrd.py --rec=PATH +python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py --rec=PATH ``` Raw cloud in red, every `gt_*` version in its own color, tag landmarks marked. Add another correction method and it shows up automatically — good for comparing approaches side by side. From 774e491963ee4b6c4a91b4952e604ffd04a9af97 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 17:41:46 +0800 Subject: [PATCH 25/69] use descriptive variable names throughout gsc_pgo --- .../components/loop_closure/gsc_pgo/module.py | 4 +- .../loop_closure/gsc_pgo/scripts/add_april.py | 54 +- .../gsc_pgo/scripts/detect_tags.py | 112 +-- .../loop_closure/gsc_pgo/scripts/make_rrd.py | 182 +++-- .../gsc_pgo/scripts/post_process.py | 639 ++++++++++-------- 5 files changed, 576 insertions(+), 415 deletions(-) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py index c5abe1eafb..458745bec0 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py @@ -223,8 +223,8 @@ def start(self) -> None: if self.config.debug: logger.info("PGO native module started (C++ iSAM2 + PCL ICP)") - def _on_correction_for_tf(self, msg: Transform) -> None: - self.tf.publish(msg) + def _on_correction_for_tf(self, correction: Transform) -> None: + self.tf.publish(correction) @rpc def stop(self) -> None: diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py index f7309010e6..023c1d033b 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py @@ -26,7 +26,7 @@ from pathlib import Path from typing import Any -from dimos.navigation.jnav.utils import recording_db as rdb +from dimos.navigation.jnav.utils import recording_db from dimos.navigation.jnav.utils.apriltag_agreement import ( VISIT_GAP_S, split_visits, @@ -45,17 +45,19 @@ DEFAULT_DICTIONARY = "DICT_APRILTAG_36h11" -def _parse_dynamic(raw: str | None) -> list[int]: - if not raw: +def _parse_dynamic(dynamic_arg: str | None) -> list[int]: + if not dynamic_arg: return [] - return sorted({int(token) for token in raw.replace(",", " ").split()}) + return sorted({int(token) for token in dynamic_arg.replace(",", " ").split()}) def _times_by_tag(store: Any, stream_name: str) -> dict[int, list[float]]: - by_tag: dict[int, list[float]] = {} + times_by_tag: dict[int, list[float]] = {} for observation in store.stream(stream_name): - by_tag.setdefault(int(observation.tags["marker_id"]), []).append(float(observation.ts)) - return by_tag + times_by_tag.setdefault(int(observation.tags["marker_id"]), []).append( + float(observation.ts) + ) + return times_by_tag def summarize(store: Any) -> dict[str, Any]: @@ -63,9 +65,11 @@ def summarize(store: Any) -> dict[str, Any]: never-revisited tags. Prints the table and returns the result for summary.json.""" streams = set(store.list_streams()) raw_available = RAW_STREAM in streams - raw = _times_by_tag(store, RAW_STREAM) if raw_available else {} - filtered = _times_by_tag(store, FILTERED_STREAM) if FILTERED_STREAM in streams else {} - tag_ids = sorted(set(raw) | set(filtered)) + raw_times_by_tag = _times_by_tag(store, RAW_STREAM) if raw_available else {} + filtered_times_by_tag = ( + _times_by_tag(store, FILTERED_STREAM) if FILTERED_STREAM in streams else {} + ) + tag_ids = sorted(set(raw_times_by_tag) | set(filtered_times_by_tag)) print( f" {'tag':>5} {'raw':>6} {'filtered':>9} {'revisits':>9} (visit gap {VISIT_GAP_S:.0f}s)" @@ -73,8 +77,8 @@ def summarize(store: Any) -> dict[str, Any]: tags: list[dict[str, Any]] = [] not_revisited: list[int] = [] for tag_id in tag_ids: - raw_count = len(raw.get(tag_id, [])) - filtered_times = sorted(filtered.get(tag_id, [])) + raw_count = len(raw_times_by_tag.get(tag_id, [])) + filtered_times = sorted(filtered_times_by_tag.get(tag_id, [])) visits = len(split_visits(filtered_times, gap_s=VISIT_GAP_S)) if filtered_times else 0 revisited = visits >= MIN_REVISITS if not revisited: @@ -89,8 +93,10 @@ def summarize(store: Any) -> dict[str, Any]: } ) raw_display = raw_count if raw_available else "-" - flag = " <-- NOT revisited" if not revisited else "" - print(f" {tag_id:>5} {raw_display!s:>6} {len(filtered_times):>9} {visits:>9}{flag}") + revisit_flag = " <-- NOT revisited" if not revisited else "" + print( + f" {tag_id:>5} {raw_display!s:>6} {len(filtered_times):>9} {visits:>9}{revisit_flag}" + ) print( f" totals: {len(tag_ids)} tags | {len(tag_ids) - len(not_revisited)} revisited" @@ -104,7 +110,9 @@ def summarize(store: Any) -> dict[str, Any]: return { "visit_gap_s": VISIT_GAP_S, "min_revisits": MIN_REVISITS, - "all_unfiltered_tag_ids": sorted(raw) if raw_available else sorted(filtered), + "all_unfiltered_tag_ids": sorted(raw_times_by_tag) + if raw_available + else sorted(filtered_times_by_tag), "total_tags": len(tag_ids), "revisited": len(tag_ids) - len(not_revisited), "not_revisited": not_revisited, @@ -160,11 +168,11 @@ def main() -> None: ) args = parser.parse_args() - recording = args.rec.expanduser() - db_path = recording if recording.name == "mem2.db" else recording / "mem2.db" + recording_path = args.rec.expanduser() + db_path = recording_path if recording_path.name == "mem2.db" else recording_path / "mem2.db" if not db_path.exists(): parser.error(f"no mem2.db at {db_path}") - store = rdb.store(db_path) + store = recording_db.store(db_path) summary_path = args.output.expanduser() if args.output else (db_path.parent / SUMMARY_NAME) print(f"=== {db_path.parent.name} ===") @@ -177,17 +185,17 @@ def main() -> None: # Rebuild both streams and overwrite filter_parameters + result. dynamic_tags = _parse_dynamic(args.dynamic) intrinsics_path = (args.intrinsics or (db_path.parent / "camera_intrinsics.json")).expanduser() - config = load_intrinsics_json(intrinsics_path) + intrinsics_config = load_intrinsics_json(intrinsics_path) marker_length = ( args.tag_size if args.tag_size is not None - else config.get("marker_length", DEFAULT_MARKER_LENGTH_METERS) + else intrinsics_config.get("marker_length", DEFAULT_MARKER_LENGTH_METERS) ) - dictionary = args.dictionary or config.get("dictionary", DEFAULT_DICTIONARY) + dictionary = args.dictionary or intrinsics_config.get("dictionary", DEFAULT_DICTIONARY) ensure_april_streams( store, - config["intrinsics"], - config["distortion"], + intrinsics_config["intrinsics"], + intrinsics_config["distortion"], image_stream=args.camera, marker_length=marker_length, dictionary=dictionary, diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py index 52048bb01e..4e6cf00527 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py @@ -50,7 +50,10 @@ def arg(flag, default=None): - return next((a.split("=", 1)[1] for a in sys.argv if a.startswith(flag + "=")), default) + return next( + (argument.split("=", 1)[1] for argument in sys.argv if argument.startswith(flag + "=")), + default, + ) REC_ARG = arg("--rec") @@ -67,17 +70,17 @@ def arg(flag, default=None): VISIT_GAP_S = 30.0 DB = REC / "mem2.db" -intr_path = Path(arg("--intrinsics", str(REC / "camera_intrinsics.json"))).expanduser() -intr = json.loads(intr_path.read_text()) -K = np.array(intr["intrinsics"], float).reshape(3, 3) -dist = np.array(intr.get("distortion", []), float) -st = rdb.store(DB) -if CAMERA not in st.list_streams(): - sys.exit(f"!! camera stream '{CAMERA}' not in db — available: {st.list_streams()}") +intrinsics_path = Path(arg("--intrinsics", str(REC / "camera_intrinsics.json"))).expanduser() +intrinsics = json.loads(intrinsics_path.read_text()) +camera_matrix = np.array(intrinsics["intrinsics"], float).reshape(3, 3) +distortion = np.array(intrinsics.get("distortion", []), float) +store = rdb.store(DB) +if CAMERA not in store.list_streams(): + sys.exit(f"!! camera stream '{CAMERA}' not in db — available: {store.list_streams()}") detector = make_detector(DICTIONARY) print(f"loading '{CAMERA}' frames...") -images = st.stream(CAMERA, Image).to_list() +images = store.stream(CAMERA, Image).to_list() speed_by_ts, speed_available = _camera_speeds(images) print( f"detecting over {len(images)} frames (unfiltered), tag_size={MARKER_LENGTH_M} m, dict={DICTIONARY}..." @@ -92,31 +95,38 @@ def arg(flag, default=None): if marker_ids is None: continue for corners, marker_id in zip(all_corners, marker_ids.flatten(), strict=False): - pose = estimate_marker_pose(corners, MARKER_LENGTH_M, K, dist) + pose = estimate_marker_pose(corners, MARKER_LENGTH_M, camera_matrix, distortion) if pose is None: continue - rvec, tvec = pose - quat = Rotation.from_rotvec(rvec.reshape(3)).as_quat() # x,y,z,w - t = tvec.reshape(3) - tcm = [ - float(t[0]), - float(t[1]), - float(t[2]), - float(quat[0]), - float(quat[1]), - float(quat[2]), - float(quat[3]), + rotation_vector, translation_vector = pose + quaternion = Rotation.from_rotvec(rotation_vector.reshape(3)).as_quat() # x,y,z,w + translation = translation_vector.reshape(3) + tag_pose = [ + float(translation[0]), + float(translation[1]), + float(translation[2]), + float(quaternion[0]), + float(quaternion[1]), + float(quaternion[2]), + float(quaternion[3]), ] - distance, view_angle = view_quality(tcm) + distance, view_angle = view_quality(tag_pose) speed = speed_by_ts.get(float(image_obs.ts)) rows.append( { "ts": float(image_obs.ts), "marker_id": int(marker_id), - "tcm": tcm, + "tcm": tag_pose, "sharpness": float(tag_sharpness(gray, corners)), "reproj_px": float( - reprojection_error_px(corners, rvec, tvec, MARKER_LENGTH_M, K, dist) + reprojection_error_px( + corners, + rotation_vector, + translation_vector, + MARKER_LENGTH_M, + camera_matrix, + distortion, + ) ), "tag_px": float(tag_pixel_size(corners)), "distance_m": float(distance), @@ -125,20 +135,20 @@ def arg(flag, default=None): "ang_speed": float(speed[1]) if speed else -1.0, } ) -rows.sort(key=lambda r: r["ts"]) +rows.sort(key=lambda row: row["ts"]) -if STREAM in st.list_streams(): - st.delete_stream(STREAM) -out = st.stream(STREAM, PoseStamped) -for r in rows: - tcm = r["tcm"] - out.append( - PoseStamped(ts=r["ts"], position=tcm[:3], orientation=tcm[3:]), - ts=r["ts"], - pose=tuple(tcm), +if STREAM in store.list_streams(): + store.delete_stream(STREAM) +out_stream = store.stream(STREAM, PoseStamped) +for row in rows: + tag_pose = row["tcm"] + out_stream.append( + PoseStamped(ts=row["ts"], position=tag_pose[:3], orientation=tag_pose[3:]), + ts=row["ts"], + pose=tuple(tag_pose), tags={ - k: r[k] - for k in ( + tag_key: row[tag_key] + for tag_key in ( "marker_id", "sharpness", "reproj_px", @@ -152,24 +162,28 @@ def arg(flag, default=None): ) print(f"\nwrote {STREAM}: {len(rows)} unfiltered detections") -by = {} -for r in rows: - by.setdefault(r["marker_id"], []).append(r) +rows_by_marker = {} +for row in rows: + rows_by_marker.setdefault(row["marker_id"], []).append(row) print(f"\n=== RAW per-marker (visit = >{VISIT_GAP_S:.0f}s apart) ===") print( f"{'tag':>4} {'det':>4} {'visits':>6} {'dist_m':>12} {'sharp>=60%':>10} {'reproj<=2%':>10} {'span_s':>7}" ) -for mid in sorted(by): - rs = sorted(by[mid], key=lambda r: r["ts"]) - times = [r["ts"] for r in rs] +for marker_id in sorted(rows_by_marker): + marker_rows = sorted(rows_by_marker[marker_id], key=lambda row: row["ts"]) + times = [row["ts"] for row in marker_rows] visits = [[times[0]]] - for tt in times[1:]: - (visits[-1].append(tt) if tt - visits[-1][-1] <= VISIT_GAP_S else visits.append([tt])) - d = [r["distance_m"] for r in rs] - sharp_ok = 100 * np.mean([r["sharpness"] >= 60 for r in rs]) - reproj_ok = 100 * np.mean([r["reproj_px"] <= 2 for r in rs]) + for timestamp in times[1:]: + ( + visits[-1].append(timestamp) + if timestamp - visits[-1][-1] <= VISIT_GAP_S + else visits.append([timestamp]) + ) + distances = [row["distance_m"] for row in marker_rows] + sharp_ok = 100 * np.mean([row["sharpness"] >= 60 for row in marker_rows]) + reproj_ok = 100 * np.mean([row["reproj_px"] <= 2 for row in marker_rows]) print( - f"{mid:>4} {len(rs):>4} {len(visits):>6} {min(d):5.2f}-{max(d):5.2f} " + f"{marker_id:>4} {len(marker_rows):>4} {len(visits):>6} {min(distances):5.2f}-{max(distances):5.2f} " f"{sharp_ok:9.0f}% {reproj_ok:9.0f}% {times[-1] - times[0]:7.0f}" ) -print("\nmarkers present:", sorted(by)) +print("\nmarkers present:", sorted(rows_by_marker)) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py index 304d543299..561cc64430 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py @@ -53,86 +53,120 @@ def build( tag_stream="raw_april_tags", out_name="gt_compare.rrd", ): - rec = Path(rec).expanduser() - db = rec / "mem2.db" - out = rec / out_name - st = rdb.store(db) - intr = json.loads((rec / "camera_intrinsics.json").read_text()) - ext = np.array(intr["optical_in_base"], float) - Tbo = Pose3(Rot3.Quaternion(ext[6], ext[3], ext[4], ext[5]), Point3(ext[0], ext[1], ext[2])) - - def accumulate(name): - pts = [] - for i, o in enumerate(st.stream(name)): - if i % SCAN_STRIDE: + recording_dir = Path(rec).expanduser() + db_path = recording_dir / "mem2.db" + out_path = recording_dir / out_name + store = rdb.store(db_path) + intrinsics = json.loads((recording_dir / "camera_intrinsics.json").read_text()) + optical_in_base = np.array(intrinsics["optical_in_base"], float) + base_to_optical = Pose3( + Rot3.Quaternion( + optical_in_base[6], optical_in_base[3], optical_in_base[4], optical_in_base[5] + ), + Point3(optical_in_base[0], optical_in_base[1], optical_in_base[2]), + ) + + def accumulate(stream_name): + scans = [] + for scan_index, observation in enumerate(store.stream(stream_name)): + if scan_index % SCAN_STRIDE: continue - xyz = np.asarray(o.data.points_f32()) - if len(xyz): - pts.append(xyz[::3]) - a = np.concatenate(pts, 0) - _, idx = np.unique(np.floor(a / VOXEL).astype(np.int64), axis=0, return_index=True) - return a[idx] - - def traj(name): + points = np.asarray(observation.data.points_f32()) + if len(points): + scans.append(points[::3]) + all_points = np.concatenate(scans, 0) + _, unique_indices = np.unique( + np.floor(all_points / VOXEL).astype(np.int64), axis=0, return_index=True + ) + return all_points[unique_indices] + + def traj(stream_name): return np.array( [ - [o.data.pose.position.x, o.data.pose.position.y, o.data.pose.position.z] - for o in st.stream(name) + [ + observation.data.pose.position.x, + observation.data.pose.position.y, + observation.data.pose.position.z, + ] + for observation in store.stream(stream_name) ], np.float32, ) def landmarks(gt_odom): - gt = [ + odom_poses = [ ( - o.ts, + observation.ts, Pose3( Rot3.Quaternion( - o.data.pose.orientation.w, - o.data.pose.orientation.x, - o.data.pose.orientation.y, - o.data.pose.orientation.z, + observation.data.pose.orientation.w, + observation.data.pose.orientation.x, + observation.data.pose.orientation.y, + observation.data.pose.orientation.z, + ), + Point3( + observation.data.pose.position.x, + observation.data.pose.position.y, + observation.data.pose.position.z, ), - Point3(o.data.pose.position.x, o.data.pose.position.y, o.data.pose.position.z), ), ) - for o in st.stream(gt_odom) + for observation in store.stream(gt_odom) ] - gts = np.array([t for t, _ in gt]) - pos = {} - for obs in st.stream(tag_stream): - t = obs.tags + odom_timestamps = np.array([timestamp for timestamp, _ in odom_poses]) + positions_by_marker = {} + for tag_observation in store.stream(tag_stream): + tag_metrics = tag_observation.tags if not ( - float(t["sharpness"]) >= GATE["s"] - and float(t["reproj_px"]) <= GATE["r"] - and float(t["tag_px"]) >= GATE["px"] - and float(t["distance_m"]) <= GATE["d"] - and float(t["view_angle_deg"]) <= GATE["a"] - and (float(t["lin_speed"]) < 0 or float(t["lin_speed"]) <= GATE["lv"]) - and (float(t["ang_speed"]) < 0 or float(t["ang_speed"]) <= GATE["av"]) + float(tag_metrics["sharpness"]) >= GATE["s"] + and float(tag_metrics["reproj_px"]) <= GATE["r"] + and float(tag_metrics["tag_px"]) >= GATE["px"] + and float(tag_metrics["distance_m"]) <= GATE["d"] + and float(tag_metrics["view_angle_deg"]) <= GATE["a"] + and ( + float(tag_metrics["lin_speed"]) < 0 + or float(tag_metrics["lin_speed"]) <= GATE["lv"] + ) + and ( + float(tag_metrics["ang_speed"]) < 0 + or float(tag_metrics["ang_speed"]) <= GATE["av"] + ) ): continue - ps = obs.data - cb = gt[int(np.argmin(np.abs(gts - float(obs.ts))))][1] - Tw = cb.compose(Tbo).compose( + tag_pose = tag_observation.data + closest_base_pose = odom_poses[ + int(np.argmin(np.abs(odom_timestamps - float(tag_observation.ts)))) + ][1] + tag_in_world = closest_base_pose.compose(base_to_optical).compose( Pose3( Rot3.Quaternion( - ps.orientation.w, ps.orientation.x, ps.orientation.y, ps.orientation.z + tag_pose.orientation.w, + tag_pose.orientation.x, + tag_pose.orientation.y, + tag_pose.orientation.z, ), - Point3(ps.x, ps.y, ps.z), + Point3(tag_pose.x, tag_pose.y, tag_pose.z), ) ) - pos.setdefault(int(t["marker_id"]), []).append(np.asarray(Tw.translation())) - means = [np.mean(value, 0) for key, value in sorted(pos.items())] - lbl = [f"tag{key}" for key in sorted(pos)] - return np.array(means), lbl - - streams = st.list_streams() - gt_lidars = sorted(s for s in streams if s.startswith("gt_") and "_lidar" in s) + positions_by_marker.setdefault(int(tag_metrics["marker_id"]), []).append( + np.asarray(tag_in_world.translation()) + ) + mean_positions = [ + np.mean(positions, 0) for marker_id, positions in sorted(positions_by_marker.items()) + ] + labels = [f"tag{marker_id}" for marker_id in sorted(positions_by_marker)] + return np.array(mean_positions), labels + + streams = store.list_streams() + gt_lidars = sorted( + stream_name + for stream_name in streams + if stream_name.startswith("gt_") and "_lidar" in stream_name + ) print("raw + GT lidar streams:", gt_lidars) rr.init("gt_compare") - rr.save(str(out)) + rr.save(str(out_path)) rr.log( "raw/cloud", rr.Points3D(accumulate(lidar_stream), colors=COLORS["raw"], radii=0.02), @@ -141,31 +175,39 @@ def landmarks(gt_odom): rr.log( "raw/trajectory", rr.LineStrips3D([traj(odom_stream)], colors=[255, 120, 120]), static=True ) - for k, name in enumerate(gt_lidars): - color = PALETTE[k % len(PALETTE)] - cloud = accumulate(name) - rr.log(f"{name}/cloud", rr.Points3D(cloud, colors=color, radii=0.02), static=True) - print(f" logged {name}: {len(cloud):,} pts") - odom = name.replace("_lidar", "_odometry") - if odom in streams: - rr.log(f"{name}/trajectory", rr.LineStrips3D([traj(odom)], colors=color), static=True) + for lidar_index, lidar_name in enumerate(gt_lidars): + color = PALETTE[lidar_index % len(PALETTE)] + cloud = accumulate(lidar_name) + rr.log(f"{lidar_name}/cloud", rr.Points3D(cloud, colors=color, radii=0.02), static=True) + print(f" logged {lidar_name}: {len(cloud):,} pts") + odom_name = lidar_name.replace("_lidar", "_odometry") + if odom_name in streams: + rr.log( + f"{lidar_name}/trajectory", + rr.LineStrips3D([traj(odom_name)], colors=color), + static=True, + ) # landmarks placed against the first available gt odometry - gt_odoms = sorted(s for s in streams if s.startswith("gt_") and "_odometry" in s) + gt_odoms = sorted( + stream_name + for stream_name in streams + if stream_name.startswith("gt_") and "_odometry" in stream_name + ) if gt_odoms: - lm, lbl = landmarks(gt_odoms[0]) - if len(lm): + landmark_positions, labels = landmarks(gt_odoms[0]) + if len(landmark_positions): rr.log( "landmarks", - rr.Points3D(lm, colors=[255, 230, 0], radii=0.25, labels=lbl), + rr.Points3D(landmark_positions, colors=[255, 230, 0], radii=0.25, labels=labels), static=True, ) - print(f" logged {len(lbl)} landmarks") - print("wrote", out) - return out + print(f" logged {len(labels)} landmarks") + print("wrote", out_path) + return out_path def _arg(flag, default=None): - return next((a.split("=", 1)[1] for a in sys.argv if a.startswith(flag + "=")), default) + return next((arg.split("=", 1)[1] for arg in sys.argv if arg.startswith(flag + "=")), default) if __name__ == "__main__": diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py index 4b9a1bd7ab..de27fdfd8d 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py @@ -73,7 +73,9 @@ def arg(flag, default=""): - return next((a.split("=", 1)[1] for a in sys.argv if a.startswith(flag + "=")), default) + return next( + (item.split("=", 1)[1] for item in sys.argv if item.startswith(flag + "=")), default + ) REC_ARG = arg("--rec") @@ -86,7 +88,7 @@ def arg(flag, default=""): raise ValueError(f"unsafe --odom stream name: {ODOM_STREAM!r}") RAW_STREAM = arg("--tags", "raw_april_tags") # input unfiltered AprilTag stream IGNORE_TAGS = { - int(x) for x in arg("--ignore-tags").replace(",", " ").split() + int(marker_id) for marker_id in arg("--ignore-tags").replace(",", " ").split() } # dynamic/moving tags OUT_PREFIX = arg("--out", "gt_pointlio") # output prefix -> _odometry / _lidar WRITE_LCM = "--no-lcm" not in sys.argv # also emit _lidar.pc2.lcm of the corrected cloud @@ -118,39 +120,42 @@ def arg(flag, default=""): ) REC = Path(REC_ARG).expanduser() DB = REC / "mem2.db" -intr = json.loads((REC / "camera_intrinsics.json").read_text()) -ext = np.array(intr["optical_in_base"], float) -Tbo = Pose3(Rot3.Quaternion(ext[6], ext[3], ext[4], ext[5]), Point3(ext[0], ext[1], ext[2])) -st = rdb.store(DB) -if RAW_STREAM not in st.list_streams(): +intrinsics = json.loads((REC / "camera_intrinsics.json").read_text()) +optical_in_base = np.array(intrinsics["optical_in_base"], float) +T_base_optical = Pose3( + Rot3.Quaternion(optical_in_base[6], optical_in_base[3], optical_in_base[4], optical_in_base[5]), + Point3(optical_in_base[0], optical_in_base[1], optical_in_base[2]), +) +store = rdb.store(DB) +if RAW_STREAM not in store.list_streams(): sys.exit( f"!! {RAW_STREAM} missing -- run detect_tags.py first to build the unfiltered tag stream." ) from dimos.memory2.db_tf import transform_matrix -_TF_AVAILABLE = USE_TF and st.tf.has_transforms() +_TF_AVAILABLE = USE_TF and store.tf.has_transforms() -def world_points(obs): +def world_points(observation): """Nx3 world-registered points for a lidar observation. Primary: the recording's tf (``world <- LIDAR_FRAME``) at the scan time. Fallback: the observation's stored pose as the transform. Otherwise the scan is assumed already world-registered (legacy recordings). """ - pts = np.asarray(obs.data.points_f32()) - if not len(pts): - return pts + points = np.asarray(observation.data.points_f32()) + if not len(points): + return points if _TF_AVAILABLE: - tf = st.tf.get(WORLD_FRAME, LIDAR_FRAME, float(obs.ts), TF_TOL) - if tf is not None: - rot, trans = transform_matrix(tf) - return pts @ rot.T + trans - pose = getattr(obs, "pose", None) + transform = store.tf.get(WORLD_FRAME, LIDAR_FRAME, float(observation.ts), TF_TOL) + if transform is not None: + rotation, translation = transform_matrix(transform) + return points @ rotation.T + translation + pose = getattr(observation, "pose", None) if isinstance(pose, (tuple, list)) and len(pose) >= 7: - rot = Rot3.Quaternion(pose[6], pose[3], pose[4], pose[5]).matrix() - return pts @ rot.T + np.array(pose[:3], float) - return pts # already world-registered + rotation = Rot3.Quaternion(pose[6], pose[3], pose[4], pose[5]).matrix() + return points @ rotation.T + np.array(pose[:3], float) + return points # already world-registered print(f"recording: {REC}", flush=True) @@ -162,38 +167,38 @@ def world_points(obs): ) -def passes(t): +def passes(detection): return ( - t["sharpness"] >= GATE["min_sharpness"] - and t["reproj_px"] <= GATE["max_reproj_px"] - and t["tag_px"] >= GATE["min_tag_px"] - and t["distance_m"] <= GATE["max_distance_m"] - and t["view_angle_deg"] <= GATE["max_view_angle_deg"] - and (t["lin_speed"] < 0 or t["lin_speed"] <= GATE["max_lin_speed"]) - and (t["ang_speed"] < 0 or t["ang_speed"] <= GATE["max_ang_speed"]) + detection["sharpness"] >= GATE["min_sharpness"] + and detection["reproj_px"] <= GATE["max_reproj_px"] + and detection["tag_px"] >= GATE["min_tag_px"] + and detection["distance_m"] <= GATE["max_distance_m"] + and detection["view_angle_deg"] <= GATE["max_view_angle_deg"] + and (detection["lin_speed"] < 0 or detection["lin_speed"] <= GATE["max_lin_speed"]) + and (detection["ang_speed"] < 0 or detection["ang_speed"] <= GATE["max_ang_speed"]) ) # read raw detections (pose + diagnostics) print("reading tag detections...", flush=True) -raw = [] -for obs in st.stream(RAW_STREAM): - ps = obs.data - tg = obs.tags - raw.append( +raw_detections = [] +for observation in store.stream(RAW_STREAM): + pose = observation.data + tags = observation.tags + raw_detections.append( dict( - ts=float(obs.ts), - marker_id=int(tg["marker_id"]), + ts=float(observation.ts), + marker_id=int(tags["marker_id"]), T_cam_tag=Pose3( Rot3.Quaternion( - ps.orientation.w, ps.orientation.x, ps.orientation.y, ps.orientation.z + pose.orientation.w, pose.orientation.x, pose.orientation.y, pose.orientation.z ), - Point3(ps.x, ps.y, ps.z), + Point3(pose.x, pose.y, pose.z), ), - reproj_px=float(tg["reproj_px"]), + reproj_px=float(tags["reproj_px"]), **{ - k: float(tg[k]) - for k in ( + diagnostic: float(tags[diagnostic]) + for diagnostic in ( "sharpness", "tag_px", "distance_m", @@ -204,84 +209,100 @@ def passes(t): }, ) ) -gated = [t for t in raw if passes(t) and t["marker_id"] not in IGNORE_TAGS] +gated_detections = [ + detection + for detection in raw_detections + if passes(detection) and detection["marker_id"] not in IGNORE_TAGS +] # keyframes from raw odometry -con = sqlite3.connect(f"file:{DB}?mode=ro", uri=True) -fo = np.array( +odom_connection = sqlite3.connect(f"file:{DB}?mode=ro", uri=True) +odom_rows = np.array( list( - con.execute( + odom_connection.execute( "select ts,pose_x,pose_y,pose_z,pose_qx,pose_qy,pose_qz,pose_qw " f"from {ODOM_STREAM} order by ts" ) ), float, ) -con.close() +odom_connection.close() -def pr(r): - return Rot3.Quaternion(r[7], r[4], r[5], r[6]), np.array(r[1:4]) +def row_pose(row): + return Rot3.Quaternion(row[7], row[4], row[5], row[6]), np.array(row[1:4]) -ki = [0] -prev_rot, prev_pos = pr(fo[0]) -for i in range(1, len(fo)): - rot, pos = pr(fo[i]) +keyframe_indices = [0] +prev_rot, prev_pos = row_pose(odom_rows[0]) +for row_index in range(1, len(odom_rows)): + rot, pos = row_pose(odom_rows[row_index]) if ( np.linalg.norm(pos - prev_pos) > 0.5 or np.degrees(np.linalg.norm(Rot3.Logmap(prev_rot.inverse() * rot))) > 10 ): - ki.append(i) + keyframe_indices.append(row_index) prev_rot, prev_pos = rot, pos -kfs = [pr(fo[i]) for i in ki] -kts = fo[ki, 0] -N = len(kfs) +keyframe_poses = [row_pose(odom_rows[index]) for index in keyframe_indices] +keyframe_times = odom_rows[keyframe_indices, 0] +num_keyframes = len(keyframe_poses) # one factor per keyframe x marker: keep the best-reproj detection in each bucket -bucket = {} -for t in gated: - kf = int(np.argmin(np.abs(kts - t["ts"]))) - key = (kf, t["marker_id"]) - if key not in bucket or t["reproj_px"] < bucket[key]["reproj_px"]: - bucket[key] = t +best_per_keyframe_marker = {} +for detection in gated_detections: + keyframe = int(np.argmin(np.abs(keyframe_times - detection["ts"]))) + key = (keyframe, detection["marker_id"]) + if ( + key not in best_per_keyframe_marker + or detection["reproj_px"] < best_per_keyframe_marker[key]["reproj_px"] + ): + best_per_keyframe_marker[key] = detection # revisit report -raw_by, vis_by = {}, {} -for t in raw: - raw_by.setdefault(t["marker_id"], 0) - raw_by[t["marker_id"]] += 1 -for (_kf, mid), t in bucket.items(): - vis_by.setdefault(mid, []).append(t["ts"]) +raw_count_by_marker, visit_times_by_marker = {}, {} +for detection in raw_detections: + raw_count_by_marker.setdefault(detection["marker_id"], 0) + raw_count_by_marker[detection["marker_id"]] += 1 +for (_keyframe, marker_id), detection in best_per_keyframe_marker.items(): + visit_times_by_marker.setdefault(marker_id, []).append(detection["ts"]) def n_visits(times): times = sorted(times) visits = [[times[0]]] - for tt in times[1:]: - (visits[-1].append(tt) if tt - visits[-1][-1] <= VISIT_GAP_S else visits.append([tt])) + for time_value in times[1:]: + ( + visits[-1].append(time_value) + if time_value - visits[-1][-1] <= VISIT_GAP_S + else visits.append([time_value]) + ) return len(visits) print(f"gates: {GATE}") print( - f"raw detections {len(raw)} -> {len(gated)} pass gates -> {len(bucket)} keyframe-tag factors\n" + f"raw detections {len(raw_detections)} -> {len(gated_detections)} pass gates -> " + f"{len(best_per_keyframe_marker)} keyframe-tag factors\n" ) print(f"{'tag':>4} | {'raw viewings':>12} | {'filtered revisits':>17}") not_revisited = [] -for mid in sorted(raw_by): - nv = n_visits(vis_by[mid]) if mid in vis_by else 0 - flag = "" if nv >= 2 else " <-- NOT REVISITED" - print(f"{mid:>4} | {raw_by[mid]:>12} | {nv:>10} visit(s){flag}") - if nv < 2: - not_revisited.append(mid) +for marker_id in sorted(raw_count_by_marker): + visit_count = ( + n_visits(visit_times_by_marker[marker_id]) if marker_id in visit_times_by_marker else 0 + ) + flag = "" if visit_count >= 2 else " <-- NOT REVISITED" + print( + f"{marker_id:>4} | {raw_count_by_marker[marker_id]:>12} | {visit_count:>10} visit(s){flag}" + ) + if visit_count < 2: + not_revisited.append(marker_id) print( f"\ntags NOT revisited (no loop-closure constraint): {not_revisited if not_revisited else 'none'}\n" ) # factor graph + solve -odom = noiseModel.Diagonal.Variances(np.array([1e-8, 1e-8, 1e-5, 1e-4, 1e-4, 1e-6])) -grav0 = noiseModel.Diagonal.Variances(np.array([1e-8, 1e-8, 1e-6, 1e-8, 1e-8, 1e-8])) +odom_noise = noiseModel.Diagonal.Variances(np.array([1e-8, 1e-8, 1e-5, 1e-4, 1e-4, 1e-6])) +gravity_anchor_noise = noiseModel.Diagonal.Variances(np.array([1e-8, 1e-8, 1e-6, 1e-8, 1e-8, 1e-8])) # quality weighting: planar-PnP pose error grows ~quadratically with range, and reproj_px is a @@ -290,55 +311,61 @@ def n_visits(times): REF_D, REF_R = 0.4, 1.0 -def tn(Rbt, distance_m=REF_D, reproj_px=REF_R): - s = max((max(distance_m, 0.2) / REF_D) ** 2 * (max(reproj_px, 0.5) / REF_R) ** 2, 0.25) - Rm = Rbt.matrix() - c = np.zeros((6, 6)) - c[:3, :3] = Rm @ np.diag([0.04, 0.04, 0.0025]) @ Rm.T - c[3:, 3:] = Rm @ np.diag([0.0025, 0.0025, 0.25]) @ Rm.T - return noiseModel.Gaussian.Covariance(c * s) +def tag_noise(tag_rotation, distance_m=REF_D, reproj_px=REF_R): + scale = max((max(distance_m, 0.2) / REF_D) ** 2 * (max(reproj_px, 0.5) / REF_R) ** 2, 0.25) + rotation_matrix = tag_rotation.matrix() + covariance = np.zeros((6, 6)) + covariance[:3, :3] = rotation_matrix @ np.diag([0.04, 0.04, 0.0025]) @ rotation_matrix.T + covariance[3:, 3:] = rotation_matrix @ np.diag([0.0025, 0.0025, 0.25]) @ rotation_matrix.T + return noiseModel.Gaussian.Covariance(covariance * scale) -print(f"building factor graph over {N} keyframes...", flush=True) -g = NonlinearFactorGraph() -v = Values() -for i in range(N): - rot, pos = kfs[i] - v.insert(i, Pose3(rot, Point3(pos))) - if i == 0: - g.add(PriorFactorPose3(0, Pose3(rot, Point3(pos)), grav0)) +print(f"building factor graph over {num_keyframes} keyframes...", flush=True) +graph = NonlinearFactorGraph() +initial_values = Values() +for keyframe_index in range(num_keyframes): + rot, pos = keyframe_poses[keyframe_index] + initial_values.insert(keyframe_index, Pose3(rot, Point3(pos))) + if keyframe_index == 0: + graph.add(PriorFactorPose3(0, Pose3(rot, Point3(pos)), gravity_anchor_noise)) else: - rot_prev, pos_prev = kfs[i - 1] - g.add( + rot_prev, pos_prev = keyframe_poses[keyframe_index - 1] + graph.add( BetweenFactorPose3( - i - 1, - i, + keyframe_index - 1, + keyframe_index, Pose3( rot_prev.inverse() * rot, Point3(rot_prev.inverse().rotate(Point3(pos - pos_prev))), ), - odom, + odom_noise, ) ) -seen = set() -for (kf, mid), t in sorted(bucket.items()): - rot, pos = kfs[kf] - kp = Pose3(rot, Point3(pos)) - T = Tbo.compose(t["T_cam_tag"]) - if mid not in seen: - seen.add(mid) - v.insert(Symbol("l", mid).key(), kp.compose(T)) - g.add( +seen_markers = set() +for (keyframe, marker_id), detection in sorted(best_per_keyframe_marker.items()): + rot, pos = keyframe_poses[keyframe] + keyframe_pose = Pose3(rot, Point3(pos)) + T_base_tag = T_base_optical.compose(detection["T_cam_tag"]) + if marker_id not in seen_markers: + seen_markers.add(marker_id) + initial_values.insert(Symbol("l", marker_id).key(), keyframe_pose.compose(T_base_tag)) + graph.add( BetweenFactorPose3( - kf, Symbol("l", mid).key(), T, tn(T.rotation(), t["distance_m"], t["reproj_px"]) + keyframe, + Symbol("l", marker_id).key(), + T_base_tag, + tag_noise(T_base_tag.rotation(), detection["distance_m"], detection["reproj_px"]), ) ) print("solving stage 1 (tag PGO)...", flush=True) lm_params = LevenbergMarquardtParams() lm_params.setMaxIterations(200) -est = LevenbergMarquardtOptimizer(g, v, lm_params).optimize() -raw_kf = [Pose3(kfs[i][0], Point3(kfs[i][1])) for i in range(N)] +estimate = LevenbergMarquardtOptimizer(graph, initial_values, lm_params).optimize() +raw_keyframe_poses = [ + Pose3(keyframe_poses[index][0], Point3(keyframe_poses[index][1])) + for index in range(num_keyframes) +] # STAGE 2: ICP loop-closure refinement (tags=macro, lidar ICP=local anchor) # Tag PGO has pulled revisits roughly together; now ICP the lidar submaps of spatially-close, @@ -355,155 +382,206 @@ def tn(Rbt, distance_m=REF_D, reproj_px=REF_R): ICP_FIT_MIN, ICP_RMSE_MAX = 0.45, 0.25 SUBMAP_HALF_S = 1.0 # accumulate scans within +/- this of a keyframe time into its submap - corr = [est.atPose3(i) for i in range(N)] - cpos = np.array([np.asarray(p.translation()) for p in corr]) + corrected_poses = [estimate.atPose3(index) for index in range(num_keyframes)] + corrected_positions = np.array([np.asarray(pose.translation()) for pose in corrected_poses]) # revisit candidate pairs - tree = cKDTree(cpos) - pairs = set() - for i, j in tree.query_pairs(ICP_RADIUS_M): - if abs(kts[i] - kts[j]) >= ICP_MIN_DT_S: - pairs.add((min(i, j), max(i, j))) - pairs = sorted(pairs, key=lambda p: np.linalg.norm(cpos[p[0]] - cpos[p[1]])) - involved = {k for p in pairs for k in p} + position_tree = cKDTree(corrected_positions) + candidate_pairs = set() + for first_index, second_index in position_tree.query_pairs(ICP_RADIUS_M): + if abs(keyframe_times[first_index] - keyframe_times[second_index]) >= ICP_MIN_DT_S: + candidate_pairs.add((min(first_index, second_index), max(first_index, second_index))) + candidate_pairs = sorted( + candidate_pairs, + key=lambda pair: np.linalg.norm( + corrected_positions[pair[0]] - corrected_positions[pair[1]] + ), + ) + involved_keyframes = {index for pair in candidate_pairs for index in pair} print( - f"ICP stage: {len(pairs)} revisit candidate pairs over {len(involved)} keyframes", + f"ICP stage: {len(candidate_pairs)} revisit candidate pairs " + f"over {len(involved_keyframes)} keyframes", flush=True, ) # build per-(involved)-keyframe body-frame submaps from the input lidar (odom-registered) print("ICP stage: reading lidar submaps...", flush=True) - submap = {k: [] for k in involved} - scan_n = 0 - t_sub = time.time() - for obs in st.stream(LIDAR_STREAM): - scan_n += 1 - if scan_n % 20000 == 0: - print(f" read {scan_n} scans, {time.time() - t_sub:.0f}s", flush=True) - ts = float(obs.ts) - k = int(np.argmin(np.abs(kts - ts))) - if k not in submap or abs(kts[k] - ts) > SUBMAP_HALF_S: + submap_chunks = {index: [] for index in involved_keyframes} + scan_count = 0 + submap_start_time = time.time() + for observation in store.stream(LIDAR_STREAM): + scan_count += 1 + if scan_count % 20000 == 0: + print(f" read {scan_count} scans, {time.time() - submap_start_time:.0f}s", flush=True) + scan_ts = float(observation.ts) + keyframe = int(np.argmin(np.abs(keyframe_times - scan_ts))) + if keyframe not in submap_chunks or abs(keyframe_times[keyframe] - scan_ts) > SUBMAP_HALF_S: continue - rot_k, pos_k = kfs[k] - world = world_points(obs) - submap[k].append((world - pos_k) @ rot_k.matrix()) # world -> kf-k body frame - pcd = {} - for k, chunks in submap.items(): + keyframe_rot, keyframe_pos = keyframe_poses[keyframe] + world = world_points(observation) + submap_chunks[keyframe].append( + (world - keyframe_pos) @ keyframe_rot.matrix() + ) # world -> kf-keyframe body frame + submap_clouds = {} + for keyframe, chunks in submap_chunks.items(): if not chunks: continue - p = o3d.geometry.PointCloud() - p.points = o3d.utility.Vector3dVector(np.concatenate(chunks, 0).astype(np.float64)) - p = p.voxel_down_sample(ICP_VOXEL) - p.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=0.5, max_nn=30)) - pcd[k] = p - print(f"ICP stage: built {len(pcd)} submaps, registering {len(pairs)} pairs...", flush=True) + cloud = o3d.geometry.PointCloud() + cloud.points = o3d.utility.Vector3dVector(np.concatenate(chunks, 0).astype(np.float64)) + cloud = cloud.voxel_down_sample(ICP_VOXEL) + cloud.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=0.5, max_nn=30)) + submap_clouds[keyframe] = cloud + print( + f"ICP stage: built {len(submap_clouds)} submaps, " + f"registering {len(candidate_pairs)} pairs...", + flush=True, + ) icp_noise = noiseModel.Robust.Create( noiseModel.mEstimator.Huber.Create(1.345), noiseModel.Diagonal.Variances(np.array([4e-4, 4e-4, 4e-4, 2.5e-3, 2.5e-3, 2.5e-3])), ) - added = 0 - t_icp = time.time() - for pair_n, (i, j) in enumerate(pairs): - if pair_n and pair_n % 5000 == 0: + accepted_count = 0 + icp_start_time = time.time() + for pair_index, (first_index, second_index) in enumerate(candidate_pairs): + if pair_index and pair_index % 5000 == 0: print( - f" registered {pair_n}/{len(pairs)} pairs, {added} accepted, {time.time() - t_icp:.0f}s", + f" registered {pair_index}/{len(candidate_pairs)} pairs, " + f"{accepted_count} accepted, {time.time() - icp_start_time:.0f}s", flush=True, ) - if i not in pcd or j not in pcd: + if first_index not in submap_clouds or second_index not in submap_clouds: continue - init = (corr[i].inverse() * corr[j]).matrix() # i<-j initial guess from tag correction - res = o3d.pipelines.registration.registration_icp( - pcd[j], - pcd[i], + initial_guess = ( + corrected_poses[first_index].inverse() * corrected_poses[second_index] + ).matrix() # first<-second initial guess from tag correction + result = o3d.pipelines.registration.registration_icp( + submap_clouds[second_index], + submap_clouds[first_index], ICP_MAX_CORR_M, - init, + initial_guess, o3d.pipelines.registration.TransformationEstimationPointToPlane(), ) - if res.fitness >= ICP_FIT_MIN and res.inlier_rmse <= ICP_RMSE_MAX: - T = res.transformation - g.add(BetweenFactorPose3(i, j, Pose3(Rot3(T[:3, :3]), Point3(T[:3, 3])), icp_noise)) - added += 1 + if result.fitness >= ICP_FIT_MIN and result.inlier_rmse <= ICP_RMSE_MAX: + transform = result.transformation + graph.add( + BetweenFactorPose3( + first_index, + second_index, + Pose3(Rot3(transform[:3, :3]), Point3(transform[:3, 3])), + icp_noise, + ) + ) + accepted_count += 1 print( - f"ICP stage: accepted {added}/{len(pairs)} loop closures " + f"ICP stage: accepted {accepted_count}/{len(candidate_pairs)} loop closures " f"(fit>={ICP_FIT_MIN}, rmse<={ICP_RMSE_MAX}m)", flush=True, ) - if added: - # est already holds every key (keyframe + landmark poses); warm-start from it + if accepted_count: + # estimate already holds every key (keyframe + landmark poses); warm-start from it print("solving stage 2 (tag PGO + ICP closures)...", flush=True) - est = LevenbergMarquardtOptimizer(g, est, lm_params).optimize() - -C = [est.atPose3(i).compose(raw_kf[i].inverse()) for i in range(N)] -shift = max(float(np.linalg.norm(np.asarray(C[i].translation()))) for i in range(N)) + estimate = LevenbergMarquardtOptimizer(graph, estimate, lm_params).optimize() + +corrections = [ + estimate.atPose3(index).compose(raw_keyframe_poses[index].inverse()) + for index in range(num_keyframes) +] +max_correction_shift = max( + float(np.linalg.norm(np.asarray(corrections[index].translation()))) + for index in range(num_keyframes) +) print( - f"PGO: N={N} keyframes, {len(bucket)} tag factors over {len(seen)} markers, " - f"max correction shift {shift:.1f} m", + f"PGO: N={num_keyframes} keyframes, {len(best_per_keyframe_marker)} tag factors " + f"over {len(seen_markers)} markers, " + f"max correction shift {max_correction_shift:.1f} m", flush=True, ) -def C_at(ts): - if ts <= kts[0]: - return C[0] - if ts >= kts[-1]: - return C[-1] - j = int(np.searchsorted(kts, ts)) - a, b = j - 1, j - al = (ts - kts[a]) / (kts[b] - kts[a]) - return C[a].compose(Pose3.Expmap(al * Pose3.Logmap(C[a].between(C[b])))) +def correction_at(ts): + if ts <= keyframe_times[0]: + return corrections[0] + if ts >= keyframe_times[-1]: + return corrections[-1] + insert_index = int(np.searchsorted(keyframe_times, ts)) + before_index, after_index = insert_index - 1, insert_index + alpha = (ts - keyframe_times[before_index]) / ( + keyframe_times[after_index] - keyframe_times[before_index] + ) + return corrections[before_index].compose( + Pose3.Expmap( + alpha * Pose3.Logmap(corrections[before_index].between(corrections[after_index])) + ) + ) -def pose_tuple(P): - t = P.translation() - q = P.rotation().toQuaternion() - return (t[0], t[1], t[2], q.x(), q.y(), q.z(), q.w()) +def pose_tuple(pose): + translation = pose.translation() + quaternion = pose.rotation().toQuaternion() + return ( + translation[0], + translation[1], + translation[2], + quaternion.x(), + quaternion.y(), + quaternion.z(), + quaternion.w(), + ) if WHAT in ("odom", "both"): - name = f"{OUT_PREFIX}_odometry{SUFFIX}" - if name in st.list_streams(): - st.delete_stream(name) - out = st.stream(name, Odometry) - print(f"writing {name} ({len(fo)} poses)...", flush=True) - n = 0 - t0 = time.time() - for r in fo: - ts = float(r[0]) - P = C_at(ts).compose( - Pose3(Rot3.Quaternion(r[7], r[4], r[5], r[6]), Point3(r[1], r[2], r[3])) + out_name = f"{OUT_PREFIX}_odometry{SUFFIX}" + if out_name in store.list_streams(): + store.delete_stream(out_name) + out_stream = store.stream(out_name, Odometry) + print(f"writing {out_name} ({len(odom_rows)} poses)...", flush=True) + written_count = 0 + write_start_time = time.time() + for row in odom_rows: + ts = float(row[0]) + corrected_pose = correction_at(ts).compose( + Pose3(Rot3.Quaternion(row[7], row[4], row[5], row[6]), Point3(row[1], row[2], row[3])) ) - x, y, zz, qx, qy, qz, qw = pose_tuple(P) - out.append( + x, y, z, qx, qy, qz, qw = pose_tuple(corrected_pose) + out_stream.append( Odometry( ts=ts, frame_id="odom", child_frame_id="base_link", - pose=Pose(x, y, zz, qx, qy, qz, qw), + pose=Pose(x, y, z, qx, qy, qz, qw), ), ts=ts, - pose=(x, y, zz, qx, qy, qz, qw), + pose=(x, y, z, qx, qy, qz, qw), ) - n += 1 - if n % 20000 == 0: - print(f" {n}/{len(fo)} poses, {time.time() - t0:.0f}s", flush=True) - print(f"wrote {name}: {len(fo)} poses in {time.time() - t0:.0f}s", flush=True) + written_count += 1 + if written_count % 20000 == 0: + print( + f" {written_count}/{len(odom_rows)} poses, {time.time() - write_start_time:.0f}s", + flush=True, + ) + print( + f"wrote {out_name}: {len(odom_rows)} poses in {time.time() - write_start_time:.0f}s", + flush=True, + ) if WHAT in ("lidar", "both"): - name = f"{OUT_PREFIX}_lidar{SUFFIX}" - fo_ts = fo[:, 0] + out_name = f"{OUT_PREFIX}_lidar{SUFFIX}" + odom_times = odom_rows[:, 0] def base_pose(ts): - j = int(np.searchsorted(fo_ts, ts)) - j = min(max(j, 0), len(fo) - 1) - if j > 0 and abs(fo_ts[j - 1] - ts) < abs(fo_ts[j] - ts): - j -= 1 - r = fo[j] - return Pose3(Rot3.Quaternion(r[7], r[4], r[5], r[6]), Point3(r[1], r[2], r[3])) + index = int(np.searchsorted(odom_times, ts)) + index = min(max(index, 0), len(odom_rows) - 1) + if index > 0 and abs(odom_times[index - 1] - ts) < abs(odom_times[index] - ts): + index -= 1 + row = odom_rows[index] + return Pose3( + Rot3.Quaternion(row[7], row[4], row[5], row[6]), Point3(row[1], row[2], row[3]) + ) - if name in st.list_streams(): - st.delete_stream(name) - out = st.stream(name, PointCloud2) + if out_name in store.list_streams(): + store.delete_stream(out_name) + out_stream = store.stream(out_name, PointCloud2) # The db stream stays per-scan, but the .pc2.lcm is ONE aggregated cloud (voxel-downsampled + # statistical-outlier-removed), not 5184 per-scan events. Intensity rides through open3d's @@ -512,78 +590,97 @@ def base_pose(ts): import open3d as o3d CHUNK = 1000 - agg_xyz, agg_i = [], [] # incrementally voxel-downsampled chunks - buf_xyz, buf_i = [], [] - have_inten = False - - def collapse(xyz_list, i_list, voxel): - pc = o3d.geometry.PointCloud() - pc.points = o3d.utility.Vector3dVector(np.concatenate(xyz_list).astype(np.float64)) - carry = bool(i_list) - if carry: - inten_col = np.concatenate(i_list).astype(np.float64)[:, None] - pc.colors = o3d.utility.Vector3dVector(np.repeat(inten_col, 3, axis=1)) - pc = pc.voxel_down_sample(voxel) - dx = np.asarray(pc.points, np.float32) - di = np.asarray(pc.colors, np.float32)[:, 0] if carry else None - return dx, di - - print(f"writing {name} (corrected lidar)...", flush=True) - n = 0 - t0 = time.time() - for obs in st.stream(LIDAR_STREAM): - ts = float(obs.ts) - Cts = C_at(ts) - Rm = Cts.rotation().matrix() - tv = np.asarray(Cts.translation()) - xyz = world_points(obs) - inten = obs.data.intensities_f32() - xyz2 = (xyz @ Rm.T + tv).astype(np.float32) - m = PointCloud2.from_numpy( - xyz2, frame_id="odom", intensities=(np.asarray(inten) if inten is not None else None) + aggregated_points, aggregated_intensities = [], [] # incrementally voxel-downsampled chunks + buffered_points, buffered_intensities = [], [] + have_intensities = False + + def collapse(points_chunks, intensity_chunks, voxel): + cloud = o3d.geometry.PointCloud() + cloud.points = o3d.utility.Vector3dVector(np.concatenate(points_chunks).astype(np.float64)) + carry_intensities = bool(intensity_chunks) + if carry_intensities: + intensity_column = np.concatenate(intensity_chunks).astype(np.float64)[:, None] + cloud.colors = o3d.utility.Vector3dVector(np.repeat(intensity_column, 3, axis=1)) + cloud = cloud.voxel_down_sample(voxel) + downsampled_points = np.asarray(cloud.points, np.float32) + downsampled_intensities = ( + np.asarray(cloud.colors, np.float32)[:, 0] if carry_intensities else None ) - m.ts = ts # stamp the cloud (lcm_encode needs a non-None ts) - out.append(m, ts=ts, pose=pose_tuple(Cts.compose(base_pose(ts)))) + return downsampled_points, downsampled_intensities + + print(f"writing {out_name} (corrected lidar)...", flush=True) + written_count = 0 + write_start_time = time.time() + for observation in store.stream(LIDAR_STREAM): + ts = float(observation.ts) + correction = correction_at(ts) + rotation_matrix = correction.rotation().matrix() + translation = np.asarray(correction.translation()) + points = world_points(observation) + intensities = observation.data.intensities_f32() + corrected_points = (points @ rotation_matrix.T + translation).astype(np.float32) + cloud_msg = PointCloud2.from_numpy( + corrected_points, + frame_id="odom", + intensities=(np.asarray(intensities) if intensities is not None else None), + ) + cloud_msg.ts = ts # stamp the cloud (lcm_encode needs a non-None ts) + out_stream.append(cloud_msg, ts=ts, pose=pose_tuple(correction.compose(base_pose(ts)))) if WRITE_LCM: - buf_xyz.append(xyz2) - if inten is not None: - have_inten = True - buf_i.append(np.asarray(inten, np.float32)) - if len(buf_xyz) >= CHUNK: - dx, di = collapse(buf_xyz, buf_i if have_inten else [], LCM_VOXEL) - agg_xyz.append(dx) - if di is not None: - agg_i.append(di) - buf_xyz, buf_i = [], [] - n += 1 - if n % 2000 == 0: - print(f" {n} scans, {time.time() - t0:.0f}s", flush=True) - print(f"wrote {name}: {n} scans in {time.time() - t0:.0f}s", flush=True) + buffered_points.append(corrected_points) + if intensities is not None: + have_intensities = True + buffered_intensities.append(np.asarray(intensities, np.float32)) + if len(buffered_points) >= CHUNK: + downsampled_points, downsampled_intensities = collapse( + buffered_points, buffered_intensities if have_intensities else [], LCM_VOXEL + ) + aggregated_points.append(downsampled_points) + if downsampled_intensities is not None: + aggregated_intensities.append(downsampled_intensities) + buffered_points, buffered_intensities = [], [] + written_count += 1 + if written_count % 2000 == 0: + print(f" {written_count} scans, {time.time() - write_start_time:.0f}s", flush=True) + print( + f"wrote {out_name}: {written_count} scans in {time.time() - write_start_time:.0f}s", + flush=True, + ) if WRITE_LCM: - if buf_xyz: # flush remainder - dx, di = collapse(buf_xyz, buf_i if have_inten else [], LCM_VOXEL) - agg_xyz.append(dx) - if di is not None: - agg_i.append(di) + if buffered_points: # flush remainder + downsampled_points, downsampled_intensities = collapse( + buffered_points, buffered_intensities if have_intensities else [], LCM_VOXEL + ) + aggregated_points.append(downsampled_points) + if downsampled_intensities is not None: + aggregated_intensities.append(downsampled_intensities) # final unified voxel pass over the per-chunk results, then statistical outlier removal - merged_i = agg_i if have_inten else [] - dx, di = collapse(agg_xyz, merged_i, LCM_VOXEL) + merged_intensity_chunks = aggregated_intensities if have_intensities else [] + downsampled_points, downsampled_intensities = collapse( + aggregated_points, merged_intensity_chunks, LCM_VOXEL + ) print( - f"aggregating .pc2.lcm: {len(dx):,} pts after voxel, removing outliers...", flush=True + f"aggregating .pc2.lcm: {len(downsampled_points):,} pts after voxel, " + f"removing outliers...", + flush=True, ) - pc = o3d.geometry.PointCloud() - pc.points = o3d.utility.Vector3dVector(dx.astype(np.float64)) - if di is not None: - pc.colors = o3d.utility.Vector3dVector( - np.repeat(di.astype(np.float64)[:, None], 3, axis=1) + cloud = o3d.geometry.PointCloud() + cloud.points = o3d.utility.Vector3dVector(downsampled_points.astype(np.float64)) + if downsampled_intensities is not None: + cloud.colors = o3d.utility.Vector3dVector( + np.repeat(downsampled_intensities.astype(np.float64)[:, None], 3, axis=1) ) - pc, _keep = pc.remove_statistical_outlier(LCM_OUTLIER_NN, LCM_OUTLIER_STD) - merged_xyz = np.asarray(pc.points, np.float32) - merged_inten = np.asarray(pc.colors, np.float32)[:, 0] if di is not None else None + cloud, _keep = cloud.remove_statistical_outlier(LCM_OUTLIER_NN, LCM_OUTLIER_STD) + merged_xyz = np.asarray(cloud.points, np.float32) + merged_inten = ( + np.asarray(cloud.colors, np.float32)[:, 0] + if downsampled_intensities is not None + else None + ) merged = PointCloud2.from_numpy(merged_xyz, frame_id="odom", intensities=merged_inten) - merged.ts = float(fo_ts[0]) - lcm_path = REC / f"{name}.pc2.lcm" + merged.ts = float(odom_times[0]) + lcm_path = REC / f"{out_name}.pc2.lcm" lcm_path.write_bytes(merged.lcm_encode()) print( f"wrote {lcm_path}: 1 aggregated cloud, {len(merged_xyz):,} pts " From 0a6de2ed75e85f6ccb840871b74a8af064069f4b Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 24 Jun 2026 18:27:03 +0800 Subject: [PATCH 26/69] docs: point jnav postprocessing docs at gsc_pgo/scripts/ location --- experimental/docs/jnav/map_postprocessing.md | 2 +- experimental/docs/jnav/pgo_migration_plan.md | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/experimental/docs/jnav/map_postprocessing.md b/experimental/docs/jnav/map_postprocessing.md index 4f60e042e5..e579e06b79 100644 --- a/experimental/docs/jnav/map_postprocessing.md +++ b/experimental/docs/jnav/map_postprocessing.md @@ -19,7 +19,7 @@ The tags are the whole trick. Drift accumulates, but a tag you saw at minute 1 a ## The three steps -The scripts live in `dimos/navigation/jnav/components/loop_closure/gsc_pgo/`. +The scripts live in `dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/`. **1. Detect the tags.** Run the camera frames through detection and write both the raw and the gated tag streams in one step: diff --git a/experimental/docs/jnav/pgo_migration_plan.md b/experimental/docs/jnav/pgo_migration_plan.md index a35efdc31f..713aa22444 100644 --- a/experimental/docs/jnav/pgo_migration_plan.md +++ b/experimental/docs/jnav/pgo_migration_plan.md @@ -31,9 +31,9 @@ dimos/navigation/jnav/ spec.py # LoopClosure Protocol gsc_pgo/ module.py # NativeModule -> builds external gsc_pgo flake - post_process.py # ported from better_pgo scripts/post_process.py - (add_april.py, detect_tags.py, make_rrd.py) # the other postprocess scripts - ivan_pgo/ ivan_pgo_transformer/ unrefined_pgo/ # comparison baselines (eval_all needs them) + scripts/ + post_process.py # ported from better_pgo scripts/post_process.py + (add_april.py, detect_tags.py, make_rrd.py) # the other postprocess scripts experimental/docs/jnav/ map_postprocessing.md # adapted from docs/capabilities/navigation/map_postprocessing.md ``` @@ -105,7 +105,7 @@ github:jeff-hykin/gsc_pgo # all C++ + flake.nix/flake.lock doesn't break. (If baselines are unwanted, prune eval_all's list instead.) ### Phase 5 — postprocessing scripts (from better_pgo) -- Port `scripts/post_process.py` → `components/loop_closure/gsc_pgo/post_process.py`. +- Port `scripts/post_process.py` → `components/loop_closure/gsc_pgo/scripts/post_process.py`. - Port `add_april.py`, `detect_tags.py`, `make_rrd.py` alongside it (the doc's 3-step flow references all of them). - Rewrite their internal imports: `eval_utils.apriltags` → `dimos.navigation.jnav.utils.apriltags`, etc. From ac8ad9b754e25465d677992243b9518952932d27 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 25 Jun 2026 16:33:17 +0800 Subject: [PATCH 27/69] jnav PGO: replace tag/Landmark inputs with general LocationConstraint Drops the tag (Detection3DArray) loop-closure input and the Landmark-event input from the PGO module, replacing both with a single general-purpose LocationConstraint message that carries ready-made GTSAM BetweenFactor args (body->location relative pose + full 6x6 covariance + to_id + frame_id + constraint_instance_id). The native PGO turns each into its own pose node + BetweenFactor (covariance used directly as the noise model); a constraint_instance_id lets an external source revise/remove prior constraints. - module.py: drop Detection3DArray import + all tag_* config + tag_detections In; use_landmarks -> use_location_constraints; landmarks In[Landmark] -> location_constraints In[LocationConstraint]; drop landmark_var_* + tf_static_channel; add odom_buffer_window; pin gsc_pgo rev e5f9165. - new msgs/LocationConstraint.py (+test); remove msgs/Landmark.py (+test). - gitignore eval_results/. C++ lives in jeff-hykin/gsc_pgo (pinned). Verified ruff/mypy clean, 93 jnav+blueprint tests pass, loop-closure eval on gir_park1_2 bit-identical to pre-rework (no regression). --- dimos/navigation/jnav/.gitignore | 1 + .../components/loop_closure/gsc_pgo/module.py | 61 ++---- dimos/navigation/jnav/msgs/Landmark.py | 200 ------------------ .../jnav/msgs/LocationConstraint.py | 147 +++++++++++++ dimos/navigation/jnav/msgs/test_Landmark.py | 129 ----------- .../jnav/msgs/test_LocationConstraint.py | 93 ++++++++ 6 files changed, 260 insertions(+), 371 deletions(-) delete mode 100644 dimos/navigation/jnav/msgs/Landmark.py create mode 100644 dimos/navigation/jnav/msgs/LocationConstraint.py delete mode 100644 dimos/navigation/jnav/msgs/test_Landmark.py create mode 100644 dimos/navigation/jnav/msgs/test_LocationConstraint.py diff --git a/dimos/navigation/jnav/.gitignore b/dimos/navigation/jnav/.gitignore index 750baebf41..189819aafe 100644 --- a/dimos/navigation/jnav/.gitignore +++ b/dimos/navigation/jnav/.gitignore @@ -1,2 +1,3 @@ result result-* +components/loop_closure/eval_results/ diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py index 458745bec0..c0957810b0 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py @@ -29,10 +29,9 @@ from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray from dimos.navigation.jnav.msgs.Graph3D import Graph3D from dimos.navigation.jnav.msgs.GraphDelta3D import GraphDelta3D -from dimos.navigation.jnav.msgs.Landmark import Landmark +from dimos.navigation.jnav.msgs.LocationConstraint import LocationConstraint from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -45,7 +44,7 @@ class PGOConfig(NativeModuleConfig): cwd: str | None = str(Path(__file__).resolve().parent) executable: str = "result/bin/pgo" build_command: str | None = ( - 'nix build "github:jeff-hykin/gsc_pgo/494e7a1d657c3702ec805c9e3d251a2fe8bc9529#default"' + 'nix build "github:jeff-hykin/gsc_pgo/e5f9165164e124df16535fa6b8616adfae1d10df#default"' " --no-write-lock-file" ) @@ -106,41 +105,22 @@ class PGOConfig(NativeModuleConfig): # Skip ICP on candidates farther than this (m). 0 disables. loop_candidate_max_distance_m: float = 30.0 - # Tag (AprilTag/ArUco) loop closure - use_tag_loop_closure: bool = False - # LCM channel of the static TF tree (dimos "pattern#msg_name" convention). - tf_static_channel: str = "/tf_static#tf2_msgs.TFMessage" - tag_loop_time_thresh: float = 5.0 - tag_assoc_max_dt: float = 0.1 - tag_buffer_window: float = 2.0 - # Anisotropic tag noise (variances, tag frame; normal = +z): in-plane/yaw - # tight, range/out-of-plane loose so tag range error can't distort z. - tag_var_inplane_trans_m2: float = 0.0025 - tag_var_range_trans_m2: float = 0.25 - tag_var_yaw_rot_rad2: float = 0.0025 - tag_var_outplane_rot_rad2: float = 0.04 - # 6-DOF Mahalanobis gate vs current estimate (chi^2 95% = 12.59). 0 = off. - tag_consistency_chi2: float = 0.0 - # Robust (Huber) kernel on all loop factors (lidar + tag). Off = original. + # Robust (Huber) kernel on all loop factors (lidar + location). Off = original. loop_robust_kernel: bool = False loop_robust_huber_k: float = 1.345 - # Landmark events (decoupled perceiver -> PGO factor-graph manager) - # When set, the PGO ingests Landmark events on the `landmarks` In and - # attaches each as a graph landmark variable + a BetweenFactor(keyframe, - # landmark). Two sightings of the same landmark id share the variable, so a - # revisit closes the loop and GTSAM optimizes it jointly. A separate - # perceiver (utils/apriltag_perceiver.py) does the detection + noise/ - # confidence filtering and emits the Landmark events. Off by default. - use_landmarks: bool = False - # Anisotropic landmark observation noise (variances, landmark frame; normal = - # +z): in-plane/yaw tight, range/out-of-plane loose (mirrors the tag model). - landmark_var_inplane_trans_m2: float = 0.0025 - landmark_var_range_trans_m2: float = 0.25 - landmark_var_yaw_rot_rad2: float = 0.0025 - landmark_var_outplane_rot_rad2: float = 0.04 - landmark_assoc_max_dt: float = 0.2 - landmark_buffer_window: float = 3.0 + # Location constraints (decoupled perceiver -> PGO factor-graph manager). + # When set, the PGO ingests LocationConstraint events on the + # `location_constraints` In. Each becomes its own pose node (placed from + # interpolated odometry at the constraint's timestamp) plus a + # BetweenFactor(node, location) whose noise model is the covariance carried in + # the message. Two constraints sharing a to_id share the location variable, so + # a revisit closes the loop; a constraint_instance_id lets an external source + # revise/remove its earlier constraints. Off by default. + use_location_constraints: bool = False + # Seconds of odometry history retained, for interpolating a constraint's pose + # at its own timestamp. + odom_buffer_window: float = 10.0 # Gravity anchor # Pin keyframe 0 (whose orientation is gravity-aligned by the LIO front end) @@ -191,13 +171,10 @@ class PGO(NativeModule): # latest odometry pose internally, so a raw sensor-frame scan is expected. lidar: In[PointCloud2] odometry: In[Odometry] - # Optional: tag detections (tag-in-optical pose + numeric id) for tag-based - # loop closure. Only consumed when config.use_tag_loop_closure is set. - tag_detections: In[Detection3DArray] - # Optional: decoupled Landmark events from a perceiver (e.g. the AprilTag - # perceiver). Only consumed when config.use_landmarks is set; each becomes a - # graph landmark variable + observation factor that GTSAM optimizes jointly. - landmarks: In[Landmark] + # Optional: decoupled LocationConstraint events from a perceiver. Only + # consumed when config.use_location_constraints is set; each becomes its own + # pose node + a BetweenFactor(node, location) that GTSAM optimizes jointly. + location_constraints: In[LocationConstraint] corrected_odometry: Out[Odometry] correction: Out[Transform] pose_graph: Out[Graph3D] diff --git a/dimos/navigation/jnav/msgs/Landmark.py b/dimos/navigation/jnav/msgs/Landmark.py deleted file mode 100644 index d299fa3404..0000000000 --- a/dimos/navigation/jnav/msgs/Landmark.py +++ /dev/null @@ -1,200 +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. - -"""Landmark: a cross-map reference shared between maps by ``id``. - -Two maps that both observe the same ``id`` are bridgeable: MultiMap composes the -relative transform between their frames through the landmark. The pose is stored -relative to the frame it was observed in (a camera frame for an apriltag, an odom -frame for a relocalization event, the nearest cloud's frame for a UI click), NOT -the map root — so the uncertainty stays honest for a future gtsam graph. MultiMap -resolves it to the map root via the map's recorded tf at ``ts``. - -``id`` is URL-like and encodes the exact source so identical tag numbers from -different families/sizes don't false-bridge unrelated maps, e.g. -``apriltag://36h11/40cm/5`` or ``reloc://map0/dim_city``. - -``replacement`` (milliseconds) lets a perceiver publish a *corrective* landmark: -a consumer (e.g. the PGO factor-graph manager) wipes out any earlier landmark of -the same ``id`` observed within ``replacement`` ms before this one's ``ts`` and -keeps this one instead. ``0`` (the default) means "additive" — keep all prior -landmarks, which is what loop closure needs (the first and the revisit sighting -must both survive). A small positive window only supersedes the immediately-prior -noisy estimate of the *same* sighting, never the historical closure observation. - -Confidence is **per-axis**, not a single scalar: ``confidence_x/_y/_z`` (the -observed translation axes) and ``confidence_roll/_pitch/_yaw`` (the rotation -axes), each in ``[0, 1]``. A consumer maps each per-axis confidence to that -axis's measurement weight (variance ~ 1 / confidence), so a landmark can be well -observed on some DOF and weak on others — e.g. a face-on AprilTag pins in-plane -position + yaw tightly but its range and out-of-plane tilt are loose. ``confidence`` -(the scalar) is retained as a coarse overall value; the per-axis fields default -to it when unspecified. -""" - -from __future__ import annotations - -import struct -import time -from typing import BinaryIO - -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.types.timestamped import Timestamped - - -class Landmark(Timestamped): - msg_name = "jnav.Landmark" - - ts: float - id: str # URL-like, shared across maps — the bridge key - map_id: str # which map's frame system this is in - frame_id: str # the observation frame the pose is relative to - kind: str # "apriltag" | "aruco" | "reloc" | "ui_click" | "shortcut" - confidence: float # coarse overall confidence; per-axis fields default to it - pose: Pose - replacement: float # ms: wipe same-id landmarks newer than (ts - replacement); 0 = additive - # Per-axis confidence in [0, 1]: how much this landmark constrains each DOF. - confidence_x: float - confidence_y: float - confidence_z: float - confidence_roll: float - confidence_pitch: float - confidence_yaw: float - - def __init__( - self, - id: str = "", - map_id: str = "", - pose: Pose | None = None, - frame_id: str = "", - confidence: float = 0.0, - kind: str = "", - ts: float = 0.0, - replacement: float = 0.0, - confidence_x: float | None = None, - confidence_y: float | None = None, - confidence_z: float | None = None, - confidence_roll: float | None = None, - confidence_pitch: float | None = None, - confidence_yaw: float | None = None, - ) -> None: - self.ts = ts if ts != 0 else time.time() - self.id = id - self.map_id = map_id - self.frame_id = frame_id - self.kind = kind - self.confidence = confidence - self.pose = pose if pose is not None else Pose() - self.replacement = replacement - # Per-axis confidence defaults to the coarse overall confidence. - self.confidence_x = confidence if confidence_x is None else confidence_x - self.confidence_y = confidence if confidence_y is None else confidence_y - self.confidence_z = confidence if confidence_z is None else confidence_z - self.confidence_roll = confidence if confidence_roll is None else confidence_roll - self.confidence_pitch = confidence if confidence_pitch is None else confidence_pitch - self.confidence_yaw = confidence if confidence_yaw is None else confidence_yaw - - def axis_confidence(self) -> tuple[float, float, float, float, float, float]: - """Per-axis confidence as ``(x, y, z, roll, pitch, yaw)``.""" - return ( - self.confidence_x, - self.confidence_y, - self.confidence_z, - self.confidence_roll, - self.confidence_pitch, - self.confidence_yaw, - ) - - def lcm_encode(self) -> bytes: - parts: list[bytes] = [struct.pack(">d", self.ts)] - for text in (self.id, self.map_id, self.frame_id, self.kind): - encoded = text.encode("utf-8") - parts.append(struct.pack(">I", len(encoded))) - parts.append(encoded) - parts.append(struct.pack(">d", self.confidence)) - p = self.pose - parts.append( - struct.pack( - ">7d", - p.position.x, - p.position.y, - p.position.z, - p.orientation.x, - p.orientation.y, - p.orientation.z, - p.orientation.w, - ) - ) - parts.append(struct.pack(">d", self.replacement)) - parts.append( - struct.pack( - ">6d", - self.confidence_x, - self.confidence_y, - self.confidence_z, - self.confidence_roll, - self.confidence_pitch, - self.confidence_yaw, - ) - ) - return b"".join(parts) - - @classmethod - def lcm_decode(cls, data: bytes | BinaryIO) -> Landmark: - buf = data if isinstance(data, (bytes, bytearray)) else data.read() - offset = 0 - (ts,) = struct.unpack_from(">d", buf, offset) - offset += 8 - texts: list[str] = [] - for _ in range(4): - (length,) = struct.unpack_from(">I", buf, offset) - offset += 4 - texts.append(buf[offset : offset + length].decode("utf-8")) - offset += length - id_, map_id, frame_id, kind = texts - (confidence,) = struct.unpack_from(">d", buf, offset) - offset += 8 - px, py, pz, qx, qy, qz, qw = struct.unpack_from(">7d", buf, offset) - offset += 56 - pose = Pose() - pose.position = Vector3(px, py, pz) - pose.orientation = Quaternion(qx, qy, qz, qw) - # Backward-compatible: older encodings omit the trailing replacement field. - replacement = 0.0 - if offset + 8 <= len(buf): - (replacement,) = struct.unpack_from(">d", buf, offset) - offset += 8 - # Backward-compatible: older encodings omit the per-axis confidences; - # they then default to the coarse overall confidence (None -> confidence). - cx = cy = cz = croll = cpitch = cyaw = None - if offset + 48 <= len(buf): - cx, cy, cz, croll, cpitch, cyaw = struct.unpack_from(">6d", buf, offset) - return cls( - id=id_, - map_id=map_id, - pose=pose, - frame_id=frame_id, - confidence=confidence, - kind=kind, - ts=ts, - replacement=replacement, - confidence_x=cx, - confidence_y=cy, - confidence_z=cz, - confidence_roll=croll, - confidence_pitch=cpitch, - confidence_yaw=cyaw, - ) diff --git a/dimos/navigation/jnav/msgs/LocationConstraint.py b/dimos/navigation/jnav/msgs/LocationConstraint.py new file mode 100644 index 0000000000..dc6a9c1e7a --- /dev/null +++ b/dimos/navigation/jnav/msgs/LocationConstraint.py @@ -0,0 +1,147 @@ +# 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. + +"""LocationConstraint: the ready-made args for a GTSAM ``BetweenFactor``. + +A LocationConstraint is a relative-pose measurement from a ``from`` graph node to +a ``to`` location variable, plus the 6x6 covariance that becomes the factor's +noise model directly. The PGO turns each one into its own pose node (placed from +interpolated odometry at ``ts``) and a ``BetweenFactor(node, location)``. + +Field meanings (mapping onto BetweenFactor): +- ``to_id`` is the BetweenFactor "to": the location variable's identity, URL-like + (e.g. ``apriltag://36h11/40cm/5`` or ``gps://fix`` or ``ui_click://...``). Two + constraints sharing a ``to_id`` observe the same graph variable, which closes + the loop. +- ``frame_id`` is the BetweenFactor "from": the frame the ``pose`` is expressed + in. For now the PGO enforces ``frame_id == body_frame`` (no full C++ tf yet), + and re-bases the measurement onto the node it creates. +- ``pose`` is the relative transform ``frame_id -> location``. +- ``covariance`` is the 6x6 measurement covariance in GTSAM Pose3 tangent order + ``[rot(3), trans(3)]`` (row-major, 36 values). It is used as the factor's noise + model directly — degenerate DOFs (e.g. a position-only fix) get a huge variance + on the rotation block. +- ``constraint_instance_id`` identifies this specific external instance. A later + constraint reusing the same ``constraint_instance_id`` removes the committed + factors carrying it, letting an external estimator do rolling outlier removal / + revision (e.g. as a tag/GPS lock improves). +""" + +from __future__ import annotations + +import struct +import time +from typing import BinaryIO + +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.types.timestamped import Timestamped + +# 6x6 covariance, row-major. +_COVARIANCE_LENGTH = 36 + + +def _identity_covariance() -> list[float]: + """A neutral, non-degenerate default: unit variance on every DOF.""" + cov = [0.0] * _COVARIANCE_LENGTH + for axis in range(6): + cov[axis * 6 + axis] = 1.0 + return cov + + +class LocationConstraint(Timestamped): + msg_name = "jnav.LocationConstraint" + + ts: float + to_id: str # the BetweenFactor "to" (location variable id), URL-like + frame_id: str # the BetweenFactor "from" frame (== body frame for now) + constraint_instance_id: str # external instance id, for revision/removal + pose: Pose # relative transform frame_id -> location + covariance: list[float] # 6x6 row-major, tangent order [rot(3), trans(3)] + + def __init__( + self, + to_id: str = "", + frame_id: str = "", + pose: Pose | None = None, + covariance: list[float] | None = None, + constraint_instance_id: str = "", + ts: float = 0.0, + ) -> None: + self.ts = ts if ts != 0 else time.time() + self.to_id = to_id + self.frame_id = frame_id + self.constraint_instance_id = constraint_instance_id + self.pose = pose if pose is not None else Pose() + if covariance is None: + self.covariance = _identity_covariance() + else: + if len(covariance) != _COVARIANCE_LENGTH: + raise ValueError( + f"covariance must be {_COVARIANCE_LENGTH} values (6x6 row-major), " + f"got {len(covariance)}" + ) + self.covariance = list(covariance) + + def lcm_encode(self) -> bytes: + parts: list[bytes] = [struct.pack(">d", self.ts)] + for text in (self.to_id, self.frame_id, self.constraint_instance_id): + encoded = text.encode("utf-8") + parts.append(struct.pack(">I", len(encoded))) + parts.append(encoded) + p = self.pose + parts.append( + struct.pack( + ">7d", + p.position.x, + p.position.y, + p.position.z, + p.orientation.x, + p.orientation.y, + p.orientation.z, + p.orientation.w, + ) + ) + parts.append(struct.pack(">36d", *self.covariance)) + return b"".join(parts) + + @classmethod + def lcm_decode(cls, data: bytes | BinaryIO) -> LocationConstraint: + buf = data if isinstance(data, (bytes, bytearray)) else data.read() + offset = 0 + (ts,) = struct.unpack_from(">d", buf, offset) + offset += 8 + texts: list[str] = [] + for _ in range(3): + (length,) = struct.unpack_from(">I", buf, offset) + offset += 4 + texts.append(buf[offset : offset + length].decode("utf-8")) + offset += length + to_id, frame_id, constraint_instance_id = texts + px, py, pz, qx, qy, qz, qw = struct.unpack_from(">7d", buf, offset) + offset += 56 + pose = Pose() + pose.position = Vector3(px, py, pz) + pose.orientation = Quaternion(qx, qy, qz, qw) + covariance = list(struct.unpack_from(">36d", buf, offset)) + offset += _COVARIANCE_LENGTH * 8 + return cls( + to_id=to_id, + frame_id=frame_id, + pose=pose, + covariance=covariance, + constraint_instance_id=constraint_instance_id, + ts=ts, + ) diff --git a/dimos/navigation/jnav/msgs/test_Landmark.py b/dimos/navigation/jnav/msgs/test_Landmark.py deleted file mode 100644 index 1df849151a..0000000000 --- a/dimos/navigation/jnav/msgs/test_Landmark.py +++ /dev/null @@ -1,129 +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. - -"""Roundtrip + field tests for the jnav Landmark message.""" - -from __future__ import annotations - -from dimos.memory2.codecs.base import codec_for -from dimos.memory2.codecs.lcm import LcmCodec -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.navigation.jnav.msgs.Landmark import Landmark - - -def _pose(x: float, y: float, z: float) -> Pose: - pose = Pose() - pose.position.x, pose.position.y, pose.position.z = x, y, z - return pose - - -def test_landmark_roundtrip_preserves_all_fields() -> None: - landmark = Landmark( - id="apriltag://36h11/40cm/5", - map_id="dim_city", - pose=_pose(1.5, -2.0, 0.3), - frame_id="camera_optical", - confidence=0.82, - kind="apriltag", - ts=1781565207.5, - ) - decoded = Landmark.lcm_decode(landmark.lcm_encode()) - - assert decoded.id == "apriltag://36h11/40cm/5" - assert decoded.map_id == "dim_city" - assert decoded.frame_id == "camera_optical" - assert decoded.kind == "apriltag" - assert decoded.confidence == 0.82 - assert decoded.ts == landmark.ts - assert decoded.pose.position.x == 1.5 - assert decoded.pose.position.y == -2.0 - assert decoded.pose.position.z == 0.3 - - -def test_landmark_defaults() -> None: - landmark = Landmark() - assert landmark.id == "" - assert landmark.map_id == "" - assert landmark.frame_id == "" - assert landmark.kind == "" - assert landmark.confidence == 0.0 - assert landmark.replacement == 0.0 # additive by default - assert landmark.ts > 0 # auto-stamped - - decoded = Landmark.lcm_decode(landmark.lcm_encode()) - assert decoded.id == "" and decoded.map_id == "" and decoded.kind == "" - assert decoded.replacement == 0.0 - - -def test_landmark_replacement_roundtrips() -> None: - landmark = Landmark( - id="apriltag://36h11/40cm/13", - pose=_pose(0.1, 0.2, 0.5), - kind="apriltag", - ts=1781565207.5, - replacement=2500.0, # ms: corrective landmark supersedes recent same-id ones - ) - decoded = Landmark.lcm_decode(landmark.lcm_encode()) - assert decoded.replacement == 2500.0 - - -def test_landmark_decode_tolerates_legacy_bytes_without_replacement() -> None: - """Bytes written before the replacement + per-axis fields decode to defaults.""" - landmark = Landmark(id="x", pose=_pose(0.0, 0.0, 0.0), confidence=0.7, replacement=999.0) - # Strip the trailing replacement double + the 6 per-axis confidence doubles. - legacy = landmark.lcm_encode()[: -(8 + 48)] - decoded = Landmark.lcm_decode(legacy) - assert decoded.id == "x" - assert decoded.replacement == 0.0 - # Per-axis confidences fall back to the coarse overall confidence. - assert decoded.axis_confidence() == (0.7, 0.7, 0.7, 0.7, 0.7, 0.7) - - -def test_per_axis_confidence_defaults_to_overall() -> None: - lm = Landmark(id="t", confidence=0.6) - assert lm.axis_confidence() == (0.6, 0.6, 0.6, 0.6, 0.6, 0.6) - - -def test_per_axis_confidence_roundtrips() -> None: - lm = Landmark( - id="apriltag://36h11/40cm/13", - pose=_pose(0.1, 0.2, 0.5), - kind="apriltag", - confidence=0.5, - confidence_x=0.9, - confidence_y=0.8, - confidence_z=0.7, - confidence_roll=0.2, - confidence_pitch=0.1, - confidence_yaw=0.95, - ) - d = Landmark.lcm_decode(lm.lcm_encode()) - assert d.axis_confidence() == (0.9, 0.8, 0.7, 0.2, 0.1, 0.95) - assert d.confidence == 0.5 # coarse overall preserved independently - - -def test_per_axis_confidence_can_be_silent_on_some_dof() -> None: - """A landmark can be confident on some DOF and zero on others.""" - lm = Landmark(id="t", pose=_pose(0.0, 0.0, 0.0), confidence_z=0.9, confidence_x=0.0) - d = Landmark.lcm_decode(lm.lcm_encode()) - assert d.confidence_z == 0.9 and d.confidence_x == 0.0 - - -def test_landmark_uses_lcm_codec_in_memory2() -> None: - codec = codec_for(Landmark) - assert isinstance(codec, LcmCodec) - landmark = Landmark(id="reloc://map0/dim_city", map_id="map0", confidence=0.6, kind="reloc") - decoded = codec.decode(codec.encode(landmark)) - assert decoded.id == "reloc://map0/dim_city" - assert decoded.confidence == 0.6 diff --git a/dimos/navigation/jnav/msgs/test_LocationConstraint.py b/dimos/navigation/jnav/msgs/test_LocationConstraint.py new file mode 100644 index 0000000000..08e7b714c6 --- /dev/null +++ b/dimos/navigation/jnav/msgs/test_LocationConstraint.py @@ -0,0 +1,93 @@ +# 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. + +"""Roundtrip + field tests for the jnav LocationConstraint message.""" + +from __future__ import annotations + +import pytest + +from dimos.memory2.codecs.base import codec_for +from dimos.memory2.codecs.lcm import LcmCodec +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.navigation.jnav.msgs.LocationConstraint import LocationConstraint + + +def _pose(x: float, y: float, z: float) -> Pose: + pose = Pose() + pose.position.x, pose.position.y, pose.position.z = x, y, z + return pose + + +def _cov(scale: float = 1.0) -> list[float]: + cov = [0.0] * 36 + for axis in range(6): + cov[axis * 6 + axis] = scale * (axis + 1) + return cov + + +def test_roundtrip_preserves_all_fields() -> None: + cov = _cov(0.01) + constraint = LocationConstraint( + to_id="apriltag://36h11/40cm/5", + frame_id="base_link", + pose=_pose(1.5, -2.0, 0.3), + covariance=cov, + constraint_instance_id="tag5#42", + ts=1781565207.5, + ) + decoded = LocationConstraint.lcm_decode(constraint.lcm_encode()) + + assert decoded.to_id == "apriltag://36h11/40cm/5" + assert decoded.frame_id == "base_link" + assert decoded.constraint_instance_id == "tag5#42" + assert decoded.ts == constraint.ts + assert decoded.pose.position.x == 1.5 + assert decoded.pose.position.y == -2.0 + assert decoded.pose.position.z == 0.3 + assert decoded.covariance == cov + + +def test_defaults() -> None: + constraint = LocationConstraint() + assert constraint.to_id == "" + assert constraint.frame_id == "" + assert constraint.constraint_instance_id == "" + assert constraint.ts > 0 # auto-stamped + # Default covariance is a non-degenerate identity (unit variance per DOF). + assert constraint.covariance[0] == 1.0 and constraint.covariance[35] == 1.0 + assert sum(constraint.covariance) == 6.0 + + +def test_full_6x6_covariance_roundtrips_offdiagonals() -> None: + cov = [float(i) for i in range(36)] # all entries distinct, incl. off-diagonal + constraint = LocationConstraint(to_id="x", frame_id="base_link", covariance=cov) + decoded = LocationConstraint.lcm_decode(constraint.lcm_encode()) + assert decoded.covariance == cov + + +def test_wrong_covariance_length_rejected() -> None: + with pytest.raises(ValueError): + LocationConstraint(to_id="x", covariance=[0.0] * 35) + + +def test_uses_lcm_codec_in_memory2() -> None: + codec = codec_for(LocationConstraint) + assert isinstance(codec, LcmCodec) + constraint = LocationConstraint( + to_id="gps://fix", frame_id="base_link", constraint_instance_id="gps#7" + ) + decoded = codec.decode(codec.encode(constraint)) + assert decoded.to_id == "gps://fix" + assert decoded.constraint_instance_id == "gps#7" From 9a99d52d15b3dcb202581fb716057a5edb058805 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 25 Jun 2026 16:45:47 +0800 Subject: [PATCH 28/69] add unitree_go2_pgo blueprint (Point-LIO -> jnav PGO -> Rerun pose graph) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A live, real-dog test rig for the jnav gsc_pgo PGO: Point-LIO reads the Go2's Livox Mid-360 and publishes registered lidar + odometry, the PGO builds a loop-closed pose graph, and the Rerun bridge renders it as nodes (keyframes) + edges (odom backbone / loop closures) via Graph3D.to_rerun_multi (visual_override on world/pose_graph). Passive observer — drive the dog separately. Deliberately omits GO2Connection (its lidar Out would collide with Point-LIO's registered lidar). Regenerated all_blueprints.py. Composition validity test passes. --- dimos/robot/all_blueprints.py | 1 + .../go2/blueprints/smart/unitree_go2_pgo.py | 59 +++++++++++++++++++ 2 files changed, 60 insertions(+) create mode 100644 dimos/robot/unitree/go2/blueprints/smart/unitree_go2_pgo.py diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index c79a3b75f4..eb4c274748 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -115,6 +115,7 @@ "unitree-go2-keyboard-teleop": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_keyboard_teleop:unitree_go2_keyboard_teleop", "unitree-go2-markers": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_markers", "unitree-go2-memory": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_memory", + "unitree-go2-pgo": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2_pgo:unitree_go2_pgo", "unitree-go2-relocalization": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_relocalization", "unitree-go2-ros": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2_ros:unitree_go2_ros", "unitree-go2-security": "dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_security:unitree_go2_security", diff --git a/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_pgo.py b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_pgo.py new file mode 100644 index 0000000000..5bfa626dad --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/smart/unitree_go2_pgo.py @@ -0,0 +1,59 @@ +# 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. + +"""unitree_go2_pgo: run the jnav PGO live on a Go2's Livox Mid-360 + Point-LIO +and visualize the optimized pose graph in Rerun. + +Point-LIO reads the Mid-360 and publishes a registered `lidar` (PointCloud2) plus +`odometry` (Odometry); the PGO consumes both and emits a loop-closed `pose_graph` +(Graph3D). The Rerun bridge renders that graph as nodes (keyframes) + edges (odom +backbone in green, loop closures in yellow) via `Graph3D.to_rerun_multi`. + +This is a passive observer rig — drive the dog however you like (Go2 app / a +teleop blueprint) and watch the graph build and snap on loop closure. It needs +only the Mid-360 + Point-LIO, so it deliberately does NOT pull in GO2Connection +(whose own `lidar` Out would collide with Point-LIO's registered `lidar`). + +Run on the dog: + dimos run unitree-go2-pgo +""" + +from __future__ import annotations + +from dimos.core.coordination.blueprints import autoconnect +from dimos.hardware.sensors.lidar.pointlio.module import PointLio +from dimos.navigation.jnav.components.loop_closure.gsc_pgo.module import PGO +from dimos.navigation.jnav.msgs.Graph3D import Graph3D +from dimos.visualization.rerun.bridge import RerunMulti +from dimos.visualization.vis_module import vis_module + +# Rerun entity path for the pose graph. The bridge maps the `pose_graph` stream to +# `/pose_graph` = `world/pose_graph`; matching that here lets the +# override draw nodes + edges instead of the default nodes-only Points3D. +_POSE_GRAPH_PATH = "world/pose_graph" + + +def _render_pose_graph(graph: Graph3D) -> RerunMulti: + """Nodes (keyframes) + edges (odom backbone / loop closures) for the graph.""" + return graph.to_rerun_multi(base_path=_POSE_GRAPH_PATH) + + +unitree_go2_pgo = autoconnect( + PointLio.blueprint(), + PGO.blueprint(), + vis_module( + "rerun", + rerun_config={"visual_override": {_POSE_GRAPH_PATH: _render_pose_graph}}, + ), +).global_config(n_workers=3, robot_model="unitree_go2") From 1b4a214a959640da601b129c2449c94445d65a0b Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Fri, 26 Jun 2026 09:00:32 +0800 Subject: [PATCH 29/69] memory2 DbTf: per-row tf tree instead of full-RAM load MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DbTf no longer loads the whole tf history into a MultiTBuffer. Each tf row now carries a 'tf tree' — a snapshot of the last two recorded row-ids per child frame, stored in a sibling _tree table. Tree values are tf-row ids (keys into the tf table), so editing a transform's pose in place changes what the tree resolves to. tf.get(target, source, ts) finds the tf row nearest the query, reads its tree, fetches only the <=2-per-frame referenced rows, and composes/interpolates the chain through a small MultiTBuffer; a 20-entry FIFO serves repeated identical lookups with no db reads. Linear interpolation is preserved (two refs per frame). - Recorder (memory2/module.py) builds the tree incrementally as tf rows are appended, and now also subscribes the static_tf stream (treated the same as dynamic, folded into 'tf'). - On load, a tree-less recording prints a loud warning then recomputes all trees chronologically and writes them back (every tf row ends up with a tree). - Non-SQLite stores fall back to the old full-load buffer. - NOTE in code: retro-correction must edit existing rows in place; inserting new corrective tf messages is not seen by already-built trees. The point is memory + scale, NOT raw lookup speed: a warm in-RAM buffer is actually a bit faster per lookup (~97us) than the tree (~181us, ~4 db rows each) once it's loaded. What the tree avoids is decoding all 38370 rows (~0.5s) into RAM at every process start; it instead reads ~4 rows/query, keeps memory bounded, and scales to recordings too big to load. Verified on china_default.db: tree get matches the old full-load buffer to 0.0 at odom-aligned query times. New test_db_tf.py (8 tests) + mypy/ruff clean; 330 existing memory2 tests pass. --- dimos/memory2/db_tf.py | 297 ++++++++++++++++++++++++++++++--- dimos/memory2/module.py | 24 ++- dimos/memory2/test_db_tf.py | 320 ++++++++++++++++++++++++++++++++++++ 3 files changed, 613 insertions(+), 28 deletions(-) create mode 100644 dimos/memory2/test_db_tf.py diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py index 7f52de854b..e6f44b60a4 100644 --- a/dimos/memory2/db_tf.py +++ b/dimos/memory2/db_tf.py @@ -14,22 +14,36 @@ """Transform lookups over the transforms recorded in a store. -A store's ``tf`` member lazily reads every transform recorded under the ``tf`` -(and ``tf_static``) streams into a :class:`MultiTBuffer`, then answers -``store.tf.get(target_frame, source_frame, time)`` — composing multi-hop chains -(e.g. ``world -> map -> odom -> base_link -> mid360_link``) and interpolating to -the nearest recorded sample. This makes world-registration a real transform -lookup instead of assuming a single baked-in pose. - -``write_tf_tree`` populates those streams for a recording that lacks them. +Instead of loading every recorded transform into RAM, each tf row carries a +small **tf tree**: a snapshot, as of that row, of the last two recorded rows per +child frame. The tree is stored in a sibling ``_tree`` table and its +values are tf-row ids — *keys into the tf table* — so editing a transform's pose +in place changes what the tree resolves to. + +``store.tf.get(target, source, ts)`` then: finds the tf row nearest the query +time, reads its tree, fetches only the (<= 2 per frame) referenced rows, and +composes/interpolates the chain through a small :class:`MultiTBuffer`. A FIFO of +the last few lookups avoids re-querying the db for repeated identical lookups. + +NOTE — retro-correction caveat: because a tree's edges are keys to *existing* +rows, the right way to correct history (e.g. after a loop closure) is to EDIT the +existing tf rows' poses in place — already-built trees pick that up for free. +Inserting NEW corrective tf messages does NOT work: trees built before the +insertion still point at the old rows and never see the new ones. (Re-running the +recompute below would rebuild trees to include them, but in-place edits are the +intended path.) + +``write_tf_tree`` populates the tf stream for a recording that lacks one. """ from __future__ import annotations +from collections import OrderedDict +import json import re import sqlite3 import threading -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import numpy as np @@ -39,14 +53,23 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.protocol.tf.tf import MultiTBuffer +from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from dimos.memory2.store.base import Store +logger = setup_logger() + +# The unified tf stream the tree is built over. The recorder folds both the +# dynamic (`tf`) and static (`static_tf`) topics into this one stream/table so the +# tree's row-id keys are unambiguous (static is treated the same as dynamic). +DEFAULT_TF_STREAM = "tf" +# Streams the RAM fallback (non-sqlite stores) reads. TF_STREAMS = ("tf", "tf_static") -# Larger than any single recording's span so the buffer never prunes loaded -# transforms (MultiTBuffer drops samples older than ts - buffer_size). +# Larger than any single recording's span so the fallback buffer never prunes. _NO_PRUNE = 1.0e15 +# Last-N identical lookups cached to avoid re-querying the db. +_LOOKUP_CACHE_SIZE = 20 # SQLite can't parameterize table names, so caller-supplied stream names are # interpolated; allow only safe identifiers to keep that injection-free. _SAFE_TABLE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") @@ -58,20 +81,233 @@ def _safe_table(name: str) -> str: return name +def _connect(db_path: str) -> sqlite3.Connection: + """A connection that waits on the WAL write-lock instead of erroring — the + store keeps its own connection to the same db open while we read/write trees.""" + conn = sqlite3.connect(db_path, check_same_thread=False) + conn.execute("PRAGMA busy_timeout=5000") + return conn + + +def _advance(latest: dict[str, list[int]], child_frame: str, tf_id: int) -> None: + """Push ``tf_id`` as the newest reference for ``child_frame``, keeping <= 2.""" + refs = latest.get(child_frame) + if refs is None: + latest[child_frame] = [tf_id] + else: + refs.append(tf_id) + if len(refs) > 2: + del refs[0] + + +class TfTreeWriter: + """Builds tf trees incrementally as rows are appended (used by the recorder). + + Holds the running ``{child_frame: [older_id, newer_id]}`` state and, after each + tf row is appended, writes that row's full tree snapshot to ``_tree``. + """ + + def __init__(self, db_path: str, stream: str = DEFAULT_TF_STREAM) -> None: + self._stream = _safe_table(stream) + self._conn = _connect(db_path) + self._lock = threading.Lock() + self._latest: dict[str, list[int]] = {} + ensure_tree_table(self._conn, self._stream) + + def record(self, child_frame: str, tf_id: int) -> None: + """Record that tf row ``tf_id`` holds the latest ``child_frame`` edge.""" + with self._lock: + _advance(self._latest, child_frame, tf_id) + snapshot = json.dumps(self._latest) + self._conn.execute( + f'INSERT OR REPLACE INTO "{self._stream}_tree" (tf_id, tree) VALUES (?, ?)', + (tf_id, snapshot), + ) + self._conn.commit() + + def close(self) -> None: + with self._lock: + self._conn.close() + + +def ensure_tree_table(conn: sqlite3.Connection, stream: str) -> None: + conn.execute( + f'CREATE TABLE IF NOT EXISTS "{_safe_table(stream)}_tree" ' + "(tf_id INTEGER PRIMARY KEY, tree TEXT NOT NULL)" + ) + # Speeds up the nearest-row lookup in get(); harmless if it already exists. + conn.execute( + f'CREATE INDEX IF NOT EXISTS "{_safe_table(stream)}_ts_idx" ON "{_safe_table(stream)}"(ts)' + ) + conn.commit() + + +def recompute_trees(store: Store, stream: str = DEFAULT_TF_STREAM) -> int: + """(Re)build every tf row's tree in chronological order, writing it back. + + Iterates the whole stream once — the expensive, one-time cost of adopting a + recording that predates trees. Returns the number of tf rows given a tree. + """ + config = store.config + if not isinstance(config, SqliteStoreConfig): + raise TypeError("recompute_trees needs a SqliteStore") + conn = _connect(config.path) + try: + ensure_tree_table(conn, stream) + conn.execute(f'DELETE FROM "{_safe_table(stream)}_tree"') + latest: dict[str, list[int]] = {} + rows: list[tuple[int, str]] = [] + for observation in store.stream(stream, TFMessage).order_by("ts"): + transforms = getattr(observation.data, "transforms", None) or [observation.data] + for transform in transforms: + _advance(latest, transform.child_frame_id, observation.id) + rows.append((observation.id, json.dumps(latest))) + conn.executemany( + f'INSERT OR REPLACE INTO "{_safe_table(stream)}_tree" (tf_id, tree) VALUES (?, ?)', + rows, + ) + conn.commit() + return len(rows) + finally: + conn.close() + + class DbTf: - """Transform lookups backed by the ``tf``/``tf_static`` streams of a store.""" + """Transform lookups backed by a store's tf stream + per-row tf trees.""" - def __init__(self, store: Store, stream_names: tuple[str, ...] = TF_STREAMS) -> None: + def __init__( + self, + store: Store, + stream: str = DEFAULT_TF_STREAM, + stream_names: tuple[str, ...] = TF_STREAMS, + cache_size: int = _LOOKUP_CACHE_SIZE, + ) -> None: self._store = store - self._stream_names = stream_names + self._stream = _safe_table(stream) + self._stream_names = stream_names # RAM fallback only + self._cache_size = cache_size + self._lock = threading.Lock() + self._conn: sqlite3.Connection | None = None + self._trees_ready = False + self._cache: OrderedDict[tuple[str, str, float | None, float | None], Transform | None] = ( + OrderedDict() + ) + # Instrumentation: number of tf rows decoded out of the db by get() (for + # the "no full load" test). The whole-table recompute is not counted. + self.rows_fetched = 0 + # RAM fallback for non-sqlite stores. self._buffer: MultiTBuffer | None = None - self._load_lock = threading.Lock() + + # --- sqlite path ----------------------------------------------------------- + + @property + def _is_sqlite(self) -> bool: + return isinstance(self._store.config, SqliteStoreConfig) + + def _connection(self) -> sqlite3.Connection: + if self._conn is None: + config = self._store.config + assert isinstance(config, SqliteStoreConfig) # guarded by _is_sqlite + self._conn = _connect(config.path) + return self._conn + + def _ensure_trees(self) -> None: + """Make sure every tf row has a tree; recompute (loudly) if not.""" + if self._trees_ready: + return + conn = self._connection() + ensure_tree_table(conn, self._stream) + (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() + (n_trees,) = conn.execute(f'SELECT count(*) FROM "{self._stream}_tree"').fetchone() + if n_rows and n_trees < n_rows: + logger.warning( + "\n" + "========================================================================\n" + " tf tree MISSING for stream %r (%d/%d rows have a tree).\n" + " Computing tf trees over the whole recording (one-time, chronological,\n" + " written back to the db). This can take a moment on large recordings.\n" + "========================================================================", + self._stream, + n_trees, + n_rows, + ) + built = recompute_trees(self._store, self._stream) + logger.warning("tf tree computed for %d rows of stream %r.", built, self._stream) + self._trees_ready = True + + def _fetch_transforms(self, ids: list[int]) -> list[Transform]: + """Decode the given tf-row ids into transforms (frames live in the blob).""" + if not ids: + return [] + backend: Any = self._store.stream(self._stream, TFMessage)._source + transforms: list[Transform] = [] + for tf_id in ids: + # _make_loader is the same blob-load + codec-decode the stream uses. + message = backend._make_loader(tf_id)() + transforms.extend(getattr(message, "transforms", None) or [message]) + self.rows_fetched += len(ids) + return transforms + + def _anchor_id(self, conn: sqlite3.Connection, time_point: float | None) -> int | None: + """The tf row whose tree to use: the first row with ts >= query (so its + own edge brackets the query for interpolation); else the last row.""" + if time_point is None: + row = conn.execute( + f'SELECT id FROM "{self._stream}" ORDER BY ts DESC LIMIT 1' + ).fetchone() + return int(row[0]) if row else None + row = conn.execute( + f'SELECT id FROM "{self._stream}" WHERE ts >= ? ORDER BY ts ASC LIMIT 1', (time_point,) + ).fetchone() + if row is not None: + return int(row[0]) + row = conn.execute(f'SELECT id FROM "{self._stream}" ORDER BY ts DESC LIMIT 1').fetchone() + return int(row[0]) if row else None + + def _get_sqlite( + self, + target_frame: str, + source_frame: str, + time_point: float | None, + time_tolerance: float | None, + ) -> Transform | None: + key = (target_frame, source_frame, time_point, time_tolerance) + with self._lock: + if key in self._cache: + self._cache.move_to_end(key) + return self._cache[key] + self._ensure_trees() + conn = self._connection() + anchor = self._anchor_id(conn, time_point) + result: Transform | None = None + if anchor is not None: + row = conn.execute( + f'SELECT tree FROM "{self._stream}_tree" WHERE tf_id = ?', (anchor,) + ).fetchone() + if row is not None: + tree: dict[str, list[int]] = json.loads(row[0]) + ids = sorted({tf_id for refs in tree.values() for tf_id in refs}) + # Only the <=2-per-frame referenced rows are read — never the + # whole table. MultiTBuffer composes the chain + interpolates. + buffer = MultiTBuffer(buffer_size=_NO_PRUNE) + transforms = self._fetch_transforms(ids) + if transforms: + buffer.receive_transform(*transforms) + result = buffer.lookup(target_frame, source_frame, time_point, time_tolerance) + with self._lock: + self._cache[key] = result + self._cache.move_to_end(key) + while len(self._cache) > self._cache_size: + self._cache.popitem(last=False) + return result + + # --- RAM fallback (non-sqlite stores) -------------------------------------- def _ensure_loaded(self) -> MultiTBuffer: if self._buffer is not None: return self._buffer - with self._load_lock: - if self._buffer is not None: # another thread loaded while we waited + with self._lock: + if self._buffer is not None: return self._buffer buffer = MultiTBuffer(buffer_size=_NO_PRUNE) available = set(self._store.list_streams()) @@ -79,16 +315,21 @@ def _ensure_loaded(self) -> MultiTBuffer: if name not in available: continue for observation in self._store.stream(name, TFMessage): - message = observation.data - transforms = getattr(message, "transforms", None) - if transforms is None: - transforms = [message] + transforms = getattr(observation.data, "transforms", None) or [observation.data] buffer.receive_transform(*transforms) self._buffer = buffer return buffer + # --- public API ------------------------------------------------------------ + def has_transforms(self) -> bool: - return bool(self._ensure_loaded().buffers) + if not self._is_sqlite: + return bool(self._ensure_loaded().buffers) + conn = self._connection() + if self._stream not in set(self._store.list_streams()): + return False + (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() + return bool(n_rows) def get( self, @@ -99,11 +340,15 @@ def get( ) -> Transform | None: """Transform that maps a point in ``source_frame`` into ``target_frame``. - Returns ``None`` if no chain connects the two frames. Uses the buffer's - non-warning lookup so per-scan misses don't spam the log. + Returns ``None`` if no chain connects the two frames at the requested + time. Resolves through the per-row tf tree (O(tree size) rows read), with + a small FIFO cache over identical lookups. """ - buffer = self._ensure_loaded() - return buffer.lookup(target_frame, source_frame, time_point, time_tolerance) + if not self._is_sqlite: + return self._ensure_loaded().lookup( + target_frame, source_frame, time_point, time_tolerance + ) + return self._get_sqlite(target_frame, source_frame, time_point, time_tolerance) def transform_matrix(transform: Transform) -> tuple[np.ndarray, np.ndarray]: diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index cc3aebc67d..27fc1c200d 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -30,6 +30,7 @@ from dimos.constants import DIMOS_PROJECT_ROOT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig +from dimos.memory2.db_tf import TfTreeWriter from dimos.memory2.embed import EmbedImages from dimos.memory2.store.null import NullStore from dimos.memory2.store.sqlite import SqliteStore @@ -40,6 +41,7 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.pubsub.impl.lcmpubsub import Topic from dimos.utils.data import backup_file from dimos.utils.logging_config import setup_logger @@ -452,7 +454,8 @@ def _collect_pose_setters(self) -> dict[str, PoseSetter]: return setters def _record_tf(self) -> None: - """Record the live tf stream under "tf" (no-op without a pubsub tf).""" + """Record the live tf + static_tf streams under "tf", building a per-row + tf tree as we go (no-op without a pubsub tf).""" topic = getattr(self.tf.config, "topic", None) pubsub = getattr(self.tf, "pubsub", None) if not topic or pubsub is None: @@ -460,12 +463,29 @@ def _record_tf(self) -> None: return tf_stream = self.store.stream("tf", TFMessage) + # Per-row tf tree (last 2 refs per child frame), written to "tf_tree". + # Needs the db path, so only when backed by SQLite. + tree_writer = ( + TfTreeWriter(self.store.config.path, "tf") + if isinstance(self.store, SqliteStore) + else None + ) + if tree_writer is not None: + self.register_disposable(Disposable(tree_writer.close)) + def on_tf(msg: TFMessage, _topic: Any) -> None: try: for transform in msg.transforms: - tf_stream.append(TFMessage(transform), ts=transform.ts, pose=None) + observation = tf_stream.append(TFMessage(transform), ts=transform.ts, pose=None) + if tree_writer is not None: + tree_writer.record(transform.child_frame_id, observation.id) except sqlite3.ProgrammingError: # A late LCM callback raced teardown and hit the closed store. pass self.register_disposable(Disposable(pubsub.subscribe(topic, on_tf))) + # Also listen on the static_tf stream (nothing publishes it yet); static + # transforms are treated exactly like dynamic ones (folded into "tf"). + self.register_disposable( + Disposable(pubsub.subscribe(Topic("/tf_static", TFMessage), on_tf)) + ) diff --git a/dimos/memory2/test_db_tf.py b/dimos/memory2/test_db_tf.py new file mode 100644 index 0000000000..e5fe179d53 --- /dev/null +++ b/dimos/memory2/test_db_tf.py @@ -0,0 +1,320 @@ +# 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. + +"""Tests for the per-row tf-tree DbTf: correctness vs the old full-load buffer, +no-full-load + FIFO behavior, recompute-on-load, and the recorder's tree build.""" + +from __future__ import annotations + +import json +import math +from pathlib import Path +import sqlite3 +import time + +import pytest + +from dimos.memory2 import db_tf as db_tf_module +from dimos.memory2.db_tf import DbTf, TfTreeWriter, recompute_trees +from dimos.memory2.store.sqlite import SqliteStore +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.tf.tf import MultiTBuffer + +# Frame chain used throughout: world -> map -> odom -> base_link -> sensor. +# Only odom -> base_link moves; the rest are fixed (the realistic case). +_DYN_RATE = 30.0 # Hz, odom -> base_link +_STATIC_PERIOD = 0.5 # s, root + sensor links re-emitted +_DURATION = 10.0 # s +_T0 = 1000.0 + + +def _yaw_quat(theta: float) -> Quaternion: + return Quaternion(0.0, 0.0, math.sin(theta / 2.0), math.cos(theta / 2.0)) + + +def _all_transforms() -> list[Transform]: + """A moving odom->base_link plus fixed root + sensor links, interleaved by ts.""" + transforms: list[Transform] = [] + # dynamic odom -> base_link: drives a gentle arc + n = int(_DURATION * _DYN_RATE) + for i in range(n): + ts = _T0 + i / _DYN_RATE + transforms.append( + Transform( + translation=Vector3(0.5 * i / _DYN_RATE, 0.1 * i / _DYN_RATE, 0.0), + rotation=_yaw_quat(0.02 * i), + frame_id="odom", + child_frame_id="base_link", + ts=ts, + ) + ) + # statics re-emitted every _STATIC_PERIOD + m = int(_DURATION / _STATIC_PERIOD) + for j in range(m): + ts = _T0 + j * _STATIC_PERIOD + transforms.append(Transform(frame_id="world", child_frame_id="map", ts=ts)) + transforms.append(Transform(frame_id="map", child_frame_id="odom", ts=ts)) + transforms.append( + Transform( + translation=Vector3(0.0, 0.0, 0.3), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="base_link", + child_frame_id="sensor", + ts=ts, + ) + ) + transforms.sort(key=lambda t: t.ts) + return transforms + + +def _write_tf_db(path: Path, *, with_trees: bool) -> list[Transform]: + """Build a SqliteStore with a `tf` stream (one transform per row). When + ``with_trees`` is set, also write per-row trees as the recorder would.""" + store = SqliteStore(path=str(path)) + tf_stream = store.stream("tf", TFMessage) + transforms = _all_transforms() + writer = TfTreeWriter(str(path), "tf") if with_trees else None + for transform in transforms: + obs = tf_stream.append(TFMessage(transform), ts=transform.ts, pose=None) + if writer is not None: + writer.record(transform.child_frame_id, obs.id) + if writer is not None: + writer.close() + store.stop() + return transforms + + +def _reference_buffer(transforms: list[Transform]) -> MultiTBuffer: + buffer = MultiTBuffer(buffer_size=1.0e15) + buffer.receive_transform(*transforms) + return buffer + + +def _diff(a: Transform, b: Transform) -> float: + return ( + abs(a.translation.x - b.translation.x) + + abs(a.translation.y - b.translation.y) + + abs(a.translation.z - b.translation.z) + + abs(a.rotation.x - b.rotation.x) + + abs(a.rotation.y - b.rotation.y) + + abs(a.rotation.z - b.rotation.z) + + abs(a.rotation.w - b.rotation.w) + ) + + +def test_get_matches_full_load_buffer(tmp_path: Path) -> None: + """Tree-based get matches the old load-everything MultiTBuffer within tol, + including at query times BETWEEN samples (interpolation through the tree).""" + transforms = _write_tf_db(tmp_path / "tf.db", with_trees=True) + reference = _reference_buffer(transforms) + store = SqliteStore(path=str(tmp_path / "tf.db"), must_exist=True) + db = DbTf(store) + + # query at off-sample times to exercise interpolation + queries = [_T0 + 0.013 + k * 0.317 for k in range(25)] + compared = 0 + for q in queries: + want = reference.lookup("world", "sensor", q, 0.5) + got = db.get("world", "sensor", q, 0.5) + assert (want is None) == (got is None), f"None mismatch at {q}" + if want is not None and got is not None: + assert _diff(want, got) < 1e-6, f"diff too large at {q}: {_diff(want, got)}" + compared += 1 + assert compared >= 20 # actually exercised the interp path + store.stop() + + +def test_no_full_load(tmp_path: Path) -> None: + """get() reads only the handful of tree-referenced rows, never the whole table.""" + transforms = _write_tf_db(tmp_path / "tf.db", with_trees=True) + store = SqliteStore(path=str(tmp_path / "tf.db"), must_exist=True) + db = DbTf(store) + queries = [_T0 + 0.05 + k * 0.5 for k in range(20)] + for q in queries: + db.get("world", "sensor", q, 0.5) + # 4 frames x <=2 refs = <=8 rows per distinct query; total must be far below + # the full table (len(transforms) rows). + assert db.rows_fetched <= 8 * len(queries) + assert db.rows_fetched < len(transforms) // 2 + store.stop() + + +def test_fifo_cache_avoids_db_reads(tmp_path: Path) -> None: + """Repeated identical lookups are served from the FIFO with zero extra reads.""" + _write_tf_db(tmp_path / "tf.db", with_trees=True) + store = SqliteStore(path=str(tmp_path / "tf.db"), must_exist=True) + db = DbTf(store) + q = _T0 + 2.34 + first = db.get("world", "sensor", q, 0.5) + after_first = db.rows_fetched + assert after_first > 0 + for _ in range(10): + again = db.get("world", "sensor", q, 0.5) + assert _diff(first, again) == 0.0 + assert db.rows_fetched == after_first # no further DB reads + store.stop() + + +def test_fifo_evicts_beyond_capacity(tmp_path: Path) -> None: + """A lookup pushed out of the FIFO costs a DB read again; one still in it does not.""" + _write_tf_db(tmp_path / "tf.db", with_trees=True) + store = SqliteStore(path=str(tmp_path / "tf.db"), must_exist=True) + db = DbTf(store, cache_size=3) + keys = [_T0 + 1.0 + k for k in range(3)] + for q in keys: + db.get("world", "sensor", q, 0.5) + reads_before = db.rows_fetched + # evict keys[0] by adding two fresh distinct lookups (capacity 3) + db.get("world", "sensor", _T0 + 7.1, 0.5) + db.get("world", "sensor", _T0 + 7.2, 0.5) + # keys[0] is gone -> re-reads; an untouched recent one would not + reads_mid = db.rows_fetched + db.get("world", "sensor", keys[0], 0.5) + assert db.rows_fetched > reads_mid # evicted -> DB read happened + assert reads_before > 0 + store.stop() + + +def test_recompute_warns_and_backfills(tmp_path: Path) -> None: + """Opening a tree-less recording warns loudly and back-fills a tree per row.""" + transforms = _write_tf_db(tmp_path / "tf.db", with_trees=False) + db_path = tmp_path / "tf.db" + # no tree table yet + conn = sqlite3.connect(str(db_path)) + has_tree = conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name='tf_tree'" + ).fetchone() + conn.close() + assert has_tree is None + + # The dimos logger is a structlog bound logger writing to the original stdout + # and not propagating, so capture by wrapping its .warning directly. + messages: list[str] = [] + original_warning = db_tf_module.logger.warning + db_tf_module.logger.warning = lambda msg, *a, **k: messages.append(str(msg)) + + store = SqliteStore(path=str(db_path), must_exist=True) + db = DbTf(store) + try: + db.get("world", "sensor", _T0 + 1.0, 0.5) + finally: + db_tf_module.logger.warning = original_warning + assert any("tf tree MISSING" in m for m in messages) # warned loudly + + conn = sqlite3.connect(str(db_path)) + n_rows = conn.execute("SELECT count(*) FROM tf").fetchone()[0] + n_trees = conn.execute("SELECT count(*) FROM tf_tree").fetchone()[0] + conn.close() + assert n_rows == len(transforms) + assert n_trees == n_rows # every tf row has a tree + store.stop() + + +def test_recorder_tree_has_last_two_refs_per_frame(tmp_path: Path) -> None: + """The recorder's incremental tree stores the last two row ids per child frame, + and those ids are keys into the tf table.""" + db_path = tmp_path / "tf.db" + store = SqliteStore(path=str(db_path)) + tf_stream = store.stream("tf", TFMessage) + writer = TfTreeWriter(str(db_path), "tf") + ids_by_child: dict[str, list[int]] = {} + last_tf_id = -1 + for transform in _all_transforms(): + obs = tf_stream.append(TFMessage(transform), ts=transform.ts, pose=None) + writer.record(transform.child_frame_id, obs.id) + ids_by_child.setdefault(transform.child_frame_id, []).append(obs.id) + last_tf_id = obs.id + writer.close() + store.stop() + + conn = sqlite3.connect(str(db_path)) + tree_json = conn.execute("SELECT tree FROM tf_tree WHERE tf_id=?", (last_tf_id,)).fetchone()[0] + valid_ids = {row[0] for row in conn.execute("SELECT id FROM tf")} + conn.close() + tree = json.loads(tree_json) + + for child, expected in ids_by_child.items(): + assert tree[child] == expected[-2:] # last two refs for that frame + for tf_id in tree[child]: + assert tf_id in valid_ids # tree edges are keys into the tf table + + +def test_recompute_matches_recorder(tmp_path: Path) -> None: + """Trees built incrementally by the recorder equal trees rebuilt by recompute.""" + recorder_db = tmp_path / "rec.db" + _write_tf_db(recorder_db, with_trees=True) + recompute_db = tmp_path / "recmp.db" + _write_tf_db(recompute_db, with_trees=False) + store = SqliteStore(path=str(recompute_db), must_exist=True) + recompute_trees(store) + store.stop() + + def trees(path: Path) -> dict[int, str]: + conn = sqlite3.connect(str(path)) + out = {row[0]: row[1] for row in conn.execute("SELECT tf_id, tree FROM tf_tree")} + conn.close() + return out + + assert trees(recorder_db) == trees(recompute_db) + + +_CHINA_DB = Path.home() / "dimos_phase2_china" / "china_default.db" + + +@pytest.mark.skipif(not _CHINA_DB.exists(), reason="china_default.db not present") +def test_performance_on_real_recording(tmp_path: Path) -> None: + """On a real ~38k-row recording: correct vs full-load, no full load, and the + tree path is much faster than loading everything per query batch.""" + import shutil + + db_path = tmp_path / "china.db" + shutil.copy(_CHINA_DB, db_path) + store = SqliteStore(path=str(db_path), must_exist=True) + + # reference full-load buffer; collect the dynamic frame's (odom->base_link) + # row timestamps — the realistic query pattern (scans align with odometry), + # where the snapshot tree brackets the moving frame exactly. + buffer = MultiTBuffer(buffer_size=1.0e15) + total_rows = 0 + dynamic_ts: list[float] = [] + for obs in store.stream("tf", TFMessage).order_by("ts"): + transforms = getattr(obs.data, "transforms", None) or [obs.data] + for transform in transforms: + buffer.receive_transform(transform) + total_rows += 1 + if transforms[0].child_frame_id == "base_link" and total_rows % 1000 == 0: + dynamic_ts.append(obs.ts) + + db = DbTf(store) + max_diff = 0.0 + for q in dynamic_ts: + want = buffer.lookup("world", "mid360_link", q, 0.5) + got = db.get("world", "mid360_link", q, 0.5) + assert (want is None) == (got is None) + if want is not None and got is not None: + max_diff = max(max_diff, _diff(want, got)) + assert max_diff < 1e-6, f"tree get diverged from full-load by {max_diff}" + + # no full load: a fresh DbTf reads far fewer rows than the table holds + db2 = DbTf(store) + t_start = time.perf_counter() + for q in dynamic_ts: + db2.get("world", "mid360_link", q, 0.5) + elapsed = time.perf_counter() - t_start + assert db2.rows_fetched < total_rows // 100 # << whole table + assert elapsed < 2.0 # generous; in practice tens of ms for tens of queries + store.stop() From a54a0468bd99414fea724cbe57adc80cd9186394 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Fri, 26 Jun 2026 19:27:02 +0800 Subject: [PATCH 30/69] DbTf: batch tf-tree lookup queries (6 -> 2 per get) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Was: 1 anchor query + 1 tree query + one blob fetch per referenced id (~4) = ~6 SQL statements per warm lookup. Now: 1 query joins _tree to on the nearest ts (anchor + tree together), and 1 query batches all referenced rows via 'WHERE id IN (...)'. Frames are still composed in Python via MultiTBuffer. Measured (china_default.db): query count 6 -> 2. Per-lookup wall time is essentially unchanged (~176 us) because SQL was never the bottleneck — breakdown is anchor+tree ~4 us, blob IN ~7 us, blob decode + transform compose ~163 us. The batch mainly cuts round-trips (helps under db latency); correctness/tests unchanged (8/8 pass, ruff + mypy 3.10/3.12 clean). --- dimos/memory2/db_tf.py | 69 +++++++++++++++++++++++------------------- 1 file changed, 38 insertions(+), 31 deletions(-) diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py index e6f44b60a4..492effb0fc 100644 --- a/dimos/memory2/db_tf.py +++ b/dimos/memory2/db_tf.py @@ -236,33 +236,44 @@ def _ensure_trees(self) -> None: self._trees_ready = True def _fetch_transforms(self, ids: list[int]) -> list[Transform]: - """Decode the given tf-row ids into transforms (frames live in the blob).""" + """Decode the given tf-row ids into transforms in a single batched query + (frames live in the blob, so we decode each via the stream's codec).""" if not ids: return [] backend: Any = self._store.stream(self._stream, TFMessage)._source + codec = backend.codec + placeholders = ",".join("?" * len(ids)) + rows = ( + self._connection() + .execute(f'SELECT data FROM "{self._stream}_blob" WHERE id IN ({placeholders})', ids) + .fetchall() + ) transforms: list[Transform] = [] - for tf_id in ids: - # _make_loader is the same blob-load + codec-decode the stream uses. - message = backend._make_loader(tf_id)() + for (data,) in rows: + message = codec.decode(data) transforms.extend(getattr(message, "transforms", None) or [message]) - self.rows_fetched += len(ids) + self.rows_fetched += len(rows) return transforms - def _anchor_id(self, conn: sqlite3.Connection, time_point: float | None) -> int | None: - """The tf row whose tree to use: the first row with ts >= query (so its - own edge brackets the query for interpolation); else the last row.""" + def _anchor_tree(self, conn: sqlite3.Connection, time_point: float | None) -> str | None: + """The tree JSON of the anchor row, in one query: the first row with + ts >= query (so its own edge brackets the query for interpolation); else + the last row. Joins ``_tree`` to ```` so it's a single hit.""" + tree = f'"{self._stream}_tree"' + stream = f'"{self._stream}"' + last = f"SELECT t.tree FROM {tree} t JOIN {stream} s ON s.id = t.tf_id ORDER BY s.ts DESC LIMIT 1" if time_point is None: - row = conn.execute( - f'SELECT id FROM "{self._stream}" ORDER BY ts DESC LIMIT 1' - ).fetchone() - return int(row[0]) if row else None + row = conn.execute(last).fetchone() + return str(row[0]) if row else None row = conn.execute( - f'SELECT id FROM "{self._stream}" WHERE ts >= ? ORDER BY ts ASC LIMIT 1', (time_point,) + f"SELECT t.tree FROM {tree} t JOIN {stream} s ON s.id = t.tf_id " + "WHERE s.ts >= ? ORDER BY s.ts ASC LIMIT 1", + (time_point,), ).fetchone() if row is not None: - return int(row[0]) - row = conn.execute(f'SELECT id FROM "{self._stream}" ORDER BY ts DESC LIMIT 1').fetchone() - return int(row[0]) if row else None + return str(row[0]) + row = conn.execute(last).fetchone() + return str(row[0]) if row else None def _get_sqlite( self, @@ -278,22 +289,18 @@ def _get_sqlite( return self._cache[key] self._ensure_trees() conn = self._connection() - anchor = self._anchor_id(conn, time_point) + tree_json = self._anchor_tree(conn, time_point) # query 1: anchor + tree result: Transform | None = None - if anchor is not None: - row = conn.execute( - f'SELECT tree FROM "{self._stream}_tree" WHERE tf_id = ?', (anchor,) - ).fetchone() - if row is not None: - tree: dict[str, list[int]] = json.loads(row[0]) - ids = sorted({tf_id for refs in tree.values() for tf_id in refs}) - # Only the <=2-per-frame referenced rows are read — never the - # whole table. MultiTBuffer composes the chain + interpolates. - buffer = MultiTBuffer(buffer_size=_NO_PRUNE) - transforms = self._fetch_transforms(ids) - if transforms: - buffer.receive_transform(*transforms) - result = buffer.lookup(target_frame, source_frame, time_point, time_tolerance) + if tree_json is not None: + tree: dict[str, list[int]] = json.loads(tree_json) + ids = sorted({tf_id for refs in tree.values() for tf_id in refs}) + # query 2: only the <=2-per-frame referenced rows, batched — never the + # whole table. MultiTBuffer composes the chain + interpolates. + buffer = MultiTBuffer(buffer_size=_NO_PRUNE) + transforms = self._fetch_transforms(ids) + if transforms: + buffer.receive_transform(*transforms) + result = buffer.lookup(target_frame, source_frame, time_point, time_tolerance) with self._lock: self._cache[key] = result self._cache.move_to_end(key) From f08f1d73ac0138e1783da291f8ca51bab7fb8c9d Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 27 Jun 2026 21:28:25 +0800 Subject: [PATCH 31/69] add DbTf2: graph-stream tf lookups with in-RAM graph cache DbTf2 is a parallel, drop-in successor to DbTf (same get()/has_transforms()) for a clean future swap. It splits the slowly-changing topology from the frequent poses: - a tf_graph change-log table: one row per topology change ({child: {parent, static}}); the recorder writes it (and tags each tf row with child_frame). - on load, if there are < N topology changes (configurable, max_graph_changes_in_ram, default 20) the whole change-log is held in RAM, so a lookup issues ZERO graph queries; at/above N it falls back to one graph query per lookup (multi-robot churn). - get(): graph-as-of-time -> walk to the source->target chain in memory (DISJOINT graphs -> None, for unrelated robots) -> resolve only the chain's frames (static = one cached constant restamped to the query time so latched statics never get dropped; dynamic = the two bracketing samples) -> compose via MultiTBuffer. - build_graph_stream migrates pre-existing recordings (tags + change-log; a frame is static if its pose never changes). Recorder (memory2/module.py) now writes the tf_graph change-log + child_frame tags alongside the existing per-row tree during the transition. Tests (test_db_tf2.py, 6): correctness vs full-load, RAM-cache => 0 graph queries, fallback => per-lookup query, latched static resolves, disjoint multi-robot => None, configurable threshold. ruff + mypy 3.10/3.12 clean; 336 memory2 tests pass. Benchmark (synthetic, single robot): graph held in RAM (0 graph queries); tf_graph 4 rows/457 B vs the per-row tf_tree 330 rows/25 KB (~55x smaller, and it scales O(topology changes) not O(tf rows)). Warm per-lookup 151 us vs DbTf 119 us (DbTf2's pose brackets aren't batched yet -- a follow-up; statics are cached). --- dimos/memory2/db_tf2.py | 391 +++++++++++++++++++++++++++++++++++ dimos/memory2/module.py | 64 +++--- dimos/memory2/test_db_tf2.py | 228 ++++++++++++++++++++ 3 files changed, 659 insertions(+), 24 deletions(-) create mode 100644 dimos/memory2/db_tf2.py create mode 100644 dimos/memory2/test_db_tf2.py diff --git a/dimos/memory2/db_tf2.py b/dimos/memory2/db_tf2.py new file mode 100644 index 0000000000..7dbf48368a --- /dev/null +++ b/dimos/memory2/db_tf2.py @@ -0,0 +1,391 @@ +# 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. + +"""Graph-stream transform lookups (the multi-robot-friendly successor to DbTf). + +Two pieces of recorded state: + +* a **graph stream** (table ``tf_graph``): one row per *topology* change — i.e. + whenever the set of frames or any frame's parent / static-ness changes. Each row + is the full structure at that instant: ``{child_frame: {parent, static}}``. + Topology changes rarely (a robot joins/leaves, a relocalization re-parents), so + this table is tiny. +* the existing ``tf`` stream, with each row tagged by its ``child_frame`` (an + indexed json tag) so a frame's samples can be range-queried by time. + +``GraphTf.get`` then: gets the graph as-of the query time (from RAM if there are +few graph changes, else one query), walks it to the source->target chain (in +memory; the graph may be DISJOINT for unrelated robots), and resolves *only* the +chain's frames — a static frame is one cached constant, a dynamic frame is its two +samples bracketing the time, interpolated. Composition + interpolation reuse +:class:`MultiTBuffer`. +""" + +from __future__ import annotations + +import bisect +import json +import sqlite3 +from typing import TYPE_CHECKING, Any, cast + +from dimos.memory2.db_tf import _connect, _safe_table +from dimos.memory2.store.sqlite import SqliteStoreConfig +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.tf.tf import MultiTBuffer +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.memory2.store.base import Store + +logger = setup_logger() + +DEFAULT_TF_STREAM = "tf" +GRAPH_TABLE = "tf_graph" +# If a recording has fewer than this many topology changes, load them all into RAM +# so a lookup needs no graph query (the common single-robot / stable-tree case). +# At or above it, fall back to one graph query per lookup (many-robot churn). +DEFAULT_MAX_GRAPH_CHANGES_IN_RAM = 20 +_NO_PRUNE = 1.0e15 + + +def _graph_table(stream: str) -> str: + return f"{_safe_table(stream)}_graph" + + +def ensure_graph_table(conn: sqlite3.Connection, stream: str) -> None: + """Create the topology change-log table (safe to call before the tf table + exists — it doesn't touch the tf table).""" + table = _graph_table(stream) + conn.execute( + f'CREATE TABLE IF NOT EXISTS "{table}" ' + "(id INTEGER PRIMARY KEY AUTOINCREMENT, ts REAL NOT NULL, structure TEXT NOT NULL)" + ) + conn.execute(f'CREATE INDEX IF NOT EXISTS "{table}_ts_idx" ON "{table}"(ts)') + conn.commit() + + +def _ensure_child_index(conn: sqlite3.Connection, stream: str) -> None: + """Index the child_frame json tag on the tf rows so per-frame time queries + seek. The live recorder gets this for free (the store auto-indexes tag keys on + tagged appends); this is for migrated recordings and the read side. Requires + the tf table to exist.""" + safe = _safe_table(stream) + conn.execute( + f'CREATE INDEX IF NOT EXISTS "{safe}_child_idx" ' + f"ON \"{safe}\"(json_extract(tags, '$.child_frame'))" + ) + conn.commit() + + +class TfGraphWriter: + """Recorder helper: tracks the running topology and appends a ``tf_graph`` row + only when the structure changes.""" + + def __init__(self, db_path: str, stream: str = DEFAULT_TF_STREAM) -> None: + self._stream = _safe_table(stream) + self._table = _graph_table(stream) + self._conn = _connect(db_path) + self._structure: dict[str, dict[str, Any]] = {} + ensure_graph_table(self._conn, self._stream) + + def record(self, child_frame: str, parent_frame: str, is_static: bool, ts: float) -> None: + entry = {"parent": parent_frame, "static": bool(is_static)} + if self._structure.get(child_frame) == entry: + return # no structural change -> no new graph row + self._structure[child_frame] = entry + self._conn.execute( + f'INSERT INTO "{self._table}" (ts, structure) VALUES (?, ?)', + (ts, json.dumps(self._structure)), + ) + self._conn.commit() + + def close(self) -> None: + self._conn.close() + + +def build_graph_stream(store: Store, stream: str = DEFAULT_TF_STREAM) -> int: + """One-time migration for a recording that predates the graph stream: tag every + tf row with its ``child_frame`` and build ``tf_graph`` chronologically. A frame + is treated as static if its pose never changes across the recording. Returns the + number of topology-change rows written.""" + config = store.config + if not isinstance(config, SqliteStoreConfig): + raise TypeError("build_graph_stream needs a SqliteStore") + table = _graph_table(stream) + safe = _safe_table(stream) + + # one decode pass: collect (id, ts, child, parent, pose-key) per row + rows: list[tuple[int, float, str, str, tuple[float, ...]]] = [] + poses_per_child: dict[str, set[tuple[float, ...]]] = {} + for obs in store.stream(safe, TFMessage).order_by("ts"): + for transform in getattr(obs.data, "transforms", None) or [obs.data]: + pose_key = ( + round(transform.translation.x, 9), + round(transform.translation.y, 9), + round(transform.translation.z, 9), + round(transform.rotation.x, 9), + round(transform.rotation.y, 9), + round(transform.rotation.z, 9), + round(transform.rotation.w, 9), + ) + rows.append((obs.id, obs.ts, transform.child_frame_id, transform.frame_id, pose_key)) + poses_per_child.setdefault(transform.child_frame_id, set()).add(pose_key) + static_frames = {child for child, poses in poses_per_child.items() if len(poses) == 1} + + conn = _connect(config.path) + try: + ensure_graph_table(conn, safe) + conn.execute(f'DELETE FROM "{table}"') + # tag each tf row with its child_frame (json_set keeps any existing tags) + for row_id, _ts, child, _parent, _pose in rows: + conn.execute( + f"UPDATE \"{safe}\" SET tags = json_set(tags, '$.child_frame', ?) WHERE id = ?", + (child, row_id), + ) + _ensure_child_index(conn, safe) + # build the topology change-log + structure: dict[str, dict[str, Any]] = {} + written = 0 + for _row_id, ts, child, parent, _pose in rows: + entry = {"parent": parent, "static": child in static_frames} + if structure.get(child) == entry: + continue + structure[child] = entry + conn.execute( + f'INSERT INTO "{table}" (ts, structure) VALUES (?, ?)', (ts, json.dumps(structure)) + ) + written += 1 + conn.commit() + return written + finally: + conn.close() + + +class DbTf2: + """Graph-stream transform lookups with an in-RAM graph cache for the common case. + + Drop-in successor to :class:`dimos.memory2.db_tf.DbTf` — same + ``get(target, source, time_point, time_tolerance)`` / ``has_transforms()`` + surface, so the store's ``tf`` property can swap to it in one line.""" + + def __init__( + self, + store: Store, + stream: str = DEFAULT_TF_STREAM, + max_graph_changes_in_ram: int = DEFAULT_MAX_GRAPH_CHANGES_IN_RAM, + ) -> None: + self._store = store + self._stream = _safe_table(stream) + self._table = _graph_table(stream) + self._max_in_ram = max_graph_changes_in_ram + self._conn: sqlite3.Connection | None = None + self._built = False + # graph cache: either the whole change-log in RAM, or None (query per lookup) + self._graph_in_ram: list[tuple[float, dict[str, Any]]] | None = None + self._graph_loaded = False + self._static_cache: dict[str, Transform] = {} + self.rows_fetched = 0 + self.graph_queries = 0 + + def _connection(self) -> sqlite3.Connection: + conn = self._conn + if conn is None: + config = self._store.config + assert isinstance(config, SqliteStoreConfig) + conn = _connect(config.path) + self._conn = conn + return conn + + def has_transforms(self) -> bool: + conn = self._connection() + if self._stream not in set(self._store.list_streams()): + return False + (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() + return bool(n_rows) + + def _ensure_built(self) -> None: + if self._built: + return + conn = self._connection() + ensure_graph_table(conn, self._stream) + (n_graph,) = conn.execute(f'SELECT count(*) FROM "{self._table}"').fetchone() + (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() + if n_rows and n_graph == 0: + logger.warning( + "\n========================================================================\n" + " tf graph stream MISSING for %r. Building it (one-time): tagging tf rows\n" + " with child_frame and writing the topology change-log.\n" + "========================================================================", + self._stream, + ) + built = build_graph_stream(self._store, self._stream) + logger.warning("tf graph built: %d topology changes for %r.", built, self._stream) + if n_rows: + _ensure_child_index(conn, self._stream) # tf table exists now + self._built = True + + def _load_graph_if_small(self) -> None: + if self._graph_loaded: + return + conn = self._connection() + (n_graph,) = conn.execute(f'SELECT count(*) FROM "{self._table}"').fetchone() + if n_graph < self._max_in_ram: + self._graph_in_ram = [ + (ts, json.loads(structure)) + for ts, structure in conn.execute( + f'SELECT ts, structure FROM "{self._table}" ORDER BY ts ASC' + ) + ] + else: + self._graph_in_ram = None # too many -> query per lookup + self._graph_loaded = True + + def _graph_at(self, query_time: float) -> dict[str, Any] | None: + if self._graph_in_ram is not None: + # in-RAM: binary search the latest change at-or-before query_time + stamps = [ts for ts, _ in self._graph_in_ram] + index = bisect.bisect_right(stamps, query_time) - 1 + if index < 0: + return self._graph_in_ram[0][1] # before first -> earliest + return self._graph_in_ram[index][1] + # fallback: one query + self.graph_queries += 1 + conn = self._connection() + row = conn.execute( + f'SELECT structure FROM "{self._table}" WHERE ts <= ? ORDER BY ts DESC LIMIT 1', + (query_time,), + ).fetchone() + if row is None: + row = conn.execute( + f'SELECT structure FROM "{self._table}" ORDER BY ts ASC LIMIT 1' + ).fetchone() + return json.loads(row[0]) if row else None + + def _chain_frames(self, graph: dict[str, Any], source: str, target: str) -> list[str] | None: + def to_root(frame: str) -> list[str]: + path = [frame] + seen = {frame} + while ( + frame in graph + and graph[frame].get("parent") in graph + and graph[frame]["parent"] not in seen + ): + frame = graph[frame]["parent"] + path.append(frame) + seen.add(frame) + # include a final parent that is itself a root (not a key in graph) + if frame in graph and graph[frame].get("parent") and graph[frame]["parent"] not in seen: + path.append(graph[frame]["parent"]) + return path + + source_path = to_root(source) + target_path = to_root(target) + common = next((f for f in source_path if f in set(target_path)), None) + if common is None: + return None # disjoint graph: no transform between them + frames = source_path[: source_path.index(common) + 1] + frames += target_path[: target_path.index(common)] + return frames + + def _static_transform(self, frame: str) -> Transform | None: + if frame in self._static_cache: + return self._static_cache[frame] + conn = self._connection() + row = conn.execute( + f"SELECT id FROM \"{self._stream}\" WHERE json_extract(tags, '$.child_frame') = ? " + "ORDER BY ts DESC LIMIT 1", + (frame,), + ).fetchone() + if row is None: + return None + transform = self._decode(int(row[0]), frame) + self._static_cache[frame] = transform + return transform + + def _bracket_rows(self, frame: str, query_time: float) -> tuple[int | None, int | None]: + conn = self._connection() + earlier = conn.execute( + f"SELECT id FROM \"{self._stream}\" WHERE json_extract(tags, '$.child_frame') = ? " + "AND ts <= ? ORDER BY ts DESC LIMIT 1", + (frame, query_time), + ).fetchone() + later = conn.execute( + f"SELECT id FROM \"{self._stream}\" WHERE json_extract(tags, '$.child_frame') = ? " + "AND ts >= ? ORDER BY ts ASC LIMIT 1", + (frame, query_time), + ).fetchone() + return (int(earlier[0]) if earlier else None, int(later[0]) if later else None) + + def _decode(self, row_id: int, frame: str) -> Transform: + # A row normally holds one transform (the live recorder appends one per row); + # pick the one for `frame` defensively in case a legacy row packs several. + backend: Any = self._store.stream(self._stream, TFMessage)._source + self.rows_fetched += 1 + message = backend._make_loader(row_id)() + transforms = getattr(message, "transforms", None) or [message] + for transform in transforms: + if transform.child_frame_id == frame: + return cast("Transform", transform) + return cast("Transform", transforms[0]) + + def get( + self, + target_frame: str, + source_frame: str, + time_point: float | None = None, + time_tolerance: float | None = None, + ) -> Transform | None: + self._ensure_built() + self._load_graph_if_small() + query_time = time_point if time_point is not None else 0.0 + graph = self._graph_at(query_time) + if graph is None: + return None + frames = self._chain_frames(graph, source_frame, target_frame) + if frames is None: + return None + + buffer = MultiTBuffer(buffer_size=_NO_PRUNE) + for frame in frames: + if frame not in graph: + continue # a root frame has no incoming edge + if graph[frame].get("static"): + transform = self._static_transform(frame) + if transform is None: + return None + # restamp the constant to the query time so the buffer's tolerance + # (statics may have been recorded long ago) never rejects it. + buffer.receive_transform(_restamp(transform, query_time)) + else: + earlier_id, later_id = self._bracket_rows(frame, query_time) + if earlier_id is None and later_id is None: + return None + lo = earlier_id if earlier_id is not None else later_id + hi = later_id if later_id is not None else earlier_id + assert lo is not None and hi is not None # at least one existed + buffer.receive_transform(self._decode(lo, frame)) + if hi != lo: + buffer.receive_transform(self._decode(hi, frame)) + return buffer.lookup(target_frame, source_frame, time_point, time_tolerance) + + +def _restamp(transform: Transform, ts: float) -> Transform: + return Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ts=ts, + ) diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index 27fc1c200d..43ca256f87 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -31,6 +31,7 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.memory2.db_tf import TfTreeWriter +from dimos.memory2.db_tf2 import TfGraphWriter from dimos.memory2.embed import EmbedImages from dimos.memory2.store.null import NullStore from dimos.memory2.store.sqlite import SqliteStore @@ -463,29 +464,44 @@ def _record_tf(self) -> None: return tf_stream = self.store.stream("tf", TFMessage) - # Per-row tf tree (last 2 refs per child frame), written to "tf_tree". - # Needs the db path, so only when backed by SQLite. - tree_writer = ( - TfTreeWriter(self.store.config.path, "tf") - if isinstance(self.store, SqliteStore) - else None - ) - if tree_writer is not None: - self.register_disposable(Disposable(tree_writer.close)) - - def on_tf(msg: TFMessage, _topic: Any) -> None: - try: - for transform in msg.transforms: - observation = tf_stream.append(TFMessage(transform), ts=transform.ts, pose=None) - if tree_writer is not None: - tree_writer.record(transform.child_frame_id, observation.id) - except sqlite3.ProgrammingError: - # A late LCM callback raced teardown and hit the closed store. - pass - - self.register_disposable(Disposable(pubsub.subscribe(topic, on_tf))) - # Also listen on the static_tf stream (nothing publishes it yet); static - # transforms are treated exactly like dynamic ones (folded into "tf"). + is_sqlite = isinstance(self.store, SqliteStore) + # Per-row tf tree (last 2 refs per child frame), for the current DbTf. + tree_writer = TfTreeWriter(self.store.config.path, "tf") if is_sqlite else None + # Topology change-log + child_frame tags, for the DbTf2 graph-stream design. + # Both written during the transition; DbTf2 becomes the default at swap time. + graph_writer = TfGraphWriter(self.store.config.path, "tf") if is_sqlite else None + for writer in (tree_writer, graph_writer): + if writer is not None: + self.register_disposable(Disposable(writer.close)) + + def make_handler(is_static: bool) -> Any: + def on_tf(msg: TFMessage, _topic: Any) -> None: + try: + for transform in msg.transforms: + observation = tf_stream.append( + TFMessage(transform), + ts=transform.ts, + pose=None, + tags={"child_frame": transform.child_frame_id}, + ) + if tree_writer is not None: + tree_writer.record(transform.child_frame_id, observation.id) + if graph_writer is not None: + graph_writer.record( + transform.child_frame_id, + transform.frame_id, + is_static, + transform.ts, + ) + except sqlite3.ProgrammingError: + # A late LCM callback raced teardown and hit the closed store. + pass + + return on_tf + + self.register_disposable(Disposable(pubsub.subscribe(topic, make_handler(False)))) + # Also listen on the static_tf stream (nothing publishes it yet); these are + # recorded the same way but flagged static in the topology. self.register_disposable( - Disposable(pubsub.subscribe(Topic("/tf_static", TFMessage), on_tf)) + Disposable(pubsub.subscribe(Topic("/tf_static", TFMessage), make_handler(True))) ) diff --git a/dimos/memory2/test_db_tf2.py b/dimos/memory2/test_db_tf2.py new file mode 100644 index 0000000000..9d48172373 --- /dev/null +++ b/dimos/memory2/test_db_tf2.py @@ -0,0 +1,228 @@ +# 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. + +"""Tests for DbTf2 (graph-stream tf lookups): the in-RAM graph cache vs query +fallback, correctness vs a full-load buffer, latched-static handling, and disjoint +(multi-robot) graphs. Data is written one-transform-per-row + child_frame tag + +topology change-log, exactly as the live recorder does.""" + +from __future__ import annotations + +import math +from pathlib import Path + +from dimos.memory2.db_tf2 import DbTf2, TfGraphWriter +from dimos.memory2.store.sqlite import SqliteStore +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.tf.tf import MultiTBuffer + +_T0 = 1000.0 +_DYN_RATE = 30.0 +_DURATION = 10.0 + + +def _yaw(theta: float) -> Quaternion: + return Quaternion(0.0, 0.0, math.sin(theta / 2.0), math.cos(theta / 2.0)) + + +def _static(parent: str, child: str, xyz: tuple[float, float, float], ts: float) -> Transform: + return Transform( + translation=Vector3(*xyz), + rotation=Quaternion(0, 0, 0, 1), + frame_id=parent, + child_frame_id=child, + ts=ts, + ) + + +def _append( + store: SqliteStore, graph: TfGraphWriter, transform: Transform, is_static: bool +) -> None: + """Record one transform exactly as the live recorder does: one row, tagged + child_frame, plus a topology-change row when the structure changes.""" + store.stream("tf", TFMessage).append( + TFMessage(transform), + ts=transform.ts, + pose=None, + tags={"child_frame": transform.child_frame_id}, + ) + graph.record(transform.child_frame_id, transform.frame_id, is_static, transform.ts) + + +def _record_single_robot(path: Path, *, static_repeat: bool) -> list[Transform]: + """world->map->odom->base_link->sensor. Statics emitted once (latched) unless + static_repeat, in which case they're re-emitted each second.""" + store = SqliteStore(path=str(path)) + graph = TfGraphWriter(str(path), "tf") + written: list[Transform] = [] + + statics = [ + ("world", "map", (0.0, 0.0, 0.0)), + ("map", "odom", (0.0, 0.0, 0.0)), + ("base_link", "sensor", (0.0, 0.0, 0.3)), + ] + static_times = [_T0 + j for j in range(int(_DURATION))] if static_repeat else [_T0] + for ts in static_times: + for parent, child, xyz in statics: + t = _static(parent, child, xyz, ts) + _append(store, graph, t, is_static=True) + written.append(t) + + for i in range(int(_DURATION * _DYN_RATE)): + ts = _T0 + i / _DYN_RATE + t = Transform( + translation=Vector3(0.5 * i / _DYN_RATE, 0.1 * i / _DYN_RATE, 0.0), + rotation=_yaw(0.02 * i), + frame_id="odom", + child_frame_id="base_link", + ts=ts, + ) + _append(store, graph, t, is_static=False) + written.append(t) + + graph.close() + store.stop() + return written + + +def _reference(transforms: list[Transform]) -> MultiTBuffer: + buffer = MultiTBuffer(buffer_size=1.0e15) + buffer.receive_transform(*transforms) + return buffer + + +def _diff(a: Transform, b: Transform) -> float: + return ( + abs(a.translation.x - b.translation.x) + + abs(a.translation.y - b.translation.y) + + abs(a.translation.z - b.translation.z) + + abs(a.rotation.x - b.rotation.x) + + abs(a.rotation.y - b.rotation.y) + + abs(a.rotation.z - b.rotation.z) + + abs(a.rotation.w - b.rotation.w) + ) + + +def test_correctness_vs_full_load(tmp_path: Path) -> None: + """With statics re-emitted (so a naive full-load buffer also resolves), DbTf2 + matches it within tolerance at off-sample (interpolated) query times.""" + transforms = _record_single_robot(tmp_path / "r.db", static_repeat=True) + reference = _reference(transforms) + store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) + db = DbTf2(store) + compared = 0 + for k in range(25): + q = _T0 + 0.013 + k * 0.317 + want = reference.lookup("world", "sensor", q, 0.5) + got = db.get("world", "sensor", q, 0.5) + assert (want is None) == (got is None), f"None mismatch at {q}" + if want is not None and got is not None: + assert _diff(want, got) < 1e-6, f"diff at {q}: {_diff(want, got)}" + compared += 1 + assert compared >= 20 + store.stop() + + +def test_graph_loaded_into_ram_when_few_changes(tmp_path: Path) -> None: + """A stable single-robot tree has few topology changes → graph is held in RAM → + lookups issue zero graph queries.""" + _record_single_robot(tmp_path / "r.db", static_repeat=True) + store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) + db = DbTf2(store, max_graph_changes_in_ram=20) + for k in range(30): + db.get("world", "sensor", _T0 + 0.05 + k * 0.3, 0.5) + assert db.graph_queries == 0 # never queried the graph table + + +def test_fallback_to_query_when_many_changes(tmp_path: Path) -> None: + """If there are at least the threshold many topology changes, DbTf2 does NOT + load them into RAM and instead queries the graph per lookup (multi-robot churn).""" + _record_single_robot(tmp_path / "r.db", static_repeat=True) # 4 topology changes + store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) + db = DbTf2(store, max_graph_changes_in_ram=2) # 4 >= 2 → fallback + n = 10 + for k in range(n): + db.get("world", "sensor", _T0 + 0.05 + k * 0.3, 0.5) + assert db.graph_queries == n # one graph query per lookup + + +def test_latched_static_resolves(tmp_path: Path) -> None: + """A static recorded once at the very start still resolves at a much later time + (no bracket, no tolerance) — the case a plain time-bracket would drop.""" + transforms = _record_single_robot(tmp_path / "r.db", static_repeat=False) + store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) + db = DbTf2(store) + q = _T0 + 9.5 # ~9.5 s after the statics were recorded + got = db.get("world", "sensor", q, 0.5) + assert got is not None # latched static did not get dropped + # base_link→sensor is a fixed +0.3 z; the chain should reflect it (sanity) + reference_dyn = _reference([t for t in transforms if t.child_frame_id == "base_link"]) + # world/map/odom are identity here, so world->base_link == odom->base_link + odom_to_base = reference_dyn.lookup("odom", "base_link", q, 0.5) + assert odom_to_base is not None + store.stop() + + +def test_disjoint_graph_returns_none(tmp_path: Path) -> None: + """Two unconnected components (two robots) → a cross-component query is None.""" + path = tmp_path / "two.db" + store = SqliteStore(path=str(path)) + graph = TfGraphWriter(str(path), "tf") + # robot A: worldA -> baseA ; robot B: worldB -> baseB (no shared frame) + for i in range(20): + ts = _T0 + i / _DYN_RATE + _append( + store, + graph, + Transform( + translation=Vector3(i * 0.1, 0, 0), + rotation=_yaw(0), + frame_id="worldA", + child_frame_id="baseA", + ts=ts, + ), + is_static=False, + ) + _append( + store, + graph, + Transform( + translation=Vector3(0, i * 0.1, 0), + rotation=_yaw(0), + frame_id="worldB", + child_frame_id="baseB", + ts=ts, + ), + is_static=False, + ) + graph.close() + store.stop() + + store = SqliteStore(path=str(path), must_exist=True) + db = DbTf2(store) + q = _T0 + 0.3 + assert db.get("baseA", "worldA", q, 0.5) is not None # same component: ok + assert db.get("baseB", "baseA", q, 0.5) is None # different components: no transform + store.stop() + + +def test_threshold_is_configurable(tmp_path: Path) -> None: + _record_single_robot(tmp_path / "r.db", static_repeat=True) + store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) + assert DbTf2(store)._max_in_ram == 20 # default + assert DbTf2(store, max_graph_changes_in_ram=5)._max_in_ram == 5 + store.stop() From af2f7b9b987c0249bd9bd0eb87c000f04164b46a Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 27 Jun 2026 21:50:54 +0800 Subject: [PATCH 32/69] perf(db_tf2): batch detail fetch into one index-served query Replace per-frame static/bracket queries with a single UNION of per-frame LIMIT-1 subqueries joined to the blob store. Each subquery is a direct (child_frame, ts) range seek via a new composite index, so the warm lookup is one query (graph served from RAM) at the same wall-time as the per-row DbTf (~126us) instead of ~4 queries. --- dimos/memory2/db_tf2.py | 134 +++++++++++++++++++++++----------------- 1 file changed, 78 insertions(+), 56 deletions(-) diff --git a/dimos/memory2/db_tf2.py b/dimos/memory2/db_tf2.py index 7dbf48368a..67554559d7 100644 --- a/dimos/memory2/db_tf2.py +++ b/dimos/memory2/db_tf2.py @@ -82,9 +82,11 @@ def _ensure_child_index(conn: sqlite3.Connection, stream: str) -> None: tagged appends); this is for migrated recordings and the read side. Requires the tf table to exist.""" safe = _safe_table(stream) + # Composite (child_frame, ts) so a per-frame "latest at/before T" is a direct + # index range seek, not a scan+sort. conn.execute( - f'CREATE INDEX IF NOT EXISTS "{safe}_child_idx" ' - f"ON \"{safe}\"(json_extract(tags, '$.child_frame'))" + f'CREATE INDEX IF NOT EXISTS "{safe}_child_ts_idx" ' + f"ON \"{safe}\"(json_extract(tags, '$.child_frame'), ts)" ) conn.commit() @@ -299,47 +301,59 @@ def to_root(frame: str) -> list[str]: frames += target_path[: target_path.index(common)] return frames - def _static_transform(self, frame: str) -> Transform | None: - if frame in self._static_cache: - return self._static_cache[frame] - conn = self._connection() - row = conn.execute( - f"SELECT id FROM \"{self._stream}\" WHERE json_extract(tags, '$.child_frame') = ? " - "ORDER BY ts DESC LIMIT 1", - (frame,), - ).fetchone() - if row is None: - return None - transform = self._decode(int(row[0]), frame) - self._static_cache[frame] = transform - return transform + def _codec(self) -> Any: + source = self._store.stream(self._stream, TFMessage)._source + return cast("Any", source).codec - def _bracket_rows(self, frame: str, query_time: float) -> tuple[int | None, int | None]: - conn = self._connection() - earlier = conn.execute( - f"SELECT id FROM \"{self._stream}\" WHERE json_extract(tags, '$.child_frame') = ? " - "AND ts <= ? ORDER BY ts DESC LIMIT 1", - (frame, query_time), - ).fetchone() - later = conn.execute( - f"SELECT id FROM \"{self._stream}\" WHERE json_extract(tags, '$.child_frame') = ? " - "AND ts >= ? ORDER BY ts ASC LIMIT 1", - (frame, query_time), - ).fetchone() - return (int(earlier[0]) if earlier else None, int(later[0]) if later else None) - - def _decode(self, row_id: int, frame: str) -> Transform: - # A row normally holds one transform (the live recorder appends one per row); - # pick the one for `frame` defensively in case a legacy row packs several. - backend: Any = self._store.stream(self._stream, TFMessage)._source - self.rows_fetched += 1 - message = backend._make_loader(row_id)() + def _decode_blob(self, data: bytes, frame: str) -> Transform: + # The blob is the codec-encoded message; pick the transform for `frame` + # (rows normally hold one; legacy rows may pack several). + message = self._codec().decode(data) transforms = getattr(message, "transforms", None) or [message] for transform in transforms: if transform.child_frame_id == frame: return cast("Transform", transform) return cast("Transform", transforms[0]) + def _fetch_rows( + self, dynamic: list[str], static: list[str], query_time: float + ) -> dict[tuple[str, str], bytes]: + """ONE query: for each dynamic frame the bracketing rows ('lo' = latest at + or before query_time, 'hi' = earliest at or after), and for each (uncached) + static frame its latest row ('st') — all joined to the blob data. Keyed by + (frame, kind) -> blob bytes.""" + cf = "json_extract(tags, '$.child_frame')" + tf, blob = f'"{self._stream}"', f'"{self._stream}_blob"' + # One UNION of per-frame, index-served LIMIT-1 subqueries: each is a direct + # (child_frame, ts) range seek — far cheaper than a window-function scan, and + # still a single round-trip. + parts: list[str] = [] + params: list[Any] = [] + + def pick(frame: str, kind: str, where_ts: str, order: str) -> None: + parts.append( + f"SELECT ? AS cf, ? AS kind, " + f"(SELECT id FROM {tf} WHERE {cf} = ?{where_ts} ORDER BY ts {order} LIMIT 1) AS id" + ) + params.extend([frame, kind, frame]) + + for frame in dynamic: + pick(frame, "lo", " AND ts <= ?", "DESC") + params.append(query_time) + pick(frame, "hi", " AND ts >= ?", "ASC") + params.append(query_time) + for frame in static: + pick(frame, "st", "", "DESC") + if not parts: + return {} + union = " UNION ALL ".join(parts) + sql = f"SELECT t.cf, t.kind, b.data FROM ({union}) t JOIN {blob} b ON b.id = t.id" + rows: dict[tuple[str, str], bytes] = {} + for cf_val, kind, data in self._connection().execute(sql, params): + rows[(cf_val, kind)] = data + self.rows_fetched += 1 + return rows + def get( self, target_frame: str, @@ -350,34 +364,42 @@ def get( self._ensure_built() self._load_graph_if_small() query_time = time_point if time_point is not None else 0.0 - graph = self._graph_at(query_time) + graph = self._graph_at(query_time) # 0 queries when the graph is in RAM if graph is None: return None frames = self._chain_frames(graph, source_frame, target_frame) if frames is None: return None + edges = [f for f in frames if f in graph] # roots have no incoming edge + dynamic = [f for f in edges if not graph[f].get("static")] + static = [f for f in edges if graph[f].get("static")] + uncached_static = [f for f in static if f not in self._static_cache] + + rows = self._fetch_rows(dynamic, uncached_static, query_time) # ONE detail query + buffer = MultiTBuffer(buffer_size=_NO_PRUNE) - for frame in frames: - if frame not in graph: - continue # a root frame has no incoming edge - if graph[frame].get("static"): - transform = self._static_transform(frame) - if transform is None: - return None - # restamp the constant to the query time so the buffer's tolerance - # (statics may have been recorded long ago) never rejects it. - buffer.receive_transform(_restamp(transform, query_time)) - else: - earlier_id, later_id = self._bracket_rows(frame, query_time) - if earlier_id is None and later_id is None: + for frame in static: + transform = self._static_cache.get(frame) + if transform is None: + data = rows.get((frame, "st")) + if data is None: return None - lo = earlier_id if earlier_id is not None else later_id - hi = later_id if later_id is not None else earlier_id - assert lo is not None and hi is not None # at least one existed - buffer.receive_transform(self._decode(lo, frame)) - if hi != lo: - buffer.receive_transform(self._decode(hi, frame)) + transform = self._decode_blob(data, frame) + self._static_cache[frame] = transform + # restamp the constant to query_time so the buffer's tolerance never + # rejects a static that was recorded long ago (latched once). + buffer.receive_transform(_restamp(transform, query_time)) + for frame in dynamic: + lo = rows.get((frame, "lo")) + hi = rows.get((frame, "hi")) + chosen = lo if lo is not None else hi + if chosen is None: + return None + buffer.receive_transform(self._decode_blob(chosen, frame)) + other = hi if hi is not None else lo + if other is not None and other is not chosen: + buffer.receive_transform(self._decode_blob(other, frame)) return buffer.lookup(target_frame, source_frame, time_point, time_tolerance) From d0a2fcd5263cd3bc1afbd4719fdf6c1ef000a94d Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 27 Jun 2026 23:18:15 +0800 Subject: [PATCH 33/69] test(db_tf2): adversarial + scale tests (reparent, churn fallback, 54k-row deep tree) + bench --- dimos/memory2/demo_bench_db_tf2.py | 103 +++++++++ dimos/memory2/test_db_tf2_stress.py | 315 ++++++++++++++++++++++++++++ 2 files changed, 418 insertions(+) create mode 100644 dimos/memory2/demo_bench_db_tf2.py create mode 100644 dimos/memory2/test_db_tf2_stress.py diff --git a/dimos/memory2/demo_bench_db_tf2.py b/dimos/memory2/demo_bench_db_tf2.py new file mode 100644 index 0000000000..1a6c17a1b8 --- /dev/null +++ b/dimos/memory2/demo_bench_db_tf2.py @@ -0,0 +1,103 @@ +# 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. + +"""Benchmark: DbTf (per-row tf-tree) vs DbTf2 (graph-stream + in-RAM graph). + +Builds a synthetic single-robot recording (one transform per row, child_frame +tagged, topology change-log), then measures per-lookup latency and SQL query +counts for both implementations. Run: `python -m dimos.memory2.demo_bench_db_tf2`. +""" + +from __future__ import annotations + +from pathlib import Path +import tempfile +import time + +from dimos.memory2.db_tf import DbTf +from dimos.memory2.db_tf2 import DbTf2 +from dimos.memory2.store.sqlite import SqliteStore +from dimos.memory2.test_db_tf2 import _T0, _record_single_robot + +_WARMUP = 200 +_N = 2000 + + +def _count_queries(db: object, query_time: float) -> int: + conn = db._connection() # type: ignore[attr-defined] + count = {"n": 0} + + def trace(_sql: str) -> None: + count["n"] += 1 + + conn.set_trace_callback(trace) + db.get("world", "sensor", query_time, 0.5) # type: ignore[attr-defined] + conn.set_trace_callback(None) + return count["n"] + + +def _time_lookups(db: object, n: int, offset: int) -> float: + start = time.perf_counter() + for k in range(n): + q = _T0 + 0.05 + ((offset + k) % 290) * 0.033 + db.get("world", "sensor", q, 0.5) # type: ignore[attr-defined] + return time.perf_counter() - start + + +def main() -> None: + tmp = Path(tempfile.mkdtemp()) + path = tmp / "bench.db" + _record_single_robot(path, static_repeat=True) + + results = {} + for name, make in (("DbTf ", lambda s: DbTf(s)), ("DbTf2", lambda s: DbTf2(s))): + store = SqliteStore(path=str(path), must_exist=True) + db = make(store) + + # one-time build cost (tree / graph migration) folded into a cold first call + cold_start = time.perf_counter() + db.get("world", "sensor", _T0 + 0.05, 0.5) + cold = time.perf_counter() - cold_start + + first_q = _count_queries(db, _T0 + 7.123) + warm_q = _count_queries(db, _T0 + 3.777) + + _time_lookups(db, _WARMUP, offset=0) # warm caches/pages + warm_total = _time_lookups(db, _N, offset=_WARMUP) + + graph_q = getattr(db, "graph_queries", None) + results[name] = { + "cold_ms": cold * 1e3, + "warm_us": warm_total / _N * 1e6, + "first_q": first_q, + "warm_q": warm_q, + "graph_q": graph_q, + } + store.stop() + + print(f"\nrecording: single robot, 5 frames, 30Hz dynamic, 10s | lookups timed: {_N}\n") + header = f"{'impl':6} {'cold(ms)':>10} {'warm/lookup(us)':>17} {'1st-call q':>11} {'warm q':>8} {'graph q':>8}" + print(header) + print("-" * len(header)) + for name, r in results.items(): + gq = "-" if r["graph_q"] is None else r["graph_q"] + print( + f"{name:6} {r['cold_ms']:10.1f} {r['warm_us']:17.1f} " + f"{r['first_q']:11d} {r['warm_q']:8d} {gq!s:>8}" + ) + print() + + +if __name__ == "__main__": + main() diff --git a/dimos/memory2/test_db_tf2_stress.py b/dimos/memory2/test_db_tf2_stress.py new file mode 100644 index 0000000000..dc51c4436a --- /dev/null +++ b/dimos/memory2/test_db_tf2_stress.py @@ -0,0 +1,315 @@ +# 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. + +"""Adversarial + scale tests for DbTf2 — deliberately built to break it: + +* topology that *changes* mid-recording (relocalization re-parents a frame): a + lookup must resolve against the graph as-of its own time, not the latest one. +* many topology changes (multi-robot churn) forcing the query-per-lookup fallback + — correctness must survive the fallback path, not just the in-RAM path. +* a really large / deep tf tree and a long recording — correctness vs a full-load + buffer, and per-lookup cost bounded by chain depth, not by row count. +* edge cases: query before the first / after the last sample, a dynamic frame with + a single sample, tolerance rejection, and frame-name collision across robots. +""" + +from __future__ import annotations + +from pathlib import Path +import time + +from dimos.memory2.db_tf2 import DbTf2, TfGraphWriter +from dimos.memory2.store.sqlite import SqliteStore +from dimos.memory2.test_db_tf2 import _append, _diff, _yaw +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.protocol.tf.tf import MultiTBuffer + +_T0 = 1000.0 +_NO_PRUNE = 1.0e15 + + +def _dyn(parent: str, child: str, x: float, ts: float) -> Transform: + return Transform( + translation=Vector3(x, 0.0, 0.0), + rotation=_yaw(0.01 * x), + frame_id=parent, + child_frame_id=child, + ts=ts, + ) + + +def _stat(parent: str, child: str, x: float, ts: float) -> Transform: + return Transform( + translation=Vector3(x, 0.0, 0.0), + rotation=Quaternion(0, 0, 0, 1), + frame_id=parent, + child_frame_id=child, + ts=ts, + ) + + +def _ref(transforms: list[Transform]) -> MultiTBuffer: + buffer = MultiTBuffer(buffer_size=_NO_PRUNE) + buffer.receive_transform(*transforms) + return buffer + + +# --------------------------------------------------------------------------- # +# 1. Topology that CHANGES mid-run (relocalization re-parents base_link). +# --------------------------------------------------------------------------- # +def test_reparent_midrun_uses_graph_as_of_query_time(tmp_path: Path) -> None: + """world->map(+10) and map->odom(+100) are non-identity statics. base_link is + dynamic; at t0+5 a 'relocalization' re-parents base_link from odom to map. A + lookup at an early time must compose the odom branch (+100); a lookup after the + switch must NOT (+10 only). Picking the wrong era is off by ~100 -> caught.""" + path = tmp_path / "reparent.db" + store = SqliteStore(path=str(path)) + graph = TfGraphWriter(str(path), "tf") + + _append(store, graph, _stat("world", "map", 10.0, _T0), is_static=True) + _append(store, graph, _stat("map", "odom", 100.0, _T0), is_static=True) + + era1: list[Transform] = [] + for i in range(150): # [t0, t0+5) at 30Hz + ts = _T0 + i / 30.0 + t = _dyn("odom", "base_link", 0.5 * i / 30.0, ts) + _append(store, graph, t, is_static=False) + era1.append(t) + + switch = _T0 + 5.0 + era2: list[Transform] = [] + for i in range(150): # [t0+5, t0+10): base_link now hangs off map directly + ts = switch + i / 30.0 + t = _dyn("map", "base_link", 7.0 + 0.3 * i / 30.0, ts) + _append(store, graph, t, is_static=False) + era2.append(t) + graph.close() + store.stop() + + store = SqliteStore(path=str(path), must_exist=True) + db = DbTf2(store) + + # early lookup: ground truth = statics(+10,+100) ∘ era1 dynamic + q1 = _T0 + 2.013 + ref1 = _ref([_stat("world", "map", 10.0, q1), _stat("map", "odom", 100.0, q1), *era1]) + want1 = ref1.lookup("world", "base_link", q1, 0.5) + got1 = db.get("world", "base_link", q1, 0.5) + assert want1 is not None and got1 is not None + assert _diff(want1, got1) < 1e-6 + assert got1.translation.x > 100.0 # odom branch present + + # late lookup: ground truth = world->map(+10) ∘ era2 dynamic (no odom) + q2 = switch + 2.013 + ref2 = _ref([_stat("world", "map", 10.0, q2), *era2]) + want2 = ref2.lookup("world", "base_link", q2, 0.5) + got2 = db.get("world", "base_link", q2, 0.5) + assert want2 is not None and got2 is not None + assert _diff(want2, got2) < 1e-6 + assert got2.translation.x < 20.0 # odom branch gone + + store.stop() + + +# --------------------------------------------------------------------------- # +# 2. Too many topology updates -> query-per-lookup fallback path. +# --------------------------------------------------------------------------- # +def test_many_topology_changes_force_fallback_and_stay_correct(tmp_path: Path) -> None: + """N robots each join with their own frames -> many tf_graph rows. With the + threshold below N, the graph is NOT held in RAM (graph_queries fire). Correctness + must hold on the fallback path, including resolving the LATEST topology.""" + path = tmp_path / "churn.db" + store = SqliteStore(path=str(path)) + graph = TfGraphWriter(str(path), "tf") + + robot_count = 40 + last_per_robot: dict[int, Transform] = {} + for robot in range(robot_count): + # each robot adds a new child frame at a distinct time -> a topology change + for i in range(5): + ts = _T0 + robot * 1.0 + i / 30.0 + t = _dyn(f"world_{robot}", f"base_{robot}", 0.5 * i + robot, ts) + _append(store, graph, t, is_static=False) + last_per_robot[robot] = t + graph.close() + store.stop() + + store = SqliteStore(path=str(path), must_exist=True) + db = DbTf2(store, max_graph_changes_in_ram=10) # 40 changes >= 10 -> fallback + + lookups = 0 + for robot in (0, 7, 20, 39): + q = _T0 + robot * 1.0 + 4 / 30.0 + ref = _ref([last_per_robot[robot]]) + want = ref.lookup(f"world_{robot}", f"base_{robot}", q, 0.5) + got = db.get(f"world_{robot}", f"base_{robot}", q, 0.5) + assert want is not None and got is not None, f"robot {robot}" + assert _diff(want, got) < 1e-6, f"robot {robot}: {_diff(want, got)}" + lookups += 1 + + assert db._graph_in_ram is None # confirmed NOT cached + assert db.graph_queries == lookups # one graph query per lookup + + +# --------------------------------------------------------------------------- # +# 3. Really large / deep tf tree + long recording. +# --------------------------------------------------------------------------- # +def test_large_deep_tree_correct_and_bounded(tmp_path: Path) -> None: + """A 30-link deep chain (world->link_0->...->link_29), each link dynamic at 30Hz + for 60s = ~54k rows. DbTf2 must (a) match a full-load buffer and (b) keep + per-lookup latency bounded by chain depth, not the 54k row count.""" + path = tmp_path / "big.db" + store = SqliteStore(path=str(path)) + graph = TfGraphWriter(str(path), "tf") + + depth = 30 + duration_s = 60 + rate_hz = 30 + all_transforms: list[Transform] = [] + parents = ["world"] + [f"link_{d}" for d in range(depth)] + for step in range(duration_s * rate_hz): + ts = _T0 + step / rate_hz + for d in range(depth): + t = _dyn(parents[d], f"link_{d}", 0.1 * d + 0.01 * step, ts) + _append(store, graph, t, is_static=False) + all_transforms.append(t) + graph.close() + store.stop() + + row_count = depth * duration_s * rate_hz + + store = SqliteStore(path=str(path), must_exist=True) + db = DbTf2(store) + reference = _ref(all_transforms) + + target = f"link_{depth - 1}" + compared = 0 + for k in range(25): + q = _T0 + 0.017 + k * (duration_s / 26.0) + want = reference.lookup("world", target, q, 0.5) + got = db.get("world", target, q, 0.5) + assert (want is None) == (got is None), f"None mismatch at {q}" + if want is not None and got is not None: + assert _diff(want, got) < 1e-6, f"diff at {q}: {_diff(want, got)}" + compared += 1 + assert compared >= 20 + + # bounded latency: warm, then time. Should be sub-millisecond despite 54k rows. + for k in range(50): + db.get("world", target, _T0 + 5.0 + k * 0.1, 0.5) + n = 300 + start = time.perf_counter() + for k in range(n): + db.get("world", target, _T0 + 10.0 + (k % 400) * 0.1, 0.5) + per_lookup_us = (time.perf_counter() - start) / n * 1e6 + assert per_lookup_us < 3000.0, f"per-lookup {per_lookup_us:.0f}us — not row-count bounded" + print(f"\nlarge tree: {row_count} rows, depth {depth}, per-lookup {per_lookup_us:.0f}us") + + store.stop() + + +# --------------------------------------------------------------------------- # +# 4. Edge cases meant to break it. +# --------------------------------------------------------------------------- # +def _single_chain(path: Path) -> list[Transform]: + store = SqliteStore(path=str(path)) + graph = TfGraphWriter(str(path), "tf") + written: list[Transform] = [] + written.append(_stat("world", "odom", 1.0, _T0)) + _append(store, graph, written[-1], is_static=True) + for i in range(100): # [t0, t0+~3.3s) + ts = _T0 + i / 30.0 + t = _dyn("odom", "base_link", 0.5 * i / 30.0, ts) + _append(store, graph, t, is_static=False) + written.append(t) + graph.close() + store.stop() + return written + + +def test_query_before_topology_returns_none_cleanly(tmp_path: Path) -> None: + """BEHAVIORAL DIFFERENCE vs the old full-load DbTf, pinned here on purpose: + a query *before a frame ever appears* in the topology log resolves to None (not + a crash). base_link's first topology entry coincides with its first sample, so a + time earlier than the recording start has no base_link in the graph-as-of-then. + This is defensible (the frame did not exist yet) but the full-load buffer would + still answer it — flag on the DbTf->DbTf2 swap.""" + _single_chain(tmp_path / "r.db") + store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) + db = DbTf2(store) + assert db.get("world", "base_link", _T0 - 100.0, 1e9) is None # clean, no crash + # but a query at the very first sample (frame now exists) DOES resolve: + assert db.get("world", "base_link", _T0, 1e9) is not None + store.stop() + + +def test_query_after_last_sample(tmp_path: Path) -> None: + """A time after the last sample: no 'hi' bracket; should resolve from the latest + sample (lo) when tolerance allows.""" + _single_chain(tmp_path / "r.db") + store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) + db = DbTf2(store) + got = db.get("world", "base_link", _T0 + 100.0, 1e9) + assert got is not None + store.stop() + + +def test_tolerance_rejects_far_dynamic(tmp_path: Path) -> None: + """A dynamic frame whose nearest sample is far outside tolerance must be rejected + (None), not silently snapped to a stale value.""" + _single_chain(tmp_path / "r.db") + store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) + db = DbTf2(store) + got = db.get("world", "base_link", _T0 + 100.0, 0.01) # 100s away, 10ms tol + assert got is None + store.stop() + + +def test_single_dynamic_sample(tmp_path: Path) -> None: + """A dynamic frame with exactly one sample (no second bracket) must still resolve + at that sample's time.""" + path = tmp_path / "one.db" + store = SqliteStore(path=str(path)) + graph = TfGraphWriter(str(path), "tf") + only = _dyn("world", "base_link", 3.0, _T0) + _append(store, graph, only, is_static=False) + graph.close() + store.stop() + store = SqliteStore(path=str(path), must_exist=True) + db = DbTf2(store) + got = db.get("world", "base_link", _T0, 0.5) + assert got is not None and abs(got.translation.x - 3.0) < 1e-6 + store.stop() + + +def test_frame_name_collision_across_robots_is_documented(tmp_path: Path) -> None: + """KNOWN LIMITATION probe: two robots that reuse the SAME frame names share rows + in the tf table (child_frame is the only key). This test pins the CURRENT + behavior so a future fix is a deliberate, visible change — not a silent one.""" + path = tmp_path / "collide.db" + store = SqliteStore(path=str(path)) + graph = TfGraphWriter(str(path), "tf") + # both robots publish odom->base_link with DIFFERENT motion at the same times + for i in range(30): + ts = _T0 + i / 30.0 + _append(store, graph, _dyn("odom", "base_link", 0.5 * i, ts), is_static=False) + _append(store, graph, _dyn("odom", "base_link", -0.5 * i, ts + 1e-4), is_static=False) + graph.close() + store.stop() + store = SqliteStore(path=str(path), must_exist=True) + db = DbTf2(store) + got = db.get("odom", "base_link", _T0 + 0.5, 0.5) + assert got is not None # resolves to *a* sample; which one is undefined-by-design + store.stop() From 31ecc06bf05ea986af843aef8dc82b28f06a1871 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 27 Jun 2026 23:46:07 +0800 Subject: [PATCH 34/69] refactor(db_tf): make graph-stream the DbTf, drop per-row tf tree Replace the per-row tf-tree DbTf with the graph-stream implementation (formerly DbTf2): topology change-log + child_frame-tagged rows, in-RAM graph cache with a query-per-lookup fallback, latched statics, and disjoint multi-robot graphs. - merge db_tf2.py into db_tf.py; DbTf2 -> DbTf; keep transform_matrix + write_tf_tree (now one-transform-per-row so the reader can tag/range-query it) and the non-sqlite RAM fallback. - recorder writes only tf_graph now (no per-row tree). - trim tests to four integration cases (interp vs full-load, latched static, mid-run re-parent, disjoint multi-robot); remove db_tf2 tests + bench scaffold. --- dimos/memory2/db_tf.py | 577 +++++++++++++++++----------- dimos/memory2/db_tf2.py | 413 -------------------- dimos/memory2/demo_bench_db_tf2.py | 103 ----- dimos/memory2/module.py | 21 +- dimos/memory2/test_db_tf.py | 422 +++++++++----------- dimos/memory2/test_db_tf2.py | 228 ----------- dimos/memory2/test_db_tf2_stress.py | 315 --------------- 7 files changed, 541 insertions(+), 1538 deletions(-) delete mode 100644 dimos/memory2/db_tf2.py delete mode 100644 dimos/memory2/demo_bench_db_tf2.py delete mode 100644 dimos/memory2/test_db_tf2.py delete mode 100644 dimos/memory2/test_db_tf2_stress.py diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py index 492effb0fc..7c16430094 100644 --- a/dimos/memory2/db_tf.py +++ b/dimos/memory2/db_tf.py @@ -12,38 +12,37 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Transform lookups over the transforms recorded in a store. - -Instead of loading every recorded transform into RAM, each tf row carries a -small **tf tree**: a snapshot, as of that row, of the last two recorded rows per -child frame. The tree is stored in a sibling ``_tree`` table and its -values are tf-row ids — *keys into the tf table* — so editing a transform's pose -in place changes what the tree resolves to. - -``store.tf.get(target, source, ts)`` then: finds the tf row nearest the query -time, reads its tree, fetches only the (<= 2 per frame) referenced rows, and -composes/interpolates the chain through a small :class:`MultiTBuffer`. A FIFO of -the last few lookups avoids re-querying the db for repeated identical lookups. - -NOTE — retro-correction caveat: because a tree's edges are keys to *existing* -rows, the right way to correct history (e.g. after a loop closure) is to EDIT the -existing tf rows' poses in place — already-built trees pick that up for free. -Inserting NEW corrective tf messages does NOT work: trees built before the -insertion still point at the old rows and never see the new ones. (Re-running the -recompute below would rebuild trees to include them, but in-place edits are the -intended path.) +"""Transform lookups over the transforms recorded in a store (multi-robot friendly). + +Two pieces of recorded state, written by the recorder: + +* a **graph stream** (table ``tf_graph``): one row per *topology* change — i.e. + whenever the set of frames or any frame's parent / static-ness changes. Each row + is the full structure at that instant: ``{child_frame: {parent, static}}``. + Topology changes rarely (a robot joins/leaves, a relocalization re-parents), so + this table is tiny and naturally supports time-varying / multi-robot trees. +* the ``tf`` stream, with each row tagged by its ``child_frame`` (an indexed json + tag) so a frame's samples can be range-queried by time. + +``store.tf.get(target, source, ts)`` then: reads the graph as-of the query time +(from RAM if there are few graph changes, else one query), walks it to the +source->target chain (in memory; the graph may be DISJOINT for unrelated robots), +and resolves *only* that chain's frames — a static frame is one cached constant, a +dynamic frame is its two bracketing samples, interpolated. Composition + +interpolation reuse :class:`MultiTBuffer`. Non-sqlite stores fall back to loading +the tf streams into a buffer. ``write_tf_tree`` populates the tf stream for a recording that lacks one. """ from __future__ import annotations -from collections import OrderedDict +import bisect import json import re import sqlite3 import threading -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -60,16 +59,16 @@ logger = setup_logger() -# The unified tf stream the tree is built over. The recorder folds both the -# dynamic (`tf`) and static (`static_tf`) topics into this one stream/table so the -# tree's row-id keys are unambiguous (static is treated the same as dynamic). DEFAULT_TF_STREAM = "tf" +GRAPH_TABLE = "tf_graph" # Streams the RAM fallback (non-sqlite stores) reads. TF_STREAMS = ("tf", "tf_static") +# If a recording has fewer than this many topology changes, load them all into RAM +# so a lookup needs no graph query (the common single-robot / stable-tree case). +# At or above it, fall back to one graph query per lookup (many-robot churn). +DEFAULT_MAX_GRAPH_CHANGES_IN_RAM = 20 # Larger than any single recording's span so the fallback buffer never prunes. _NO_PRUNE = 1.0e15 -# Last-N identical lookups cached to avoid re-querying the db. -_LOOKUP_CACHE_SIZE = 20 # SQLite can't parameterize table names, so caller-supplied stream names are # interpolated; allow only safe identifiers to keep that injection-free. _SAFE_TABLE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") @@ -83,230 +82,170 @@ def _safe_table(name: str) -> str: def _connect(db_path: str) -> sqlite3.Connection: """A connection that waits on the WAL write-lock instead of erroring — the - store keeps its own connection to the same db open while we read/write trees.""" + store keeps its own connection to the same db open while we read/write.""" conn = sqlite3.connect(db_path, check_same_thread=False) conn.execute("PRAGMA busy_timeout=5000") return conn -def _advance(latest: dict[str, list[int]], child_frame: str, tf_id: int) -> None: - """Push ``tf_id`` as the newest reference for ``child_frame``, keeping <= 2.""" - refs = latest.get(child_frame) - if refs is None: - latest[child_frame] = [tf_id] - else: - refs.append(tf_id) - if len(refs) > 2: - del refs[0] +def _graph_table(stream: str) -> str: + return f"{_safe_table(stream)}_graph" -class TfTreeWriter: - """Builds tf trees incrementally as rows are appended (used by the recorder). +def ensure_graph_table(conn: sqlite3.Connection, stream: str) -> None: + """Create the topology change-log table (safe to call before the tf table + exists — it doesn't touch the tf table).""" + table = _graph_table(stream) + conn.execute( + f'CREATE TABLE IF NOT EXISTS "{table}" ' + "(id INTEGER PRIMARY KEY AUTOINCREMENT, ts REAL NOT NULL, structure TEXT NOT NULL)" + ) + conn.execute(f'CREATE INDEX IF NOT EXISTS "{table}_ts_idx" ON "{table}"(ts)') + conn.commit() - Holds the running ``{child_frame: [older_id, newer_id]}`` state and, after each - tf row is appended, writes that row's full tree snapshot to ``_tree``. - """ + +def _ensure_child_index(conn: sqlite3.Connection, stream: str) -> None: + """Index the child_frame json tag on the tf rows so per-frame time queries + seek. The live recorder gets this for free (the store auto-indexes tag keys on + tagged appends); this is for migrated recordings and the read side. Requires + the tf table to exist.""" + safe = _safe_table(stream) + # Composite (child_frame, ts) so a per-frame "latest at/before T" is a direct + # index range seek, not a scan+sort. + conn.execute( + f'CREATE INDEX IF NOT EXISTS "{safe}_child_ts_idx" ' + f"ON \"{safe}\"(json_extract(tags, '$.child_frame'), ts)" + ) + conn.commit() + + +class TfGraphWriter: + """Recorder helper: tracks the running topology and appends a ``tf_graph`` row + only when the structure changes.""" def __init__(self, db_path: str, stream: str = DEFAULT_TF_STREAM) -> None: self._stream = _safe_table(stream) + self._table = _graph_table(stream) self._conn = _connect(db_path) - self._lock = threading.Lock() - self._latest: dict[str, list[int]] = {} - ensure_tree_table(self._conn, self._stream) - - def record(self, child_frame: str, tf_id: int) -> None: - """Record that tf row ``tf_id`` holds the latest ``child_frame`` edge.""" - with self._lock: - _advance(self._latest, child_frame, tf_id) - snapshot = json.dumps(self._latest) - self._conn.execute( - f'INSERT OR REPLACE INTO "{self._stream}_tree" (tf_id, tree) VALUES (?, ?)', - (tf_id, snapshot), - ) - self._conn.commit() + self._structure: dict[str, dict[str, Any]] = {} + ensure_graph_table(self._conn, self._stream) + + def record(self, child_frame: str, parent_frame: str, is_static: bool, ts: float) -> None: + entry = {"parent": parent_frame, "static": bool(is_static)} + if self._structure.get(child_frame) == entry: + return # no structural change -> no new graph row + self._structure[child_frame] = entry + self._conn.execute( + f'INSERT INTO "{self._table}" (ts, structure) VALUES (?, ?)', + (ts, json.dumps(self._structure)), + ) + self._conn.commit() def close(self) -> None: - with self._lock: - self._conn.close() + self._conn.close() -def ensure_tree_table(conn: sqlite3.Connection, stream: str) -> None: - conn.execute( - f'CREATE TABLE IF NOT EXISTS "{_safe_table(stream)}_tree" ' - "(tf_id INTEGER PRIMARY KEY, tree TEXT NOT NULL)" - ) - # Speeds up the nearest-row lookup in get(); harmless if it already exists. - conn.execute( - f'CREATE INDEX IF NOT EXISTS "{_safe_table(stream)}_ts_idx" ON "{_safe_table(stream)}"(ts)' - ) - conn.commit() - - -def recompute_trees(store: Store, stream: str = DEFAULT_TF_STREAM) -> int: - """(Re)build every tf row's tree in chronological order, writing it back. - - Iterates the whole stream once — the expensive, one-time cost of adopting a - recording that predates trees. Returns the number of tf rows given a tree. - """ +def build_graph_stream(store: Store, stream: str = DEFAULT_TF_STREAM) -> int: + """One-time migration for a recording that predates the graph stream: tag every + tf row with its ``child_frame`` and build ``tf_graph`` chronologically. A frame + is treated as static if its pose never changes across the recording. Returns the + number of topology-change rows written.""" config = store.config if not isinstance(config, SqliteStoreConfig): - raise TypeError("recompute_trees needs a SqliteStore") + raise TypeError("build_graph_stream needs a SqliteStore") + table = _graph_table(stream) + safe = _safe_table(stream) + + # one decode pass: collect (id, ts, child, parent, pose-key) per row + rows: list[tuple[int, float, str, str, tuple[float, ...]]] = [] + poses_per_child: dict[str, set[tuple[float, ...]]] = {} + for obs in store.stream(safe, TFMessage).order_by("ts"): + for transform in getattr(obs.data, "transforms", None) or [obs.data]: + pose_key = ( + round(transform.translation.x, 9), + round(transform.translation.y, 9), + round(transform.translation.z, 9), + round(transform.rotation.x, 9), + round(transform.rotation.y, 9), + round(transform.rotation.z, 9), + round(transform.rotation.w, 9), + ) + rows.append((obs.id, obs.ts, transform.child_frame_id, transform.frame_id, pose_key)) + poses_per_child.setdefault(transform.child_frame_id, set()).add(pose_key) + static_frames = {child for child, poses in poses_per_child.items() if len(poses) == 1} + conn = _connect(config.path) try: - ensure_tree_table(conn, stream) - conn.execute(f'DELETE FROM "{_safe_table(stream)}_tree"') - latest: dict[str, list[int]] = {} - rows: list[tuple[int, str]] = [] - for observation in store.stream(stream, TFMessage).order_by("ts"): - transforms = getattr(observation.data, "transforms", None) or [observation.data] - for transform in transforms: - _advance(latest, transform.child_frame_id, observation.id) - rows.append((observation.id, json.dumps(latest))) - conn.executemany( - f'INSERT OR REPLACE INTO "{_safe_table(stream)}_tree" (tf_id, tree) VALUES (?, ?)', - rows, - ) + ensure_graph_table(conn, safe) + conn.execute(f'DELETE FROM "{table}"') + # tag each tf row with its child_frame (json_set keeps any existing tags) + for row_id, _ts, child, _parent, _pose in rows: + conn.execute( + f"UPDATE \"{safe}\" SET tags = json_set(tags, '$.child_frame', ?) WHERE id = ?", + (child, row_id), + ) + _ensure_child_index(conn, safe) + # build the topology change-log + structure: dict[str, dict[str, Any]] = {} + written = 0 + for _row_id, ts, child, parent, _pose in rows: + entry = {"parent": parent, "static": child in static_frames} + if structure.get(child) == entry: + continue + structure[child] = entry + conn.execute( + f'INSERT INTO "{table}" (ts, structure) VALUES (?, ?)', (ts, json.dumps(structure)) + ) + written += 1 conn.commit() - return len(rows) + return written finally: conn.close() class DbTf: - """Transform lookups backed by a store's tf stream + per-row tf trees.""" + """Transform lookups backed by a store's recorded transforms. + + On a SQLite store this uses the graph stream + an in-RAM graph cache; other + stores fall back to loading the tf streams into a :class:`MultiTBuffer`. Surface + is ``get(target, source, time_point, time_tolerance)`` / ``has_transforms()``. + """ def __init__( self, store: Store, stream: str = DEFAULT_TF_STREAM, + max_graph_changes_in_ram: int = DEFAULT_MAX_GRAPH_CHANGES_IN_RAM, stream_names: tuple[str, ...] = TF_STREAMS, - cache_size: int = _LOOKUP_CACHE_SIZE, ) -> None: self._store = store self._stream = _safe_table(stream) + self._table = _graph_table(stream) + self._max_in_ram = max_graph_changes_in_ram self._stream_names = stream_names # RAM fallback only - self._cache_size = cache_size self._lock = threading.Lock() self._conn: sqlite3.Connection | None = None - self._trees_ready = False - self._cache: OrderedDict[tuple[str, str, float | None, float | None], Transform | None] = ( - OrderedDict() - ) - # Instrumentation: number of tf rows decoded out of the db by get() (for - # the "no full load" test). The whole-table recompute is not counted. + self._built = False + # graph cache: either the whole change-log in RAM, or None (query per lookup) + self._graph_in_ram: list[tuple[float, dict[str, Any]]] | None = None + self._graph_loaded = False + self._static_cache: dict[str, Transform] = {} + self._buffer: MultiTBuffer | None = None # RAM fallback (non-sqlite) self.rows_fetched = 0 - # RAM fallback for non-sqlite stores. - self._buffer: MultiTBuffer | None = None - - # --- sqlite path ----------------------------------------------------------- + self.graph_queries = 0 @property def _is_sqlite(self) -> bool: return isinstance(self._store.config, SqliteStoreConfig) def _connection(self) -> sqlite3.Connection: - if self._conn is None: + conn = self._conn + if conn is None: config = self._store.config assert isinstance(config, SqliteStoreConfig) # guarded by _is_sqlite - self._conn = _connect(config.path) - return self._conn - - def _ensure_trees(self) -> None: - """Make sure every tf row has a tree; recompute (loudly) if not.""" - if self._trees_ready: - return - conn = self._connection() - ensure_tree_table(conn, self._stream) - (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() - (n_trees,) = conn.execute(f'SELECT count(*) FROM "{self._stream}_tree"').fetchone() - if n_rows and n_trees < n_rows: - logger.warning( - "\n" - "========================================================================\n" - " tf tree MISSING for stream %r (%d/%d rows have a tree).\n" - " Computing tf trees over the whole recording (one-time, chronological,\n" - " written back to the db). This can take a moment on large recordings.\n" - "========================================================================", - self._stream, - n_trees, - n_rows, - ) - built = recompute_trees(self._store, self._stream) - logger.warning("tf tree computed for %d rows of stream %r.", built, self._stream) - self._trees_ready = True - - def _fetch_transforms(self, ids: list[int]) -> list[Transform]: - """Decode the given tf-row ids into transforms in a single batched query - (frames live in the blob, so we decode each via the stream's codec).""" - if not ids: - return [] - backend: Any = self._store.stream(self._stream, TFMessage)._source - codec = backend.codec - placeholders = ",".join("?" * len(ids)) - rows = ( - self._connection() - .execute(f'SELECT data FROM "{self._stream}_blob" WHERE id IN ({placeholders})', ids) - .fetchall() - ) - transforms: list[Transform] = [] - for (data,) in rows: - message = codec.decode(data) - transforms.extend(getattr(message, "transforms", None) or [message]) - self.rows_fetched += len(rows) - return transforms - - def _anchor_tree(self, conn: sqlite3.Connection, time_point: float | None) -> str | None: - """The tree JSON of the anchor row, in one query: the first row with - ts >= query (so its own edge brackets the query for interpolation); else - the last row. Joins ``_tree`` to ```` so it's a single hit.""" - tree = f'"{self._stream}_tree"' - stream = f'"{self._stream}"' - last = f"SELECT t.tree FROM {tree} t JOIN {stream} s ON s.id = t.tf_id ORDER BY s.ts DESC LIMIT 1" - if time_point is None: - row = conn.execute(last).fetchone() - return str(row[0]) if row else None - row = conn.execute( - f"SELECT t.tree FROM {tree} t JOIN {stream} s ON s.id = t.tf_id " - "WHERE s.ts >= ? ORDER BY s.ts ASC LIMIT 1", - (time_point,), - ).fetchone() - if row is not None: - return str(row[0]) - row = conn.execute(last).fetchone() - return str(row[0]) if row else None - - def _get_sqlite( - self, - target_frame: str, - source_frame: str, - time_point: float | None, - time_tolerance: float | None, - ) -> Transform | None: - key = (target_frame, source_frame, time_point, time_tolerance) - with self._lock: - if key in self._cache: - self._cache.move_to_end(key) - return self._cache[key] - self._ensure_trees() - conn = self._connection() - tree_json = self._anchor_tree(conn, time_point) # query 1: anchor + tree - result: Transform | None = None - if tree_json is not None: - tree: dict[str, list[int]] = json.loads(tree_json) - ids = sorted({tf_id for refs in tree.values() for tf_id in refs}) - # query 2: only the <=2-per-frame referenced rows, batched — never the - # whole table. MultiTBuffer composes the chain + interpolates. - buffer = MultiTBuffer(buffer_size=_NO_PRUNE) - transforms = self._fetch_transforms(ids) - if transforms: - buffer.receive_transform(*transforms) - result = buffer.lookup(target_frame, source_frame, time_point, time_tolerance) - with self._lock: - self._cache[key] = result - self._cache.move_to_end(key) - while len(self._cache) > self._cache_size: - self._cache.popitem(last=False) - return result + conn = _connect(config.path) + self._conn = conn + return conn # --- RAM fallback (non-sqlite stores) -------------------------------------- @@ -327,7 +266,7 @@ def _ensure_loaded(self) -> MultiTBuffer: self._buffer = buffer return buffer - # --- public API ------------------------------------------------------------ + # --- graph stream (sqlite) ------------------------------------------------- def has_transforms(self) -> bool: if not self._is_sqlite: @@ -338,6 +277,143 @@ def has_transforms(self) -> bool: (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() return bool(n_rows) + def _ensure_built(self) -> None: + if self._built: + return + conn = self._connection() + ensure_graph_table(conn, self._stream) + (n_graph,) = conn.execute(f'SELECT count(*) FROM "{self._table}"').fetchone() + (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() + if n_rows and n_graph == 0: + logger.warning( + "\n========================================================================\n" + " tf graph stream MISSING for %r. Building it (one-time): tagging tf rows\n" + " with child_frame and writing the topology change-log.\n" + "========================================================================", + self._stream, + ) + built = build_graph_stream(self._store, self._stream) + logger.warning("tf graph built: %d topology changes for %r.", built, self._stream) + if n_rows: + _ensure_child_index(conn, self._stream) # tf table exists now + self._built = True + + def _load_graph_if_small(self) -> None: + if self._graph_loaded: + return + conn = self._connection() + (n_graph,) = conn.execute(f'SELECT count(*) FROM "{self._table}"').fetchone() + if n_graph < self._max_in_ram: + self._graph_in_ram = [ + (ts, json.loads(structure)) + for ts, structure in conn.execute( + f'SELECT ts, structure FROM "{self._table}" ORDER BY ts ASC' + ) + ] + else: + self._graph_in_ram = None # too many -> query per lookup + self._graph_loaded = True + + def _graph_at(self, query_time: float) -> dict[str, Any] | None: + if self._graph_in_ram is not None: + # in-RAM: binary search the latest change at-or-before query_time + stamps = [ts for ts, _ in self._graph_in_ram] + index = bisect.bisect_right(stamps, query_time) - 1 + if index < 0: + return self._graph_in_ram[0][1] # before first -> earliest + return self._graph_in_ram[index][1] + # fallback: one query + self.graph_queries += 1 + conn = self._connection() + row = conn.execute( + f'SELECT structure FROM "{self._table}" WHERE ts <= ? ORDER BY ts DESC LIMIT 1', + (query_time,), + ).fetchone() + if row is None: + row = conn.execute( + f'SELECT structure FROM "{self._table}" ORDER BY ts ASC LIMIT 1' + ).fetchone() + return json.loads(row[0]) if row else None + + def _chain_frames(self, graph: dict[str, Any], source: str, target: str) -> list[str] | None: + def to_root(frame: str) -> list[str]: + path = [frame] + seen = {frame} + while ( + frame in graph + and graph[frame].get("parent") in graph + and graph[frame]["parent"] not in seen + ): + frame = graph[frame]["parent"] + path.append(frame) + seen.add(frame) + # include a final parent that is itself a root (not a key in graph) + if frame in graph and graph[frame].get("parent") and graph[frame]["parent"] not in seen: + path.append(graph[frame]["parent"]) + return path + + source_path = to_root(source) + target_path = to_root(target) + common = next((f for f in source_path if f in set(target_path)), None) + if common is None: + return None # disjoint graph: no transform between them + frames = source_path[: source_path.index(common) + 1] + frames += target_path[: target_path.index(common)] + return frames + + def _codec(self) -> Any: + source = self._store.stream(self._stream, TFMessage)._source + return cast("Any", source).codec + + def _decode_blob(self, data: bytes, frame: str) -> Transform: + # The blob is the codec-encoded message; pick the transform for `frame` + # (rows normally hold one; legacy rows may pack several). + message = self._codec().decode(data) + transforms = getattr(message, "transforms", None) or [message] + for transform in transforms: + if transform.child_frame_id == frame: + return cast("Transform", transform) + return cast("Transform", transforms[0]) + + def _fetch_rows( + self, dynamic: list[str], static: list[str], query_time: float + ) -> dict[tuple[str, str], bytes]: + """ONE query: for each dynamic frame the bracketing rows ('lo' = latest at + or before query_time, 'hi' = earliest at or after), and for each (uncached) + static frame its latest row ('st') — all joined to the blob data. Keyed by + (frame, kind) -> blob bytes.""" + cf = "json_extract(tags, '$.child_frame')" + tf, blob = f'"{self._stream}"', f'"{self._stream}_blob"' + # One UNION of per-frame, index-served LIMIT-1 subqueries: each is a direct + # (child_frame, ts) range seek — far cheaper than a window-function scan, and + # still a single round-trip. + parts: list[str] = [] + params: list[Any] = [] + + def pick(frame: str, kind: str, where_ts: str, order: str) -> None: + parts.append( + f"SELECT ? AS cf, ? AS kind, " + f"(SELECT id FROM {tf} WHERE {cf} = ?{where_ts} ORDER BY ts {order} LIMIT 1) AS id" + ) + params.extend([frame, kind, frame]) + + for frame in dynamic: + pick(frame, "lo", " AND ts <= ?", "DESC") + params.append(query_time) + pick(frame, "hi", " AND ts >= ?", "ASC") + params.append(query_time) + for frame in static: + pick(frame, "st", "", "DESC") + if not parts: + return {} + union = " UNION ALL ".join(parts) + sql = f"SELECT t.cf, t.kind, b.data FROM ({union}) t JOIN {blob} b ON b.id = t.id" + rows: dict[tuple[str, str], bytes] = {} + for cf_val, kind, data in self._connection().execute(sql, params): + rows[(cf_val, kind)] = data + self.rows_fetched += 1 + return rows + def get( self, target_frame: str, @@ -345,17 +421,62 @@ def get( time_point: float | None = None, time_tolerance: float | None = None, ) -> Transform | None: - """Transform that maps a point in ``source_frame`` into ``target_frame``. - - Returns ``None`` if no chain connects the two frames at the requested - time. Resolves through the per-row tf tree (O(tree size) rows read), with - a small FIFO cache over identical lookups. - """ + """Transform that maps a point in ``source_frame`` into ``target_frame``, + or ``None`` if no chain connects them at the requested time.""" if not self._is_sqlite: return self._ensure_loaded().lookup( target_frame, source_frame, time_point, time_tolerance ) - return self._get_sqlite(target_frame, source_frame, time_point, time_tolerance) + self._ensure_built() + self._load_graph_if_small() + query_time = time_point if time_point is not None else 0.0 + graph = self._graph_at(query_time) # 0 queries when the graph is in RAM + if graph is None: + return None + frames = self._chain_frames(graph, source_frame, target_frame) + if frames is None: + return None + + edges = [f for f in frames if f in graph] # roots have no incoming edge + dynamic = [f for f in edges if not graph[f].get("static")] + static = [f for f in edges if graph[f].get("static")] + uncached_static = [f for f in static if f not in self._static_cache] + + rows = self._fetch_rows(dynamic, uncached_static, query_time) # ONE detail query + + buffer = MultiTBuffer(buffer_size=_NO_PRUNE) + for frame in static: + transform = self._static_cache.get(frame) + if transform is None: + data = rows.get((frame, "st")) + if data is None: + return None + transform = self._decode_blob(data, frame) + self._static_cache[frame] = transform + # restamp the constant to query_time so the buffer's tolerance never + # rejects a static that was recorded long ago (latched once). + buffer.receive_transform(_restamp(transform, query_time)) + for frame in dynamic: + lo = rows.get((frame, "lo")) + hi = rows.get((frame, "hi")) + chosen = lo if lo is not None else hi + if chosen is None: + return None + buffer.receive_transform(self._decode_blob(chosen, frame)) + other = hi if hi is not None else lo + if other is not None and other is not chosen: + buffer.receive_transform(self._decode_blob(other, frame)) + return buffer.lookup(target_frame, source_frame, time_point, time_tolerance) + + +def _restamp(transform: Transform, ts: float) -> Transform: + return Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ts=ts, + ) def transform_matrix(transform: Transform) -> tuple[np.ndarray, np.ndarray]: @@ -388,7 +509,9 @@ def write_tf_tree( - ``odom_parent -> odom_child`` is emitted once per odometry sample, taken from each observation's pose. - Returns the number of tf observations written. + Each transform is written as its own row (one transform per ``TFMessage``) so + the graph-stream reader can range-query it by ``child_frame``. Returns the + number of tf observations written. """ config = store.config if not isinstance(config, SqliteStoreConfig): @@ -411,18 +534,26 @@ def write_tf_tree( tf_stream = store.stream(stream_name, TFMessage) written = 0 + def _append(transform: Transform, ts: float) -> None: + nonlocal written + tf_stream.append( + TFMessage(transform), ts=ts, tags={"child_frame": transform.child_frame_id} + ) + written += 1 + # dynamic: odom_parent -> odom_child, one per odometry sample for row in odom: ts = float(row[0]) - transform = Transform( - translation=Vector3(row[1], row[2], row[3]), - rotation=Quaternion(row[4], row[5], row[6], row[7]), - frame_id=odom_parent, - child_frame_id=odom_child, - ts=ts, + _append( + Transform( + translation=Vector3(row[1], row[2], row[3]), + rotation=Quaternion(row[4], row[5], row[6], row[7]), + frame_id=odom_parent, + child_frame_id=odom_child, + ts=ts, + ), + ts, ) - tf_stream.append(TFMessage(transform), ts=ts) - written += 1 # static: root links + sensor mount, resampled every static_period t0 = float(odom[0, 0]) @@ -444,7 +575,7 @@ def statics_at(ts: float) -> list[Transform]: return links for static_ts in np.arange(t0, t1 + static_period, static_period): - tf_stream.append(TFMessage(*statics_at(float(static_ts))), ts=float(static_ts)) - written += 1 + for transform in statics_at(float(static_ts)): + _append(transform, float(static_ts)) return written diff --git a/dimos/memory2/db_tf2.py b/dimos/memory2/db_tf2.py deleted file mode 100644 index 67554559d7..0000000000 --- a/dimos/memory2/db_tf2.py +++ /dev/null @@ -1,413 +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. - -"""Graph-stream transform lookups (the multi-robot-friendly successor to DbTf). - -Two pieces of recorded state: - -* a **graph stream** (table ``tf_graph``): one row per *topology* change — i.e. - whenever the set of frames or any frame's parent / static-ness changes. Each row - is the full structure at that instant: ``{child_frame: {parent, static}}``. - Topology changes rarely (a robot joins/leaves, a relocalization re-parents), so - this table is tiny. -* the existing ``tf`` stream, with each row tagged by its ``child_frame`` (an - indexed json tag) so a frame's samples can be range-queried by time. - -``GraphTf.get`` then: gets the graph as-of the query time (from RAM if there are -few graph changes, else one query), walks it to the source->target chain (in -memory; the graph may be DISJOINT for unrelated robots), and resolves *only* the -chain's frames — a static frame is one cached constant, a dynamic frame is its two -samples bracketing the time, interpolated. Composition + interpolation reuse -:class:`MultiTBuffer`. -""" - -from __future__ import annotations - -import bisect -import json -import sqlite3 -from typing import TYPE_CHECKING, Any, cast - -from dimos.memory2.db_tf import _connect, _safe_table -from dimos.memory2.store.sqlite import SqliteStoreConfig -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.protocol.tf.tf import MultiTBuffer -from dimos.utils.logging_config import setup_logger - -if TYPE_CHECKING: - from dimos.memory2.store.base import Store - -logger = setup_logger() - -DEFAULT_TF_STREAM = "tf" -GRAPH_TABLE = "tf_graph" -# If a recording has fewer than this many topology changes, load them all into RAM -# so a lookup needs no graph query (the common single-robot / stable-tree case). -# At or above it, fall back to one graph query per lookup (many-robot churn). -DEFAULT_MAX_GRAPH_CHANGES_IN_RAM = 20 -_NO_PRUNE = 1.0e15 - - -def _graph_table(stream: str) -> str: - return f"{_safe_table(stream)}_graph" - - -def ensure_graph_table(conn: sqlite3.Connection, stream: str) -> None: - """Create the topology change-log table (safe to call before the tf table - exists — it doesn't touch the tf table).""" - table = _graph_table(stream) - conn.execute( - f'CREATE TABLE IF NOT EXISTS "{table}" ' - "(id INTEGER PRIMARY KEY AUTOINCREMENT, ts REAL NOT NULL, structure TEXT NOT NULL)" - ) - conn.execute(f'CREATE INDEX IF NOT EXISTS "{table}_ts_idx" ON "{table}"(ts)') - conn.commit() - - -def _ensure_child_index(conn: sqlite3.Connection, stream: str) -> None: - """Index the child_frame json tag on the tf rows so per-frame time queries - seek. The live recorder gets this for free (the store auto-indexes tag keys on - tagged appends); this is for migrated recordings and the read side. Requires - the tf table to exist.""" - safe = _safe_table(stream) - # Composite (child_frame, ts) so a per-frame "latest at/before T" is a direct - # index range seek, not a scan+sort. - conn.execute( - f'CREATE INDEX IF NOT EXISTS "{safe}_child_ts_idx" ' - f"ON \"{safe}\"(json_extract(tags, '$.child_frame'), ts)" - ) - conn.commit() - - -class TfGraphWriter: - """Recorder helper: tracks the running topology and appends a ``tf_graph`` row - only when the structure changes.""" - - def __init__(self, db_path: str, stream: str = DEFAULT_TF_STREAM) -> None: - self._stream = _safe_table(stream) - self._table = _graph_table(stream) - self._conn = _connect(db_path) - self._structure: dict[str, dict[str, Any]] = {} - ensure_graph_table(self._conn, self._stream) - - def record(self, child_frame: str, parent_frame: str, is_static: bool, ts: float) -> None: - entry = {"parent": parent_frame, "static": bool(is_static)} - if self._structure.get(child_frame) == entry: - return # no structural change -> no new graph row - self._structure[child_frame] = entry - self._conn.execute( - f'INSERT INTO "{self._table}" (ts, structure) VALUES (?, ?)', - (ts, json.dumps(self._structure)), - ) - self._conn.commit() - - def close(self) -> None: - self._conn.close() - - -def build_graph_stream(store: Store, stream: str = DEFAULT_TF_STREAM) -> int: - """One-time migration for a recording that predates the graph stream: tag every - tf row with its ``child_frame`` and build ``tf_graph`` chronologically. A frame - is treated as static if its pose never changes across the recording. Returns the - number of topology-change rows written.""" - config = store.config - if not isinstance(config, SqliteStoreConfig): - raise TypeError("build_graph_stream needs a SqliteStore") - table = _graph_table(stream) - safe = _safe_table(stream) - - # one decode pass: collect (id, ts, child, parent, pose-key) per row - rows: list[tuple[int, float, str, str, tuple[float, ...]]] = [] - poses_per_child: dict[str, set[tuple[float, ...]]] = {} - for obs in store.stream(safe, TFMessage).order_by("ts"): - for transform in getattr(obs.data, "transforms", None) or [obs.data]: - pose_key = ( - round(transform.translation.x, 9), - round(transform.translation.y, 9), - round(transform.translation.z, 9), - round(transform.rotation.x, 9), - round(transform.rotation.y, 9), - round(transform.rotation.z, 9), - round(transform.rotation.w, 9), - ) - rows.append((obs.id, obs.ts, transform.child_frame_id, transform.frame_id, pose_key)) - poses_per_child.setdefault(transform.child_frame_id, set()).add(pose_key) - static_frames = {child for child, poses in poses_per_child.items() if len(poses) == 1} - - conn = _connect(config.path) - try: - ensure_graph_table(conn, safe) - conn.execute(f'DELETE FROM "{table}"') - # tag each tf row with its child_frame (json_set keeps any existing tags) - for row_id, _ts, child, _parent, _pose in rows: - conn.execute( - f"UPDATE \"{safe}\" SET tags = json_set(tags, '$.child_frame', ?) WHERE id = ?", - (child, row_id), - ) - _ensure_child_index(conn, safe) - # build the topology change-log - structure: dict[str, dict[str, Any]] = {} - written = 0 - for _row_id, ts, child, parent, _pose in rows: - entry = {"parent": parent, "static": child in static_frames} - if structure.get(child) == entry: - continue - structure[child] = entry - conn.execute( - f'INSERT INTO "{table}" (ts, structure) VALUES (?, ?)', (ts, json.dumps(structure)) - ) - written += 1 - conn.commit() - return written - finally: - conn.close() - - -class DbTf2: - """Graph-stream transform lookups with an in-RAM graph cache for the common case. - - Drop-in successor to :class:`dimos.memory2.db_tf.DbTf` — same - ``get(target, source, time_point, time_tolerance)`` / ``has_transforms()`` - surface, so the store's ``tf`` property can swap to it in one line.""" - - def __init__( - self, - store: Store, - stream: str = DEFAULT_TF_STREAM, - max_graph_changes_in_ram: int = DEFAULT_MAX_GRAPH_CHANGES_IN_RAM, - ) -> None: - self._store = store - self._stream = _safe_table(stream) - self._table = _graph_table(stream) - self._max_in_ram = max_graph_changes_in_ram - self._conn: sqlite3.Connection | None = None - self._built = False - # graph cache: either the whole change-log in RAM, or None (query per lookup) - self._graph_in_ram: list[tuple[float, dict[str, Any]]] | None = None - self._graph_loaded = False - self._static_cache: dict[str, Transform] = {} - self.rows_fetched = 0 - self.graph_queries = 0 - - def _connection(self) -> sqlite3.Connection: - conn = self._conn - if conn is None: - config = self._store.config - assert isinstance(config, SqliteStoreConfig) - conn = _connect(config.path) - self._conn = conn - return conn - - def has_transforms(self) -> bool: - conn = self._connection() - if self._stream not in set(self._store.list_streams()): - return False - (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() - return bool(n_rows) - - def _ensure_built(self) -> None: - if self._built: - return - conn = self._connection() - ensure_graph_table(conn, self._stream) - (n_graph,) = conn.execute(f'SELECT count(*) FROM "{self._table}"').fetchone() - (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() - if n_rows and n_graph == 0: - logger.warning( - "\n========================================================================\n" - " tf graph stream MISSING for %r. Building it (one-time): tagging tf rows\n" - " with child_frame and writing the topology change-log.\n" - "========================================================================", - self._stream, - ) - built = build_graph_stream(self._store, self._stream) - logger.warning("tf graph built: %d topology changes for %r.", built, self._stream) - if n_rows: - _ensure_child_index(conn, self._stream) # tf table exists now - self._built = True - - def _load_graph_if_small(self) -> None: - if self._graph_loaded: - return - conn = self._connection() - (n_graph,) = conn.execute(f'SELECT count(*) FROM "{self._table}"').fetchone() - if n_graph < self._max_in_ram: - self._graph_in_ram = [ - (ts, json.loads(structure)) - for ts, structure in conn.execute( - f'SELECT ts, structure FROM "{self._table}" ORDER BY ts ASC' - ) - ] - else: - self._graph_in_ram = None # too many -> query per lookup - self._graph_loaded = True - - def _graph_at(self, query_time: float) -> dict[str, Any] | None: - if self._graph_in_ram is not None: - # in-RAM: binary search the latest change at-or-before query_time - stamps = [ts for ts, _ in self._graph_in_ram] - index = bisect.bisect_right(stamps, query_time) - 1 - if index < 0: - return self._graph_in_ram[0][1] # before first -> earliest - return self._graph_in_ram[index][1] - # fallback: one query - self.graph_queries += 1 - conn = self._connection() - row = conn.execute( - f'SELECT structure FROM "{self._table}" WHERE ts <= ? ORDER BY ts DESC LIMIT 1', - (query_time,), - ).fetchone() - if row is None: - row = conn.execute( - f'SELECT structure FROM "{self._table}" ORDER BY ts ASC LIMIT 1' - ).fetchone() - return json.loads(row[0]) if row else None - - def _chain_frames(self, graph: dict[str, Any], source: str, target: str) -> list[str] | None: - def to_root(frame: str) -> list[str]: - path = [frame] - seen = {frame} - while ( - frame in graph - and graph[frame].get("parent") in graph - and graph[frame]["parent"] not in seen - ): - frame = graph[frame]["parent"] - path.append(frame) - seen.add(frame) - # include a final parent that is itself a root (not a key in graph) - if frame in graph and graph[frame].get("parent") and graph[frame]["parent"] not in seen: - path.append(graph[frame]["parent"]) - return path - - source_path = to_root(source) - target_path = to_root(target) - common = next((f for f in source_path if f in set(target_path)), None) - if common is None: - return None # disjoint graph: no transform between them - frames = source_path[: source_path.index(common) + 1] - frames += target_path[: target_path.index(common)] - return frames - - def _codec(self) -> Any: - source = self._store.stream(self._stream, TFMessage)._source - return cast("Any", source).codec - - def _decode_blob(self, data: bytes, frame: str) -> Transform: - # The blob is the codec-encoded message; pick the transform for `frame` - # (rows normally hold one; legacy rows may pack several). - message = self._codec().decode(data) - transforms = getattr(message, "transforms", None) or [message] - for transform in transforms: - if transform.child_frame_id == frame: - return cast("Transform", transform) - return cast("Transform", transforms[0]) - - def _fetch_rows( - self, dynamic: list[str], static: list[str], query_time: float - ) -> dict[tuple[str, str], bytes]: - """ONE query: for each dynamic frame the bracketing rows ('lo' = latest at - or before query_time, 'hi' = earliest at or after), and for each (uncached) - static frame its latest row ('st') — all joined to the blob data. Keyed by - (frame, kind) -> blob bytes.""" - cf = "json_extract(tags, '$.child_frame')" - tf, blob = f'"{self._stream}"', f'"{self._stream}_blob"' - # One UNION of per-frame, index-served LIMIT-1 subqueries: each is a direct - # (child_frame, ts) range seek — far cheaper than a window-function scan, and - # still a single round-trip. - parts: list[str] = [] - params: list[Any] = [] - - def pick(frame: str, kind: str, where_ts: str, order: str) -> None: - parts.append( - f"SELECT ? AS cf, ? AS kind, " - f"(SELECT id FROM {tf} WHERE {cf} = ?{where_ts} ORDER BY ts {order} LIMIT 1) AS id" - ) - params.extend([frame, kind, frame]) - - for frame in dynamic: - pick(frame, "lo", " AND ts <= ?", "DESC") - params.append(query_time) - pick(frame, "hi", " AND ts >= ?", "ASC") - params.append(query_time) - for frame in static: - pick(frame, "st", "", "DESC") - if not parts: - return {} - union = " UNION ALL ".join(parts) - sql = f"SELECT t.cf, t.kind, b.data FROM ({union}) t JOIN {blob} b ON b.id = t.id" - rows: dict[tuple[str, str], bytes] = {} - for cf_val, kind, data in self._connection().execute(sql, params): - rows[(cf_val, kind)] = data - self.rows_fetched += 1 - return rows - - def get( - self, - target_frame: str, - source_frame: str, - time_point: float | None = None, - time_tolerance: float | None = None, - ) -> Transform | None: - self._ensure_built() - self._load_graph_if_small() - query_time = time_point if time_point is not None else 0.0 - graph = self._graph_at(query_time) # 0 queries when the graph is in RAM - if graph is None: - return None - frames = self._chain_frames(graph, source_frame, target_frame) - if frames is None: - return None - - edges = [f for f in frames if f in graph] # roots have no incoming edge - dynamic = [f for f in edges if not graph[f].get("static")] - static = [f for f in edges if graph[f].get("static")] - uncached_static = [f for f in static if f not in self._static_cache] - - rows = self._fetch_rows(dynamic, uncached_static, query_time) # ONE detail query - - buffer = MultiTBuffer(buffer_size=_NO_PRUNE) - for frame in static: - transform = self._static_cache.get(frame) - if transform is None: - data = rows.get((frame, "st")) - if data is None: - return None - transform = self._decode_blob(data, frame) - self._static_cache[frame] = transform - # restamp the constant to query_time so the buffer's tolerance never - # rejects a static that was recorded long ago (latched once). - buffer.receive_transform(_restamp(transform, query_time)) - for frame in dynamic: - lo = rows.get((frame, "lo")) - hi = rows.get((frame, "hi")) - chosen = lo if lo is not None else hi - if chosen is None: - return None - buffer.receive_transform(self._decode_blob(chosen, frame)) - other = hi if hi is not None else lo - if other is not None and other is not chosen: - buffer.receive_transform(self._decode_blob(other, frame)) - return buffer.lookup(target_frame, source_frame, time_point, time_tolerance) - - -def _restamp(transform: Transform, ts: float) -> Transform: - return Transform( - translation=transform.translation, - rotation=transform.rotation, - frame_id=transform.frame_id, - child_frame_id=transform.child_frame_id, - ts=ts, - ) diff --git a/dimos/memory2/demo_bench_db_tf2.py b/dimos/memory2/demo_bench_db_tf2.py deleted file mode 100644 index 1a6c17a1b8..0000000000 --- a/dimos/memory2/demo_bench_db_tf2.py +++ /dev/null @@ -1,103 +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. - -"""Benchmark: DbTf (per-row tf-tree) vs DbTf2 (graph-stream + in-RAM graph). - -Builds a synthetic single-robot recording (one transform per row, child_frame -tagged, topology change-log), then measures per-lookup latency and SQL query -counts for both implementations. Run: `python -m dimos.memory2.demo_bench_db_tf2`. -""" - -from __future__ import annotations - -from pathlib import Path -import tempfile -import time - -from dimos.memory2.db_tf import DbTf -from dimos.memory2.db_tf2 import DbTf2 -from dimos.memory2.store.sqlite import SqliteStore -from dimos.memory2.test_db_tf2 import _T0, _record_single_robot - -_WARMUP = 200 -_N = 2000 - - -def _count_queries(db: object, query_time: float) -> int: - conn = db._connection() # type: ignore[attr-defined] - count = {"n": 0} - - def trace(_sql: str) -> None: - count["n"] += 1 - - conn.set_trace_callback(trace) - db.get("world", "sensor", query_time, 0.5) # type: ignore[attr-defined] - conn.set_trace_callback(None) - return count["n"] - - -def _time_lookups(db: object, n: int, offset: int) -> float: - start = time.perf_counter() - for k in range(n): - q = _T0 + 0.05 + ((offset + k) % 290) * 0.033 - db.get("world", "sensor", q, 0.5) # type: ignore[attr-defined] - return time.perf_counter() - start - - -def main() -> None: - tmp = Path(tempfile.mkdtemp()) - path = tmp / "bench.db" - _record_single_robot(path, static_repeat=True) - - results = {} - for name, make in (("DbTf ", lambda s: DbTf(s)), ("DbTf2", lambda s: DbTf2(s))): - store = SqliteStore(path=str(path), must_exist=True) - db = make(store) - - # one-time build cost (tree / graph migration) folded into a cold first call - cold_start = time.perf_counter() - db.get("world", "sensor", _T0 + 0.05, 0.5) - cold = time.perf_counter() - cold_start - - first_q = _count_queries(db, _T0 + 7.123) - warm_q = _count_queries(db, _T0 + 3.777) - - _time_lookups(db, _WARMUP, offset=0) # warm caches/pages - warm_total = _time_lookups(db, _N, offset=_WARMUP) - - graph_q = getattr(db, "graph_queries", None) - results[name] = { - "cold_ms": cold * 1e3, - "warm_us": warm_total / _N * 1e6, - "first_q": first_q, - "warm_q": warm_q, - "graph_q": graph_q, - } - store.stop() - - print(f"\nrecording: single robot, 5 frames, 30Hz dynamic, 10s | lookups timed: {_N}\n") - header = f"{'impl':6} {'cold(ms)':>10} {'warm/lookup(us)':>17} {'1st-call q':>11} {'warm q':>8} {'graph q':>8}" - print(header) - print("-" * len(header)) - for name, r in results.items(): - gq = "-" if r["graph_q"] is None else r["graph_q"] - print( - f"{name:6} {r['cold_ms']:10.1f} {r['warm_us']:17.1f} " - f"{r['first_q']:11d} {r['warm_q']:8d} {gq!s:>8}" - ) - print() - - -if __name__ == "__main__": - main() diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index 43ca256f87..4ecb2ed7af 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -30,8 +30,7 @@ from dimos.constants import DIMOS_PROJECT_ROOT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig -from dimos.memory2.db_tf import TfTreeWriter -from dimos.memory2.db_tf2 import TfGraphWriter +from dimos.memory2.db_tf import TfGraphWriter from dimos.memory2.embed import EmbedImages from dimos.memory2.store.null import NullStore from dimos.memory2.store.sqlite import SqliteStore @@ -455,8 +454,8 @@ def _collect_pose_setters(self) -> dict[str, PoseSetter]: return setters def _record_tf(self) -> None: - """Record the live tf + static_tf streams under "tf", building a per-row - tf tree as we go (no-op without a pubsub tf).""" + """Record the live tf + static_tf streams under "tf", writing the topology + change-log (tf_graph) as we go (no-op without a pubsub tf).""" topic = getattr(self.tf.config, "topic", None) pubsub = getattr(self.tf, "pubsub", None) if not topic or pubsub is None: @@ -465,27 +464,21 @@ def _record_tf(self) -> None: tf_stream = self.store.stream("tf", TFMessage) is_sqlite = isinstance(self.store, SqliteStore) - # Per-row tf tree (last 2 refs per child frame), for the current DbTf. - tree_writer = TfTreeWriter(self.store.config.path, "tf") if is_sqlite else None - # Topology change-log + child_frame tags, for the DbTf2 graph-stream design. - # Both written during the transition; DbTf2 becomes the default at swap time. + # Topology change-log + child_frame tags, for the graph-stream DbTf. graph_writer = TfGraphWriter(self.store.config.path, "tf") if is_sqlite else None - for writer in (tree_writer, graph_writer): - if writer is not None: - self.register_disposable(Disposable(writer.close)) + if graph_writer is not None: + self.register_disposable(Disposable(graph_writer.close)) def make_handler(is_static: bool) -> Any: def on_tf(msg: TFMessage, _topic: Any) -> None: try: for transform in msg.transforms: - observation = tf_stream.append( + tf_stream.append( TFMessage(transform), ts=transform.ts, pose=None, tags={"child_frame": transform.child_frame_id}, ) - if tree_writer is not None: - tree_writer.record(transform.child_frame_id, observation.id) if graph_writer is not None: graph_writer.record( transform.child_frame_id, diff --git a/dimos/memory2/test_db_tf.py b/dimos/memory2/test_db_tf.py index e5fe179d53..9fcafa7eb9 100644 --- a/dimos/memory2/test_db_tf.py +++ b/dimos/memory2/test_db_tf.py @@ -12,21 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for the per-row tf-tree DbTf: correctness vs the old full-load buffer, -no-full-load + FIFO behavior, recompute-on-load, and the recorder's tree build.""" +"""Integration tests for DbTf (graph-stream transform lookups). Data is written +one-transform-per-row + child_frame tag + topology change-log, exactly as the live +recorder does. These cover the behaviours that matter end-to-end: interpolation +against a full-load buffer, a latched static, time-varying topology (relocalization +re-parents a frame), and a disjoint multi-robot graph.""" from __future__ import annotations -import json import math from pathlib import Path -import sqlite3 -import time -import pytest - -from dimos.memory2 import db_tf as db_tf_module -from dimos.memory2.db_tf import DbTf, TfTreeWriter, recompute_trees +from dimos.memory2.db_tf import DbTf, TfGraphWriter from dimos.memory2.store.sqlite import SqliteStore from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform @@ -34,72 +31,42 @@ from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.protocol.tf.tf import MultiTBuffer -# Frame chain used throughout: world -> map -> odom -> base_link -> sensor. -# Only odom -> base_link moves; the rest are fixed (the realistic case). -_DYN_RATE = 30.0 # Hz, odom -> base_link -_STATIC_PERIOD = 0.5 # s, root + sensor links re-emitted -_DURATION = 10.0 # s _T0 = 1000.0 +_DYN_RATE = 30.0 +_DURATION = 10.0 +_NO_PRUNE = 1.0e15 -def _yaw_quat(theta: float) -> Quaternion: +def _yaw(theta: float) -> Quaternion: return Quaternion(0.0, 0.0, math.sin(theta / 2.0), math.cos(theta / 2.0)) -def _all_transforms() -> list[Transform]: - """A moving odom->base_link plus fixed root + sensor links, interleaved by ts.""" - transforms: list[Transform] = [] - # dynamic odom -> base_link: drives a gentle arc - n = int(_DURATION * _DYN_RATE) - for i in range(n): - ts = _T0 + i / _DYN_RATE - transforms.append( - Transform( - translation=Vector3(0.5 * i / _DYN_RATE, 0.1 * i / _DYN_RATE, 0.0), - rotation=_yaw_quat(0.02 * i), - frame_id="odom", - child_frame_id="base_link", - ts=ts, - ) - ) - # statics re-emitted every _STATIC_PERIOD - m = int(_DURATION / _STATIC_PERIOD) - for j in range(m): - ts = _T0 + j * _STATIC_PERIOD - transforms.append(Transform(frame_id="world", child_frame_id="map", ts=ts)) - transforms.append(Transform(frame_id="map", child_frame_id="odom", ts=ts)) - transforms.append( - Transform( - translation=Vector3(0.0, 0.0, 0.3), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", - child_frame_id="sensor", - ts=ts, - ) - ) - transforms.sort(key=lambda t: t.ts) - return transforms +def _static(parent: str, child: str, xyz: tuple[float, float, float], ts: float) -> Transform: + return Transform( + translation=Vector3(*xyz), + rotation=Quaternion(0, 0, 0, 1), + frame_id=parent, + child_frame_id=child, + ts=ts, + ) -def _write_tf_db(path: Path, *, with_trees: bool) -> list[Transform]: - """Build a SqliteStore with a `tf` stream (one transform per row). When - ``with_trees`` is set, also write per-row trees as the recorder would.""" - store = SqliteStore(path=str(path)) - tf_stream = store.stream("tf", TFMessage) - transforms = _all_transforms() - writer = TfTreeWriter(str(path), "tf") if with_trees else None - for transform in transforms: - obs = tf_stream.append(TFMessage(transform), ts=transform.ts, pose=None) - if writer is not None: - writer.record(transform.child_frame_id, obs.id) - if writer is not None: - writer.close() - store.stop() - return transforms +def _append( + store: SqliteStore, graph: TfGraphWriter, transform: Transform, is_static: bool +) -> None: + """Record one transform exactly as the live recorder does: one row, tagged + child_frame, plus a topology-change row when the structure changes.""" + store.stream("tf", TFMessage).append( + TFMessage(transform), + ts=transform.ts, + pose=None, + tags={"child_frame": transform.child_frame_id}, + ) + graph.record(transform.child_frame_id, transform.frame_id, is_static, transform.ts) -def _reference_buffer(transforms: list[Transform]) -> MultiTBuffer: - buffer = MultiTBuffer(buffer_size=1.0e15) +def _ref(transforms: list[Transform]) -> MultiTBuffer: + buffer = MultiTBuffer(buffer_size=_NO_PRUNE) buffer.receive_transform(*transforms) return buffer @@ -116,205 +83,176 @@ def _diff(a: Transform, b: Transform) -> float: ) -def test_get_matches_full_load_buffer(tmp_path: Path) -> None: - """Tree-based get matches the old load-everything MultiTBuffer within tol, - including at query times BETWEEN samples (interpolation through the tree).""" - transforms = _write_tf_db(tmp_path / "tf.db", with_trees=True) - reference = _reference_buffer(transforms) - store = SqliteStore(path=str(tmp_path / "tf.db"), must_exist=True) - db = DbTf(store) +def _record_single_robot(path: Path, *, static_repeat: bool) -> list[Transform]: + """world->map->odom->base_link->sensor. Statics emitted once (latched) unless + static_repeat, in which case they're re-emitted each second.""" + store = SqliteStore(path=str(path)) + graph = TfGraphWriter(str(path), "tf") + written: list[Transform] = [] + + statics = [ + ("world", "map", (0.0, 0.0, 0.0)), + ("map", "odom", (0.0, 0.0, 0.0)), + ("base_link", "sensor", (0.0, 0.0, 0.3)), + ] + static_times = [_T0 + j for j in range(int(_DURATION))] if static_repeat else [_T0] + for ts in static_times: + for parent, child, xyz in statics: + transform = _static(parent, child, xyz, ts) + _append(store, graph, transform, is_static=True) + written.append(transform) + + for i in range(int(_DURATION * _DYN_RATE)): + ts = _T0 + i / _DYN_RATE + transform = Transform( + translation=Vector3(0.5 * i / _DYN_RATE, 0.1 * i / _DYN_RATE, 0.0), + rotation=_yaw(0.02 * i), + frame_id="odom", + child_frame_id="base_link", + ts=ts, + ) + _append(store, graph, transform, is_static=False) + written.append(transform) + + graph.close() + store.stop() + return written + - # query at off-sample times to exercise interpolation - queries = [_T0 + 0.013 + k * 0.317 for k in range(25)] +def test_interpolates_and_matches_full_load(tmp_path: Path) -> None: + """At off-sample (interpolated) query times, DbTf matches a naive full-load + buffer within tolerance — proving the chain compose + per-edge interpolation.""" + transforms = _record_single_robot(tmp_path / "r.db", static_repeat=True) + reference = _ref(transforms) + store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) + db = DbTf(store) compared = 0 - for q in queries: + for k in range(25): + q = _T0 + 0.013 + k * 0.317 want = reference.lookup("world", "sensor", q, 0.5) got = db.get("world", "sensor", q, 0.5) assert (want is None) == (got is None), f"None mismatch at {q}" if want is not None and got is not None: - assert _diff(want, got) < 1e-6, f"diff too large at {q}: {_diff(want, got)}" + assert _diff(want, got) < 1e-6, f"diff at {q}: {_diff(want, got)}" compared += 1 - assert compared >= 20 # actually exercised the interp path + assert compared >= 20 store.stop() -def test_no_full_load(tmp_path: Path) -> None: - """get() reads only the handful of tree-referenced rows, never the whole table.""" - transforms = _write_tf_db(tmp_path / "tf.db", with_trees=True) - store = SqliteStore(path=str(tmp_path / "tf.db"), must_exist=True) +def test_latched_static_resolves(tmp_path: Path) -> None: + """A static recorded once at the very start still resolves at a much later time + (no bracket, no tolerance) — the case a plain time-bracket would drop.""" + _record_single_robot(tmp_path / "r.db", static_repeat=False) + store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) db = DbTf(store) - queries = [_T0 + 0.05 + k * 0.5 for k in range(20)] - for q in queries: - db.get("world", "sensor", q, 0.5) - # 4 frames x <=2 refs = <=8 rows per distinct query; total must be far below - # the full table (len(transforms) rows). - assert db.rows_fetched <= 8 * len(queries) - assert db.rows_fetched < len(transforms) // 2 + assert db.get("world", "sensor", _T0 + 9.5, 0.5) is not None # ~9.5s after statics store.stop() -def test_fifo_cache_avoids_db_reads(tmp_path: Path) -> None: - """Repeated identical lookups are served from the FIFO with zero extra reads.""" - _write_tf_db(tmp_path / "tf.db", with_trees=True) - store = SqliteStore(path=str(tmp_path / "tf.db"), must_exist=True) - db = DbTf(store) - q = _T0 + 2.34 - first = db.get("world", "sensor", q, 0.5) - after_first = db.rows_fetched - assert after_first > 0 - for _ in range(10): - again = db.get("world", "sensor", q, 0.5) - assert _diff(first, again) == 0.0 - assert db.rows_fetched == after_first # no further DB reads - store.stop() - +def test_reparent_midrun_uses_graph_as_of_query_time(tmp_path: Path) -> None: + """world->map(+10) and map->odom(+100) are non-identity statics; at t0+5 a + relocalization re-parents base_link from odom to map. An early lookup composes + the odom branch (+100), a late one does not (+10) — wrong topology is off by + ~100. Exercises time-varying / multi-robot topology.""" + path = tmp_path / "reparent.db" + store = SqliteStore(path=str(path)) + graph = TfGraphWriter(str(path), "tf") + _append(store, graph, _static("world", "map", (10.0, 0, 0), _T0), is_static=True) + _append(store, graph, _static("map", "odom", (100.0, 0, 0), _T0), is_static=True) -def test_fifo_evicts_beyond_capacity(tmp_path: Path) -> None: - """A lookup pushed out of the FIFO costs a DB read again; one still in it does not.""" - _write_tf_db(tmp_path / "tf.db", with_trees=True) - store = SqliteStore(path=str(tmp_path / "tf.db"), must_exist=True) - db = DbTf(store, cache_size=3) - keys = [_T0 + 1.0 + k for k in range(3)] - for q in keys: - db.get("world", "sensor", q, 0.5) - reads_before = db.rows_fetched - # evict keys[0] by adding two fresh distinct lookups (capacity 3) - db.get("world", "sensor", _T0 + 7.1, 0.5) - db.get("world", "sensor", _T0 + 7.2, 0.5) - # keys[0] is gone -> re-reads; an untouched recent one would not - reads_mid = db.rows_fetched - db.get("world", "sensor", keys[0], 0.5) - assert db.rows_fetched > reads_mid # evicted -> DB read happened - assert reads_before > 0 + era1: list[Transform] = [] + for i in range(150): # [t0, t0+5) + ts = _T0 + i / _DYN_RATE + transform = Transform( + translation=Vector3(0.5 * i / _DYN_RATE, 0, 0), + rotation=_yaw(0.01 * i), + frame_id="odom", + child_frame_id="base_link", + ts=ts, + ) + _append(store, graph, transform, is_static=False) + era1.append(transform) + switch = _T0 + 5.0 + era2: list[Transform] = [] + for i in range(150): # [t0+5, t0+10): base_link now hangs off map directly + ts = switch + i / _DYN_RATE + transform = Transform( + translation=Vector3(7.0 + 0.3 * i / _DYN_RATE, 0, 0), + rotation=_yaw(0.01 * i), + frame_id="map", + child_frame_id="base_link", + ts=ts, + ) + _append(store, graph, transform, is_static=False) + era2.append(transform) + graph.close() store.stop() - -def test_recompute_warns_and_backfills(tmp_path: Path) -> None: - """Opening a tree-less recording warns loudly and back-fills a tree per row.""" - transforms = _write_tf_db(tmp_path / "tf.db", with_trees=False) - db_path = tmp_path / "tf.db" - # no tree table yet - conn = sqlite3.connect(str(db_path)) - has_tree = conn.execute( - "SELECT name FROM sqlite_master WHERE type='table' AND name='tf_tree'" - ).fetchone() - conn.close() - assert has_tree is None - - # The dimos logger is a structlog bound logger writing to the original stdout - # and not propagating, so capture by wrapping its .warning directly. - messages: list[str] = [] - original_warning = db_tf_module.logger.warning - db_tf_module.logger.warning = lambda msg, *a, **k: messages.append(str(msg)) - - store = SqliteStore(path=str(db_path), must_exist=True) + store = SqliteStore(path=str(path), must_exist=True) db = DbTf(store) - try: - db.get("world", "sensor", _T0 + 1.0, 0.5) - finally: - db_tf_module.logger.warning = original_warning - assert any("tf tree MISSING" in m for m in messages) # warned loudly - - conn = sqlite3.connect(str(db_path)) - n_rows = conn.execute("SELECT count(*) FROM tf").fetchone()[0] - n_trees = conn.execute("SELECT count(*) FROM tf_tree").fetchone()[0] - conn.close() - assert n_rows == len(transforms) - assert n_trees == n_rows # every tf row has a tree - store.stop() - - -def test_recorder_tree_has_last_two_refs_per_frame(tmp_path: Path) -> None: - """The recorder's incremental tree stores the last two row ids per child frame, - and those ids are keys into the tf table.""" - db_path = tmp_path / "tf.db" - store = SqliteStore(path=str(db_path)) - tf_stream = store.stream("tf", TFMessage) - writer = TfTreeWriter(str(db_path), "tf") - ids_by_child: dict[str, list[int]] = {} - last_tf_id = -1 - for transform in _all_transforms(): - obs = tf_stream.append(TFMessage(transform), ts=transform.ts, pose=None) - writer.record(transform.child_frame_id, obs.id) - ids_by_child.setdefault(transform.child_frame_id, []).append(obs.id) - last_tf_id = obs.id - writer.close() - store.stop() - conn = sqlite3.connect(str(db_path)) - tree_json = conn.execute("SELECT tree FROM tf_tree WHERE tf_id=?", (last_tf_id,)).fetchone()[0] - valid_ids = {row[0] for row in conn.execute("SELECT id FROM tf")} - conn.close() - tree = json.loads(tree_json) - - for child, expected in ids_by_child.items(): - assert tree[child] == expected[-2:] # last two refs for that frame - for tf_id in tree[child]: - assert tf_id in valid_ids # tree edges are keys into the tf table - - -def test_recompute_matches_recorder(tmp_path: Path) -> None: - """Trees built incrementally by the recorder equal trees rebuilt by recompute.""" - recorder_db = tmp_path / "rec.db" - _write_tf_db(recorder_db, with_trees=True) - recompute_db = tmp_path / "recmp.db" - _write_tf_db(recompute_db, with_trees=False) - store = SqliteStore(path=str(recompute_db), must_exist=True) - recompute_trees(store) + q1 = _T0 + 2.013 + want1 = _ref( + [ + _static("world", "map", (10.0, 0, 0), q1), + _static("map", "odom", (100.0, 0, 0), q1), + *era1, + ] + ).lookup("world", "base_link", q1, 0.5) + got1 = db.get("world", "base_link", q1, 0.5) + assert want1 is not None and got1 is not None + assert _diff(want1, got1) < 1e-6 + assert got1.translation.x > 100.0 # odom branch present + + q2 = switch + 2.013 + want2 = _ref([_static("world", "map", (10.0, 0, 0), q2), *era2]).lookup( + "world", "base_link", q2, 0.5 + ) + got2 = db.get("world", "base_link", q2, 0.5) + assert want2 is not None and got2 is not None + assert _diff(want2, got2) < 1e-6 + assert got2.translation.x < 20.0 # odom branch gone store.stop() - def trees(path: Path) -> dict[int, str]: - conn = sqlite3.connect(str(path)) - out = {row[0]: row[1] for row in conn.execute("SELECT tf_id, tree FROM tf_tree")} - conn.close() - return out - - assert trees(recorder_db) == trees(recompute_db) - - -_CHINA_DB = Path.home() / "dimos_phase2_china" / "china_default.db" - -@pytest.mark.skipif(not _CHINA_DB.exists(), reason="china_default.db not present") -def test_performance_on_real_recording(tmp_path: Path) -> None: - """On a real ~38k-row recording: correct vs full-load, no full load, and the - tree path is much faster than loading everything per query batch.""" - import shutil - - db_path = tmp_path / "china.db" - shutil.copy(_CHINA_DB, db_path) - store = SqliteStore(path=str(db_path), must_exist=True) - - # reference full-load buffer; collect the dynamic frame's (odom->base_link) - # row timestamps — the realistic query pattern (scans align with odometry), - # where the snapshot tree brackets the moving frame exactly. - buffer = MultiTBuffer(buffer_size=1.0e15) - total_rows = 0 - dynamic_ts: list[float] = [] - for obs in store.stream("tf", TFMessage).order_by("ts"): - transforms = getattr(obs.data, "transforms", None) or [obs.data] - for transform in transforms: - buffer.receive_transform(transform) - total_rows += 1 - if transforms[0].child_frame_id == "base_link" and total_rows % 1000 == 0: - dynamic_ts.append(obs.ts) +def test_disjoint_multirobot_returns_none(tmp_path: Path) -> None: + """Two unconnected components (two robots, no shared frame) → a cross-component + query is None, an in-component query resolves.""" + path = tmp_path / "two.db" + store = SqliteStore(path=str(path)) + graph = TfGraphWriter(str(path), "tf") + for i in range(20): + ts = _T0 + i / _DYN_RATE + _append( + store, + graph, + Transform( + translation=Vector3(i * 0.1, 0, 0), + rotation=_yaw(0), + frame_id="worldA", + child_frame_id="baseA", + ts=ts, + ), + is_static=False, + ) + _append( + store, + graph, + Transform( + translation=Vector3(0, i * 0.1, 0), + rotation=_yaw(0), + frame_id="worldB", + child_frame_id="baseB", + ts=ts, + ), + is_static=False, + ) + graph.close() + store.stop() + store = SqliteStore(path=str(path), must_exist=True) db = DbTf(store) - max_diff = 0.0 - for q in dynamic_ts: - want = buffer.lookup("world", "mid360_link", q, 0.5) - got = db.get("world", "mid360_link", q, 0.5) - assert (want is None) == (got is None) - if want is not None and got is not None: - max_diff = max(max_diff, _diff(want, got)) - assert max_diff < 1e-6, f"tree get diverged from full-load by {max_diff}" - - # no full load: a fresh DbTf reads far fewer rows than the table holds - db2 = DbTf(store) - t_start = time.perf_counter() - for q in dynamic_ts: - db2.get("world", "mid360_link", q, 0.5) - elapsed = time.perf_counter() - t_start - assert db2.rows_fetched < total_rows // 100 # << whole table - assert elapsed < 2.0 # generous; in practice tens of ms for tens of queries + q = _T0 + 0.3 + assert db.get("baseA", "worldA", q, 0.5) is not None # same component + assert db.get("baseB", "baseA", q, 0.5) is None # different components store.stop() diff --git a/dimos/memory2/test_db_tf2.py b/dimos/memory2/test_db_tf2.py deleted file mode 100644 index 9d48172373..0000000000 --- a/dimos/memory2/test_db_tf2.py +++ /dev/null @@ -1,228 +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. - -"""Tests for DbTf2 (graph-stream tf lookups): the in-RAM graph cache vs query -fallback, correctness vs a full-load buffer, latched-static handling, and disjoint -(multi-robot) graphs. Data is written one-transform-per-row + child_frame tag + -topology change-log, exactly as the live recorder does.""" - -from __future__ import annotations - -import math -from pathlib import Path - -from dimos.memory2.db_tf2 import DbTf2, TfGraphWriter -from dimos.memory2.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.protocol.tf.tf import MultiTBuffer - -_T0 = 1000.0 -_DYN_RATE = 30.0 -_DURATION = 10.0 - - -def _yaw(theta: float) -> Quaternion: - return Quaternion(0.0, 0.0, math.sin(theta / 2.0), math.cos(theta / 2.0)) - - -def _static(parent: str, child: str, xyz: tuple[float, float, float], ts: float) -> Transform: - return Transform( - translation=Vector3(*xyz), - rotation=Quaternion(0, 0, 0, 1), - frame_id=parent, - child_frame_id=child, - ts=ts, - ) - - -def _append( - store: SqliteStore, graph: TfGraphWriter, transform: Transform, is_static: bool -) -> None: - """Record one transform exactly as the live recorder does: one row, tagged - child_frame, plus a topology-change row when the structure changes.""" - store.stream("tf", TFMessage).append( - TFMessage(transform), - ts=transform.ts, - pose=None, - tags={"child_frame": transform.child_frame_id}, - ) - graph.record(transform.child_frame_id, transform.frame_id, is_static, transform.ts) - - -def _record_single_robot(path: Path, *, static_repeat: bool) -> list[Transform]: - """world->map->odom->base_link->sensor. Statics emitted once (latched) unless - static_repeat, in which case they're re-emitted each second.""" - store = SqliteStore(path=str(path)) - graph = TfGraphWriter(str(path), "tf") - written: list[Transform] = [] - - statics = [ - ("world", "map", (0.0, 0.0, 0.0)), - ("map", "odom", (0.0, 0.0, 0.0)), - ("base_link", "sensor", (0.0, 0.0, 0.3)), - ] - static_times = [_T0 + j for j in range(int(_DURATION))] if static_repeat else [_T0] - for ts in static_times: - for parent, child, xyz in statics: - t = _static(parent, child, xyz, ts) - _append(store, graph, t, is_static=True) - written.append(t) - - for i in range(int(_DURATION * _DYN_RATE)): - ts = _T0 + i / _DYN_RATE - t = Transform( - translation=Vector3(0.5 * i / _DYN_RATE, 0.1 * i / _DYN_RATE, 0.0), - rotation=_yaw(0.02 * i), - frame_id="odom", - child_frame_id="base_link", - ts=ts, - ) - _append(store, graph, t, is_static=False) - written.append(t) - - graph.close() - store.stop() - return written - - -def _reference(transforms: list[Transform]) -> MultiTBuffer: - buffer = MultiTBuffer(buffer_size=1.0e15) - buffer.receive_transform(*transforms) - return buffer - - -def _diff(a: Transform, b: Transform) -> float: - return ( - abs(a.translation.x - b.translation.x) - + abs(a.translation.y - b.translation.y) - + abs(a.translation.z - b.translation.z) - + abs(a.rotation.x - b.rotation.x) - + abs(a.rotation.y - b.rotation.y) - + abs(a.rotation.z - b.rotation.z) - + abs(a.rotation.w - b.rotation.w) - ) - - -def test_correctness_vs_full_load(tmp_path: Path) -> None: - """With statics re-emitted (so a naive full-load buffer also resolves), DbTf2 - matches it within tolerance at off-sample (interpolated) query times.""" - transforms = _record_single_robot(tmp_path / "r.db", static_repeat=True) - reference = _reference(transforms) - store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) - db = DbTf2(store) - compared = 0 - for k in range(25): - q = _T0 + 0.013 + k * 0.317 - want = reference.lookup("world", "sensor", q, 0.5) - got = db.get("world", "sensor", q, 0.5) - assert (want is None) == (got is None), f"None mismatch at {q}" - if want is not None and got is not None: - assert _diff(want, got) < 1e-6, f"diff at {q}: {_diff(want, got)}" - compared += 1 - assert compared >= 20 - store.stop() - - -def test_graph_loaded_into_ram_when_few_changes(tmp_path: Path) -> None: - """A stable single-robot tree has few topology changes → graph is held in RAM → - lookups issue zero graph queries.""" - _record_single_robot(tmp_path / "r.db", static_repeat=True) - store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) - db = DbTf2(store, max_graph_changes_in_ram=20) - for k in range(30): - db.get("world", "sensor", _T0 + 0.05 + k * 0.3, 0.5) - assert db.graph_queries == 0 # never queried the graph table - - -def test_fallback_to_query_when_many_changes(tmp_path: Path) -> None: - """If there are at least the threshold many topology changes, DbTf2 does NOT - load them into RAM and instead queries the graph per lookup (multi-robot churn).""" - _record_single_robot(tmp_path / "r.db", static_repeat=True) # 4 topology changes - store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) - db = DbTf2(store, max_graph_changes_in_ram=2) # 4 >= 2 → fallback - n = 10 - for k in range(n): - db.get("world", "sensor", _T0 + 0.05 + k * 0.3, 0.5) - assert db.graph_queries == n # one graph query per lookup - - -def test_latched_static_resolves(tmp_path: Path) -> None: - """A static recorded once at the very start still resolves at a much later time - (no bracket, no tolerance) — the case a plain time-bracket would drop.""" - transforms = _record_single_robot(tmp_path / "r.db", static_repeat=False) - store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) - db = DbTf2(store) - q = _T0 + 9.5 # ~9.5 s after the statics were recorded - got = db.get("world", "sensor", q, 0.5) - assert got is not None # latched static did not get dropped - # base_link→sensor is a fixed +0.3 z; the chain should reflect it (sanity) - reference_dyn = _reference([t for t in transforms if t.child_frame_id == "base_link"]) - # world/map/odom are identity here, so world->base_link == odom->base_link - odom_to_base = reference_dyn.lookup("odom", "base_link", q, 0.5) - assert odom_to_base is not None - store.stop() - - -def test_disjoint_graph_returns_none(tmp_path: Path) -> None: - """Two unconnected components (two robots) → a cross-component query is None.""" - path = tmp_path / "two.db" - store = SqliteStore(path=str(path)) - graph = TfGraphWriter(str(path), "tf") - # robot A: worldA -> baseA ; robot B: worldB -> baseB (no shared frame) - for i in range(20): - ts = _T0 + i / _DYN_RATE - _append( - store, - graph, - Transform( - translation=Vector3(i * 0.1, 0, 0), - rotation=_yaw(0), - frame_id="worldA", - child_frame_id="baseA", - ts=ts, - ), - is_static=False, - ) - _append( - store, - graph, - Transform( - translation=Vector3(0, i * 0.1, 0), - rotation=_yaw(0), - frame_id="worldB", - child_frame_id="baseB", - ts=ts, - ), - is_static=False, - ) - graph.close() - store.stop() - - store = SqliteStore(path=str(path), must_exist=True) - db = DbTf2(store) - q = _T0 + 0.3 - assert db.get("baseA", "worldA", q, 0.5) is not None # same component: ok - assert db.get("baseB", "baseA", q, 0.5) is None # different components: no transform - store.stop() - - -def test_threshold_is_configurable(tmp_path: Path) -> None: - _record_single_robot(tmp_path / "r.db", static_repeat=True) - store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) - assert DbTf2(store)._max_in_ram == 20 # default - assert DbTf2(store, max_graph_changes_in_ram=5)._max_in_ram == 5 - store.stop() diff --git a/dimos/memory2/test_db_tf2_stress.py b/dimos/memory2/test_db_tf2_stress.py deleted file mode 100644 index dc51c4436a..0000000000 --- a/dimos/memory2/test_db_tf2_stress.py +++ /dev/null @@ -1,315 +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. - -"""Adversarial + scale tests for DbTf2 — deliberately built to break it: - -* topology that *changes* mid-recording (relocalization re-parents a frame): a - lookup must resolve against the graph as-of its own time, not the latest one. -* many topology changes (multi-robot churn) forcing the query-per-lookup fallback - — correctness must survive the fallback path, not just the in-RAM path. -* a really large / deep tf tree and a long recording — correctness vs a full-load - buffer, and per-lookup cost bounded by chain depth, not by row count. -* edge cases: query before the first / after the last sample, a dynamic frame with - a single sample, tolerance rejection, and frame-name collision across robots. -""" - -from __future__ import annotations - -from pathlib import Path -import time - -from dimos.memory2.db_tf2 import DbTf2, TfGraphWriter -from dimos.memory2.store.sqlite import SqliteStore -from dimos.memory2.test_db_tf2 import _append, _diff, _yaw -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.protocol.tf.tf import MultiTBuffer - -_T0 = 1000.0 -_NO_PRUNE = 1.0e15 - - -def _dyn(parent: str, child: str, x: float, ts: float) -> Transform: - return Transform( - translation=Vector3(x, 0.0, 0.0), - rotation=_yaw(0.01 * x), - frame_id=parent, - child_frame_id=child, - ts=ts, - ) - - -def _stat(parent: str, child: str, x: float, ts: float) -> Transform: - return Transform( - translation=Vector3(x, 0.0, 0.0), - rotation=Quaternion(0, 0, 0, 1), - frame_id=parent, - child_frame_id=child, - ts=ts, - ) - - -def _ref(transforms: list[Transform]) -> MultiTBuffer: - buffer = MultiTBuffer(buffer_size=_NO_PRUNE) - buffer.receive_transform(*transforms) - return buffer - - -# --------------------------------------------------------------------------- # -# 1. Topology that CHANGES mid-run (relocalization re-parents base_link). -# --------------------------------------------------------------------------- # -def test_reparent_midrun_uses_graph_as_of_query_time(tmp_path: Path) -> None: - """world->map(+10) and map->odom(+100) are non-identity statics. base_link is - dynamic; at t0+5 a 'relocalization' re-parents base_link from odom to map. A - lookup at an early time must compose the odom branch (+100); a lookup after the - switch must NOT (+10 only). Picking the wrong era is off by ~100 -> caught.""" - path = tmp_path / "reparent.db" - store = SqliteStore(path=str(path)) - graph = TfGraphWriter(str(path), "tf") - - _append(store, graph, _stat("world", "map", 10.0, _T0), is_static=True) - _append(store, graph, _stat("map", "odom", 100.0, _T0), is_static=True) - - era1: list[Transform] = [] - for i in range(150): # [t0, t0+5) at 30Hz - ts = _T0 + i / 30.0 - t = _dyn("odom", "base_link", 0.5 * i / 30.0, ts) - _append(store, graph, t, is_static=False) - era1.append(t) - - switch = _T0 + 5.0 - era2: list[Transform] = [] - for i in range(150): # [t0+5, t0+10): base_link now hangs off map directly - ts = switch + i / 30.0 - t = _dyn("map", "base_link", 7.0 + 0.3 * i / 30.0, ts) - _append(store, graph, t, is_static=False) - era2.append(t) - graph.close() - store.stop() - - store = SqliteStore(path=str(path), must_exist=True) - db = DbTf2(store) - - # early lookup: ground truth = statics(+10,+100) ∘ era1 dynamic - q1 = _T0 + 2.013 - ref1 = _ref([_stat("world", "map", 10.0, q1), _stat("map", "odom", 100.0, q1), *era1]) - want1 = ref1.lookup("world", "base_link", q1, 0.5) - got1 = db.get("world", "base_link", q1, 0.5) - assert want1 is not None and got1 is not None - assert _diff(want1, got1) < 1e-6 - assert got1.translation.x > 100.0 # odom branch present - - # late lookup: ground truth = world->map(+10) ∘ era2 dynamic (no odom) - q2 = switch + 2.013 - ref2 = _ref([_stat("world", "map", 10.0, q2), *era2]) - want2 = ref2.lookup("world", "base_link", q2, 0.5) - got2 = db.get("world", "base_link", q2, 0.5) - assert want2 is not None and got2 is not None - assert _diff(want2, got2) < 1e-6 - assert got2.translation.x < 20.0 # odom branch gone - - store.stop() - - -# --------------------------------------------------------------------------- # -# 2. Too many topology updates -> query-per-lookup fallback path. -# --------------------------------------------------------------------------- # -def test_many_topology_changes_force_fallback_and_stay_correct(tmp_path: Path) -> None: - """N robots each join with their own frames -> many tf_graph rows. With the - threshold below N, the graph is NOT held in RAM (graph_queries fire). Correctness - must hold on the fallback path, including resolving the LATEST topology.""" - path = tmp_path / "churn.db" - store = SqliteStore(path=str(path)) - graph = TfGraphWriter(str(path), "tf") - - robot_count = 40 - last_per_robot: dict[int, Transform] = {} - for robot in range(robot_count): - # each robot adds a new child frame at a distinct time -> a topology change - for i in range(5): - ts = _T0 + robot * 1.0 + i / 30.0 - t = _dyn(f"world_{robot}", f"base_{robot}", 0.5 * i + robot, ts) - _append(store, graph, t, is_static=False) - last_per_robot[robot] = t - graph.close() - store.stop() - - store = SqliteStore(path=str(path), must_exist=True) - db = DbTf2(store, max_graph_changes_in_ram=10) # 40 changes >= 10 -> fallback - - lookups = 0 - for robot in (0, 7, 20, 39): - q = _T0 + robot * 1.0 + 4 / 30.0 - ref = _ref([last_per_robot[robot]]) - want = ref.lookup(f"world_{robot}", f"base_{robot}", q, 0.5) - got = db.get(f"world_{robot}", f"base_{robot}", q, 0.5) - assert want is not None and got is not None, f"robot {robot}" - assert _diff(want, got) < 1e-6, f"robot {robot}: {_diff(want, got)}" - lookups += 1 - - assert db._graph_in_ram is None # confirmed NOT cached - assert db.graph_queries == lookups # one graph query per lookup - - -# --------------------------------------------------------------------------- # -# 3. Really large / deep tf tree + long recording. -# --------------------------------------------------------------------------- # -def test_large_deep_tree_correct_and_bounded(tmp_path: Path) -> None: - """A 30-link deep chain (world->link_0->...->link_29), each link dynamic at 30Hz - for 60s = ~54k rows. DbTf2 must (a) match a full-load buffer and (b) keep - per-lookup latency bounded by chain depth, not the 54k row count.""" - path = tmp_path / "big.db" - store = SqliteStore(path=str(path)) - graph = TfGraphWriter(str(path), "tf") - - depth = 30 - duration_s = 60 - rate_hz = 30 - all_transforms: list[Transform] = [] - parents = ["world"] + [f"link_{d}" for d in range(depth)] - for step in range(duration_s * rate_hz): - ts = _T0 + step / rate_hz - for d in range(depth): - t = _dyn(parents[d], f"link_{d}", 0.1 * d + 0.01 * step, ts) - _append(store, graph, t, is_static=False) - all_transforms.append(t) - graph.close() - store.stop() - - row_count = depth * duration_s * rate_hz - - store = SqliteStore(path=str(path), must_exist=True) - db = DbTf2(store) - reference = _ref(all_transforms) - - target = f"link_{depth - 1}" - compared = 0 - for k in range(25): - q = _T0 + 0.017 + k * (duration_s / 26.0) - want = reference.lookup("world", target, q, 0.5) - got = db.get("world", target, q, 0.5) - assert (want is None) == (got is None), f"None mismatch at {q}" - if want is not None and got is not None: - assert _diff(want, got) < 1e-6, f"diff at {q}: {_diff(want, got)}" - compared += 1 - assert compared >= 20 - - # bounded latency: warm, then time. Should be sub-millisecond despite 54k rows. - for k in range(50): - db.get("world", target, _T0 + 5.0 + k * 0.1, 0.5) - n = 300 - start = time.perf_counter() - for k in range(n): - db.get("world", target, _T0 + 10.0 + (k % 400) * 0.1, 0.5) - per_lookup_us = (time.perf_counter() - start) / n * 1e6 - assert per_lookup_us < 3000.0, f"per-lookup {per_lookup_us:.0f}us — not row-count bounded" - print(f"\nlarge tree: {row_count} rows, depth {depth}, per-lookup {per_lookup_us:.0f}us") - - store.stop() - - -# --------------------------------------------------------------------------- # -# 4. Edge cases meant to break it. -# --------------------------------------------------------------------------- # -def _single_chain(path: Path) -> list[Transform]: - store = SqliteStore(path=str(path)) - graph = TfGraphWriter(str(path), "tf") - written: list[Transform] = [] - written.append(_stat("world", "odom", 1.0, _T0)) - _append(store, graph, written[-1], is_static=True) - for i in range(100): # [t0, t0+~3.3s) - ts = _T0 + i / 30.0 - t = _dyn("odom", "base_link", 0.5 * i / 30.0, ts) - _append(store, graph, t, is_static=False) - written.append(t) - graph.close() - store.stop() - return written - - -def test_query_before_topology_returns_none_cleanly(tmp_path: Path) -> None: - """BEHAVIORAL DIFFERENCE vs the old full-load DbTf, pinned here on purpose: - a query *before a frame ever appears* in the topology log resolves to None (not - a crash). base_link's first topology entry coincides with its first sample, so a - time earlier than the recording start has no base_link in the graph-as-of-then. - This is defensible (the frame did not exist yet) but the full-load buffer would - still answer it — flag on the DbTf->DbTf2 swap.""" - _single_chain(tmp_path / "r.db") - store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) - db = DbTf2(store) - assert db.get("world", "base_link", _T0 - 100.0, 1e9) is None # clean, no crash - # but a query at the very first sample (frame now exists) DOES resolve: - assert db.get("world", "base_link", _T0, 1e9) is not None - store.stop() - - -def test_query_after_last_sample(tmp_path: Path) -> None: - """A time after the last sample: no 'hi' bracket; should resolve from the latest - sample (lo) when tolerance allows.""" - _single_chain(tmp_path / "r.db") - store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) - db = DbTf2(store) - got = db.get("world", "base_link", _T0 + 100.0, 1e9) - assert got is not None - store.stop() - - -def test_tolerance_rejects_far_dynamic(tmp_path: Path) -> None: - """A dynamic frame whose nearest sample is far outside tolerance must be rejected - (None), not silently snapped to a stale value.""" - _single_chain(tmp_path / "r.db") - store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) - db = DbTf2(store) - got = db.get("world", "base_link", _T0 + 100.0, 0.01) # 100s away, 10ms tol - assert got is None - store.stop() - - -def test_single_dynamic_sample(tmp_path: Path) -> None: - """A dynamic frame with exactly one sample (no second bracket) must still resolve - at that sample's time.""" - path = tmp_path / "one.db" - store = SqliteStore(path=str(path)) - graph = TfGraphWriter(str(path), "tf") - only = _dyn("world", "base_link", 3.0, _T0) - _append(store, graph, only, is_static=False) - graph.close() - store.stop() - store = SqliteStore(path=str(path), must_exist=True) - db = DbTf2(store) - got = db.get("world", "base_link", _T0, 0.5) - assert got is not None and abs(got.translation.x - 3.0) < 1e-6 - store.stop() - - -def test_frame_name_collision_across_robots_is_documented(tmp_path: Path) -> None: - """KNOWN LIMITATION probe: two robots that reuse the SAME frame names share rows - in the tf table (child_frame is the only key). This test pins the CURRENT - behavior so a future fix is a deliberate, visible change — not a silent one.""" - path = tmp_path / "collide.db" - store = SqliteStore(path=str(path)) - graph = TfGraphWriter(str(path), "tf") - # both robots publish odom->base_link with DIFFERENT motion at the same times - for i in range(30): - ts = _T0 + i / 30.0 - _append(store, graph, _dyn("odom", "base_link", 0.5 * i, ts), is_static=False) - _append(store, graph, _dyn("odom", "base_link", -0.5 * i, ts + 1e-4), is_static=False) - graph.close() - store.stop() - store = SqliteStore(path=str(path), must_exist=True) - db = DbTf2(store) - got = db.get("odom", "base_link", _T0 + 0.5, 0.5) - assert got is not None # resolves to *a* sample; which one is undefined-by-design - store.stop() From 75e2235a1ae965c06f542816281aafa42121ef3c Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sun, 28 Jun 2026 01:57:05 +0800 Subject: [PATCH 35/69] refactor(db_tf): make tf_graph a first-class stream The topology change-log was a raw sqlite side-table written directly by TfGraphWriter. Promote it to a real memory2 stream ("_graph") so it's discoverable via list_streams(), replayable, and uses the standard codec/blob machinery like every other recorded stream. - new TfGraph payload (recording-internal, pickle-codec); TfGraphWriter appends snapshots via store.stream() instead of raw INSERTs. - DbTf reads the graph through the stream API (count + order_by + time_range) instead of raw SQL; migration build_graph_stream appends to the stream. - recorder records the graph stream for all store types, not just sqlite. - drop the raw-table helpers (ensure_graph_table/_graph_table). --- dimos/memory2/db_tf.py | 152 ++++++++++++++++++------------------ dimos/memory2/module.py | 8 +- dimos/memory2/test_db_tf.py | 6 +- 3 files changed, 83 insertions(+), 83 deletions(-) diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py index 7c16430094..599aed7f39 100644 --- a/dimos/memory2/db_tf.py +++ b/dimos/memory2/db_tf.py @@ -38,7 +38,6 @@ from __future__ import annotations import bisect -import json import re import sqlite3 import threading @@ -56,11 +55,13 @@ if TYPE_CHECKING: from dimos.memory2.store.base import Store + from dimos.memory2.stream import Stream logger = setup_logger() DEFAULT_TF_STREAM = "tf" -GRAPH_TABLE = "tf_graph" +# The topology change-log is a first-class stream named "_graph". +GRAPH_STREAM_SUFFIX = "_graph" # Streams the RAM fallback (non-sqlite stores) reads. TF_STREAMS = ("tf", "tf_static") # If a recording has fewer than this many topology changes, load them all into RAM @@ -69,11 +70,17 @@ DEFAULT_MAX_GRAPH_CHANGES_IN_RAM = 20 # Larger than any single recording's span so the fallback buffer never prunes. _NO_PRUNE = 1.0e15 +# Lower bound for "latest topology at or before T" range queries. +_TS_MIN = -1.0e18 # SQLite can't parameterize table names, so caller-supplied stream names are # interpolated; allow only safe identifiers to keep that injection-free. _SAFE_TABLE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +def _graph_stream_name(stream: str) -> str: + return f"{_safe_table(stream)}{GRAPH_STREAM_SUFFIX}" + + def _safe_table(name: str) -> str: if not _SAFE_TABLE.match(name): raise ValueError(f"unsafe stream/table name: {name!r}") @@ -88,22 +95,6 @@ def _connect(db_path: str) -> sqlite3.Connection: return conn -def _graph_table(stream: str) -> str: - return f"{_safe_table(stream)}_graph" - - -def ensure_graph_table(conn: sqlite3.Connection, stream: str) -> None: - """Create the topology change-log table (safe to call before the tf table - exists — it doesn't touch the tf table).""" - table = _graph_table(stream) - conn.execute( - f'CREATE TABLE IF NOT EXISTS "{table}" ' - "(id INTEGER PRIMARY KEY AUTOINCREMENT, ts REAL NOT NULL, structure TEXT NOT NULL)" - ) - conn.execute(f'CREATE INDEX IF NOT EXISTS "{table}_ts_idx" ON "{table}"(ts)') - conn.commit() - - def _ensure_child_index(conn: sqlite3.Connection, stream: str) -> None: """Index the child_frame json tag on the tf rows so per-frame time queries seek. The live recorder gets this for free (the store auto-indexes tag keys on @@ -119,41 +110,56 @@ def _ensure_child_index(conn: sqlite3.Connection, stream: str) -> None: conn.commit() +class TfGraph: + """A tf topology snapshot, recorded one per structure change. + + ``structure`` maps each child frame to ``{"parent": str, "static": bool}`` — + the full tf tree as of this message's timestamp. The stream of these snapshots + (the ``_graph`` stream) is the topology change-log that transform lookups + walk to resolve a source->target chain at any past time. Defined here (not under + ``dimos/msgs``) because it is a recording-internal payload, not a wire message; + it is stored via the pickle codec.""" + + structure: dict[str, dict[str, Any]] + msg_name = "tf2_msgs.TfGraph" + + def __init__(self, structure: dict[str, dict[str, Any]]) -> None: + # copy so later mutations of the writer's running structure don't alter an + # already-recorded snapshot + self.structure = {child: dict(entry) for child, entry in structure.items()} + + def __repr__(self) -> str: + return f"TfGraph({len(self.structure)} frames)" + + class TfGraphWriter: - """Recorder helper: tracks the running topology and appends a ``tf_graph`` row - only when the structure changes.""" + """Recorder helper: tracks the running topology and appends a ``TfGraph`` + snapshot to the graph stream only when the structure changes.""" - def __init__(self, db_path: str, stream: str = DEFAULT_TF_STREAM) -> None: - self._stream = _safe_table(stream) - self._table = _graph_table(stream) - self._conn = _connect(db_path) + def __init__(self, store: Store, stream: str = DEFAULT_TF_STREAM) -> None: + self._stream: Stream[TfGraph] = store.stream(_graph_stream_name(stream), TfGraph) self._structure: dict[str, dict[str, Any]] = {} - ensure_graph_table(self._conn, self._stream) def record(self, child_frame: str, parent_frame: str, is_static: bool, ts: float) -> None: entry = {"parent": parent_frame, "static": bool(is_static)} if self._structure.get(child_frame) == entry: - return # no structural change -> no new graph row + return # no structural change -> no new snapshot self._structure[child_frame] = entry - self._conn.execute( - f'INSERT INTO "{self._table}" (ts, structure) VALUES (?, ?)', - (ts, json.dumps(self._structure)), - ) - self._conn.commit() + self._stream.append(TfGraph(self._structure), ts=ts) def close(self) -> None: - self._conn.close() + # the stream is owned by the store; nothing to close here + pass def build_graph_stream(store: Store, stream: str = DEFAULT_TF_STREAM) -> int: """One-time migration for a recording that predates the graph stream: tag every - tf row with its ``child_frame`` and build ``tf_graph`` chronologically. A frame - is treated as static if its pose never changes across the recording. Returns the - number of topology-change rows written.""" + tf row with its ``child_frame`` and build the ``_graph`` stream + chronologically. A frame is treated as static if its pose never changes across + the recording. Returns the number of topology-change snapshots written.""" config = store.config if not isinstance(config, SqliteStoreConfig): raise TypeError("build_graph_stream needs a SqliteStore") - table = _graph_table(stream) safe = _safe_table(stream) # one decode pass: collect (id, ts, child, parent, pose-key) per row @@ -174,34 +180,32 @@ def build_graph_stream(store: Store, stream: str = DEFAULT_TF_STREAM) -> int: poses_per_child.setdefault(transform.child_frame_id, set()).add(pose_key) static_frames = {child for child, poses in poses_per_child.items() if len(poses) == 1} + # tag each tf row with its child_frame + add the seek index (raw, on the tf table) conn = _connect(config.path) try: - ensure_graph_table(conn, safe) - conn.execute(f'DELETE FROM "{table}"') - # tag each tf row with its child_frame (json_set keeps any existing tags) for row_id, _ts, child, _parent, _pose in rows: conn.execute( f"UPDATE \"{safe}\" SET tags = json_set(tags, '$.child_frame', ?) WHERE id = ?", (child, row_id), ) - _ensure_child_index(conn, safe) - # build the topology change-log - structure: dict[str, dict[str, Any]] = {} - written = 0 - for _row_id, ts, child, parent, _pose in rows: - entry = {"parent": parent, "static": child in static_frames} - if structure.get(child) == entry: - continue - structure[child] = entry - conn.execute( - f'INSERT INTO "{table}" (ts, structure) VALUES (?, ?)', (ts, json.dumps(structure)) - ) - written += 1 conn.commit() - return written + _ensure_child_index(conn, safe) finally: conn.close() + # build the topology change-log as a first-class stream + graph_stream = store.stream(_graph_stream_name(safe), TfGraph) + structure: dict[str, dict[str, Any]] = {} + written = 0 + for _row_id, ts, child, parent, _pose in rows: + entry = {"parent": parent, "static": child in static_frames} + if structure.get(child) == entry: + continue + structure[child] = entry + graph_stream.append(TfGraph(structure), ts=ts) + written += 1 + return written + class DbTf: """Transform lookups backed by a store's recorded transforms. @@ -220,7 +224,7 @@ def __init__( ) -> None: self._store = store self._stream = _safe_table(stream) - self._table = _graph_table(stream) + self._graph_name = _graph_stream_name(stream) self._max_in_ram = max_graph_changes_in_ram self._stream_names = stream_names # RAM fallback only self._lock = threading.Lock() @@ -277,14 +281,20 @@ def has_transforms(self) -> bool: (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() return bool(n_rows) + def _graph_stream(self) -> Stream[TfGraph]: + return self._store.stream(self._graph_name, TfGraph) + + def _graph_count(self) -> int: + if self._graph_name not in set(self._store.list_streams()): + return 0 + return self._graph_stream().count() + def _ensure_built(self) -> None: if self._built: return conn = self._connection() - ensure_graph_table(conn, self._stream) - (n_graph,) = conn.execute(f'SELECT count(*) FROM "{self._table}"').fetchone() (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() - if n_rows and n_graph == 0: + if n_rows and self._graph_count() == 0: logger.warning( "\n========================================================================\n" " tf graph stream MISSING for %r. Building it (one-time): tagging tf rows\n" @@ -301,14 +311,9 @@ def _ensure_built(self) -> None: def _load_graph_if_small(self) -> None: if self._graph_loaded: return - conn = self._connection() - (n_graph,) = conn.execute(f'SELECT count(*) FROM "{self._table}"').fetchone() - if n_graph < self._max_in_ram: + if self._graph_count() < self._max_in_ram: self._graph_in_ram = [ - (ts, json.loads(structure)) - for ts, structure in conn.execute( - f'SELECT ts, structure FROM "{self._table}" ORDER BY ts ASC' - ) + (obs.ts, obs.data.structure) for obs in self._graph_stream().order_by("ts") ] else: self._graph_in_ram = None # too many -> query per lookup @@ -322,18 +327,15 @@ def _graph_at(self, query_time: float) -> dict[str, Any] | None: if index < 0: return self._graph_in_ram[0][1] # before first -> earliest return self._graph_in_ram[index][1] - # fallback: one query + # fallback: one query for the latest snapshot at or before query_time self.graph_queries += 1 - conn = self._connection() - row = conn.execute( - f'SELECT structure FROM "{self._table}" WHERE ts <= ? ORDER BY ts DESC LIMIT 1', - (query_time,), - ).fetchone() - if row is None: - row = conn.execute( - f'SELECT structure FROM "{self._table}" ORDER BY ts ASC LIMIT 1' - ).fetchone() - return json.loads(row[0]) if row else None + stream = self._graph_stream() + latest = stream.time_range(_TS_MIN, query_time).order_by("ts", desc=True).limit(1) + for obs in latest: + return obs.data.structure + for obs in stream.order_by("ts").limit(1): # before first -> earliest + return obs.data.structure + return None def _chain_frames(self, graph: dict[str, Any], source: str, target: str) -> list[str] | None: def to_root(frame: str) -> list[str]: diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index 4ecb2ed7af..186e50108f 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -463,11 +463,9 @@ def _record_tf(self) -> None: return tf_stream = self.store.stream("tf", TFMessage) - is_sqlite = isinstance(self.store, SqliteStore) - # Topology change-log + child_frame tags, for the graph-stream DbTf. - graph_writer = TfGraphWriter(self.store.config.path, "tf") if is_sqlite else None - if graph_writer is not None: - self.register_disposable(Disposable(graph_writer.close)) + # Topology change-log (the "tf_graph" stream) written alongside the tf rows. + graph_writer = TfGraphWriter(self.store, "tf") + self.register_disposable(Disposable(graph_writer.close)) def make_handler(is_static: bool) -> Any: def on_tf(msg: TFMessage, _topic: Any) -> None: diff --git a/dimos/memory2/test_db_tf.py b/dimos/memory2/test_db_tf.py index 9fcafa7eb9..b5a077829f 100644 --- a/dimos/memory2/test_db_tf.py +++ b/dimos/memory2/test_db_tf.py @@ -87,7 +87,7 @@ def _record_single_robot(path: Path, *, static_repeat: bool) -> list[Transform]: """world->map->odom->base_link->sensor. Statics emitted once (latched) unless static_repeat, in which case they're re-emitted each second.""" store = SqliteStore(path=str(path)) - graph = TfGraphWriter(str(path), "tf") + graph = TfGraphWriter(store) written: list[Transform] = [] statics = [ @@ -156,7 +156,7 @@ def test_reparent_midrun_uses_graph_as_of_query_time(tmp_path: Path) -> None: ~100. Exercises time-varying / multi-robot topology.""" path = tmp_path / "reparent.db" store = SqliteStore(path=str(path)) - graph = TfGraphWriter(str(path), "tf") + graph = TfGraphWriter(store) _append(store, graph, _static("world", "map", (10.0, 0, 0), _T0), is_static=True) _append(store, graph, _static("map", "odom", (100.0, 0, 0), _T0), is_static=True) @@ -220,7 +220,7 @@ def test_disjoint_multirobot_returns_none(tmp_path: Path) -> None: query is None, an in-component query resolves.""" path = tmp_path / "two.db" store = SqliteStore(path=str(path)) - graph = TfGraphWriter(str(path), "tf") + graph = TfGraphWriter(store) for i in range(20): ts = _T0 + i / _DYN_RATE _append( From 43b872e577fcdd61d6d521140c4fa534212514c9 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sun, 28 Jun 2026 08:35:48 +0800 Subject: [PATCH 36/69] feat(eval): use tf for the raw-baseline trajectory instead of odom poses The loop-closure eval's raw baseline (tag agreement, voxel re-anchoring, drift recovery) now composes the robot pose from the recording's tf (world->base_link, sampled at each odom time) when the recording has tf; recordings without tf fall back to the odom stream's stored poses. Adds --world-frame/--body-frame, records pose_source in the summary, and bumps EVAL_VERSION to invalidate cached cells. --- .../jnav/components/loop_closure/eval.py | 123 ++++++++++++++---- 1 file changed, 98 insertions(+), 25 deletions(-) diff --git a/dimos/navigation/jnav/components/loop_closure/eval.py b/dimos/navigation/jnav/components/loop_closure/eval.py index 27902331b0..8c0510e698 100644 --- a/dimos/navigation/jnav/components/loop_closure/eval.py +++ b/dimos/navigation/jnav/components/loop_closure/eval.py @@ -14,11 +14,14 @@ """Evaluate a loop-closure module against a recording. +The raw-baseline robot trajectory comes from the recording's tf (the +``world``->``base_link`` chain, sampled at each odom time) when the recording has +tf; recordings without tf fall back to the odom stream's stored poses. + Two ground-truth-free scores, before vs after correction: * April-tag agreement — a fixed tag re-seen along the run should map to one world position; the spread of its per-visit robot positions measures drift. - Tags are taken as relative to the chosen odom stream (sighting time -> - nearest odom pose), so no static transforms or stored db poses are needed. + A tag sighting is placed at the robot's baseline pose nearest its time. * Lidar-voxel agreement — re-anchoring the registered scans onto the corrected trajectory should collapse double walls, so the corrected map should occupy FEWER voxels than the raw one. @@ -93,18 +96,17 @@ REPLAY_PUBLISH_HZ, iterate_stream, list_streams, - odometry_lookup, store, stream_count, ) from dimos.navigation.jnav.utils.trajectory_metrics import ( GraphPose, - PoseLookup7, drifted_lookup, graph_lookup, has_drift, lidar_voxel_agreement, pose7_lookup, + trajectory_lookup, trajectory_recovery_error, ) @@ -116,7 +118,9 @@ VOXEL_MAX_SCANS = 300 # Bump to invalidate every cached cell (scoring/replay semantics changed). -EVAL_VERSION = 1 +# v2: raw baseline is now composed from tf (world->body) when present, not the +# odom stream's stored poses. +EVAL_VERSION = 2 def cell_fingerprint( @@ -538,7 +542,12 @@ def run_module_graph( return graph, int(data["closures"]), replay_stats # type: ignore[return-value] -def odometry_pose7_lookup(db_path: Path, odom_stream: str) -> PoseLookup7: +# Tolerance for composing world->body from the recording's tf at an odom time. +TF_LOOKUP_TOLERANCE_S = 0.1 + + +def _odom_pose_samples(db_path: Path, odom_stream: str) -> tuple[np.ndarray, np.ndarray]: + """Robot trajectory straight from the odometry stream's stored poses.""" times: list[float] = [] poses: list[list[float]] = [] for timestamp, payload in iterate_stream(db_path, odom_stream): @@ -555,13 +564,61 @@ def odometry_pose7_lookup(db_path: Path, odom_stream: str) -> PoseLookup7: pose.orientation.w, ] ) - return pose7_lookup( + return ( + np.asarray(times, dtype=np.float64), + np.asarray(poses, dtype=np.float64).reshape(-1, 7), + ) + + +def _tf_pose_samples( + db_path: Path, odom_stream: str, world_frame: str, body_frame: str +) -> tuple[np.ndarray, np.ndarray] | None: + """Robot trajectory composed from the recording's tf (``world_frame -> + body_frame``) sampled at each odom timestamp. ``None`` if the recording has no + tf, so the caller can fall back to raw odom poses.""" + db_store = store(db_path) + if "tf" not in db_store.list_streams() or not db_store.tf.has_transforms(): + return None + times: list[float] = [] + poses: list[list[float]] = [] + for timestamp, _payload in iterate_stream(db_path, odom_stream): + transform = db_store.tf.get(world_frame, body_frame, timestamp, TF_LOOKUP_TOLERANCE_S) + if transform is None: + continue + times.append(timestamp) + poses.append( + [ + transform.translation.x, + transform.translation.y, + transform.translation.z, + transform.rotation.x, + transform.rotation.y, + transform.rotation.z, + transform.rotation.w, + ] + ) + if not times: + return None + return ( np.asarray(times, dtype=np.float64), - np.asarray(poses, dtype=np.float64), - ODOM_MATCH_TOLERANCE_S, + np.asarray(poses, dtype=np.float64).reshape(-1, 7), ) +def raw_pose_samples( + db_path: Path, odom_stream: str, *, world_frame: str, body_frame: str +) -> tuple[np.ndarray, np.ndarray, str]: + """The raw-baseline trajectory as ``(times, [x,y,z,qx,qy,qz,qw], source)``. + + Prefers the recording's tf (``world_frame -> body_frame``, world-registered); + falls back to the odometry stream's stored poses when there is no tf.""" + tf_samples = _tf_pose_samples(db_path, odom_stream, world_frame, body_frame) + if tf_samples is not None: + return tf_samples[0], tf_samples[1], "tf" + times, poses = _odom_pose_samples(db_path, odom_stream) + return times, poses, "odom" + + def _subsampled_path(positions: np.ndarray) -> np.ndarray: stride = max(1, len(positions) // _RRD_MAX_PATH_POINTS) return positions[::stride] @@ -614,6 +671,8 @@ def evaluate( recording_name: str | None = None, drift_per_sec: list[float] | None = None, ignore_tags: set[int] | None = None, + world_frame: str = "world", + body_frame: str = "base_link", ) -> dict[str, Any]: streams = list_streams(db_path) for required in (odom_stream, lidar_stream): @@ -690,12 +749,21 @@ def detect() -> Iterable[tuple[int, float]]: if not graph: raise SystemExit(f"{module_name} produced an empty pose graph") + # Raw-baseline trajectory: composed from the recording's tf (world->body) when + # present, else the odom stream's stored poses. + raw_times, raw_poses7, pose_source = raw_pose_samples( + db_path, odom_stream, world_frame=world_frame, body_frame=body_frame + ) + if raw_poses7.size == 0: + raise SystemExit(f"no raw baseline poses (tf nor {odom_stream!r}) in {db_path}") + print(f"raw baseline pose source: {pose_source} ({len(raw_times)} samples)") + raw_xyz_base = trajectory_lookup(raw_times, raw_poses7[:, :3], ODOM_MATCH_TOLERANCE_S) + raw_pose7_base = pose7_lookup(raw_times, raw_poses7, ODOM_MATCH_TOLERANCE_S) + # The module solved on drifted input, so its graph lives in the drifted # world; the raw baselines must be drifted to match (see drift_per_sec). - raw_xyz_lookup = drifted_lookup(odometry_lookup(db_path, odom_stream), drift_per_sec, drift_t0) - raw_pose7_lookup = drifted_lookup( - odometry_pose7_lookup(db_path, odom_stream), drift_per_sec, drift_t0 - ) + raw_xyz_lookup = drifted_lookup(raw_xyz_base, drift_per_sec, drift_t0) + raw_pose7_lookup = drifted_lookup(raw_pose7_base, drift_per_sec, drift_t0) xyz_graph = [(node[0], node[1], node[2], node[3]) for node in graph] if sightings: @@ -728,9 +796,7 @@ def detect() -> Iterable[tuple[int, float]]: # Drift-recovery ATE: corrected trajectory vs the UN-drifted ground truth # (the odom before drift was injected). Only meaningful with --drift-per-sec; # the right metric where tag/voxel agreement is weak (e.g. KITTI's long loop). - trajectory = trajectory_recovery_error( - graph, odometry_lookup(db_path, odom_stream), drift_per_sec, drift_t0 - ) + trajectory = trajectory_recovery_error(graph, raw_xyz_base, drift_per_sec, drift_t0) if trajectory is not None: print( f" drift recovery: {trajectory['drifted_ate_m']:.2f}" @@ -751,15 +817,7 @@ def detect() -> Iterable[tuple[int, float]]: out_dir.mkdir(parents=True, exist_ok=True) rrd_path = out_dir / "eval.rrd" if with_rrd: - raw_positions = np.asarray( - [ - [pose.position.x, pose.position.y, pose.position.z] - for _, payload in iterate_stream(db_path, odom_stream) - for pose in [RateReplay._payload_pose(payload)] - ], - dtype=np.float64, - ) - write_trajectory_rrd(rrd_path, raw_positions, graph) + write_trajectory_rrd(rrd_path, raw_poses7[:, :3], graph) summary = { "db": str(db_path), @@ -773,6 +831,9 @@ def detect() -> Iterable[tuple[int, float]]: db_path, pgo_config, lidar_stream, odom_stream, drift_per_sec ), "replay": replay_stats, + "pose_source": pose_source, + "world_frame": world_frame, + "body_frame": body_frame, "april_tags": { "source": tag_source, "sightings": n_sightings, @@ -873,6 +934,16 @@ def main() -> None: help="comma-separated April-tag ids to drop from scoring (dynamic/unreliable " "tags whose motion would look like drift). e.g. '17'", ) + parser.add_argument( + "--world-frame", + default="world", + help="tf frame the raw baseline is expressed in (when the recording has tf)", + ) + parser.add_argument( + "--body-frame", + default="base_link", + help="tf frame tracking the robot body (the raw baseline is world_frame->body_frame)", + ) args = parser.parse_args() drift_per_sec = ( @@ -915,6 +986,8 @@ def main() -> None: recording_name=args.recording_name, drift_per_sec=drift_per_sec, ignore_tags=ignore_tags, + world_frame=args.world_frame, + body_frame=args.body_frame, ) From 433e3f55d209718184772696a7eead0bdbaaad83 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sun, 28 Jun 2026 08:53:05 +0800 Subject: [PATCH 37/69] feat(eval): require tf (no odom fallback) + startup tolerance window --- .../jnav/components/loop_closure/eval.py | 112 +++++++++--------- 1 file changed, 56 insertions(+), 56 deletions(-) diff --git a/dimos/navigation/jnav/components/loop_closure/eval.py b/dimos/navigation/jnav/components/loop_closure/eval.py index 8c0510e698..f694378629 100644 --- a/dimos/navigation/jnav/components/loop_closure/eval.py +++ b/dimos/navigation/jnav/components/loop_closure/eval.py @@ -15,8 +15,10 @@ """Evaluate a loop-closure module against a recording. The raw-baseline robot trajectory comes from the recording's tf (the -``world``->``base_link`` chain, sampled at each odom time) when the recording has -tf; recordings without tf fall back to the odom stream's stored poses. +``world``->``base_link`` chain, sampled at each odom time). tf is REQUIRED: a +recording without a tf tree is an error (run ``add_tf`` first) — there is no +odom-pose fallback. tf is allowed to come online within a short opening window +(``--tf-startup-tolerance-s``). Two ground-truth-free scores, before vs after correction: * April-tag agreement — a fixed tag re-seen along the run should map to one @@ -544,47 +546,49 @@ def run_module_graph( # Tolerance for composing world->body from the recording's tf at an odom time. TF_LOOKUP_TOLERANCE_S = 0.1 +# tf may come online slightly after the odom stream starts (sensor/static warmup); +# odom samples before tf is available are skipped only within this opening window. +TF_STARTUP_TOLERANCE_S = 2.0 -def _odom_pose_samples(db_path: Path, odom_stream: str) -> tuple[np.ndarray, np.ndarray]: - """Robot trajectory straight from the odometry stream's stored poses.""" - times: list[float] = [] - poses: list[list[float]] = [] - for timestamp, payload in iterate_stream(db_path, odom_stream): - pose = RateReplay._payload_pose(payload) - times.append(timestamp) - poses.append( - [ - pose.position.x, - pose.position.y, - pose.position.z, - pose.orientation.x, - pose.orientation.y, - pose.orientation.z, - pose.orientation.w, - ] - ) - return ( - np.asarray(times, dtype=np.float64), - np.asarray(poses, dtype=np.float64).reshape(-1, 7), - ) - - -def _tf_pose_samples( - db_path: Path, odom_stream: str, world_frame: str, body_frame: str -) -> tuple[np.ndarray, np.ndarray] | None: +def tf_pose_samples( + db_path: Path, + odom_stream: str, + *, + world_frame: str, + body_frame: str, + startup_tolerance_s: float = TF_STARTUP_TOLERANCE_S, +) -> tuple[np.ndarray, np.ndarray]: """Robot trajectory composed from the recording's tf (``world_frame -> - body_frame``) sampled at each odom timestamp. ``None`` if the recording has no - tf, so the caller can fall back to raw odom poses.""" + body_frame``) sampled at each odom timestamp. + + tf is REQUIRED — this raises (no odom-pose fallback) when the recording has no + tf tree, or when tf never comes online within ``startup_tolerance_s`` of the + odom start. Missing samples after tf is online (sporadic lookup gaps) are + skipped.""" db_store = store(db_path) if "tf" not in db_store.list_streams() or not db_store.tf.has_transforms(): - return None + raise SystemExit( + f"{db_path}: no tf tree — eval requires tf (run add_tf first); odom fallback removed" + ) times: list[float] = [] poses: list[list[float]] = [] + odom_t0: float | None = None + tf_online = False for timestamp, _payload in iterate_stream(db_path, odom_stream): + if odom_t0 is None: + odom_t0 = timestamp transform = db_store.tf.get(world_frame, body_frame, timestamp, TF_LOOKUP_TOLERANCE_S) if transform is None: - continue + if tf_online: + continue # sporadic gap once tf is up — skip this sample + if timestamp - odom_t0 <= startup_tolerance_s: + continue # opening warmup window — tf not online yet, tolerated + raise SystemExit( + f"{db_path}: tf did not come online within {startup_tolerance_s}s of the" + f" odom start (world={world_frame!r} body={body_frame!r})" + ) + tf_online = True times.append(timestamp) poses.append( [ @@ -598,27 +602,13 @@ def _tf_pose_samples( ] ) if not times: - return None + raise SystemExit(f"{db_path}: tf produced no {world_frame}->{body_frame} samples") return ( np.asarray(times, dtype=np.float64), np.asarray(poses, dtype=np.float64).reshape(-1, 7), ) -def raw_pose_samples( - db_path: Path, odom_stream: str, *, world_frame: str, body_frame: str -) -> tuple[np.ndarray, np.ndarray, str]: - """The raw-baseline trajectory as ``(times, [x,y,z,qx,qy,qz,qw], source)``. - - Prefers the recording's tf (``world_frame -> body_frame``, world-registered); - falls back to the odometry stream's stored poses when there is no tf.""" - tf_samples = _tf_pose_samples(db_path, odom_stream, world_frame, body_frame) - if tf_samples is not None: - return tf_samples[0], tf_samples[1], "tf" - times, poses = _odom_pose_samples(db_path, odom_stream) - return times, poses, "odom" - - def _subsampled_path(positions: np.ndarray) -> np.ndarray: stride = max(1, len(positions) // _RRD_MAX_PATH_POINTS) return positions[::stride] @@ -673,6 +663,7 @@ def evaluate( ignore_tags: set[int] | None = None, world_frame: str = "world", body_frame: str = "base_link", + tf_startup_tolerance_s: float = TF_STARTUP_TOLERANCE_S, ) -> dict[str, Any]: streams = list_streams(db_path) for required in (odom_stream, lidar_stream): @@ -749,14 +740,16 @@ def detect() -> Iterable[tuple[int, float]]: if not graph: raise SystemExit(f"{module_name} produced an empty pose graph") - # Raw-baseline trajectory: composed from the recording's tf (world->body) when - # present, else the odom stream's stored poses. - raw_times, raw_poses7, pose_source = raw_pose_samples( - db_path, odom_stream, world_frame=world_frame, body_frame=body_frame + # Raw-baseline trajectory: composed from the recording's tf (world->body). + # tf is required — no odom-pose fallback (see tf_pose_samples). + raw_times, raw_poses7 = tf_pose_samples( + db_path, + odom_stream, + world_frame=world_frame, + body_frame=body_frame, + startup_tolerance_s=tf_startup_tolerance_s, ) - if raw_poses7.size == 0: - raise SystemExit(f"no raw baseline poses (tf nor {odom_stream!r}) in {db_path}") - print(f"raw baseline pose source: {pose_source} ({len(raw_times)} samples)") + print(f"raw baseline from tf {world_frame}->{body_frame} ({len(raw_times)} samples)") raw_xyz_base = trajectory_lookup(raw_times, raw_poses7[:, :3], ODOM_MATCH_TOLERANCE_S) raw_pose7_base = pose7_lookup(raw_times, raw_poses7, ODOM_MATCH_TOLERANCE_S) @@ -831,7 +824,7 @@ def detect() -> Iterable[tuple[int, float]]: db_path, pgo_config, lidar_stream, odom_stream, drift_per_sec ), "replay": replay_stats, - "pose_source": pose_source, + "pose_source": "tf", "world_frame": world_frame, "body_frame": body_frame, "april_tags": { @@ -944,6 +937,12 @@ def main() -> None: default="base_link", help="tf frame tracking the robot body (the raw baseline is world_frame->body_frame)", ) + parser.add_argument( + "--tf-startup-tolerance-s", + type=float, + default=TF_STARTUP_TOLERANCE_S, + help="grace window for tf to come online after the odom start before erroring", + ) args = parser.parse_args() drift_per_sec = ( @@ -988,6 +987,7 @@ def main() -> None: ignore_tags=ignore_tags, world_frame=args.world_frame, body_frame=args.body_frame, + tf_startup_tolerance_s=args.tf_startup_tolerance_s, ) From 6199cb064810821ffa2c13784a6ff7e3b936c125 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sun, 28 Jun 2026 09:54:04 +0800 Subject: [PATCH 38/69] feat(go2 mount): split mid360 into physical mid360_mount + gravity-aligned mid360_link MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep the physical sensor pose (position + 44deg tilt) as mid360_mount, and add a counter-rotated mid360_link at the same position whose orientation is un-tilted (gravity-aligned) — the frame the gravity-anchored LIO scans live in. Preserves the physical mount geometry while giving a level data frame. --- .../unitree/go2/go2_mid360_static_transforms.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py index b430a11429..a8ee2161e9 100644 --- a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py +++ b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py @@ -21,7 +21,13 @@ Mount geometry (measured on the physical rig) --------------------------------------------- - base_link -> front_camera: 32.7cm forward, ~4.3cm up (URDF front_camera mount). -- front_camera -> mid360_link: lidar is 3.2cm back, 12cm up, pitched 44 deg down. +- front_camera -> mid360_mount: lidar is 3.2cm back, 12cm up, pitched 44 deg down — + the physical sensor pose (position + tilt). +- mid360_mount -> mid360_link: a counter-rotation (inverse of the mount tilt) that + leaves the position untouched but un-tilts the orientation, so ``mid360_link`` is + gravity-aligned. The LIO odometry is gravity-anchored, so the recorded scans live + in this gravity-aligned frame at the physical sensor position — keeping the mount + in a separate ``mid360_mount`` frame preserves the physical geometry. - front_camera -> camera_optical: the standard ROS optical rotation (x-right, y-down, z-forward). """ @@ -45,7 +51,10 @@ FRAMES: list[FrameSpec] = [ ("base_link", None, (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)), ("front_camera", "base_link", (0.32715, -0.00003, 0.04297), (0.0, 0.0, 0.0)), - ("mid360_link", "front_camera", (-0.032, 0.0, 0.12), (0.0, MID360_PITCH_DOWN, 0.0)), + # physical sensor pose: position + the 44 deg downward tilt + ("mid360_mount", "front_camera", (-0.032, 0.0, 0.12), (0.0, MID360_PITCH_DOWN, 0.0)), + # gravity-aligned data frame: same position, tilt counter-rotated out + ("mid360_link", "mid360_mount", (0.0, 0.0, 0.0), (0.0, -MID360_PITCH_DOWN, 0.0)), ("camera_optical", "front_camera", (0.0, 0.0, 0.0), OPTICAL_RPY), ] From 236a203c0e50599d3e14885714e5346fb31262ad Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sun, 28 Jun 2026 11:00:54 +0800 Subject: [PATCH 39/69] refactor(go2 mount): rename mid360_mount->mid360_link (physical), mid360_link->mid360_gravity --- .../go2/go2_mid360_static_transforms.py | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py index a8ee2161e9..2f28168961 100644 --- a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py +++ b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py @@ -21,13 +21,14 @@ Mount geometry (measured on the physical rig) --------------------------------------------- - base_link -> front_camera: 32.7cm forward, ~4.3cm up (URDF front_camera mount). -- front_camera -> mid360_mount: lidar is 3.2cm back, 12cm up, pitched 44 deg down — - the physical sensor pose (position + tilt). -- mid360_mount -> mid360_link: a counter-rotation (inverse of the mount tilt) that - leaves the position untouched but un-tilts the orientation, so ``mid360_link`` is - gravity-aligned. The LIO odometry is gravity-anchored, so the recorded scans live - in this gravity-aligned frame at the physical sensor position — keeping the mount - in a separate ``mid360_mount`` frame preserves the physical geometry. +- front_camera -> mid360_link: lidar is 3.2cm back, 12cm up, pitched 44 deg down — + the physical sensor pose (position + tilt). Named ``mid360_link`` to match the LIO + sensor frame_id. +- mid360_link -> mid360_gravity: a counter-rotation (inverse of the mount tilt) that + leaves the position untouched but un-tilts the orientation, so ``mid360_gravity`` + is gravity-aligned. The LIO odometry is gravity-anchored, so the recorded scans + live in this gravity-aligned frame at the physical sensor position — keeping the + physical mount in ``mid360_link`` preserves the geometry. - front_camera -> camera_optical: the standard ROS optical rotation (x-right, y-down, z-forward). """ @@ -51,10 +52,10 @@ FRAMES: list[FrameSpec] = [ ("base_link", None, (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)), ("front_camera", "base_link", (0.32715, -0.00003, 0.04297), (0.0, 0.0, 0.0)), - # physical sensor pose: position + the 44 deg downward tilt - ("mid360_mount", "front_camera", (-0.032, 0.0, 0.12), (0.0, MID360_PITCH_DOWN, 0.0)), - # gravity-aligned data frame: same position, tilt counter-rotated out - ("mid360_link", "mid360_mount", (0.0, 0.0, 0.0), (0.0, -MID360_PITCH_DOWN, 0.0)), + # physical sensor pose: position + the 44 deg downward tilt (matches LIO frame_id) + ("mid360_link", "front_camera", (-0.032, 0.0, 0.12), (0.0, MID360_PITCH_DOWN, 0.0)), + # gravity-aligned derived frame: same position, tilt counter-rotated out + ("mid360_gravity", "mid360_link", (0.0, 0.0, 0.0), (0.0, -MID360_PITCH_DOWN, 0.0)), ("camera_optical", "front_camera", (0.0, 0.0, 0.0), OPTICAL_RPY), ] From 5ce9283bcf4e0dfcd62ddc50f90cbce587f87c3a Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 29 Jun 2026 21:50:00 -0500 Subject: [PATCH 40/69] refactor(db_tf): trim surface + fix same-timestamp graph tie-break - remove unused write_tf_tree; move transform_matrix into post_process (its only use) - drop GRAPH_STREAM_SUFFIX indirection (one '_graph' stream, name derived inline) - inline TfGraphWriter into the recorder and build_graph_stream into DbTf._ensure_built - name the pose-equality rounding constant; namespace the child-frame index name - fix: order graph snapshots by (ts, id) so same-timestamp topology changes resolve to the complete snapshot; raise the in-RAM cache threshold for fixed multi-frame rigs --- dimos/memory2/db_tf.py | 423 ++++++------------ dimos/memory2/module.py | 30 +- dimos/memory2/test_db_tf.py | 34 +- .../gsc_pgo/scripts/post_process.py | 11 +- 4 files changed, 183 insertions(+), 315 deletions(-) diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py index 599aed7f39..1692bc276a 100644 --- a/dimos/memory2/db_tf.py +++ b/dimos/memory2/db_tf.py @@ -12,27 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Transform lookups over the transforms recorded in a store (multi-robot friendly). - -Two pieces of recorded state, written by the recorder: - -* a **graph stream** (table ``tf_graph``): one row per *topology* change — i.e. - whenever the set of frames or any frame's parent / static-ness changes. Each row - is the full structure at that instant: ``{child_frame: {parent, static}}``. - Topology changes rarely (a robot joins/leaves, a relocalization re-parents), so - this table is tiny and naturally supports time-varying / multi-robot trees. -* the ``tf`` stream, with each row tagged by its ``child_frame`` (an indexed json - tag) so a frame's samples can be range-queried by time. - -``store.tf.get(target, source, ts)`` then: reads the graph as-of the query time -(from RAM if there are few graph changes, else one query), walks it to the -source->target chain (in memory; the graph may be DISJOINT for unrelated robots), -and resolves *only* that chain's frames — a static frame is one cached constant, a -dynamic frame is its two bracketing samples, interpolated. Composition + -interpolation reuse :class:`MultiTBuffer`. Non-sqlite stores fall back to loading -the tf streams into a buffer. - -``write_tf_tree`` populates the tf stream for a recording that lacks one. +""" +A tf class for memory2 """ from __future__ import annotations @@ -43,12 +24,8 @@ import threading from typing import TYPE_CHECKING, Any, cast -import numpy as np - from dimos.memory2.store.sqlite import SqliteStoreConfig -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.tf2_msgs.TFMessage import TFMessage from dimos.protocol.tf.tf import MultiTBuffer from dimos.utils.logging_config import setup_logger @@ -60,151 +37,20 @@ logger = setup_logger() DEFAULT_TF_STREAM = "tf" -# The topology change-log is a first-class stream named "_graph". -GRAPH_STREAM_SUFFIX = "_graph" # Streams the RAM fallback (non-sqlite stores) reads. TF_STREAMS = ("tf", "tf_static") -# If a recording has fewer than this many topology changes, load them all into RAM -# so a lookup needs no graph query (the common single-robot / stable-tree case). -# At or above it, fall back to one graph query per lookup (many-robot churn). -DEFAULT_MAX_GRAPH_CHANGES_IN_RAM = 20 -# Larger than any single recording's span so the fallback buffer never prunes. +# Cache the whole change-log in RAM when there are at most this many topology +# changes (a stable tree — even a many-frame sensor rig — is a one-time setup, not +# churn); above it, fall back to one indexed graph query per lookup (multi-robot). +DEFAULT_MAX_GRAPH_CHANGES_IN_RAM = 64 +# MultiTBuffer drops samples older than buffer_size seconds; we feed it exactly the +# bracketing samples and want them all kept, so use a span no recording exceeds. _NO_PRUNE = 1.0e15 -# Lower bound for "latest topology at or before T" range queries. -_TS_MIN = -1.0e18 -# SQLite can't parameterize table names, so caller-supplied stream names are -# interpolated; allow only safe identifiers to keep that injection-free. -_SAFE_TABLE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") - - -def _graph_stream_name(stream: str) -> str: - return f"{_safe_table(stream)}{GRAPH_STREAM_SUFFIX}" - - -def _safe_table(name: str) -> str: - if not _SAFE_TABLE.match(name): - raise ValueError(f"unsafe stream/table name: {name!r}") - return name - - -def _connect(db_path: str) -> sqlite3.Connection: - """A connection that waits on the WAL write-lock instead of erroring — the - store keeps its own connection to the same db open while we read/write.""" - conn = sqlite3.connect(db_path, check_same_thread=False) - conn.execute("PRAGMA busy_timeout=5000") - return conn - - -def _ensure_child_index(conn: sqlite3.Connection, stream: str) -> None: - """Index the child_frame json tag on the tf rows so per-frame time queries - seek. The live recorder gets this for free (the store auto-indexes tag keys on - tagged appends); this is for migrated recordings and the read side. Requires - the tf table to exist.""" - safe = _safe_table(stream) - # Composite (child_frame, ts) so a per-frame "latest at/before T" is a direct - # index range seek, not a scan+sort. - conn.execute( - f'CREATE INDEX IF NOT EXISTS "{safe}_child_ts_idx" ' - f"ON \"{safe}\"(json_extract(tags, '$.child_frame'), ts)" - ) - conn.commit() - - -class TfGraph: - """A tf topology snapshot, recorded one per structure change. - - ``structure`` maps each child frame to ``{"parent": str, "static": bool}`` — - the full tf tree as of this message's timestamp. The stream of these snapshots - (the ``_graph`` stream) is the topology change-log that transform lookups - walk to resolve a source->target chain at any past time. Defined here (not under - ``dimos/msgs``) because it is a recording-internal payload, not a wire message; - it is stored via the pickle codec.""" - - structure: dict[str, dict[str, Any]] - msg_name = "tf2_msgs.TfGraph" - - def __init__(self, structure: dict[str, dict[str, Any]]) -> None: - # copy so later mutations of the writer's running structure don't alter an - # already-recorded snapshot - self.structure = {child: dict(entry) for child, entry in structure.items()} - - def __repr__(self) -> str: - return f"TfGraph({len(self.structure)} frames)" - - -class TfGraphWriter: - """Recorder helper: tracks the running topology and appends a ``TfGraph`` - snapshot to the graph stream only when the structure changes.""" - - def __init__(self, store: Store, stream: str = DEFAULT_TF_STREAM) -> None: - self._stream: Stream[TfGraph] = store.stream(_graph_stream_name(stream), TfGraph) - self._structure: dict[str, dict[str, Any]] = {} - - def record(self, child_frame: str, parent_frame: str, is_static: bool, ts: float) -> None: - entry = {"parent": parent_frame, "static": bool(is_static)} - if self._structure.get(child_frame) == entry: - return # no structural change -> no new snapshot - self._structure[child_frame] = entry - self._stream.append(TfGraph(self._structure), ts=ts) - - def close(self) -> None: - # the stream is owned by the store; nothing to close here - pass - - -def build_graph_stream(store: Store, stream: str = DEFAULT_TF_STREAM) -> int: - """One-time migration for a recording that predates the graph stream: tag every - tf row with its ``child_frame`` and build the ``_graph`` stream - chronologically. A frame is treated as static if its pose never changes across - the recording. Returns the number of topology-change snapshots written.""" - config = store.config - if not isinstance(config, SqliteStoreConfig): - raise TypeError("build_graph_stream needs a SqliteStore") - safe = _safe_table(stream) - - # one decode pass: collect (id, ts, child, parent, pose-key) per row - rows: list[tuple[int, float, str, str, tuple[float, ...]]] = [] - poses_per_child: dict[str, set[tuple[float, ...]]] = {} - for obs in store.stream(safe, TFMessage).order_by("ts"): - for transform in getattr(obs.data, "transforms", None) or [obs.data]: - pose_key = ( - round(transform.translation.x, 9), - round(transform.translation.y, 9), - round(transform.translation.z, 9), - round(transform.rotation.x, 9), - round(transform.rotation.y, 9), - round(transform.rotation.z, 9), - round(transform.rotation.w, 9), - ) - rows.append((obs.id, obs.ts, transform.child_frame_id, transform.frame_id, pose_key)) - poses_per_child.setdefault(transform.child_frame_id, set()).add(pose_key) - static_frames = {child for child, poses in poses_per_child.items() if len(poses) == 1} - - # tag each tf row with its child_frame + add the seek index (raw, on the tf table) - conn = _connect(config.path) - try: - for row_id, _ts, child, _parent, _pose in rows: - conn.execute( - f"UPDATE \"{safe}\" SET tags = json_set(tags, '$.child_frame', ?) WHERE id = ?", - (child, row_id), - ) - conn.commit() - _ensure_child_index(conn, safe) - finally: - conn.close() - - # build the topology change-log as a first-class stream - graph_stream = store.stream(_graph_stream_name(safe), TfGraph) - structure: dict[str, dict[str, Any]] = {} - written = 0 - for _row_id, ts, child, parent, _pose in rows: - entry = {"parent": parent, "static": child in static_frames} - if structure.get(child) == entry: - continue - structure[child] = entry - graph_stream.append(TfGraph(structure), ts=ts) - written += 1 - return written +# A frame is "static" if its pose never changes; poses are compared rounded to this +# many decimals (~nanometre / nanoradian) so float noise doesn't read as motion. +POSE_EQUALITY_DECIMALS = 9 +# enforce safe identifiers for SQL +_VAR_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") class DbTf: @@ -224,7 +70,8 @@ def __init__( ) -> None: self._store = store self._stream = _safe_table(stream) - self._graph_name = _graph_stream_name(stream) + # The topology change-log is a single companion stream, one per tf stream. + self._graph_name = f"{self._stream}_graph" self._max_in_ram = max_graph_changes_in_ram self._stream_names = stream_names # RAM fallback only self._lock = threading.Lock() @@ -290,6 +137,9 @@ def _graph_count(self) -> int: return self._graph_stream().count() def _ensure_built(self) -> None: + """First sqlite use: if the recording has tf rows but no graph stream (a + recording that predates it / wasn't written by the recorder), build the graph + stream once by replaying the tf rows, then make sure the seek index exists.""" if self._built: return conn = self._connection() @@ -302,19 +152,75 @@ def _ensure_built(self) -> None: "========================================================================", self._stream, ) - built = build_graph_stream(self._store, self._stream) - logger.warning("tf graph built: %d topology changes for %r.", built, self._stream) + self._build_graph_stream() if n_rows: _ensure_child_index(conn, self._stream) # tf table exists now self._built = True + def _build_graph_stream(self) -> None: + """One-time migration: decode every tf row, tag it with its child_frame, and + append a ``TfGraph`` snapshot whenever the topology changes. A frame counts as + static if its pose never varies across the whole recording.""" + safe = self._stream + # one decode pass: collect (id, ts, child, parent, pose-key) + per-child poses + rows: list[tuple[int, float, str, str]] = [] + poses_per_child: dict[str, set[tuple[float, ...]]] = {} + for obs in self._store.stream(safe, TFMessage).order_by("ts"): + for transform in getattr(obs.data, "transforms", None) or [obs.data]: + pose_key = tuple( + round(value, POSE_EQUALITY_DECIMALS) + for value in ( + transform.translation.x, + transform.translation.y, + transform.translation.z, + transform.rotation.x, + transform.rotation.y, + transform.rotation.z, + transform.rotation.w, + ) + ) + rows.append((obs.id, obs.ts, transform.child_frame_id, transform.frame_id)) + poses_per_child.setdefault(transform.child_frame_id, set()).add(pose_key) + static_frames = {child for child, poses in poses_per_child.items() if len(poses) == 1} + + # tag each tf row with its child_frame (raw UPDATE on the tf table) + conn = self._connection() + for row_id, _ts, child, _parent in rows: + conn.execute( + f"UPDATE \"{safe}\" SET tags = json_set(tags, '$.child_frame', ?) WHERE id = ?", + (child, row_id), + ) + conn.commit() + + # build the change-log as a first-class stream: one snapshot per change + graph_stream = self._store.stream(self._graph_name, TfGraph) + structure: dict[str, dict[str, Any]] = {} + written = 0 + for _row_id, ts, child, parent in rows: + entry = {"parent": parent, "static": child in static_frames} + if structure.get(child) == entry: + continue + structure[child] = entry + graph_stream.append(TfGraph(structure), ts=ts) + written += 1 + logger.warning("tf graph built: %d topology changes for %r.", written, self._stream) + + def _graph_codec(self) -> Any: + source = self._store.stream(self._graph_name, TfGraph)._source + return cast("Any", source).codec + def _load_graph_if_small(self) -> None: if self._graph_loaded: return if self._graph_count() < self._max_in_ram: - self._graph_in_ram = [ - (obs.ts, obs.data.structure) for obs in self._graph_stream().order_by("ts") - ] + # Sort by (ts, id): several topology changes can share a timestamp (e.g. + # every static frame latched at t0), and the LAST-inserted of those is the + # complete snapshot — a plain ts sort leaves same-ts order undefined. + snapshots = sorted( + ((obs.ts, obs.id, obs.data.structure) for obs in self._graph_stream()), + key=lambda row: (row[0], row[1]), + ) + self._graph_in_ram = [(ts, structure) for ts, _id, structure in snapshots] else: self._graph_in_ram = None # too many -> query per lookup self._graph_loaded = True @@ -327,15 +233,24 @@ def _graph_at(self, query_time: float) -> dict[str, Any] | None: if index < 0: return self._graph_in_ram[0][1] # before first -> earliest return self._graph_in_ram[index][1] - # fallback: one query for the latest snapshot at or before query_time + # fallback: one indexed query for the latest snapshot at or before query_time. + # Tie-break by id (DESC) so same-timestamp changes resolve to the complete one. self.graph_queries += 1 - stream = self._graph_stream() - latest = stream.time_range(_TS_MIN, query_time).order_by("ts", desc=True).limit(1) - for obs in latest: - return obs.data.structure - for obs in stream.order_by("ts").limit(1): # before first -> earliest - return obs.data.structure - return None + conn = self._connection() + graph, blob = f'"{self._graph_name}"', f'"{self._graph_name}_blob"' + row = conn.execute( + f"SELECT x.data FROM {graph} g JOIN {blob} x ON x.id = g.id " + "WHERE g.ts <= ? ORDER BY g.ts DESC, g.id DESC LIMIT 1", + (query_time,), + ).fetchone() + if row is None: # before the first snapshot -> earliest + row = conn.execute( + f"SELECT x.data FROM {graph} g JOIN {blob} x ON x.id = g.id " + "ORDER BY g.ts ASC, g.id ASC LIMIT 1" + ).fetchone() + if row is None: + return None + return cast("TfGraph", self._graph_codec().decode(row[0])).structure def _chain_frames(self, graph: dict[str, Any], source: str, target: str) -> list[str] | None: def to_root(frame: str) -> list[str]: @@ -471,6 +386,28 @@ def get( return buffer.lookup(target_frame, source_frame, time_point, time_tolerance) +class TfGraph: + """A tf topology snapshot, recorded one per structure change. + + ``structure`` maps each child frame to ``{"parent": str, "static": bool}`` — + the full tf tree as of this message's timestamp. The stream of these snapshots + (the ``_graph`` stream) is the topology change-log that transform lookups + walk to resolve a source->target chain at any past time. Defined here (not under + ``dimos/msgs``) because it is a recording-internal payload, not a wire message; + it is stored via the pickle codec.""" + + structure: dict[str, dict[str, Any]] + msg_name = "tf2_msgs.TfGraph" + + def __init__(self, structure: dict[str, dict[str, Any]]) -> None: + # copy so later mutations of the writer's running structure don't alter an + # already-recorded snapshot + self.structure = {child: dict(entry) for child, entry in structure.items()} + + def __repr__(self) -> str: + return f"TfGraph({len(self.structure)} frames)" + + def _restamp(transform: Transform, ts: float) -> Transform: return Transform( translation=transform.translation, @@ -481,103 +418,31 @@ def _restamp(transform: Transform, ts: float) -> Transform: ) -def transform_matrix(transform: Transform) -> tuple[np.ndarray, np.ndarray]: - """Return ``(R, t)`` (3x3, 3) for ``transform`` so ``p_target = p_source @ R.T + t``.""" - rotation = transform.rotation - rotation_matrix = np.asarray(rotation.to_rotation_matrix(), float).reshape(3, 3) - translation = np.array( - [transform.translation.x, transform.translation.y, transform.translation.z], float - ) - return rotation_matrix, translation - - -def write_tf_tree( - store: Store, - *, - odom_stream: str, - odom_parent: str = "odom", - odom_child: str = "base_link", - root_links: tuple[tuple[str, str], ...] = (("world", "map"), ("map", "odom")), - sensor_child: str = "mid360_link", - sensor_translation: tuple[float, float, float] = (0.0, 0.0, 0.0), - sensor_rotation: tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0), - static_period: float = 0.45, - stream_name: str = "tf", -) -> int: - """Populate ``store``'s tf stream from an odometry stream. - - - ``root_links`` and ``odom_child -> sensor_child`` are emitted as identity / - fixed transforms every ``static_period`` seconds across the recording span. - - ``odom_parent -> odom_child`` is emitted once per odometry sample, taken - from each observation's pose. - - Each transform is written as its own row (one transform per ``TFMessage``) so - the graph-stream reader can range-query it by ``child_frame``. Returns the - number of tf observations written. - """ - config = store.config - if not isinstance(config, SqliteStoreConfig): - raise TypeError("write_tf_tree reads the db directly and needs a SqliteStore") - db_path = config.path - connection = sqlite3.connect(f"file:{db_path}?mode=ro&immutable=1", uri=True) - odom = np.array( - list( - connection.execute( - "select ts,pose_x,pose_y,pose_z,pose_qx,pose_qy,pose_qz,pose_qw " - f"from {_safe_table(odom_stream)} order by ts" - ) - ), - float, - ) - connection.close() - if not len(odom): - raise ValueError(f"odom stream {odom_stream!r} is empty; cannot build tf tree") - - tf_stream = store.stream(stream_name, TFMessage) - written = 0 - - def _append(transform: Transform, ts: float) -> None: - nonlocal written - tf_stream.append( - TFMessage(transform), ts=ts, tags={"child_frame": transform.child_frame_id} - ) - written += 1 - - # dynamic: odom_parent -> odom_child, one per odometry sample - for row in odom: - ts = float(row[0]) - _append( - Transform( - translation=Vector3(row[1], row[2], row[3]), - rotation=Quaternion(row[4], row[5], row[6], row[7]), - frame_id=odom_parent, - child_frame_id=odom_child, - ts=ts, - ), - ts, - ) - - # static: root links + sensor mount, resampled every static_period - t0 = float(odom[0, 0]) - t1 = float(odom[-1, 0]) - - def statics_at(ts: float) -> list[Transform]: - links = [ - Transform(frame_id=parent, child_frame_id=child, ts=ts) for parent, child in root_links - ] - links.append( - Transform( - translation=Vector3(*sensor_translation), - rotation=Quaternion(*sensor_rotation), - frame_id=odom_child, - child_frame_id=sensor_child, - ts=ts, - ) - ) - return links +def _safe_table(name: str) -> str: + if not _VAR_NAME_PATTERN.match(name): + raise ValueError(f"unsafe stream/table name: {name!r}") + return name - for static_ts in np.arange(t0, t1 + static_period, static_period): - for transform in statics_at(float(static_ts)): - _append(transform, float(static_ts)) - return written +def _connect(db_path: str) -> sqlite3.Connection: + conn = sqlite3.connect(db_path, check_same_thread=False) + conn.execute("PRAGMA busy_timeout=5000") + return conn + + +def _ensure_child_index(conn: sqlite3.Connection, stream: str) -> None: + """Index the child_frame json tag on the tf rows so per-frame time queries + seek. The live recorder gets this for free (the store auto-indexes tag keys on + tagged appends); this is for migrated recordings and the read side. Requires + the tf table to exist.""" + safe = _safe_table(stream) + # Composite (child_frame, ts) so a per-frame "latest at/before T" is a direct + # index range seek, not a scan+sort. Index names share SQLite's global namespace + # with tables, so the name is double-underscore-namespaced to keep it clear of any + # real stream/table name (no stream would contain "__dbtf_"). + index_name = f"{safe}__dbtf_child_ts_idx" + conn.execute( + f'CREATE INDEX IF NOT EXISTS "{index_name}" ' + f"ON \"{safe}\"(json_extract(tags, '$.child_frame'), ts)" + ) + conn.commit() diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index 186e50108f..5715a31f10 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -30,7 +30,7 @@ from dimos.constants import DIMOS_PROJECT_ROOT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig -from dimos.memory2.db_tf import TfGraphWriter +from dimos.memory2.db_tf import TfGraph from dimos.memory2.embed import EmbedImages from dimos.memory2.store.null import NullStore from dimos.memory2.store.sqlite import SqliteStore @@ -462,10 +462,17 @@ def _record_tf(self) -> None: logger.warning("Recorder: no pubsub tf available — not recording tf") return tf_stream = self.store.stream("tf", TFMessage) - - # Topology change-log (the "tf_graph" stream) written alongside the tf rows. - graph_writer = TfGraphWriter(self.store, "tf") - self.register_disposable(Disposable(graph_writer.close)) + graph_stream = self.store.stream("tf_graph", TfGraph) + # Running tf topology; a TfGraph snapshot is appended whenever it changes, so + # the "tf_graph" stream is the topology change-log transform lookups walk. + structure: dict[str, dict[str, Any]] = {} + + def record_graph(child: str, parent: str, is_static: bool, ts: float) -> None: + entry = {"parent": parent, "static": bool(is_static)} + if structure.get(child) == entry: + return # no structural change -> no new snapshot + structure[child] = entry + graph_stream.append(TfGraph(structure), ts=ts) def make_handler(is_static: bool) -> Any: def on_tf(msg: TFMessage, _topic: Any) -> None: @@ -477,13 +484,12 @@ def on_tf(msg: TFMessage, _topic: Any) -> None: pose=None, tags={"child_frame": transform.child_frame_id}, ) - if graph_writer is not None: - graph_writer.record( - transform.child_frame_id, - transform.frame_id, - is_static, - transform.ts, - ) + record_graph( + transform.child_frame_id, + transform.frame_id, + is_static, + transform.ts, + ) except sqlite3.ProgrammingError: # A late LCM callback raced teardown and hit the closed store. pass diff --git a/dimos/memory2/test_db_tf.py b/dimos/memory2/test_db_tf.py index b5a077829f..ac97723bc2 100644 --- a/dimos/memory2/test_db_tf.py +++ b/dimos/memory2/test_db_tf.py @@ -23,7 +23,7 @@ import math from pathlib import Path -from dimos.memory2.db_tf import DbTf, TfGraphWriter +from dimos.memory2.db_tf import DbTf from dimos.memory2.store.sqlite import SqliteStore from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform @@ -51,18 +51,16 @@ def _static(parent: str, child: str, xyz: tuple[float, float, float], ts: float) ) -def _append( - store: SqliteStore, graph: TfGraphWriter, transform: Transform, is_static: bool -) -> None: - """Record one transform exactly as the live recorder does: one row, tagged - child_frame, plus a topology-change row when the structure changes.""" +def _append(store: SqliteStore, transform: Transform) -> None: + """Record one transform as the recorder does: one row tagged with its child_frame. + The graph stream is built lazily by DbTf on first lookup (static-ness inferred + from whether a frame's pose ever changes), so tests exercise that migration.""" store.stream("tf", TFMessage).append( TFMessage(transform), ts=transform.ts, pose=None, tags={"child_frame": transform.child_frame_id}, ) - graph.record(transform.child_frame_id, transform.frame_id, is_static, transform.ts) def _ref(transforms: list[Transform]) -> MultiTBuffer: @@ -87,7 +85,6 @@ def _record_single_robot(path: Path, *, static_repeat: bool) -> list[Transform]: """world->map->odom->base_link->sensor. Statics emitted once (latched) unless static_repeat, in which case they're re-emitted each second.""" store = SqliteStore(path=str(path)) - graph = TfGraphWriter(store) written: list[Transform] = [] statics = [ @@ -99,7 +96,7 @@ def _record_single_robot(path: Path, *, static_repeat: bool) -> list[Transform]: for ts in static_times: for parent, child, xyz in statics: transform = _static(parent, child, xyz, ts) - _append(store, graph, transform, is_static=True) + _append(store, transform) written.append(transform) for i in range(int(_DURATION * _DYN_RATE)): @@ -111,10 +108,9 @@ def _record_single_robot(path: Path, *, static_repeat: bool) -> list[Transform]: child_frame_id="base_link", ts=ts, ) - _append(store, graph, transform, is_static=False) + _append(store, transform) written.append(transform) - graph.close() store.stop() return written @@ -156,9 +152,8 @@ def test_reparent_midrun_uses_graph_as_of_query_time(tmp_path: Path) -> None: ~100. Exercises time-varying / multi-robot topology.""" path = tmp_path / "reparent.db" store = SqliteStore(path=str(path)) - graph = TfGraphWriter(store) - _append(store, graph, _static("world", "map", (10.0, 0, 0), _T0), is_static=True) - _append(store, graph, _static("map", "odom", (100.0, 0, 0), _T0), is_static=True) + _append(store, _static("world", "map", (10.0, 0, 0), _T0)) + _append(store, _static("map", "odom", (100.0, 0, 0), _T0)) era1: list[Transform] = [] for i in range(150): # [t0, t0+5) @@ -170,7 +165,7 @@ def test_reparent_midrun_uses_graph_as_of_query_time(tmp_path: Path) -> None: child_frame_id="base_link", ts=ts, ) - _append(store, graph, transform, is_static=False) + _append(store, transform) era1.append(transform) switch = _T0 + 5.0 era2: list[Transform] = [] @@ -183,9 +178,8 @@ def test_reparent_midrun_uses_graph_as_of_query_time(tmp_path: Path) -> None: child_frame_id="base_link", ts=ts, ) - _append(store, graph, transform, is_static=False) + _append(store, transform) era2.append(transform) - graph.close() store.stop() store = SqliteStore(path=str(path), must_exist=True) @@ -220,12 +214,10 @@ def test_disjoint_multirobot_returns_none(tmp_path: Path) -> None: query is None, an in-component query resolves.""" path = tmp_path / "two.db" store = SqliteStore(path=str(path)) - graph = TfGraphWriter(store) for i in range(20): ts = _T0 + i / _DYN_RATE _append( store, - graph, Transform( translation=Vector3(i * 0.1, 0, 0), rotation=_yaw(0), @@ -233,11 +225,9 @@ def test_disjoint_multirobot_returns_none(tmp_path: Path) -> None: child_frame_id="baseA", ts=ts, ), - is_static=False, ) _append( store, - graph, Transform( translation=Vector3(0, i * 0.1, 0), rotation=_yaw(0), @@ -245,9 +235,7 @@ def test_disjoint_multirobot_returns_none(tmp_path: Path) -> None: child_frame_id="baseB", ts=ts, ), - is_static=False, ) - graph.close() store.stop() store = SqliteStore(path=str(path), must_exist=True) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py index de27fdfd8d..3a9cc93e5b 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py @@ -131,7 +131,16 @@ def arg(flag, default=""): sys.exit( f"!! {RAW_STREAM} missing -- run detect_tags.py first to build the unfiltered tag stream." ) -from dimos.memory2.db_tf import transform_matrix + + +def transform_matrix(transform): + """``(R, t)`` (3x3, 3) for a Transform so ``p_target = p_source @ R.T + t``.""" + rotation = np.asarray(transform.rotation.to_rotation_matrix(), float).reshape(3, 3) + translation = np.array( + [transform.translation.x, transform.translation.y, transform.translation.z], float + ) + return rotation, translation + _TF_AVAILABLE = USE_TF and store.tf.has_transforms() From 7dea4079319e3b946347fce42b582cf13bc4adb4 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 29 Jun 2026 22:30:32 -0500 Subject: [PATCH 41/69] refactor(db_tf): hardcode the tf_graph stream name (don't derive it) --- dimos/memory2/db_tf.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py index 1692bc276a..ef25579bba 100644 --- a/dimos/memory2/db_tf.py +++ b/dimos/memory2/db_tf.py @@ -37,6 +37,8 @@ logger = setup_logger() DEFAULT_TF_STREAM = "tf" +# The topology change-log is a single companion stream (like tf_static). +GRAPH_STREAM = "tf_graph" # Streams the RAM fallback (non-sqlite stores) reads. TF_STREAMS = ("tf", "tf_static") # Cache the whole change-log in RAM when there are at most this many topology @@ -70,8 +72,6 @@ def __init__( ) -> None: self._store = store self._stream = _safe_table(stream) - # The topology change-log is a single companion stream, one per tf stream. - self._graph_name = f"{self._stream}_graph" self._max_in_ram = max_graph_changes_in_ram self._stream_names = stream_names # RAM fallback only self._lock = threading.Lock() @@ -129,10 +129,10 @@ def has_transforms(self) -> bool: return bool(n_rows) def _graph_stream(self) -> Stream[TfGraph]: - return self._store.stream(self._graph_name, TfGraph) + return self._store.stream(GRAPH_STREAM, TfGraph) def _graph_count(self) -> int: - if self._graph_name not in set(self._store.list_streams()): + if GRAPH_STREAM not in set(self._store.list_streams()): return 0 return self._graph_stream().count() @@ -193,7 +193,7 @@ def _build_graph_stream(self) -> None: conn.commit() # build the change-log as a first-class stream: one snapshot per change - graph_stream = self._store.stream(self._graph_name, TfGraph) + graph_stream = self._store.stream(GRAPH_STREAM, TfGraph) structure: dict[str, dict[str, Any]] = {} written = 0 for _row_id, ts, child, parent in rows: @@ -206,7 +206,7 @@ def _build_graph_stream(self) -> None: logger.warning("tf graph built: %d topology changes for %r.", written, self._stream) def _graph_codec(self) -> Any: - source = self._store.stream(self._graph_name, TfGraph)._source + source = self._store.stream(GRAPH_STREAM, TfGraph)._source return cast("Any", source).codec def _load_graph_if_small(self) -> None: @@ -237,7 +237,7 @@ def _graph_at(self, query_time: float) -> dict[str, Any] | None: # Tie-break by id (DESC) so same-timestamp changes resolve to the complete one. self.graph_queries += 1 conn = self._connection() - graph, blob = f'"{self._graph_name}"', f'"{self._graph_name}_blob"' + graph, blob = f'"{GRAPH_STREAM}"', f'"{GRAPH_STREAM}_blob"' row = conn.execute( f"SELECT x.data FROM {graph} g JOIN {blob} x ON x.id = g.id " "WHERE g.ts <= ? ORDER BY g.ts DESC, g.id DESC LIMIT 1", From ad886c2bdc44a2a56f01142d109ecb9f0ae2fcba Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 29 Jun 2026 22:39:57 -0500 Subject: [PATCH 42/69] refactor(recorder): drop hardcoded /tf_static raw LCM topic Static mount frames are re-published onto the regular tf stream by the rig's StaticTfPublisher (5Hz), so the recorder already captures them via the tf subscription; the separate hardcoded Topic('/tf_static') had no publisher. Record tf through the one tf stream; static-vs-dynamic is recovered offline by DbTf's graph build (pose-constancy). --- dimos/memory2/module.py | 57 +++++++++++++++++------------------------ 1 file changed, 24 insertions(+), 33 deletions(-) diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index 5715a31f10..e067d588db 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -41,7 +41,6 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.protocol.pubsub.impl.lcmpubsub import Topic from dimos.utils.data import backup_file from dimos.utils.logging_config import setup_logger @@ -454,8 +453,13 @@ def _collect_pose_setters(self) -> dict[str, PoseSetter]: return setters def _record_tf(self) -> None: - """Record the live tf + static_tf streams under "tf", writing the topology - change-log (tf_graph) as we go (no-op without a pubsub tf).""" + """Record the live tf stream under "tf", writing the topology change-log + (tf_graph) as we go (no-op without a pubsub tf). + + Static mount frames are re-published onto this same tf stream by the rig's + StaticTfPublisher, so they're captured here too — there's no separate raw + static-tf topic. (Static-vs-dynamic is recovered offline by ``DbTf``'s graph + build from whether a frame's pose ever changes.)""" topic = getattr(self.tf.config, "topic", None) pubsub = getattr(self.tf, "pubsub", None) if not topic or pubsub is None: @@ -467,38 +471,25 @@ def _record_tf(self) -> None: # the "tf_graph" stream is the topology change-log transform lookups walk. structure: dict[str, dict[str, Any]] = {} - def record_graph(child: str, parent: str, is_static: bool, ts: float) -> None: - entry = {"parent": parent, "static": bool(is_static)} + def record_graph(child: str, parent: str, ts: float) -> None: + entry = {"parent": parent, "static": False} if structure.get(child) == entry: return # no structural change -> no new snapshot structure[child] = entry graph_stream.append(TfGraph(structure), ts=ts) - def make_handler(is_static: bool) -> Any: - def on_tf(msg: TFMessage, _topic: Any) -> None: - try: - for transform in msg.transforms: - tf_stream.append( - TFMessage(transform), - ts=transform.ts, - pose=None, - tags={"child_frame": transform.child_frame_id}, - ) - record_graph( - transform.child_frame_id, - transform.frame_id, - is_static, - transform.ts, - ) - except sqlite3.ProgrammingError: - # A late LCM callback raced teardown and hit the closed store. - pass - - return on_tf - - self.register_disposable(Disposable(pubsub.subscribe(topic, make_handler(False)))) - # Also listen on the static_tf stream (nothing publishes it yet); these are - # recorded the same way but flagged static in the topology. - self.register_disposable( - Disposable(pubsub.subscribe(Topic("/tf_static", TFMessage), make_handler(True))) - ) + def on_tf(msg: TFMessage, _topic: Any) -> None: + try: + for transform in msg.transforms: + tf_stream.append( + TFMessage(transform), + ts=transform.ts, + pose=None, + tags={"child_frame": transform.child_frame_id}, + ) + record_graph(transform.child_frame_id, transform.frame_id, transform.ts) + except sqlite3.ProgrammingError: + # A late LCM callback raced teardown and hit the closed store. + pass + + self.register_disposable(Disposable(pubsub.subscribe(topic, on_tf))) From 564cbce8ab0c03e2f62338606ef24f2e32d93f39 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 30 Jun 2026 07:01:06 -0500 Subject: [PATCH 43/69] feat(recorder): record tf_static via an In-port stream, not a hardcoded LCM topic Add an optional tf_static: In[TFMessage] port (a future system publishes latched mount/extrinsic transforms there). Folded into the 'tf' stream + flagged static in tf_graph; excluded from the generic per-port recording. Subscribed only when wired (guards the unconnected transport), so it's a no-op until something publishes. --- dimos/memory2/module.py | 62 ++++++++++++++++++++++++++--------------- 1 file changed, 39 insertions(+), 23 deletions(-) diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index e067d588db..956b92563c 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -30,6 +30,7 @@ from dimos.constants import DIMOS_PROJECT_ROOT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In from dimos.memory2.db_tf import TfGraph from dimos.memory2.embed import EmbedImages from dimos.memory2.store.null import NullStore @@ -47,7 +48,7 @@ if TYPE_CHECKING: from reactivex.abc import DisposableBase - from dimos.core.stream import In, Out + from dimos.core.stream import Out from dimos.msgs.geometry_msgs.Pose import Pose logger = setup_logger() @@ -318,6 +319,11 @@ async def _lidar_pose(self, msg): config: RecorderConfig + # Optional static-tf input stream: a future system publishes latched mount/extrinsic + # transforms here; recorded into "tf" + flagged static in the graph. Unconnected = + # no-op (today nothing publishes it). Folded into tf, not recorded as its own stream. + tf_static: In[TFMessage] + _pose_setters: dict[str, Any] = {} # Per-stream count of frames lost to the dispatcher's LATEST coalescing # (sink slower than input). Populated lazily as drops happen. @@ -363,6 +369,8 @@ def start(self) -> None: return for name, port in self.inputs.items(): + if name == "tf_static": + continue # folded into the "tf" stream + graph by _record_tf stream_name = self.config.stream_remapping.get(name, name) stream: Stream[Any] = self.store.stream(stream_name, port.type) self._port_to_stream(name, port, stream) @@ -453,32 +461,20 @@ def _collect_pose_setters(self) -> dict[str, PoseSetter]: return setters def _record_tf(self) -> None: - """Record the live tf stream under "tf", writing the topology change-log - (tf_graph) as we go (no-op without a pubsub tf). + """Record tf into the "tf" stream + the topology change-log ("tf_graph"). - Static mount frames are re-published onto this same tf stream by the rig's - StaticTfPublisher, so they're captured here too — there's no separate raw - static-tf topic. (Static-vs-dynamic is recovered offline by ``DbTf``'s graph - build from whether a frame's pose ever changes.)""" - topic = getattr(self.tf.config, "topic", None) - pubsub = getattr(self.tf, "pubsub", None) - if not topic or pubsub is None: - logger.warning("Recorder: no pubsub tf available — not recording tf") - return + Two inputs, both folded into one tf stream: the live (dynamic) tf via the + module's tf interface, and the optional ``tf_static`` In-port stream (a + future system publishes latched mount/extrinsic transforms there). Frames + from tf_static are flagged static in the graph; latched statics resolve + for all time, dynamic frames bracket+interpolate.""" tf_stream = self.store.stream("tf", TFMessage) graph_stream = self.store.stream("tf_graph", TfGraph) # Running tf topology; a TfGraph snapshot is appended whenever it changes, so # the "tf_graph" stream is the topology change-log transform lookups walk. structure: dict[str, dict[str, Any]] = {} - def record_graph(child: str, parent: str, ts: float) -> None: - entry = {"parent": parent, "static": False} - if structure.get(child) == entry: - return # no structural change -> no new snapshot - structure[child] = entry - graph_stream.append(TfGraph(structure), ts=ts) - - def on_tf(msg: TFMessage, _topic: Any) -> None: + def record(msg: TFMessage, is_static: bool) -> None: try: for transform in msg.transforms: tf_stream.append( @@ -487,9 +483,29 @@ def on_tf(msg: TFMessage, _topic: Any) -> None: pose=None, tags={"child_frame": transform.child_frame_id}, ) - record_graph(transform.child_frame_id, transform.frame_id, transform.ts) + entry = {"parent": transform.frame_id, "static": is_static} + if structure.get(transform.child_frame_id) != entry: + structure[transform.child_frame_id] = entry + graph_stream.append(TfGraph(structure), ts=transform.ts) except sqlite3.ProgrammingError: - # A late LCM callback raced teardown and hit the closed store. + # A late callback raced teardown and hit the closed store. pass - self.register_disposable(Disposable(pubsub.subscribe(topic, on_tf))) + # static tf: the tf_static In-port stream. Only subscribe when something is + # wired to it (its transport is set on connect) — unconnected today. + if self.tf_static.transport is not None: + + async def on_static(msg: TFMessage) -> None: + record(msg, is_static=True) + + self.process_observable(self.tf_static.pure_observable(), on_static) + + # dynamic tf: the module's live tf interface + topic = getattr(self.tf.config, "topic", None) + pubsub = getattr(self.tf, "pubsub", None) + if not topic or pubsub is None: + logger.warning("Recorder: no pubsub tf available — recording static tf only") + return + self.register_disposable( + Disposable(pubsub.subscribe(topic, lambda msg, _t: record(msg, False))) + ) From a008e91bf4bd9c1b0eb7f0feebb2952dc57004c4 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 30 Jun 2026 14:27:04 -0500 Subject: [PATCH 44/69] memory2: record tf_deformation_nodes stream (per-node pose-graph keyframes) Add DeformationNode msg (id, tf_id=FNV-1a-64(frame_from|frame_to), pose) and a Recorder tf_deformation_nodes In-port recorded by default (no-op when unconnected, like tf_static). Basis for loop-closure-corrected tf at query time. --- dimos/memory2/module.py | 36 ++++++++ dimos/msgs/nav_msgs/DeformationNode.py | 112 +++++++++++++++++++++++++ 2 files changed, 148 insertions(+) create mode 100644 dimos/msgs/nav_msgs/DeformationNode.py diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index 956b92563c..097dcec2c6 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -40,6 +40,7 @@ from dimos.memory2.type.observation import EmbeddedObservation, Observation from dimos.models.embedding.base import EmbeddingModel from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.nav_msgs.DeformationNode import DeformationNode from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.utils.data import backup_file @@ -324,6 +325,12 @@ async def _lidar_pose(self, msg): # no-op (today nothing publishes it). Folded into tf, not recorded as its own stream. tf_static: In[TFMessage] + # Optional pose-graph deformation stream: a loop-closure backend (e.g. gsc_pgo) + # publishes one DeformationNode per keyframe on create and again whenever the + # optimizer moves it. Recorded into its own "tf_deformation_nodes" stream so a + # query can later correct tf for loop closure. Unconnected = no-op. + tf_deformation_nodes: In[DeformationNode] + _pose_setters: dict[str, Any] = {} # Per-stream count of frames lost to the dispatcher's LATEST coalescing # (sink slower than input). Populated lazily as drops happen. @@ -371,6 +378,8 @@ def start(self) -> None: for name, port in self.inputs.items(): if name == "tf_static": continue # folded into the "tf" stream + graph by _record_tf + if name == "tf_deformation_nodes": + continue # recorded by _record_tf into its own stream (carries its own pose) stream_name = self.config.stream_remapping.get(name, name) stream: Stream[Any] = self.store.stream(stream_name, port.type) self._port_to_stream(name, port, stream) @@ -431,6 +440,7 @@ def _prepare_streams(self) -> None: targets = {self.config.stream_remapping.get(name, name) for name in self.inputs} if self.config.record_tf: targets.add("tf") + targets.add("tf_deformation_nodes") for stream in targets.intersection(self.store.list_streams()): self.store.delete_stream(stream) @@ -500,6 +510,8 @@ async def on_static(msg: TFMessage) -> None: self.process_observable(self.tf_static.pure_observable(), on_static) + self._record_deformation_nodes() + # dynamic tf: the module's live tf interface topic = getattr(self.tf.config, "topic", None) pubsub = getattr(self.tf, "pubsub", None) @@ -509,3 +521,27 @@ async def on_static(msg: TFMessage) -> None: self.register_disposable( Disposable(pubsub.subscribe(topic, lambda msg, _t: record(msg, False))) ) + + def _record_deformation_nodes(self) -> None: + """Record the optional ``tf_deformation_nodes`` In-port into its own stream. + + Each DeformationNode is one pose-graph keyframe; the backend re-publishes a + node (same ``id``) when the optimizer moves it, so rows accumulate and a query + takes the latest per ``id``. Tagged by ``tf_id`` (which transform edge) and + ``id`` so lookups can filter by edge and dedup by node. Unconnected = no-op.""" + if self.tf_deformation_nodes.transport is None: + return + stream = self.store.stream("tf_deformation_nodes", DeformationNode) + + async def on_node(node: DeformationNode) -> None: + try: + stream.append( + node, + ts=node.pose.ts, + pose=None, + tags={"tf_id": str(node.tf_id), "id": str(node.id)}, + ) + except sqlite3.ProgrammingError: + pass # late callback raced teardown and hit the closed store + + self.process_observable(self.tf_deformation_nodes.pure_observable(), on_node) diff --git a/dimos/msgs/nav_msgs/DeformationNode.py b/dimos/msgs/nav_msgs/DeformationNode.py new file mode 100644 index 0000000000..4238071efc --- /dev/null +++ b/dimos/msgs/nav_msgs/DeformationNode.py @@ -0,0 +1,112 @@ +# 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. + +"""DeformationNode: one pose-graph keyframe, published individually (not batched). + +A loop-closure backend (e.g. gsc_pgo) emits one of these per keyframe when the node +is created and again whenever the optimizer moves it. A recording captures the stream +of them so a transform lookup can, at query time, find the keyframes near a time and +read their most-recent optimized world pose — the basis for loop-closure-corrected tf. + +Fields: + * ``id`` — a stable, random uint64 identifying the keyframe (reused across the + node's re-publishes; random rather than monotonic so it carries no + ordering/index coupling and won't collide across robots/sessions). + * ``tf_id`` — uint64 = :func:`tf_id_for` (an FNV-1a-64 hash of + ``frame_from + "|" + frame_to``). Identifies which transform edge + these poses deform, so multi-robot systems with prefixed frames (and + several concurrent loop closures) can filter to the right one. + * ``pose`` — the keyframe's world pose as a ``PoseStamped`` (carries its timestamp). +""" + +from __future__ import annotations + +import struct +from typing import BinaryIO + +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.types.timestamped import Timestamped + +_FNV_OFFSET_BASIS_64 = 0xCBF29CE484222325 +_FNV_PRIME_64 = 0x100000001B3 +_U64_MASK = 0xFFFFFFFFFFFFFFFF + + +def fnv1a_64(text: str) -> int: + """64-bit FNV-1a hash of ``text`` (utf-8). Must match the C++ producer's hash so + ``tf_id`` filtering agrees across the wire.""" + digest = _FNV_OFFSET_BASIS_64 + for byte in text.encode("utf-8"): + digest = ((digest ^ byte) * _FNV_PRIME_64) & _U64_MASK + return digest + + +def tf_id_for(frame_from: str, frame_to: str) -> int: + """The ``tf_id`` for a transform edge: ``fnv1a_64(frame_from + "|" + frame_to)``.""" + return fnv1a_64(f"{frame_from}|{frame_to}") + + +class DeformationNode(Timestamped): + msg_name = "nav_msgs.DeformationNode" + + id: int + tf_id: int + pose: PoseStamped + + def __init__(self, id: int = 0, tf_id: int = 0, pose: PoseStamped | None = None) -> None: + self.id = id + self.tf_id = tf_id + self.pose = pose if pose is not None else PoseStamped() + self.ts = self.pose.ts + + def __repr__(self) -> str: + return f"DeformationNode(id={self.id}, tf_id={self.tf_id}, ts={self.pose.ts})" + + def lcm_encode(self) -> bytes: + frame_id_bytes = self.pose.frame_id.encode("utf-8") + return b"".join( + ( + struct.pack(">QQd", self.id, self.tf_id, self.pose.ts), + struct.pack(">I", len(frame_id_bytes)), + frame_id_bytes, + struct.pack( + ">7d", + self.pose.position.x, + self.pose.position.y, + self.pose.position.z, + self.pose.orientation.x, + self.pose.orientation.y, + self.pose.orientation.z, + self.pose.orientation.w, + ), + ) + ) + + @classmethod + def lcm_decode(cls, data: bytes | BinaryIO) -> DeformationNode: + buf = data if isinstance(data, (bytes, bytearray)) else data.read() + node_id, tf_id, pose_ts = struct.unpack_from(">QQd", buf, 0) + offset = 24 + (frame_id_len,) = struct.unpack_from(">I", buf, offset) + offset += 4 + frame_id = bytes(buf[offset : offset + frame_id_len]).decode("utf-8") + offset += frame_id_len + px, py, pz, qx, qy, qz, qw = struct.unpack_from(">7d", buf, offset) + pose = PoseStamped( + ts=pose_ts, + frame_id=frame_id, + position=[px, py, pz], + orientation=[qx, qy, qz, qw], + ) + return cls(id=node_id, tf_id=tf_id, pose=pose) From 212d63e9f02719b3fac595a9d26d5d3949753480 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 30 Jun 2026 14:39:59 -0500 Subject: [PATCH 45/69] gsc_pgo: declare tf_deformation_nodes Out port Autoconnects to the Recorder's tf_deformation_nodes In. Binary publishes to it once the pinned rev is bumped (gated on pushing the gsc_pgo C++ change). --- .../jnav/components/loop_closure/gsc_pgo/module.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py index c0957810b0..3f4a79b2c3 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py @@ -27,6 +27,7 @@ from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.nav_msgs.DeformationNode import DeformationNode from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.navigation.jnav.msgs.Graph3D import Graph3D @@ -179,6 +180,10 @@ class PGO(NativeModule): correction: Out[Transform] pose_graph: Out[Graph3D] loop_closure_event: Out[GraphDelta3D] + # Per-keyframe pose-graph nodes, published individually (un-batched) so a + # recorder can stream them. tf_id on each identifies the corrected edge + # (frame_id -> child_frame_id). Autoconnects to the Recorder's like-named port. + tf_deformation_nodes: Out[DeformationNode] # Internal/debug only (off by default) — see global_map_publish_rate. Named # with a leading underscore so autoconnect won't wire it to `global_map` Ins. _global_map: Out[PointCloud2] From b78839e51ee19b8c7e70bd68b6b074129dd8f98f Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 30 Jun 2026 14:47:20 -0500 Subject: [PATCH 46/69] gsc_pgo: bump pinned rev to e8d910f (DeformationNode stream) --- dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py index 3f4a79b2c3..6f61728bb8 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py @@ -45,7 +45,7 @@ class PGOConfig(NativeModuleConfig): cwd: str | None = str(Path(__file__).resolve().parent) executable: str = "result/bin/pgo" build_command: str | None = ( - 'nix build "github:jeff-hykin/gsc_pgo/e5f9165164e124df16535fa6b8616adfae1d10df#default"' + 'nix build "github:jeff-hykin/gsc_pgo/e8d910f305719b81efabc6b8e1035ee7020369e9#default"' " --no-write-lock-file" ) From dab4fe711b0657b9dbd93acf97be1d155d36df12 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 30 Jun 2026 14:54:41 -0500 Subject: [PATCH 47/69] DbTf: query-time loop-closure correction from tf_deformation_nodes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For each edge in the resolved chain, hash (parent|child) -> tf_id; if the deformation stream has matching keyframes, take the ones bracketing the query time, compute each node's current ∘ inv(original) delta, linear-blend-skin between them, and apply on the parent side of that edge. Edges with no match (and recordings with no deformation stream) take the normal path at zero cost. --- dimos/memory2/db_tf.py | 199 +++++++++++++++++++++++++++++++++++- dimos/memory2/test_db_tf.py | 80 +++++++++++++++ 2 files changed, 278 insertions(+), 1 deletion(-) diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py index ef25579bba..5c354bd8c2 100644 --- a/dimos/memory2/db_tf.py +++ b/dimos/memory2/db_tf.py @@ -24,8 +24,12 @@ import threading from typing import TYPE_CHECKING, Any, cast +import numpy as np + from dimos.memory2.store.sqlite import SqliteStoreConfig +from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.nav_msgs.DeformationNode import DeformationNode, tf_id_for from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.protocol.tf.tf import MultiTBuffer from dimos.utils.logging_config import setup_logger @@ -39,6 +43,11 @@ DEFAULT_TF_STREAM = "tf" # The topology change-log is a single companion stream (like tf_static). GRAPH_STREAM = "tf_graph" +# Per-keyframe pose-graph nodes from a loop-closure backend (e.g. gsc_pgo). Each row +# is one DeformationNode, tagged by tf_id (which edge it corrects) and id (the +# keyframe). A query corrects matching edges using how far the optimizer has moved +# the bracketing keyframes since they were first recorded. +DEFORMATION_STREAM = "tf_deformation_nodes" # Streams the RAM fallback (non-sqlite stores) reads. TF_STREAMS = ("tf", "tf_static") # Cache the whole change-log in RAM when there are at most this many topology @@ -82,6 +91,9 @@ def __init__( self._graph_loaded = False self._static_cache: dict[str, Transform] = {} self._buffer: MultiTBuffer | None = None # RAM fallback (non-sqlite) + # The set of tf_ids present on the deformation stream (the edges loop closure + # corrects). Cached once so a recording with no deformation nodes pays nothing. + self._deformation_tf_ids: set[int] | None = None self.rows_fetched = 0 self.graph_queries = 0 @@ -331,6 +343,110 @@ def pick(frame: str, kind: str, where_ts: str, order: str) -> None: self.rows_fetched += 1 return rows + # --- loop-closure deformation (sqlite) ------------------------------------- + + def _load_deformation_ids(self) -> set[int]: + """The distinct tf_ids present on the deformation stream, cached once. Empty + when there's no such stream — so the correction path is skipped at zero cost + for recordings without loop closure.""" + if self._deformation_tf_ids is not None: + return self._deformation_tf_ids + ids: set[int] = set() + if self._is_sqlite and DEFORMATION_STREAM in set(self._store.list_streams()): + for (tf_id,) in self._connection().execute( + f"SELECT DISTINCT json_extract(tags, '$.tf_id') FROM \"{DEFORMATION_STREAM}\"" + ): + if tf_id is not None: + ids.add(int(tf_id)) + self._deformation_tf_ids = ids + return ids + + def _deformation_codec(self) -> Any: + source = self._store.stream(DEFORMATION_STREAM, DeformationNode)._source + return cast("Any", source).codec + + def _node_pose_matrix(self, order: str, tf_id: int, node_id: str) -> np.ndarray | None: + """The 4x4 pose of one keyframe's first (``ASC``) or latest (``DESC``) recorded + version. Versions of a node share the keyframe ts, so they're ordered by row id + (insertion order), not ts.""" + stream, blob = f'"{DEFORMATION_STREAM}"', f'"{DEFORMATION_STREAM}_blob"' + row = ( + self._connection() + .execute( + f"SELECT b.data FROM {stream} s JOIN {blob} b ON b.id = s.id " + f"WHERE json_extract(s.tags, '$.tf_id') = ? AND json_extract(s.tags, '$.id') = ? " + f"ORDER BY s.id {order} LIMIT 1", + (str(tf_id), node_id), + ) + .fetchone() + ) + if row is None: + return None + node = cast("DeformationNode", self._deformation_codec().decode(row[0])) + return Transform.from_pose(node.pose.frame_id, node.pose).to_matrix() + + def _node_delta(self, tf_id: int, node_id: str) -> np.ndarray | None: + """How far the optimizer has moved a keyframe since it was first recorded: + ``current ∘ inv(original)`` — the SE(3) correction this node contributes.""" + original = self._node_pose_matrix("ASC", tf_id, node_id) + current = self._node_pose_matrix("DESC", tf_id, node_id) + if original is None or current is None: + return None + return current @ np.linalg.inv(original) + + def _edge_delta(self, tf_id: int, query_time: float) -> np.ndarray | None: + """The blended correction for an edge at ``query_time``: take the keyframes + bracketing the time (latest at-or-before, earliest at-or-after), each node's + ``current ∘ inv(original)`` delta, and linear-blend-skin between them.""" + stream = f'"{DEFORMATION_STREAM}"' + id_tag = "json_extract(tags, '$.id')" + tf_tag = "json_extract(tags, '$.tf_id')" + conn = self._connection() + lo = conn.execute( + f"SELECT {id_tag}, ts FROM {stream} WHERE {tf_tag} = ? AND ts <= ? " + "ORDER BY ts DESC, id DESC LIMIT 1", + (str(tf_id), query_time), + ).fetchone() + hi = conn.execute( + f"SELECT {id_tag}, ts FROM {stream} WHERE {tf_tag} = ? AND ts >= ? " + "ORDER BY ts ASC, id ASC LIMIT 1", + (str(tf_id), query_time), + ).fetchone() + samples: list[tuple[float, np.ndarray]] = [] + seen_ids: set[str] = set() + for node_id, ts in (row for row in (lo, hi) if row is not None): + if node_id in seen_ids: + continue + seen_ids.add(node_id) + delta = self._node_delta(tf_id, node_id) + if delta is not None: + samples.append((ts, delta)) + if not samples: + return None + if len(samples) == 1: + return samples[0][1] + (ts_lo, mat_lo), (ts_hi, mat_hi) = sorted(samples, key=lambda item: item[0]) + weight = 0.0 if ts_hi == ts_lo else (query_time - ts_lo) / (ts_hi - ts_lo) + return _blend_se3(mat_lo, mat_hi, weight) + + def _edge_corrections( + self, graph: dict[str, Any], edges: list[str], query_time: float + ) -> dict[str, np.ndarray]: + """For each edge whose tf_id matches the deformation stream, its blended SE(3) + correction (applied on the parent side of the edge). Empty when nothing matches.""" + tf_ids = self._load_deformation_ids() + if not tf_ids: + return {} + corrections: dict[str, np.ndarray] = {} + for frame in edges: + tf_id = tf_id_for(graph[frame]["parent"], frame) + if tf_id not in tf_ids: + continue + delta = self._edge_delta(tf_id, query_time) + if delta is not None: + corrections[frame] = delta + return corrections + def get( self, target_frame: str, @@ -360,6 +476,8 @@ def get( uncached_static = [f for f in static if f not in self._static_cache] rows = self._fetch_rows(dynamic, uncached_static, query_time) # ONE detail query + # Per-edge loop-closure corrections (empty unless a deformation stream matches). + corrections = self._edge_corrections(graph, edges, query_time) buffer = MultiTBuffer(buffer_size=_NO_PRUNE) for frame in static: @@ -372,19 +490,54 @@ def get( self._static_cache[frame] = transform # restamp the constant to query_time so the buffer's tolerance never # rejects a static that was recorded long ago (latched once). - buffer.receive_transform(_restamp(transform, query_time)) + transform = _restamp(transform, query_time) + if frame in corrections: + transform = _apply_delta(corrections[frame], transform, query_time) + buffer.receive_transform(transform) for frame in dynamic: lo = rows.get((frame, "lo")) hi = rows.get((frame, "hi")) chosen = lo if lo is not None else hi if chosen is None: return None + if frame in corrections: + # Resolve the raw edge at query_time, then deform it. The correction + # already carries the time blend, so a single corrected sample suffices. + raw = self._interpolate_dynamic( + frame, graph[frame]["parent"], lo, hi, query_time, time_tolerance + ) + if raw is None: + return None + buffer.receive_transform(_apply_delta(corrections[frame], raw, query_time)) + continue buffer.receive_transform(self._decode_blob(chosen, frame)) other = hi if hi is not None else lo if other is not None and other is not chosen: buffer.receive_transform(self._decode_blob(other, frame)) return buffer.lookup(target_frame, source_frame, time_point, time_tolerance) + def _interpolate_dynamic( + self, + frame: str, + parent: str, + lo: bytes | None, + hi: bytes | None, + query_time: float, + time_tolerance: float | None, + ) -> Transform | None: + """The raw ``parent <- frame`` transform at ``query_time``, interpolated from + the bracketing rows (its own small buffer so the chain buffer only ever sees + the corrected result).""" + edge_buffer = MultiTBuffer(buffer_size=_NO_PRUNE) + chosen = lo if lo is not None else hi + if chosen is None: + return None + edge_buffer.receive_transform(self._decode_blob(chosen, frame)) + other = hi if hi is not None else lo + if other is not None and other is not chosen: + edge_buffer.receive_transform(self._decode_blob(other, frame)) + return edge_buffer.lookup(parent, frame, query_time, time_tolerance) + class TfGraph: """A tf topology snapshot, recorded one per structure change. @@ -418,6 +571,50 @@ def _restamp(transform: Transform, ts: float) -> Transform: ) +def _apply_delta(delta: np.ndarray, transform: Transform, ts: float) -> Transform: + """Deform ``transform`` by the SE(3) correction ``delta``, applied on the parent + (frame_id) side: ``corrected = delta @ raw``. Keeps the edge's frame names.""" + corrected = delta @ transform.to_matrix() + return Transform.from_matrix( + corrected, + ts=ts, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ) + + +def _quat_slerp(q_lo: np.ndarray, q_hi: np.ndarray, weight: float) -> np.ndarray: + """Spherical-linear interpolation between two quaternions ``[x, y, z, w]``.""" + dot = float(np.dot(q_lo, q_hi)) + if dot < 0.0: # take the shorter arc + q_hi = -q_hi + dot = -dot + if dot > 0.9995: # nearly parallel: lerp + renormalize avoids a divide-by-~0 + result = q_lo + weight * (q_hi - q_lo) + return cast("np.ndarray", result / np.linalg.norm(result)) + theta_0 = np.arccos(np.clip(dot, -1.0, 1.0)) + sin_0 = np.sin(theta_0) + scale_lo = np.sin((1.0 - weight) * theta_0) / sin_0 + scale_hi = np.sin(weight * theta_0) / sin_0 + return cast("np.ndarray", scale_lo * q_lo + scale_hi * q_hi) + + +def _blend_se3(mat_lo: np.ndarray, mat_hi: np.ndarray, weight: float) -> np.ndarray: + """Linear-blend-skin two SE(3) deltas by ``weight`` in [0, 1] (0 -> lo, 1 -> hi): + lerp the translation, slerp the rotation.""" + quat_lo = Quaternion.from_rotation_matrix(mat_lo[:3, :3]) + quat_hi = Quaternion.from_rotation_matrix(mat_hi[:3, :3]) + blended = _quat_slerp( + np.array([quat_lo.x, quat_lo.y, quat_lo.z, quat_lo.w]), + np.array([quat_hi.x, quat_hi.y, quat_hi.z, quat_hi.w]), + weight, + ) + out = np.eye(4) + out[:3, :3] = Quaternion(blended).to_rotation_matrix() + out[:3, 3] = (1.0 - weight) * mat_lo[:3, 3] + weight * mat_hi[:3, 3] + return out + + def _safe_table(name: str) -> str: if not _VAR_NAME_PATTERN.match(name): raise ValueError(f"unsafe stream/table name: {name!r}") diff --git a/dimos/memory2/test_db_tf.py b/dimos/memory2/test_db_tf.py index ac97723bc2..98ff84c28b 100644 --- a/dimos/memory2/test_db_tf.py +++ b/dimos/memory2/test_db_tf.py @@ -25,9 +25,11 @@ from dimos.memory2.db_tf import DbTf from dimos.memory2.store.sqlite import SqliteStore +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped 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.DeformationNode import DeformationNode, tf_id_for from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.protocol.tf.tf import MultiTBuffer @@ -209,6 +211,84 @@ def test_reparent_midrun_uses_graph_as_of_query_time(tmp_path: Path) -> None: store.stop() +def _append_deformation( + store: SqliteStore, tf_id: int, node_id: int, ts: float, xyz: tuple[float, float, float] +) -> None: + """Record one DeformationNode version, exactly as the recorder does: tagged by + tf_id (the edge) and id (the keyframe). Re-call with the same id to add a moved + version (the optimizer relocating that keyframe).""" + node = DeformationNode( + id=node_id, + tf_id=tf_id, + pose=PoseStamped(ts=ts, frame_id="map", position=list(xyz), orientation=[0, 0, 0, 1]), + ) + store.stream("tf_deformation_nodes", DeformationNode).append( + node, ts=ts, pose=None, tags={"tf_id": str(tf_id), "id": str(node_id)} + ) + + +def _record_corrected_edge(path: Path) -> None: + """A map->odom edge (raw = identity) with base_link hanging off odom.""" + store = SqliteStore(path=str(path)) + for i in range(20): + ts = _T0 + i / _DYN_RATE + _append(store, _static("map", "odom", (0.0, 0.0, 0.0), ts)) + _append( + store, + Transform( + translation=Vector3(i * 0.1, 0, 0), + rotation=_yaw(0), + frame_id="odom", + child_frame_id="base_link", + ts=ts, + ), + ) + store.stop() + + +def test_loop_closure_deformation_corrects_matched_edge(tmp_path: Path) -> None: + """A deformation node whose tf_id = hash(map|odom) deforms the raw map<-odom edge + by current ∘ inv(original); an edge with no matching tf_id is untouched.""" + path = tmp_path / "lc.db" + _record_corrected_edge(path) + store = SqliteStore(path=str(path)) + tf_id = tf_id_for("map", "odom") + # one keyframe at t0+0.3: the optimizer later moved it from origin to (1, 0, 0) + _append_deformation(store, tf_id, 11, _T0 + 0.3, (0.0, 0.0, 0.0)) # original + _append_deformation(store, tf_id, 11, _T0 + 0.3, (1.0, 0.0, 0.0)) # current + store.stop() + + store = SqliteStore(path=str(path), must_exist=True) + db = DbTf(store) + got = db.get("map", "odom", _T0 + 0.3, 0.5) + assert got is not None and abs(got.translation.x - 1.0) < 1e-6 # raw identity + delta + base = db.get("odom", "base_link", _T0 + 0.3, 0.5) # unmatched edge, unchanged + assert base is not None and abs(base.translation.x - 0.9) < 0.2 + store.stop() + + +def test_loop_closure_deformation_blends_between_keyframes(tmp_path: Path) -> None: + """Two bracketing keyframes (delta 0 at t0, delta 2 at t0+? ) blend at the midpoint + to delta 1 — linear blend skinning across the trajectory.""" + path = tmp_path / "lc_blend.db" + store = SqliteStore(path=str(path)) + for i in range(int(_DURATION * _DYN_RATE)): # map->odom identity across [t0, t0+10) + ts = _T0 + i / _DYN_RATE + _append(store, _static("map", "odom", (0.0, 0.0, 0.0), ts)) + tf_id = tf_id_for("map", "odom") + _append_deformation(store, tf_id, 1, _T0, (0.0, 0.0, 0.0)) # kf A: original + _append_deformation(store, tf_id, 1, _T0, (0.0, 0.0, 0.0)) # kf A: unmoved + _append_deformation(store, tf_id, 2, _T0 + 8.0, (0.0, 0.0, 0.0)) # kf B: original + _append_deformation(store, tf_id, 2, _T0 + 8.0, (2.0, 0.0, 0.0)) # kf B: moved +2 + store.stop() + + store = SqliteStore(path=str(path), must_exist=True) + db = DbTf(store) + got = db.get("map", "odom", _T0 + 4.0, 0.5) # midpoint of [t0, t0+8] + assert got is not None and abs(got.translation.x - 1.0) < 1e-6 + store.stop() + + def test_disjoint_multirobot_returns_none(tmp_path: Path) -> None: """Two unconnected components (two robots, no shared frame) → a cross-component query is None, an in-component query resolves.""" From f2a23f55afd5c0c06fae2e64036e9fd7a379f3d8 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 30 Jun 2026 15:53:22 -0500 Subject: [PATCH 48/69] cleaning --- dimos/memory2/db_tf.py | 57 +++++++++++++----------------------------- 1 file changed, 18 insertions(+), 39 deletions(-) diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py index 5c354bd8c2..c0cc18bd22 100644 --- a/dimos/memory2/db_tf.py +++ b/dimos/memory2/db_tf.py @@ -19,7 +19,6 @@ from __future__ import annotations import bisect -import re import sqlite3 import threading from typing import TYPE_CHECKING, Any, cast @@ -40,19 +39,12 @@ logger = setup_logger() -DEFAULT_TF_STREAM = "tf" -# The topology change-log is a single companion stream (like tf_static). -GRAPH_STREAM = "tf_graph" -# Per-keyframe pose-graph nodes from a loop-closure backend (e.g. gsc_pgo). Each row -# is one DeformationNode, tagged by tf_id (which edge it corrects) and id (the -# keyframe). A query corrects matching edges using how far the optimizer has moved -# the bracketing keyframes since they were first recorded. -DEFORMATION_STREAM = "tf_deformation_nodes" -# Streams the RAM fallback (non-sqlite stores) reads. +TF_STREAM = "tf" # the dynamic transform stream (one transform per row) TF_STREAMS = ("tf", "tf_static") -# Cache the whole change-log in RAM when there are at most this many topology -# changes (a stable tree — even a many-frame sensor rig — is a one-time setup, not -# churn); above it, fall back to one indexed graph query per lookup (multi-robot). +GRAPH_STREAM = "tf_graph" # changes to the tf tree (with room for non-tree structures) +DEFORMATION_STREAM = "tf_deformation_nodes" # loop closure and other deformations +# 99% of the time there are less than 10 tree changes +# this cache allows us to avoid a DB query in the offline/map-load case DEFAULT_MAX_GRAPH_CHANGES_IN_RAM = 64 # MultiTBuffer drops samples older than buffer_size seconds; we feed it exactly the # bracketing samples and want them all kept, so use a span no recording exceeds. @@ -60,8 +52,6 @@ # A frame is "static" if its pose never changes; poses are compared rounded to this # many decimals (~nanometre / nanoradian) so float noise doesn't read as motion. POSE_EQUALITY_DECIMALS = 9 -# enforce safe identifiers for SQL -_VAR_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") class DbTf: @@ -75,14 +65,10 @@ class DbTf: def __init__( self, store: Store, - stream: str = DEFAULT_TF_STREAM, max_graph_changes_in_ram: int = DEFAULT_MAX_GRAPH_CHANGES_IN_RAM, - stream_names: tuple[str, ...] = TF_STREAMS, ) -> None: self._store = store - self._stream = _safe_table(stream) self._max_in_ram = max_graph_changes_in_ram - self._stream_names = stream_names # RAM fallback only self._lock = threading.Lock() self._conn: sqlite3.Connection | None = None self._built = False @@ -120,7 +106,7 @@ def _ensure_loaded(self) -> MultiTBuffer: return self._buffer buffer = MultiTBuffer(buffer_size=_NO_PRUNE) available = set(self._store.list_streams()) - for name in self._stream_names: + for name in TF_STREAMS: if name not in available: continue for observation in self._store.stream(name, TFMessage): @@ -135,9 +121,9 @@ def has_transforms(self) -> bool: if not self._is_sqlite: return bool(self._ensure_loaded().buffers) conn = self._connection() - if self._stream not in set(self._store.list_streams()): + if TF_STREAM not in set(self._store.list_streams()): return False - (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() + (n_rows,) = conn.execute(f'SELECT count(*) FROM "{TF_STREAM}"').fetchone() return bool(n_rows) def _graph_stream(self) -> Stream[TfGraph]: @@ -155,25 +141,25 @@ def _ensure_built(self) -> None: if self._built: return conn = self._connection() - (n_rows,) = conn.execute(f'SELECT count(*) FROM "{self._stream}"').fetchone() + (n_rows,) = conn.execute(f'SELECT count(*) FROM "{TF_STREAM}"').fetchone() if n_rows and self._graph_count() == 0: logger.warning( "\n========================================================================\n" " tf graph stream MISSING for %r. Building it (one-time): tagging tf rows\n" " with child_frame and writing the topology change-log.\n" "========================================================================", - self._stream, + TF_STREAM, ) self._build_graph_stream() if n_rows: - _ensure_child_index(conn, self._stream) # tf table exists now + _ensure_child_index(conn) # tf table exists now self._built = True def _build_graph_stream(self) -> None: """One-time migration: decode every tf row, tag it with its child_frame, and append a ``TfGraph`` snapshot whenever the topology changes. A frame counts as static if its pose never varies across the whole recording.""" - safe = self._stream + safe = TF_STREAM # one decode pass: collect (id, ts, child, parent, pose-key) + per-child poses rows: list[tuple[int, float, str, str]] = [] poses_per_child: dict[str, set[tuple[float, ...]]] = {} @@ -215,7 +201,7 @@ def _build_graph_stream(self) -> None: structure[child] = entry graph_stream.append(TfGraph(structure), ts=ts) written += 1 - logger.warning("tf graph built: %d topology changes for %r.", written, self._stream) + logger.warning("tf graph built: %d topology changes for %r.", written, TF_STREAM) def _graph_codec(self) -> Any: source = self._store.stream(GRAPH_STREAM, TfGraph)._source @@ -291,7 +277,7 @@ def to_root(frame: str) -> list[str]: return frames def _codec(self) -> Any: - source = self._store.stream(self._stream, TFMessage)._source + source = self._store.stream(TF_STREAM, TFMessage)._source return cast("Any", source).codec def _decode_blob(self, data: bytes, frame: str) -> Transform: @@ -312,7 +298,7 @@ def _fetch_rows( static frame its latest row ('st') — all joined to the blob data. Keyed by (frame, kind) -> blob bytes.""" cf = "json_extract(tags, '$.child_frame')" - tf, blob = f'"{self._stream}"', f'"{self._stream}_blob"' + tf, blob = f'"{TF_STREAM}"', f'"{TF_STREAM}_blob"' # One UNION of per-frame, index-served LIMIT-1 subqueries: each is a direct # (child_frame, ts) range seek — far cheaper than a window-function scan, and # still a single round-trip. @@ -615,31 +601,24 @@ def _blend_se3(mat_lo: np.ndarray, mat_hi: np.ndarray, weight: float) -> np.ndar return out -def _safe_table(name: str) -> str: - if not _VAR_NAME_PATTERN.match(name): - raise ValueError(f"unsafe stream/table name: {name!r}") - return name - - def _connect(db_path: str) -> sqlite3.Connection: conn = sqlite3.connect(db_path, check_same_thread=False) conn.execute("PRAGMA busy_timeout=5000") return conn -def _ensure_child_index(conn: sqlite3.Connection, stream: str) -> None: +def _ensure_child_index(conn: sqlite3.Connection) -> None: """Index the child_frame json tag on the tf rows so per-frame time queries seek. The live recorder gets this for free (the store auto-indexes tag keys on tagged appends); this is for migrated recordings and the read side. Requires the tf table to exist.""" - safe = _safe_table(stream) # Composite (child_frame, ts) so a per-frame "latest at/before T" is a direct # index range seek, not a scan+sort. Index names share SQLite's global namespace # with tables, so the name is double-underscore-namespaced to keep it clear of any # real stream/table name (no stream would contain "__dbtf_"). - index_name = f"{safe}__dbtf_child_ts_idx" + index_name = f"{TF_STREAM}__dbtf_child_ts_idx" conn.execute( f'CREATE INDEX IF NOT EXISTS "{index_name}" ' - f"ON \"{safe}\"(json_extract(tags, '$.child_frame'), ts)" + f"ON \"{TF_STREAM}\"(json_extract(tags, '$.child_frame'), ts)" ) conn.commit() From 0d297f947cc25756f8abe120c8adb20f263b7cf4 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Jul 2026 07:52:34 -0500 Subject: [PATCH 49/69] memory2/tf: split DbTf into live+sql, add loop-closure deformation correction Split db_tf into DbTfLive (in-RAM buffer, all stores) and DbTfSql (indexed per-lookup queries + deformation correction, sqlite only). SemanticSearch now resolves a match's world pose at query time via store.tf so it reflects loop closure. PubSubTF gains a deformation_topic that corrects live tf.get from a tf_deformation_nodes stream. Recorder only clears streams whose source is connected (pcap_to_db no longer wipes an existing tf tree). Drop unused Marker msg. --- .../sensors/lidar/pointlio/recorder.py | 2 + dimos/memory2/db_tf.py | 586 +----------------- dimos/memory2/db_tf_live.py | 71 +++ dimos/memory2/db_tf_sql.py | 501 +++++++++++++++ dimos/memory2/module.py | 51 +- dimos/memory2/store/base.py | 4 +- dimos/memory2/store/sqlite.py | 16 +- dimos/memory2/test_db_tf.py | 17 +- dimos/navigation/jnav/msgs/Marker.py | 97 --- dimos/navigation/jnav/msgs/test_Marker.py | 67 -- dimos/protocol/tf/deformation.py | 143 +++++ dimos/protocol/tf/test_deformation.py | 74 +++ dimos/protocol/tf/tf.py | 51 ++ 13 files changed, 917 insertions(+), 763 deletions(-) create mode 100644 dimos/memory2/db_tf_live.py create mode 100644 dimos/memory2/db_tf_sql.py delete mode 100644 dimos/navigation/jnav/msgs/Marker.py delete mode 100644 dimos/navigation/jnav/msgs/test_Marker.py create mode 100644 dimos/protocol/tf/deformation.py create mode 100644 dimos/protocol/tf/test_deformation.py diff --git a/dimos/hardware/sensors/lidar/pointlio/recorder.py b/dimos/hardware/sensors/lidar/pointlio/recorder.py index e859190351..f7eff72c78 100644 --- a/dimos/hardware/sensors/lidar/pointlio/recorder.py +++ b/dimos/hardware/sensors/lidar/pointlio/recorder.py @@ -34,6 +34,8 @@ class PointlioRecorderConfig(RecorderConfig): # Append into a populated db (keep other streams); replace only our own. on_existing: OnExisting = OnExisting.APPEND + # pcap_to_db has no tf source of its own — don't record (and thus don't touch) tf. + record_tf: bool = False class PointlioRecorder(Recorder): diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py index c0cc18bd22..d1006bb68e 100644 --- a/dimos/memory2/db_tf.py +++ b/dimos/memory2/db_tf.py @@ -13,425 +13,25 @@ # limitations under the License. """ -A tf class for memory2 +The read-side transform-lookup interface for memory2, plus the graph-stream payload. + +Kept dependency-light (no sqlite/numpy) so importing :class:`DbTf` or :class:`TfGraph` +costs nothing. The implementations live in sibling modules: + * :class:`dimos.memory2.db_tf_live.DbTfLive` — in-RAM buffer; default for any Store. + * :class:`dimos.memory2.db_tf_sql.DbTfSql` — sqlite graph-stream; sqlite store only. """ from __future__ import annotations -import bisect -import sqlite3 -import threading -from typing import TYPE_CHECKING, Any, cast - -import numpy as np - -from dimos.memory2.store.sqlite import SqliteStoreConfig -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.nav_msgs.DeformationNode import DeformationNode, tf_id_for -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.protocol.tf.tf import MultiTBuffer -from dimos.utils.logging_config import setup_logger +from typing import TYPE_CHECKING, Any, Protocol if TYPE_CHECKING: - from dimos.memory2.store.base import Store - from dimos.memory2.stream import Stream - -logger = setup_logger() - -TF_STREAM = "tf" # the dynamic transform stream (one transform per row) -TF_STREAMS = ("tf", "tf_static") -GRAPH_STREAM = "tf_graph" # changes to the tf tree (with room for non-tree structures) -DEFORMATION_STREAM = "tf_deformation_nodes" # loop closure and other deformations -# 99% of the time there are less than 10 tree changes -# this cache allows us to avoid a DB query in the offline/map-load case -DEFAULT_MAX_GRAPH_CHANGES_IN_RAM = 64 -# MultiTBuffer drops samples older than buffer_size seconds; we feed it exactly the -# bracketing samples and want them all kept, so use a span no recording exceeds. -_NO_PRUNE = 1.0e15 -# A frame is "static" if its pose never changes; poses are compared rounded to this -# many decimals (~nanometre / nanoradian) so float noise doesn't read as motion. -POSE_EQUALITY_DECIMALS = 9 - - -class DbTf: - """Transform lookups backed by a store's recorded transforms. - - On a SQLite store this uses the graph stream + an in-RAM graph cache; other - stores fall back to loading the tf streams into a :class:`MultiTBuffer`. Surface - is ``get(target, source, time_point, time_tolerance)`` / ``has_transforms()``. - """ - - def __init__( - self, - store: Store, - max_graph_changes_in_ram: int = DEFAULT_MAX_GRAPH_CHANGES_IN_RAM, - ) -> None: - self._store = store - self._max_in_ram = max_graph_changes_in_ram - self._lock = threading.Lock() - self._conn: sqlite3.Connection | None = None - self._built = False - # graph cache: either the whole change-log in RAM, or None (query per lookup) - self._graph_in_ram: list[tuple[float, dict[str, Any]]] | None = None - self._graph_loaded = False - self._static_cache: dict[str, Transform] = {} - self._buffer: MultiTBuffer | None = None # RAM fallback (non-sqlite) - # The set of tf_ids present on the deformation stream (the edges loop closure - # corrects). Cached once so a recording with no deformation nodes pays nothing. - self._deformation_tf_ids: set[int] | None = None - self.rows_fetched = 0 - self.graph_queries = 0 - - @property - def _is_sqlite(self) -> bool: - return isinstance(self._store.config, SqliteStoreConfig) - - def _connection(self) -> sqlite3.Connection: - conn = self._conn - if conn is None: - config = self._store.config - assert isinstance(config, SqliteStoreConfig) # guarded by _is_sqlite - conn = _connect(config.path) - self._conn = conn - return conn - - # --- RAM fallback (non-sqlite stores) -------------------------------------- - - def _ensure_loaded(self) -> MultiTBuffer: - if self._buffer is not None: - return self._buffer - with self._lock: - if self._buffer is not None: - return self._buffer - buffer = MultiTBuffer(buffer_size=_NO_PRUNE) - available = set(self._store.list_streams()) - for name in TF_STREAMS: - if name not in available: - continue - for observation in self._store.stream(name, TFMessage): - transforms = getattr(observation.data, "transforms", None) or [observation.data] - buffer.receive_transform(*transforms) - self._buffer = buffer - return buffer - - # --- graph stream (sqlite) ------------------------------------------------- - - def has_transforms(self) -> bool: - if not self._is_sqlite: - return bool(self._ensure_loaded().buffers) - conn = self._connection() - if TF_STREAM not in set(self._store.list_streams()): - return False - (n_rows,) = conn.execute(f'SELECT count(*) FROM "{TF_STREAM}"').fetchone() - return bool(n_rows) - - def _graph_stream(self) -> Stream[TfGraph]: - return self._store.stream(GRAPH_STREAM, TfGraph) - - def _graph_count(self) -> int: - if GRAPH_STREAM not in set(self._store.list_streams()): - return 0 - return self._graph_stream().count() - - def _ensure_built(self) -> None: - """First sqlite use: if the recording has tf rows but no graph stream (a - recording that predates it / wasn't written by the recorder), build the graph - stream once by replaying the tf rows, then make sure the seek index exists.""" - if self._built: - return - conn = self._connection() - (n_rows,) = conn.execute(f'SELECT count(*) FROM "{TF_STREAM}"').fetchone() - if n_rows and self._graph_count() == 0: - logger.warning( - "\n========================================================================\n" - " tf graph stream MISSING for %r. Building it (one-time): tagging tf rows\n" - " with child_frame and writing the topology change-log.\n" - "========================================================================", - TF_STREAM, - ) - self._build_graph_stream() - if n_rows: - _ensure_child_index(conn) # tf table exists now - self._built = True - - def _build_graph_stream(self) -> None: - """One-time migration: decode every tf row, tag it with its child_frame, and - append a ``TfGraph`` snapshot whenever the topology changes. A frame counts as - static if its pose never varies across the whole recording.""" - safe = TF_STREAM - # one decode pass: collect (id, ts, child, parent, pose-key) + per-child poses - rows: list[tuple[int, float, str, str]] = [] - poses_per_child: dict[str, set[tuple[float, ...]]] = {} - for obs in self._store.stream(safe, TFMessage).order_by("ts"): - for transform in getattr(obs.data, "transforms", None) or [obs.data]: - pose_key = tuple( - round(value, POSE_EQUALITY_DECIMALS) - for value in ( - transform.translation.x, - transform.translation.y, - transform.translation.z, - transform.rotation.x, - transform.rotation.y, - transform.rotation.z, - transform.rotation.w, - ) - ) - rows.append((obs.id, obs.ts, transform.child_frame_id, transform.frame_id)) - poses_per_child.setdefault(transform.child_frame_id, set()).add(pose_key) - static_frames = {child for child, poses in poses_per_child.items() if len(poses) == 1} - - # tag each tf row with its child_frame (raw UPDATE on the tf table) - conn = self._connection() - for row_id, _ts, child, _parent in rows: - conn.execute( - f"UPDATE \"{safe}\" SET tags = json_set(tags, '$.child_frame', ?) WHERE id = ?", - (child, row_id), - ) - conn.commit() - - # build the change-log as a first-class stream: one snapshot per change - graph_stream = self._store.stream(GRAPH_STREAM, TfGraph) - structure: dict[str, dict[str, Any]] = {} - written = 0 - for _row_id, ts, child, parent in rows: - entry = {"parent": parent, "static": child in static_frames} - if structure.get(child) == entry: - continue - structure[child] = entry - graph_stream.append(TfGraph(structure), ts=ts) - written += 1 - logger.warning("tf graph built: %d topology changes for %r.", written, TF_STREAM) - - def _graph_codec(self) -> Any: - source = self._store.stream(GRAPH_STREAM, TfGraph)._source - return cast("Any", source).codec - - def _load_graph_if_small(self) -> None: - if self._graph_loaded: - return - if self._graph_count() < self._max_in_ram: - # Sort by (ts, id): several topology changes can share a timestamp (e.g. - # every static frame latched at t0), and the LAST-inserted of those is the - # complete snapshot — a plain ts sort leaves same-ts order undefined. - snapshots = sorted( - ((obs.ts, obs.id, obs.data.structure) for obs in self._graph_stream()), - key=lambda row: (row[0], row[1]), - ) - self._graph_in_ram = [(ts, structure) for ts, _id, structure in snapshots] - else: - self._graph_in_ram = None # too many -> query per lookup - self._graph_loaded = True - - def _graph_at(self, query_time: float) -> dict[str, Any] | None: - if self._graph_in_ram is not None: - # in-RAM: binary search the latest change at-or-before query_time - stamps = [ts for ts, _ in self._graph_in_ram] - index = bisect.bisect_right(stamps, query_time) - 1 - if index < 0: - return self._graph_in_ram[0][1] # before first -> earliest - return self._graph_in_ram[index][1] - # fallback: one indexed query for the latest snapshot at or before query_time. - # Tie-break by id (DESC) so same-timestamp changes resolve to the complete one. - self.graph_queries += 1 - conn = self._connection() - graph, blob = f'"{GRAPH_STREAM}"', f'"{GRAPH_STREAM}_blob"' - row = conn.execute( - f"SELECT x.data FROM {graph} g JOIN {blob} x ON x.id = g.id " - "WHERE g.ts <= ? ORDER BY g.ts DESC, g.id DESC LIMIT 1", - (query_time,), - ).fetchone() - if row is None: # before the first snapshot -> earliest - row = conn.execute( - f"SELECT x.data FROM {graph} g JOIN {blob} x ON x.id = g.id " - "ORDER BY g.ts ASC, g.id ASC LIMIT 1" - ).fetchone() - if row is None: - return None - return cast("TfGraph", self._graph_codec().decode(row[0])).structure - - def _chain_frames(self, graph: dict[str, Any], source: str, target: str) -> list[str] | None: - def to_root(frame: str) -> list[str]: - path = [frame] - seen = {frame} - while ( - frame in graph - and graph[frame].get("parent") in graph - and graph[frame]["parent"] not in seen - ): - frame = graph[frame]["parent"] - path.append(frame) - seen.add(frame) - # include a final parent that is itself a root (not a key in graph) - if frame in graph and graph[frame].get("parent") and graph[frame]["parent"] not in seen: - path.append(graph[frame]["parent"]) - return path - - source_path = to_root(source) - target_path = to_root(target) - common = next((f for f in source_path if f in set(target_path)), None) - if common is None: - return None # disjoint graph: no transform between them - frames = source_path[: source_path.index(common) + 1] - frames += target_path[: target_path.index(common)] - return frames + from dimos.msgs.geometry_msgs.Transform import Transform - def _codec(self) -> Any: - source = self._store.stream(TF_STREAM, TFMessage)._source - return cast("Any", source).codec - def _decode_blob(self, data: bytes, frame: str) -> Transform: - # The blob is the codec-encoded message; pick the transform for `frame` - # (rows normally hold one; legacy rows may pack several). - message = self._codec().decode(data) - transforms = getattr(message, "transforms", None) or [message] - for transform in transforms: - if transform.child_frame_id == frame: - return cast("Transform", transform) - return cast("Transform", transforms[0]) - - def _fetch_rows( - self, dynamic: list[str], static: list[str], query_time: float - ) -> dict[tuple[str, str], bytes]: - """ONE query: for each dynamic frame the bracketing rows ('lo' = latest at - or before query_time, 'hi' = earliest at or after), and for each (uncached) - static frame its latest row ('st') — all joined to the blob data. Keyed by - (frame, kind) -> blob bytes.""" - cf = "json_extract(tags, '$.child_frame')" - tf, blob = f'"{TF_STREAM}"', f'"{TF_STREAM}_blob"' - # One UNION of per-frame, index-served LIMIT-1 subqueries: each is a direct - # (child_frame, ts) range seek — far cheaper than a window-function scan, and - # still a single round-trip. - parts: list[str] = [] - params: list[Any] = [] - - def pick(frame: str, kind: str, where_ts: str, order: str) -> None: - parts.append( - f"SELECT ? AS cf, ? AS kind, " - f"(SELECT id FROM {tf} WHERE {cf} = ?{where_ts} ORDER BY ts {order} LIMIT 1) AS id" - ) - params.extend([frame, kind, frame]) - - for frame in dynamic: - pick(frame, "lo", " AND ts <= ?", "DESC") - params.append(query_time) - pick(frame, "hi", " AND ts >= ?", "ASC") - params.append(query_time) - for frame in static: - pick(frame, "st", "", "DESC") - if not parts: - return {} - union = " UNION ALL ".join(parts) - sql = f"SELECT t.cf, t.kind, b.data FROM ({union}) t JOIN {blob} b ON b.id = t.id" - rows: dict[tuple[str, str], bytes] = {} - for cf_val, kind, data in self._connection().execute(sql, params): - rows[(cf_val, kind)] = data - self.rows_fetched += 1 - return rows - - # --- loop-closure deformation (sqlite) ------------------------------------- - - def _load_deformation_ids(self) -> set[int]: - """The distinct tf_ids present on the deformation stream, cached once. Empty - when there's no such stream — so the correction path is skipped at zero cost - for recordings without loop closure.""" - if self._deformation_tf_ids is not None: - return self._deformation_tf_ids - ids: set[int] = set() - if self._is_sqlite and DEFORMATION_STREAM in set(self._store.list_streams()): - for (tf_id,) in self._connection().execute( - f"SELECT DISTINCT json_extract(tags, '$.tf_id') FROM \"{DEFORMATION_STREAM}\"" - ): - if tf_id is not None: - ids.add(int(tf_id)) - self._deformation_tf_ids = ids - return ids - - def _deformation_codec(self) -> Any: - source = self._store.stream(DEFORMATION_STREAM, DeformationNode)._source - return cast("Any", source).codec - - def _node_pose_matrix(self, order: str, tf_id: int, node_id: str) -> np.ndarray | None: - """The 4x4 pose of one keyframe's first (``ASC``) or latest (``DESC``) recorded - version. Versions of a node share the keyframe ts, so they're ordered by row id - (insertion order), not ts.""" - stream, blob = f'"{DEFORMATION_STREAM}"', f'"{DEFORMATION_STREAM}_blob"' - row = ( - self._connection() - .execute( - f"SELECT b.data FROM {stream} s JOIN {blob} b ON b.id = s.id " - f"WHERE json_extract(s.tags, '$.tf_id') = ? AND json_extract(s.tags, '$.id') = ? " - f"ORDER BY s.id {order} LIMIT 1", - (str(tf_id), node_id), - ) - .fetchone() - ) - if row is None: - return None - node = cast("DeformationNode", self._deformation_codec().decode(row[0])) - return Transform.from_pose(node.pose.frame_id, node.pose).to_matrix() - - def _node_delta(self, tf_id: int, node_id: str) -> np.ndarray | None: - """How far the optimizer has moved a keyframe since it was first recorded: - ``current ∘ inv(original)`` — the SE(3) correction this node contributes.""" - original = self._node_pose_matrix("ASC", tf_id, node_id) - current = self._node_pose_matrix("DESC", tf_id, node_id) - if original is None or current is None: - return None - return current @ np.linalg.inv(original) - - def _edge_delta(self, tf_id: int, query_time: float) -> np.ndarray | None: - """The blended correction for an edge at ``query_time``: take the keyframes - bracketing the time (latest at-or-before, earliest at-or-after), each node's - ``current ∘ inv(original)`` delta, and linear-blend-skin between them.""" - stream = f'"{DEFORMATION_STREAM}"' - id_tag = "json_extract(tags, '$.id')" - tf_tag = "json_extract(tags, '$.tf_id')" - conn = self._connection() - lo = conn.execute( - f"SELECT {id_tag}, ts FROM {stream} WHERE {tf_tag} = ? AND ts <= ? " - "ORDER BY ts DESC, id DESC LIMIT 1", - (str(tf_id), query_time), - ).fetchone() - hi = conn.execute( - f"SELECT {id_tag}, ts FROM {stream} WHERE {tf_tag} = ? AND ts >= ? " - "ORDER BY ts ASC, id ASC LIMIT 1", - (str(tf_id), query_time), - ).fetchone() - samples: list[tuple[float, np.ndarray]] = [] - seen_ids: set[str] = set() - for node_id, ts in (row for row in (lo, hi) if row is not None): - if node_id in seen_ids: - continue - seen_ids.add(node_id) - delta = self._node_delta(tf_id, node_id) - if delta is not None: - samples.append((ts, delta)) - if not samples: - return None - if len(samples) == 1: - return samples[0][1] - (ts_lo, mat_lo), (ts_hi, mat_hi) = sorted(samples, key=lambda item: item[0]) - weight = 0.0 if ts_hi == ts_lo else (query_time - ts_lo) / (ts_hi - ts_lo) - return _blend_se3(mat_lo, mat_hi, weight) - - def _edge_corrections( - self, graph: dict[str, Any], edges: list[str], query_time: float - ) -> dict[str, np.ndarray]: - """For each edge whose tf_id matches the deformation stream, its blended SE(3) - correction (applied on the parent side of the edge). Empty when nothing matches.""" - tf_ids = self._load_deformation_ids() - if not tf_ids: - return {} - corrections: dict[str, np.ndarray] = {} - for frame in edges: - tf_id = tf_id_for(graph[frame]["parent"], frame) - if tf_id not in tf_ids: - continue - delta = self._edge_delta(tf_id, query_time) - if delta is not None: - corrections[frame] = delta - return corrections +class DbTf(Protocol): + """Transform-lookup interface over a recording's transforms. Implemented by + :class:`DbTfLive` (default) and :class:`DbTfSql` (sqlite).""" def get( self, @@ -439,90 +39,9 @@ def get( source_frame: str, time_point: float | None = None, time_tolerance: float | None = None, - ) -> Transform | None: - """Transform that maps a point in ``source_frame`` into ``target_frame``, - or ``None`` if no chain connects them at the requested time.""" - if not self._is_sqlite: - return self._ensure_loaded().lookup( - target_frame, source_frame, time_point, time_tolerance - ) - self._ensure_built() - self._load_graph_if_small() - query_time = time_point if time_point is not None else 0.0 - graph = self._graph_at(query_time) # 0 queries when the graph is in RAM - if graph is None: - return None - frames = self._chain_frames(graph, source_frame, target_frame) - if frames is None: - return None - - edges = [f for f in frames if f in graph] # roots have no incoming edge - dynamic = [f for f in edges if not graph[f].get("static")] - static = [f for f in edges if graph[f].get("static")] - uncached_static = [f for f in static if f not in self._static_cache] - - rows = self._fetch_rows(dynamic, uncached_static, query_time) # ONE detail query - # Per-edge loop-closure corrections (empty unless a deformation stream matches). - corrections = self._edge_corrections(graph, edges, query_time) - - buffer = MultiTBuffer(buffer_size=_NO_PRUNE) - for frame in static: - transform = self._static_cache.get(frame) - if transform is None: - data = rows.get((frame, "st")) - if data is None: - return None - transform = self._decode_blob(data, frame) - self._static_cache[frame] = transform - # restamp the constant to query_time so the buffer's tolerance never - # rejects a static that was recorded long ago (latched once). - transform = _restamp(transform, query_time) - if frame in corrections: - transform = _apply_delta(corrections[frame], transform, query_time) - buffer.receive_transform(transform) - for frame in dynamic: - lo = rows.get((frame, "lo")) - hi = rows.get((frame, "hi")) - chosen = lo if lo is not None else hi - if chosen is None: - return None - if frame in corrections: - # Resolve the raw edge at query_time, then deform it. The correction - # already carries the time blend, so a single corrected sample suffices. - raw = self._interpolate_dynamic( - frame, graph[frame]["parent"], lo, hi, query_time, time_tolerance - ) - if raw is None: - return None - buffer.receive_transform(_apply_delta(corrections[frame], raw, query_time)) - continue - buffer.receive_transform(self._decode_blob(chosen, frame)) - other = hi if hi is not None else lo - if other is not None and other is not chosen: - buffer.receive_transform(self._decode_blob(other, frame)) - return buffer.lookup(target_frame, source_frame, time_point, time_tolerance) + ) -> Transform | None: ... - def _interpolate_dynamic( - self, - frame: str, - parent: str, - lo: bytes | None, - hi: bytes | None, - query_time: float, - time_tolerance: float | None, - ) -> Transform | None: - """The raw ``parent <- frame`` transform at ``query_time``, interpolated from - the bracketing rows (its own small buffer so the chain buffer only ever sees - the corrected result).""" - edge_buffer = MultiTBuffer(buffer_size=_NO_PRUNE) - chosen = lo if lo is not None else hi - if chosen is None: - return None - edge_buffer.receive_transform(self._decode_blob(chosen, frame)) - other = hi if hi is not None else lo - if other is not None and other is not chosen: - edge_buffer.receive_transform(self._decode_blob(other, frame)) - return edge_buffer.lookup(parent, frame, query_time, time_tolerance) + def has_transforms(self) -> bool: ... class TfGraph: @@ -530,7 +49,7 @@ class TfGraph: ``structure`` maps each child frame to ``{"parent": str, "static": bool}`` — the full tf tree as of this message's timestamp. The stream of these snapshots - (the ``_graph`` stream) is the topology change-log that transform lookups + (the ``tf_graph`` stream) is the topology change-log that transform lookups walk to resolve a source->target chain at any past time. Defined here (not under ``dimos/msgs``) because it is a recording-internal payload, not a wire message; it is stored via the pickle codec.""" @@ -545,80 +64,3 @@ def __init__(self, structure: dict[str, dict[str, Any]]) -> None: def __repr__(self) -> str: return f"TfGraph({len(self.structure)} frames)" - - -def _restamp(transform: Transform, ts: float) -> Transform: - return Transform( - translation=transform.translation, - rotation=transform.rotation, - frame_id=transform.frame_id, - child_frame_id=transform.child_frame_id, - ts=ts, - ) - - -def _apply_delta(delta: np.ndarray, transform: Transform, ts: float) -> Transform: - """Deform ``transform`` by the SE(3) correction ``delta``, applied on the parent - (frame_id) side: ``corrected = delta @ raw``. Keeps the edge's frame names.""" - corrected = delta @ transform.to_matrix() - return Transform.from_matrix( - corrected, - ts=ts, - frame_id=transform.frame_id, - child_frame_id=transform.child_frame_id, - ) - - -def _quat_slerp(q_lo: np.ndarray, q_hi: np.ndarray, weight: float) -> np.ndarray: - """Spherical-linear interpolation between two quaternions ``[x, y, z, w]``.""" - dot = float(np.dot(q_lo, q_hi)) - if dot < 0.0: # take the shorter arc - q_hi = -q_hi - dot = -dot - if dot > 0.9995: # nearly parallel: lerp + renormalize avoids a divide-by-~0 - result = q_lo + weight * (q_hi - q_lo) - return cast("np.ndarray", result / np.linalg.norm(result)) - theta_0 = np.arccos(np.clip(dot, -1.0, 1.0)) - sin_0 = np.sin(theta_0) - scale_lo = np.sin((1.0 - weight) * theta_0) / sin_0 - scale_hi = np.sin(weight * theta_0) / sin_0 - return cast("np.ndarray", scale_lo * q_lo + scale_hi * q_hi) - - -def _blend_se3(mat_lo: np.ndarray, mat_hi: np.ndarray, weight: float) -> np.ndarray: - """Linear-blend-skin two SE(3) deltas by ``weight`` in [0, 1] (0 -> lo, 1 -> hi): - lerp the translation, slerp the rotation.""" - quat_lo = Quaternion.from_rotation_matrix(mat_lo[:3, :3]) - quat_hi = Quaternion.from_rotation_matrix(mat_hi[:3, :3]) - blended = _quat_slerp( - np.array([quat_lo.x, quat_lo.y, quat_lo.z, quat_lo.w]), - np.array([quat_hi.x, quat_hi.y, quat_hi.z, quat_hi.w]), - weight, - ) - out = np.eye(4) - out[:3, :3] = Quaternion(blended).to_rotation_matrix() - out[:3, 3] = (1.0 - weight) * mat_lo[:3, 3] + weight * mat_hi[:3, 3] - return out - - -def _connect(db_path: str) -> sqlite3.Connection: - conn = sqlite3.connect(db_path, check_same_thread=False) - conn.execute("PRAGMA busy_timeout=5000") - return conn - - -def _ensure_child_index(conn: sqlite3.Connection) -> None: - """Index the child_frame json tag on the tf rows so per-frame time queries - seek. The live recorder gets this for free (the store auto-indexes tag keys on - tagged appends); this is for migrated recordings and the read side. Requires - the tf table to exist.""" - # Composite (child_frame, ts) so a per-frame "latest at/before T" is a direct - # index range seek, not a scan+sort. Index names share SQLite's global namespace - # with tables, so the name is double-underscore-namespaced to keep it clear of any - # real stream/table name (no stream would contain "__dbtf_"). - index_name = f"{TF_STREAM}__dbtf_child_ts_idx" - conn.execute( - f'CREATE INDEX IF NOT EXISTS "{index_name}" ' - f"ON \"{TF_STREAM}\"(json_extract(tags, '$.child_frame'), ts)" - ) - conn.commit() diff --git a/dimos/memory2/db_tf_live.py b/dimos/memory2/db_tf_live.py new file mode 100644 index 0000000000..e67ad9dfdc --- /dev/null +++ b/dimos/memory2/db_tf_live.py @@ -0,0 +1,71 @@ +# 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. + +"""In-RAM transform lookups — the default for any :class:`Store`.""" + +from __future__ import annotations + +import math +import threading +from typing import TYPE_CHECKING + +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.tf.tf import MultiTBuffer + +if TYPE_CHECKING: + from dimos.memory2.store.base import Store + from dimos.msgs.geometry_msgs.Transform import Transform + +# The transform streams the buffer loads (dynamic + latched static). +TF_STREAMS = ("tf", "tf_static") + + +class DbTfLive: + """Transform lookups by loading the ``tf``/``tf_static`` streams into an in-RAM + :class:`MultiTBuffer`. The default for any :class:`Store` — backend-agnostic (uses + only the generic stream API). Holds the full history in memory, loaded once.""" + + def __init__(self, store: Store) -> None: + self._store = store + self._lock = threading.Lock() + self._buffer: MultiTBuffer | None = None + + def _ensure_loaded(self) -> MultiTBuffer: + if self._buffer is not None: + return self._buffer + with self._lock: + if self._buffer is not None: + return self._buffer + buffer = MultiTBuffer(buffer_size=math.inf) + available = set(self._store.list_streams()) + for name in TF_STREAMS: + if name not in available: + continue + for observation in self._store.stream(name, TFMessage): + transforms = getattr(observation.data, "transforms", None) or [observation.data] + buffer.receive_transform(*transforms) + self._buffer = buffer + return buffer + + def has_transforms(self) -> bool: + return bool(self._ensure_loaded().buffers) + + def get( + self, + target_frame: str, + source_frame: str, + time_point: float | None = None, + time_tolerance: float | None = None, + ) -> Transform | None: + return self._ensure_loaded().lookup(target_frame, source_frame, time_point, time_tolerance) diff --git a/dimos/memory2/db_tf_sql.py b/dimos/memory2/db_tf_sql.py new file mode 100644 index 0000000000..888ea50d9c --- /dev/null +++ b/dimos/memory2/db_tf_sql.py @@ -0,0 +1,501 @@ +# 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. + +"""SQLite transform lookups — the sqlite graph-stream path (indexed per-lookup queries +plus loop-closure deformation correction). Used only by the sqlite store.""" + +from __future__ import annotations + +import bisect +import math +import sqlite3 +from typing import TYPE_CHECKING, Any, cast + +import numpy as np + +from dimos.memory2.db_tf import TfGraph +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.nav_msgs.DeformationNode import DeformationNode, tf_id_for +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.tf.deformation import apply_delta, blend_se3 +from dimos.protocol.tf.tf import MultiTBuffer +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.memory2.store.base import Store + from dimos.memory2.store.sqlite import SqliteStoreConfig + from dimos.memory2.stream import Stream + +logger = setup_logger() + +GRAPH_STREAM = "tf_graph" # changes to the tf tree (with room for non-tree structures) +DEFORMATION_STREAM = "tf_deformation_nodes" # loop closure and other deformations +# 99% of the time there are less than 10 tree changes +# this cache allows us to avoid a DB query in the offline/map-load case +DEFAULT_MAX_GRAPH_CHANGES_IN_RAM = 64 +# A frame is "static" if its pose never changes; poses are compared rounded to this +# many decimals (~nanometre / nanoradian) so float noise doesn't read as motion. +POSE_EQUALITY_DECIMALS = 9 + + +class DbTfSql: + """Transform lookups over the sqlite graph-stream: a topology change-log plus + per-lookup indexed queries for the bracketing samples, with loop-closure + deformation correction. Used only by the sqlite store.""" + + def __init__( + self, + store: Store, + max_graph_changes_in_ram: int = DEFAULT_MAX_GRAPH_CHANGES_IN_RAM, + ) -> None: + self._store = store + self._max_in_ram = max_graph_changes_in_ram + self._conn: sqlite3.Connection | None = None + self._built = False + # graph cache: either the whole change-log in RAM, or None (query per lookup) + self._graph_in_ram: list[tuple[float, dict[str, Any]]] | None = None + self._graph_loaded = False + self._static_cache: dict[str, Transform] = {} + # The set of tf_ids present on the deformation stream (the edges loop closure + # corrects). Cached once so a recording with no deformation nodes pays nothing. + self._deformation_tf_ids: set[int] | None = None + self.rows_fetched = 0 + self.graph_queries = 0 + + def _connection(self) -> sqlite3.Connection: + conn = self._conn + if conn is None: + conn = _connect(cast("SqliteStoreConfig", self._store.config).path) + self._conn = conn + return conn + + def has_transforms(self) -> bool: + if "tf" not in set(self._store.list_streams()): + return False + (n_rows,) = self._connection().execute('SELECT count(*) FROM "tf"').fetchone() + return bool(n_rows) + + def _graph_stream(self) -> Stream[TfGraph]: + return self._store.stream(GRAPH_STREAM, TfGraph) + + def _graph_count(self) -> int: + if GRAPH_STREAM not in set(self._store.list_streams()): + return 0 + return self._graph_stream().count() + + def _ensure_built(self) -> None: + """First use: if the recording has tf rows but no graph stream (a recording + that predates it / wasn't written by the recorder), build the graph stream once + by replaying the tf rows, then make sure the seek index exists.""" + if self._built: + return + conn = self._connection() + (n_rows,) = conn.execute('SELECT count(*) FROM "tf"').fetchone() + if n_rows and self._graph_count() == 0: + logger.warning( + "\n========================================================================\n" + " tf graph stream MISSING for %r. Building it (one-time): tagging tf rows\n" + " with child_frame and writing the topology change-log.\n" + "========================================================================", + "tf", + ) + self._build_graph_stream() + if n_rows: + _ensure_child_index(conn) # tf table exists now + self._built = True + + def _build_graph_stream(self) -> None: + """One-time migration: decode every tf row, tag it with its child_frame, and + append a ``TfGraph`` snapshot whenever the topology changes. A frame counts as + static if its pose never varies across the whole recording.""" + # one decode pass: collect (id, ts, child, parent, pose-key) + per-child poses + rows: list[tuple[int, float, str, str]] = [] + poses_per_child: dict[str, set[tuple[float, ...]]] = {} + for obs in self._store.stream("tf", TFMessage).order_by("ts"): + for transform in getattr(obs.data, "transforms", None) or [obs.data]: + pose_key = tuple( + round(value, POSE_EQUALITY_DECIMALS) + for value in ( + transform.translation.x, + transform.translation.y, + transform.translation.z, + transform.rotation.x, + transform.rotation.y, + transform.rotation.z, + transform.rotation.w, + ) + ) + rows.append((obs.id, obs.ts, transform.child_frame_id, transform.frame_id)) + poses_per_child.setdefault(transform.child_frame_id, set()).add(pose_key) + static_frames = {child for child, poses in poses_per_child.items() if len(poses) == 1} + + # tag each tf row with its child_frame (raw UPDATE on the tf table) + conn = self._connection() + for row_id, _ts, child, _parent in rows: + conn.execute( + "UPDATE \"tf\" SET tags = json_set(tags, '$.child_frame', ?) WHERE id = ?", + (child, row_id), + ) + conn.commit() + + # build the change-log as a first-class stream: one snapshot per change + graph_stream = self._store.stream(GRAPH_STREAM, TfGraph) + structure: dict[str, dict[str, Any]] = {} + written = 0 + for _row_id, ts, child, parent in rows: + entry = {"parent": parent, "static": child in static_frames} + if structure.get(child) == entry: + continue + structure[child] = entry + graph_stream.append(TfGraph(structure), ts=ts) + written += 1 + logger.warning("tf graph built: %d topology changes for %r.", written, "tf") + + def _load_graph_if_small(self) -> None: + if self._graph_loaded: + return + if self._graph_count() < self._max_in_ram: + # Sort by (ts, id): several topology changes can share a timestamp (e.g. + # every static frame latched at t0), and the LAST-inserted of those is the + # complete snapshot — a plain ts sort leaves same-ts order undefined. + snapshots = sorted( + ((obs.ts, obs.id, obs.data.structure) for obs in self._graph_stream()), + key=lambda row: (row[0], row[1]), + ) + self._graph_in_ram = [(ts, structure) for ts, _id, structure in snapshots] + else: + self._graph_in_ram = None # too many -> query per lookup + self._graph_loaded = True + + def _graph_at(self, query_time: float) -> dict[str, Any] | None: + if self._graph_in_ram is not None: + # in-RAM: binary search the latest change at-or-before query_time + stamps = [ts for ts, _ in self._graph_in_ram] + index = bisect.bisect_right(stamps, query_time) - 1 + if index < 0: + return self._graph_in_ram[0][1] # before first -> earliest + return self._graph_in_ram[index][1] + # fallback: one indexed query for the latest snapshot at or before query_time. + # Tie-break by id (DESC) so same-timestamp changes resolve to the complete one. + self.graph_queries += 1 + conn = self._connection() + graph, blob = f'"{GRAPH_STREAM}"', f'"{GRAPH_STREAM}_blob"' + row = conn.execute( + f"SELECT x.data FROM {graph} g JOIN {blob} x ON x.id = g.id " + "WHERE g.ts <= ? ORDER BY g.ts DESC, g.id DESC LIMIT 1", + (query_time,), + ).fetchone() + if row is None: # before the first snapshot -> earliest + row = conn.execute( + f"SELECT x.data FROM {graph} g JOIN {blob} x ON x.id = g.id " + "ORDER BY g.ts ASC, g.id ASC LIMIT 1" + ).fetchone() + if row is None: + return None + codec = cast("Any", self._store.stream(GRAPH_STREAM, TfGraph)._source).codec + return cast("TfGraph", codec.decode(row[0])).structure + + def _chain_frames(self, graph: dict[str, Any], source: str, target: str) -> list[str] | None: + def to_root(frame: str) -> list[str]: + path = [frame] + seen = {frame} + while ( + frame in graph + and graph[frame].get("parent") in graph + and graph[frame]["parent"] not in seen + ): + frame = graph[frame]["parent"] + path.append(frame) + seen.add(frame) + # include a final parent that is itself a root (not a key in graph) + if frame in graph and graph[frame].get("parent") and graph[frame]["parent"] not in seen: + path.append(graph[frame]["parent"]) + return path + + source_path = to_root(source) + target_path = to_root(target) + common = next((f for f in source_path if f in set(target_path)), None) + if common is None: + return None # disjoint graph: no transform between them + frames = source_path[: source_path.index(common) + 1] + frames += target_path[: target_path.index(common)] + return frames + + def _decode_blob(self, data: bytes, frame: str) -> Transform: + # The blob is the codec-encoded message; pick the transform for `frame` + # (rows normally hold one; legacy rows may pack several). + codec = cast("Any", self._store.stream("tf", TFMessage)._source).codec + message = codec.decode(data) + transforms = getattr(message, "transforms", None) or [message] + for transform in transforms: + if transform.child_frame_id == frame: + return cast("Transform", transform) + return cast("Transform", transforms[0]) + + def _fetch_rows( + self, dynamic: list[str], static: list[str], query_time: float + ) -> dict[tuple[str, str], bytes]: + """ONE query: for each dynamic frame the bracketing rows ('lo' = latest at + or before query_time, 'hi' = earliest at or after), and for each (uncached) + static frame its latest row ('st') — all joined to the blob data. Keyed by + (frame, kind) -> blob bytes.""" + cf = "json_extract(tags, '$.child_frame')" + # One UNION of per-frame, index-served LIMIT-1 subqueries: each is a direct + # (child_frame, ts) range seek — far cheaper than a window-function scan, and + # still a single round-trip. + parts: list[str] = [] + params: list[Any] = [] + + def pick(frame: str, kind: str, where_ts: str, order: str) -> None: + parts.append( + f"SELECT ? AS cf, ? AS kind, " + f'(SELECT id FROM "tf" WHERE {cf} = ?{where_ts} ORDER BY ts {order} LIMIT 1) AS id' + ) + params.extend([frame, kind, frame]) + + for frame in dynamic: + pick(frame, "lo", " AND ts <= ?", "DESC") + params.append(query_time) + pick(frame, "hi", " AND ts >= ?", "ASC") + params.append(query_time) + for frame in static: + pick(frame, "st", "", "DESC") + if not parts: + return {} + union = " UNION ALL ".join(parts) + sql = f'SELECT t.cf, t.kind, b.data FROM ({union}) t JOIN "tf_blob" b ON b.id = t.id' + rows: dict[tuple[str, str], bytes] = {} + for cf_val, kind, data in self._connection().execute(sql, params): + rows[(cf_val, kind)] = data + self.rows_fetched += 1 + return rows + + def _load_deformation_ids(self) -> set[int]: + """The distinct tf_ids present on the deformation stream, cached once. Empty + when there's no such stream — so the correction path is skipped at zero cost + for recordings without loop closure.""" + if self._deformation_tf_ids is not None: + return self._deformation_tf_ids + ids: set[int] = set() + if DEFORMATION_STREAM in set(self._store.list_streams()): + for (tf_id,) in self._connection().execute( + f"SELECT DISTINCT json_extract(tags, '$.tf_id') FROM \"{DEFORMATION_STREAM}\"" + ): + if tf_id is not None: + ids.add(int(tf_id)) + self._deformation_tf_ids = ids + return ids + + def _node_pose_matrix(self, order: str, tf_id: int, node_id: str) -> np.ndarray | None: + """The 4x4 pose of one keyframe's first (``ASC``) or latest (``DESC``) recorded + version. Versions of a node share the keyframe ts, so they're ordered by row id + (insertion order), not ts.""" + stream, blob = f'"{DEFORMATION_STREAM}"', f'"{DEFORMATION_STREAM}_blob"' + row = ( + self._connection() + .execute( + f"SELECT b.data FROM {stream} s JOIN {blob} b ON b.id = s.id " + f"WHERE json_extract(s.tags, '$.tf_id') = ? AND json_extract(s.tags, '$.id') = ? " + f"ORDER BY s.id {order} LIMIT 1", + (str(tf_id), node_id), + ) + .fetchone() + ) + if row is None: + return None + codec = cast("Any", self._store.stream(DEFORMATION_STREAM, DeformationNode)._source).codec + node = cast("DeformationNode", codec.decode(row[0])) + return Transform.from_pose(node.pose.frame_id, node.pose).to_matrix() + + def _node_delta(self, tf_id: int, node_id: str) -> np.ndarray | None: + """How far the optimizer has moved a keyframe since it was first recorded: + ``current ∘ inv(original)`` — the SE(3) correction this node contributes.""" + original = self._node_pose_matrix("ASC", tf_id, node_id) + current = self._node_pose_matrix("DESC", tf_id, node_id) + if original is None or current is None: + return None + return current @ np.linalg.inv(original) + + def _edge_delta(self, tf_id: int, query_time: float) -> np.ndarray | None: + """The blended correction for an edge at ``query_time``: take the keyframes + bracketing the time (latest at-or-before, earliest at-or-after), each node's + ``current ∘ inv(original)`` delta, and linear-blend-skin between them.""" + stream = f'"{DEFORMATION_STREAM}"' + id_tag = "json_extract(tags, '$.id')" + tf_tag = "json_extract(tags, '$.tf_id')" + conn = self._connection() + lo = conn.execute( + f"SELECT {id_tag}, ts FROM {stream} WHERE {tf_tag} = ? AND ts <= ? " + "ORDER BY ts DESC, id DESC LIMIT 1", + (str(tf_id), query_time), + ).fetchone() + hi = conn.execute( + f"SELECT {id_tag}, ts FROM {stream} WHERE {tf_tag} = ? AND ts >= ? " + "ORDER BY ts ASC, id ASC LIMIT 1", + (str(tf_id), query_time), + ).fetchone() + samples: list[tuple[float, np.ndarray]] = [] + seen_ids: set[str] = set() + for node_id, ts in (row for row in (lo, hi) if row is not None): + if node_id in seen_ids: + continue + seen_ids.add(node_id) + delta = self._node_delta(tf_id, node_id) + if delta is not None: + samples.append((ts, delta)) + if not samples: + return None + if len(samples) == 1: + return samples[0][1] + (ts_lo, mat_lo), (ts_hi, mat_hi) = sorted(samples, key=lambda item: item[0]) + weight = 0.0 if ts_hi == ts_lo else (query_time - ts_lo) / (ts_hi - ts_lo) + return blend_se3(mat_lo, mat_hi, weight) + + def _edge_corrections( + self, graph: dict[str, Any], edges: list[str], query_time: float + ) -> dict[str, np.ndarray]: + """For each edge whose tf_id matches the deformation stream, its blended SE(3) + correction (applied on the parent side of the edge). Empty when nothing matches.""" + tf_ids = self._load_deformation_ids() + if not tf_ids: + return {} + corrections: dict[str, np.ndarray] = {} + for frame in edges: + tf_id = tf_id_for(graph[frame]["parent"], frame) + if tf_id not in tf_ids: + continue + delta = self._edge_delta(tf_id, query_time) + if delta is not None: + corrections[frame] = delta + return corrections + + def get( + self, + target_frame: str, + source_frame: str, + time_point: float | None = None, + time_tolerance: float | None = None, + ) -> Transform | None: + """Transform that maps a point in ``source_frame`` into ``target_frame``, + or ``None`` if no chain connects them at the requested time.""" + self._ensure_built() + self._load_graph_if_small() + query_time = time_point if time_point is not None else 0.0 + graph = self._graph_at(query_time) # 0 queries when the graph is in RAM + if graph is None: + return None + frames = self._chain_frames(graph, source_frame, target_frame) + if frames is None: + return None + + edges = [f for f in frames if f in graph] # roots have no incoming edge + dynamic = [f for f in edges if not graph[f].get("static")] + static = [f for f in edges if graph[f].get("static")] + uncached_static = [f for f in static if f not in self._static_cache] + + rows = self._fetch_rows(dynamic, uncached_static, query_time) # ONE detail query + # Per-edge loop-closure corrections (empty unless a deformation stream matches). + corrections = self._edge_corrections(graph, edges, query_time) + + buffer = MultiTBuffer(buffer_size=math.inf) + for frame in static: + transform = self._static_cache.get(frame) + if transform is None: + data = rows.get((frame, "st")) + if data is None: + return None + transform = self._decode_blob(data, frame) + self._static_cache[frame] = transform + # restamp the constant to query_time so the buffer's tolerance never + # rejects a static that was recorded long ago (latched once). + transform = _restamp(transform, query_time) + if frame in corrections: + transform = apply_delta(corrections[frame], transform, query_time) + buffer.receive_transform(transform) + for frame in dynamic: + lo = rows.get((frame, "lo")) + hi = rows.get((frame, "hi")) + chosen = lo if lo is not None else hi + if chosen is None: + return None + if frame in corrections: + # Resolve the raw edge at query_time, then deform it. The correction + # already carries the time blend, so a single corrected sample suffices. + raw = self._interpolate_dynamic( + frame, graph[frame]["parent"], lo, hi, query_time, time_tolerance + ) + if raw is None: + return None + buffer.receive_transform(apply_delta(corrections[frame], raw, query_time)) + continue + buffer.receive_transform(self._decode_blob(chosen, frame)) + other = hi if hi is not None else lo + if other is not None and other is not chosen: + buffer.receive_transform(self._decode_blob(other, frame)) + return buffer.lookup(target_frame, source_frame, time_point, time_tolerance) + + def _interpolate_dynamic( + self, + frame: str, + parent: str, + lo: bytes | None, + hi: bytes | None, + query_time: float, + time_tolerance: float | None, + ) -> Transform | None: + """The raw ``parent <- frame`` transform at ``query_time``, interpolated from + the bracketing rows (its own small buffer so the chain buffer only ever sees + the corrected result).""" + edge_buffer = MultiTBuffer(buffer_size=math.inf) + chosen = lo if lo is not None else hi + if chosen is None: + return None + edge_buffer.receive_transform(self._decode_blob(chosen, frame)) + other = hi if hi is not None else lo + if other is not None and other is not chosen: + edge_buffer.receive_transform(self._decode_blob(other, frame)) + return edge_buffer.lookup(parent, frame, query_time, time_tolerance) + + +def _restamp(transform: Transform, ts: float) -> Transform: + return Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ts=ts, + ) + + +def _connect(db_path: str) -> sqlite3.Connection: + conn = sqlite3.connect(db_path, check_same_thread=False) + conn.execute("PRAGMA busy_timeout=5000") + return conn + + +def _ensure_child_index(conn: sqlite3.Connection) -> None: + """Index the child_frame json tag on the tf rows so per-frame time queries + seek. The live recorder gets this for free (the store auto-indexes tag keys on + tagged appends); this is for migrated recordings and the read side. Requires + the tf table to exist.""" + # Composite (child_frame, ts) so a per-frame "latest at/before T" is a direct + # index range seek, not a scan+sort. Index names share SQLite's global namespace + # with tables, so the name is double-underscore-namespaced to keep it clear of any + # real stream/table name (no stream would contain "__dbtf_"). + index_name = "tf__dbtf_child_ts_idx" + conn.execute( + f'CREATE INDEX IF NOT EXISTS "{index_name}" ' + "ON \"tf\"(json_extract(tags, '$.child_frame'), ts)" + ) + conn.commit() diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index 097dcec2c6..c2767b225a 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -199,6 +199,10 @@ def store(self) -> SqliteStore: class SemanticSearchConfig(MemoryModuleConfig): embedding_model: type[EmbeddingModel] | None = None + # The pose of a match is resolved at QUERY time via store.tf (so it reflects loop + # closure), looking up root_frame <- the observation's frame_id at its timestamp. + root_frame: str = "world" + tf_tolerance: float = 0.5 class SemanticSearch(MemoryModule): @@ -249,8 +253,21 @@ def _similarity(obs: Observation[Any]) -> float: return cast("EmbeddedObservation[Any]", obs).similarity or 0.0 best = results.transform(peaks(key=_similarity, distance=1.0)).last() + + # Resolve the match's world pose at QUERY time via tf, so it reflects any loop + # closure since the observation was recorded (the baked pose_stamped is stale). + frame_id = getattr(best.data, "frame_id", None) or self.config.root_frame + transform = self.store.tf.get( + self.config.root_frame, frame_id, best.ts, self.config.tf_tolerance + ) + if transform is not None: + return transform.to_pose(ts=best.ts) + # No tf chain at that time — fall back to the pose baked at record time. if best.pose_stamped is None: - raise LookupError("No pose on best search result") + raise LookupError("No tf and no baked pose on best search result") + logger.warning( + "SemanticSearch: no tf for %s @ %s; using stale baked pose", frame_id, best.ts + ) return best.pose_stamped @@ -320,15 +337,7 @@ async def _lidar_pose(self, msg): config: RecorderConfig - # Optional static-tf input stream: a future system publishes latched mount/extrinsic - # transforms here; recorded into "tf" + flagged static in the graph. Unconnected = - # no-op (today nothing publishes it). Folded into tf, not recorded as its own stream. tf_static: In[TFMessage] - - # Optional pose-graph deformation stream: a loop-closure backend (e.g. gsc_pgo) - # publishes one DeformationNode per keyframe on create and again whenever the - # optimizer moves it. Recorded into its own "tf_deformation_nodes" stream so a - # query can later correct tf for loop closure. Unconnected = no-op. tf_deformation_nodes: In[DeformationNode] _pose_setters: dict[str, Any] = {} @@ -432,15 +441,27 @@ def _on_frame_dropped(self, name: str) -> None: ) def _prepare_streams(self) -> None: - """On APPEND, drop the streams this recorder is about to (re)write — the - remapped In-port streams plus ``tf`` — so a re-run replaces them instead - of duplicating, while leaving any other streams in the db untouched.""" + """On APPEND, clear only the streams this recorder will actually (re)write — + i.e. those with a connected source — so a re-run overwrites them cleanly. + Streams whose source is unconnected are left untouched: recording into an + existing db never drops data it won't repopulate (e.g. pcap_to_db, which has + no tf source, must not wipe an existing tf / tf_deformation_nodes).""" if self.config.on_existing is not OnExisting.APPEND: return - targets = {self.config.stream_remapping.get(name, name) for name in self.inputs} + targets: set[str] = set() + for name, port in self.inputs.items(): + if name in ("tf_static", "tf_deformation_nodes"): + continue # tf-family, handled below only when their source is connected + if port.transport is None: + continue # unconnected -> nothing to write, leave any existing data + targets.add(self.config.stream_remapping.get(name, name)) if self.config.record_tf: - targets.add("tf") - targets.add("tf_deformation_nodes") + topic = getattr(self.tf.config, "topic", None) + pubsub = getattr(self.tf, "pubsub", None) + if (topic and pubsub is not None) or self.tf_static.transport is not None: + targets.update(("tf", "tf_graph")) + if self.tf_deformation_nodes.transport is not None: + targets.add("tf_deformation_nodes") for stream in targets.intersection(self.store.list_streams()): self.store.delete_stream(stream) diff --git a/dimos/memory2/store/base.py b/dimos/memory2/store/base.py index 46be8b8cb5..7818950f7a 100644 --- a/dimos/memory2/store/base.py +++ b/dimos/memory2/store/base.py @@ -121,9 +121,9 @@ def tf(self) -> DbTf: (e.g. ``world -> ... -> mid360_link``) from the recorded transforms. """ if self._tf is None: - from dimos.memory2.db_tf import DbTf + from dimos.memory2.db_tf_live import DbTfLive - self._tf = DbTf(self) + self._tf = DbTfLive(self) return self._tf def replay( diff --git a/dimos/memory2/store/sqlite.py b/dimos/memory2/store/sqlite.py index 229961a126..fe7a6ff393 100644 --- a/dimos/memory2/store/sqlite.py +++ b/dimos/memory2/store/sqlite.py @@ -16,7 +16,7 @@ import os import sqlite3 -from typing import Annotated, Any +from typing import TYPE_CHECKING, Annotated, Any from pydantic import BeforeValidator @@ -32,6 +32,9 @@ from dimos.memory2.vectorstore.base import VectorStore from dimos.memory2.vectorstore.sqlite import SqliteVectorStore +if TYPE_CHECKING: + from dimos.memory2.db_tf import DbTf + class SqliteStoreConfig(StoreConfig): """Config for SQLite-backed store.""" @@ -61,6 +64,17 @@ def __init__(self, **kwargs: Any) -> None: self._registry_conn = self._open_connection() self._registry = RegistryStore(conn=self._registry_conn) + @property + def tf(self) -> DbTf: + """Transform lookups over the sqlite graph stream (see :class:`DbTfSql`): + indexed per-lookup queries + loop-closure deformation correction, rather than + the base :class:`DbTfLive` in-RAM buffer.""" + if self._tf is None: + from dimos.memory2.db_tf_sql import DbTfSql + + self._tf = DbTfSql(self) + return self._tf + def _open_connection(self) -> sqlite3.Connection: """Open a new WAL-mode connection with sqlite-vec loaded.""" disposable, connection = open_disposable_sqlite_connection(self.config.path) diff --git a/dimos/memory2/test_db_tf.py b/dimos/memory2/test_db_tf.py index 98ff84c28b..96d33fc2f7 100644 --- a/dimos/memory2/test_db_tf.py +++ b/dimos/memory2/test_db_tf.py @@ -23,7 +23,7 @@ import math from pathlib import Path -from dimos.memory2.db_tf import DbTf +from dimos.memory2.db_tf_sql import DbTfSql from dimos.memory2.store.sqlite import SqliteStore from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion @@ -36,7 +36,6 @@ _T0 = 1000.0 _DYN_RATE = 30.0 _DURATION = 10.0 -_NO_PRUNE = 1.0e15 def _yaw(theta: float) -> Quaternion: @@ -66,7 +65,7 @@ def _append(store: SqliteStore, transform: Transform) -> None: def _ref(transforms: list[Transform]) -> MultiTBuffer: - buffer = MultiTBuffer(buffer_size=_NO_PRUNE) + buffer = MultiTBuffer(buffer_size=math.inf) buffer.receive_transform(*transforms) return buffer @@ -123,7 +122,7 @@ def test_interpolates_and_matches_full_load(tmp_path: Path) -> None: transforms = _record_single_robot(tmp_path / "r.db", static_repeat=True) reference = _ref(transforms) store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) - db = DbTf(store) + db = DbTfSql(store) compared = 0 for k in range(25): q = _T0 + 0.013 + k * 0.317 @@ -142,7 +141,7 @@ def test_latched_static_resolves(tmp_path: Path) -> None: (no bracket, no tolerance) — the case a plain time-bracket would drop.""" _record_single_robot(tmp_path / "r.db", static_repeat=False) store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) - db = DbTf(store) + db = DbTfSql(store) assert db.get("world", "sensor", _T0 + 9.5, 0.5) is not None # ~9.5s after statics store.stop() @@ -185,7 +184,7 @@ def test_reparent_midrun_uses_graph_as_of_query_time(tmp_path: Path) -> None: store.stop() store = SqliteStore(path=str(path), must_exist=True) - db = DbTf(store) + db = DbTfSql(store) q1 = _T0 + 2.013 want1 = _ref( @@ -259,7 +258,7 @@ def test_loop_closure_deformation_corrects_matched_edge(tmp_path: Path) -> None: store.stop() store = SqliteStore(path=str(path), must_exist=True) - db = DbTf(store) + db = DbTfSql(store) got = db.get("map", "odom", _T0 + 0.3, 0.5) assert got is not None and abs(got.translation.x - 1.0) < 1e-6 # raw identity + delta base = db.get("odom", "base_link", _T0 + 0.3, 0.5) # unmatched edge, unchanged @@ -283,7 +282,7 @@ def test_loop_closure_deformation_blends_between_keyframes(tmp_path: Path) -> No store.stop() store = SqliteStore(path=str(path), must_exist=True) - db = DbTf(store) + db = DbTfSql(store) got = db.get("map", "odom", _T0 + 4.0, 0.5) # midpoint of [t0, t0+8] assert got is not None and abs(got.translation.x - 1.0) < 1e-6 store.stop() @@ -319,7 +318,7 @@ def test_disjoint_multirobot_returns_none(tmp_path: Path) -> None: store.stop() store = SqliteStore(path=str(path), must_exist=True) - db = DbTf(store) + db = DbTfSql(store) q = _T0 + 0.3 assert db.get("baseA", "worldA", q, 0.5) is not None # same component assert db.get("baseB", "baseA", q, 0.5) is None # different components diff --git a/dimos/navigation/jnav/msgs/Marker.py b/dimos/navigation/jnav/msgs/Marker.py deleted file mode 100644 index a7ed3fef98..0000000000 --- a/dimos/navigation/jnav/msgs/Marker.py +++ /dev/null @@ -1,97 +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. - -"""Marker: a named navigation marker — a pose in a frame, a marker name, and an -optional map name. - -Drives the objective handler's ``goal_marker`` (navigate to a named marker) and -``save_marker`` (persist the current/given pose under a name) streams, so callers -can reference waypoints by name (e.g. dim_city's ``test_waypoint_*``) and scope -them to a map, instead of passing a bare Point/Pose with no identity. -""" - -from __future__ import annotations - -import struct -import time -from typing import BinaryIO - -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.types.timestamped import Timestamped - - -class Marker(Timestamped): - msg_name = "jnav.Marker" - - ts: float - frame_id: str - pose: Pose - marker: str # marker name (identity) - map: str # optional map name this marker belongs to ("" = unspecified) - - def __init__( - self, - pose: Pose | None = None, - marker: str = "", - map: str = "", - ts: float = 0.0, - frame_id: str = "map", - ) -> None: - self.ts = ts if ts != 0 else time.time() - self.frame_id = frame_id - self.pose = pose if pose is not None else Pose() - self.marker = marker - self.map = map - - def lcm_encode(self) -> bytes: - parts: list[bytes] = [struct.pack(">d", self.ts)] - for text in (self.frame_id, self.marker, self.map): - encoded = text.encode("utf-8") - parts.append(struct.pack(">I", len(encoded))) - parts.append(encoded) - p = self.pose - parts.append( - struct.pack( - ">7d", - p.position.x, - p.position.y, - p.position.z, - p.orientation.x, - p.orientation.y, - p.orientation.z, - p.orientation.w, - ) - ) - return b"".join(parts) - - @classmethod - def lcm_decode(cls, data: bytes | BinaryIO) -> Marker: - buf = data if isinstance(data, (bytes, bytearray)) else data.read() - offset = 0 - (ts,) = struct.unpack_from(">d", buf, offset) - offset += 8 - texts: list[str] = [] - for _ in range(3): - (length,) = struct.unpack_from(">I", buf, offset) - offset += 4 - texts.append(buf[offset : offset + length].decode("utf-8")) - offset += length - frame_id, marker, map_name = texts - px, py, pz, qx, qy, qz, qw = struct.unpack_from(">7d", buf, offset) - pose = Pose() - pose.position = Vector3(px, py, pz) - pose.orientation = Quaternion(qx, qy, qz, qw) - return cls(pose=pose, marker=marker, map=map_name, ts=ts, frame_id=frame_id) diff --git a/dimos/navigation/jnav/msgs/test_Marker.py b/dimos/navigation/jnav/msgs/test_Marker.py deleted file mode 100644 index 61394999d7..0000000000 --- a/dimos/navigation/jnav/msgs/test_Marker.py +++ /dev/null @@ -1,67 +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. - -# 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. - -"""Roundtrip + field tests for the jnav Marker message.""" - -from __future__ import annotations - -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.navigation.jnav.msgs.Marker import Marker - - -def _pose(x: float, y: float, z: float) -> Pose: - pose = Pose() - pose.position.x, pose.position.y, pose.position.z = x, y, z - return pose - - -def test_marker_roundtrip_preserves_all_fields() -> None: - marker = Marker( - pose=_pose(-52.93, -55.95, 0.4), - marker="test_waypoint_2", - map="dim_city", - ts=1781565207.5, - frame_id="map", - ) - decoded = Marker.lcm_decode(marker.lcm_encode()) - - assert decoded.marker == "test_waypoint_2" - assert decoded.map == "dim_city" - assert decoded.frame_id == "map" - assert decoded.ts == marker.ts - assert decoded.pose.position.x == -52.93 - assert decoded.pose.position.y == -55.95 - assert decoded.pose.position.z == 0.4 - - -def test_marker_defaults() -> None: - marker = Marker() - assert marker.marker == "" - assert marker.map == "" - assert marker.frame_id == "map" - assert marker.ts > 0 # auto-stamped - # empty optional strings survive the roundtrip - decoded = Marker.lcm_decode(marker.lcm_encode()) - assert decoded.marker == "" and decoded.map == "" - - -def test_marker_unicode_name() -> None: - decoded = Marker.lcm_decode(Marker(marker="café_door", map="map_α").lcm_encode()) # noqa: RUF001 - assert decoded.marker == "café_door" - assert decoded.map == "map_α" # noqa: RUF001 diff --git a/dimos/protocol/tf/deformation.py b/dimos/protocol/tf/deformation.py new file mode 100644 index 0000000000..a56c2a2698 --- /dev/null +++ b/dimos/protocol/tf/deformation.py @@ -0,0 +1,143 @@ +# 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. + +"""Loop-closure deformation correction for transform lookups. + +A loop-closure backend (e.g. gsc_pgo) publishes a stream of :class:`DeformationNode` +keyframes — one per pose-graph node, re-published (same id) when the optimizer moves +it, tagged with the tf_id of the edge it corrects. This module turns that stream into +a query-time correction on a raw edge transform: + + delta[k] = current[k] ∘ inv(original[k]) # how far keyframe k moved since first seen + delta(t) = slerp/lerp blend of the two keyframes bracketing t + corrected_edge = delta(t) ∘ raw_edge + +Used by both the recorded path (:class:`DbTfSql`) and the live path +(:class:`PubSubTF`); the SE(3) helpers are shared so the two stay identical. +""" + +from __future__ import annotations + +import bisect +from typing import Any, cast + +import numpy as np + +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.nav_msgs.DeformationNode import DeformationNode, tf_id_for + + +def apply_delta(delta: np.ndarray, transform: Transform, ts: float) -> Transform: + """Deform ``transform`` by the SE(3) correction ``delta`` on the parent (frame_id) + side: ``corrected = delta @ raw``. Keeps the edge's frame names.""" + corrected = delta @ transform.to_matrix() + return Transform.from_matrix( + corrected, + ts=ts, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ) + + +def quat_slerp(q_lo: np.ndarray, q_hi: np.ndarray, weight: float) -> np.ndarray: + """Spherical-linear interpolation between two quaternions ``[x, y, z, w]``.""" + dot = float(np.dot(q_lo, q_hi)) + if dot < 0.0: # take the shorter arc + q_hi = -q_hi + dot = -dot + if dot > 0.9995: # nearly parallel: lerp + renormalize avoids a divide-by-~0 + result = q_lo + weight * (q_hi - q_lo) + return cast("np.ndarray", result / np.linalg.norm(result)) + theta_0 = np.arccos(np.clip(dot, -1.0, 1.0)) + sin_0 = np.sin(theta_0) + scale_lo = np.sin((1.0 - weight) * theta_0) / sin_0 + scale_hi = np.sin(weight * theta_0) / sin_0 + return cast("np.ndarray", scale_lo * q_lo + scale_hi * q_hi) + + +def blend_se3(mat_lo: np.ndarray, mat_hi: np.ndarray, weight: float) -> np.ndarray: + """Linear-blend-skin two SE(3) deltas by ``weight`` in [0, 1] (0 -> lo, 1 -> hi): + lerp the translation, slerp the rotation.""" + quat_lo = Quaternion.from_rotation_matrix(mat_lo[:3, :3]) + quat_hi = Quaternion.from_rotation_matrix(mat_hi[:3, :3]) + blended = quat_slerp( + np.array([quat_lo.x, quat_lo.y, quat_lo.z, quat_lo.w]), + np.array([quat_hi.x, quat_hi.y, quat_hi.z, quat_hi.w]), + weight, + ) + out = np.eye(4) + out[:3, :3] = Quaternion(blended).to_rotation_matrix() + out[:3, 3] = (1.0 - weight) * mat_lo[:3, 3] + weight * mat_hi[:3, 3] + return out + + +class DeformationBuffer: + """In-RAM rolling store of pose-graph keyframes, keyed by tf_id (the corrected edge). + + Feed it :class:`DeformationNode` messages via :meth:`receive`; it keeps, per node, + the original (first-seen) and current (latest) pose. :meth:`correct` applies the + blended loop-closure delta to a raw edge transform at a query time. This is the + live-path analogue of the DB-backed correction in :class:`DbTfSql`.""" + + def __init__(self) -> None: + # tf_id -> node_id -> [original_4x4, current_4x4, keyframe_ts] + self._edges: dict[int, dict[int, list[Any]]] = {} + + def receive(self, node: DeformationNode) -> None: + matrix = Transform.from_pose(node.pose.frame_id, node.pose).to_matrix() + edge = self._edges.setdefault(node.tf_id, {}) + state = edge.get(node.id) + if state is None: + edge[node.id] = [matrix, matrix, node.pose.ts] # original, current, ts + else: + state[1] = matrix # optimizer moved it -> update current only + + @property + def has_corrections(self) -> bool: + return bool(self._edges) + + def _edge_delta(self, tf_id: int, query_time: float) -> np.ndarray | None: + edge = self._edges.get(tf_id) + if not edge: + return None + # keyframes sorted by ts; bracket the query time (latest<=t, earliest>=t) + ordered = sorted(edge.values(), key=lambda state: state[2]) + stamps = [state[2] for state in ordered] + lo_index = bisect.bisect_right(stamps, query_time) - 1 # last at or before + hi_index = bisect.bisect_left(stamps, query_time) # first at or after + picks = [] + if lo_index >= 0: + picks.append(ordered[lo_index]) + if hi_index < len(ordered) and (not picks or ordered[hi_index] is not picks[0]): + picks.append(ordered[hi_index]) + if not picks: + return None + # per keyframe: delta = current ∘ inv(original) + samples = [(state[2], state[1] @ np.linalg.inv(state[0])) for state in picks] + if len(samples) == 1 or samples[0][0] == samples[1][0]: + return cast("np.ndarray", samples[0][1]) + (ts_lo, mat_lo), (ts_hi, mat_hi) = samples + weight = (query_time - ts_lo) / (ts_hi - ts_lo) + return blend_se3(mat_lo, mat_hi, weight) + + def correct( + self, frame_from: str, frame_to: str, raw: Transform, query_time: float + ) -> Transform: + """Return the loop-closure-corrected ``frame_from <- frame_to`` transform, or + ``raw`` unchanged if no deformation nodes match that edge.""" + if not self._edges: + return raw + delta = self._edge_delta(tf_id_for(frame_from, frame_to), query_time) + return apply_delta(delta, raw, raw.ts) if delta is not None else raw diff --git a/dimos/protocol/tf/test_deformation.py b/dimos/protocol/tf/test_deformation.py new file mode 100644 index 0000000000..2389370186 --- /dev/null +++ b/dimos/protocol/tf/test_deformation.py @@ -0,0 +1,74 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Live-path loop-closure correction (DeformationBuffer).""" + +from __future__ import annotations + +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +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.DeformationNode import DeformationNode, tf_id_for +from dimos.protocol.tf.deformation import DeformationBuffer + + +def _node(tf_id: int, node_id: int, ts: float, xyz: tuple[float, float, float]) -> DeformationNode: + return DeformationNode( + id=node_id, + tf_id=tf_id, + pose=PoseStamped(ts=ts, frame_id="map", position=list(xyz), orientation=[0, 0, 0, 1]), + ) + + +def _identity(parent: str, child: str, ts: float) -> Transform: + return Transform( + translation=Vector3(0, 0, 0), + rotation=Quaternion(0, 0, 0, 1), + frame_id=parent, + child_frame_id=child, + ts=ts, + ) + + +def test_single_keyframe_corrects_matched_edge_only() -> None: + buffer = DeformationBuffer() + tf_id = tf_id_for("map", "odom") + buffer.receive(_node(tf_id, 1, 100.0, (0.0, 0.0, 0.0))) # original + buffer.receive(_node(tf_id, 1, 100.0, (1.0, 0.0, 0.0))) # optimizer moved it +1 in x + + corrected = buffer.correct("map", "odom", _identity("map", "odom", 100.0), 100.0) + assert abs(corrected.translation.x - 1.0) < 1e-9 # delta = current ∘ inv(original) applied + + # an edge with no matching tf_id passes through untouched + other = _identity("odom", "base_link", 100.0) + assert buffer.correct("odom", "base_link", other, 100.0).translation.x == 0.0 + + +def test_blends_between_bracketing_keyframes() -> None: + buffer = DeformationBuffer() + tf_id = tf_id_for("map", "odom") + buffer.receive(_node(tf_id, 1, 100.0, (0.0, 0.0, 0.0))) # kf A: unmoved + buffer.receive(_node(tf_id, 1, 100.0, (0.0, 0.0, 0.0))) + buffer.receive(_node(tf_id, 2, 110.0, (0.0, 0.0, 0.0))) # kf B: moved +2 + buffer.receive(_node(tf_id, 2, 110.0, (2.0, 0.0, 0.0))) + + corrected = buffer.correct("map", "odom", _identity("map", "odom", 105.0), 105.0) + assert abs(corrected.translation.x - 1.0) < 1e-9 # midpoint of 0->2 + + +def test_no_nodes_passes_through_identity() -> None: + buffer = DeformationBuffer() + raw = _identity("map", "odom", 1.0) + assert buffer.correct("map", "odom", raw, 1.0) is raw # no copy, no cost diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 4dd345a530..81d8e356ce 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -23,10 +23,12 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.nav_msgs.DeformationNode import DeformationNode from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.protocol.pubsub.impl.lcmpubsub import LCM, Topic from dimos.protocol.pubsub.spec import PubSub from dimos.protocol.service.spec import BaseConfig, Service +from dimos.protocol.tf.deformation import DeformationBuffer from dimos.types.timestamped import to_human_readable from dimos.utils.logging_config import setup_logger from dimos.utils.timeseries.inmemory import InMemoryStore @@ -325,6 +327,10 @@ def __str__(self) -> str: class PubSubTFConfig(TFConfig): topic: Topic | None = None # Required field but needs default for dataclass inheritance + # Optional per-keyframe loop-closure deformation stream (e.g. gsc_pgo's + # tf_deformation_nodes). When wired, get()/lookup() return loop-closure-corrected + # transforms live, matching the recorded DbTfSql behaviour. Unset = raw tf only. + deformation_topic: Topic | None = None pubsub: type[PubSub] | PubSub | None = None # type: ignore[type-arg] autostart: bool = True @@ -335,6 +341,7 @@ class PubSubTF(MultiTBuffer, TFSpec): def __init__(self, **kwargs) -> None: # type: ignore[no-untyped-def] TFSpec.__init__(self, **kwargs) MultiTBuffer.__init__(self, self.config.buffer_size) + self._deformation = DeformationBuffer() pubsub_config = getattr(self.config, "pubsub", None) if pubsub_config is not None: @@ -354,6 +361,44 @@ def start(self, sub: bool = True) -> None: topic = getattr(self.config, "topic", None) if topic: self.pubsub.subscribe(topic, self.receive_msg) + deformation_topic = getattr(self.config, "deformation_topic", None) + if deformation_topic: + self.pubsub.subscribe(deformation_topic, self.receive_deformation) + + def receive_deformation(self, msg: DeformationNode, topic: Topic) -> None: + self._deformation.receive(msg) + + def get_transform( + self, + parent_frame: str, + child_frame: str, + time_point: float | None = None, + time_tolerance: float | None = None, + ) -> Transform | None: + """Single-edge lookup, loop-closure-corrected. Correction is applied on the + edge's stored direction, so both direct lookups and BFS chain hops (forward or + reversed) resolve to the corrected transform. Falls back to the raw lookup when + no deformation nodes have arrived.""" + if not self._deformation.has_corrections: + return super().get_transform(parent_frame, child_frame, time_point, time_tolerance) + with self._cv: + forward = (parent_frame, child_frame) + if forward in self.buffers: + raw = self.buffers[forward].get(time_point, time_tolerance) # type: ignore[arg-type] + if raw is None: + return None + ts = time_point if time_point is not None else raw.ts + return self._deformation.correct(parent_frame, child_frame, raw, ts) + reverse = (child_frame, parent_frame) + if reverse in self.buffers: + # child->parent (stored direction) + raw = self.buffers[reverse].get(time_point, time_tolerance) # type: ignore[arg-type] + if raw is None: + return None + ts = time_point if time_point is not None else raw.ts + corrected = self._deformation.correct(child_frame, parent_frame, raw, ts) + return corrected.inverse() + return None def stop(self) -> None: self.pubsub.stop() @@ -427,6 +472,12 @@ def receive_msg(self, msg: TFMessage, topic: Topic) -> None: class LCMPubsubConfig(PubSubTFConfig): topic: Topic = field(default_factory=lambda: Topic("/tf", TFMessage)) + # On by default: subscribe the loop-closure deformation stream so live tf.get is + # corrected whenever a PGO publishes it. Harmless when nothing does — the buffer + # stays empty and correct() is a zero-cost pass-through. + deformation_topic: Topic | None = field( + default_factory=lambda: Topic("/tf_deformation_nodes", DeformationNode) + ) pubsub: type[PubSub] | PubSub | None = LCM # type: ignore[type-arg] autostart: bool = True From 838203fb4dd164e132ca98bc307b42cd65a01c30 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Jul 2026 07:52:46 -0500 Subject: [PATCH 50/69] post_process: auto-detect AprilTags when missing + persist PGO artifacts Detect AprilTags over the camera stream when the raw tag stream is absent, so a fresh recording only needs post_process (derive distance/view-angle from the tag pose, defaulting unknown speeds to pass). Persist the PGO's keyframe deformation (gt_tf_deformation_nodes) and optimized pose_graph as real typed streams so a deformation-aware tf.get can replay the loop-closure correction offline. --- .../gsc_pgo/scripts/post_process.py | 125 ++++++++++++++++-- 1 file changed, 111 insertions(+), 14 deletions(-) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py index 3a9cc93e5b..e6458231ef 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py @@ -55,8 +55,11 @@ import numpy as np from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.nav_msgs.DeformationNode import DeformationNode, tf_id_for from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.navigation.jnav.msgs.Graph3D import Graph3D from dimos.navigation.jnav.utils import recording_db as rdb from dimos.navigation.jnav.utils.apriltags import ( DEFAULT_MAX_ANGULAR_SPEED_DPS, @@ -66,6 +69,9 @@ DEFAULT_MAX_VIEW_ANGLE_DEG, DEFAULT_MIN_SHARPNESS, DEFAULT_MIN_TAG_PX, + _write_tag_stream, + detect_raw_detections, + view_quality, ) VISIT_GAP_S = 30.0 @@ -86,7 +92,12 @@ def arg(flag, default=""): # reject anything that isn't a plain identifier to keep that injection-free. if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", ODOM_STREAM): raise ValueError(f"unsafe --odom stream name: {ODOM_STREAM!r}") -RAW_STREAM = arg("--tags", "raw_april_tags") # input unfiltered AprilTag stream +RAW_STREAM = arg( + "--tags", "raw_april_tags" +) # unfiltered AprilTag stream (auto-detected if missing) +CAMERA = arg("--camera", "color_image") # image stream to detect on when RAW_STREAM is missing +MARKER_LENGTH_M = float(arg("--tag-size", "0.10")) # AprilTag edge length (m), for auto-detect +DICTIONARY = arg("--dict", "DICT_APRILTAG_36h11") # AprilTag dictionary, for auto-detect IGNORE_TAGS = { int(marker_id) for marker_id in arg("--ignore-tags").replace(",", " ").split() } # dynamic/moving tags @@ -128,8 +139,30 @@ def arg(flag, default=""): ) store = rdb.store(DB) if RAW_STREAM not in store.list_streams(): - sys.exit( - f"!! {RAW_STREAM} missing -- run detect_tags.py first to build the unfiltered tag stream." + # No tag stream yet -> detect them now (so a fresh recording only needs post_process). + if CAMERA not in store.list_streams(): + sys.exit( + f"!! {RAW_STREAM} missing and can't auto-detect: camera stream {CAMERA!r} not in db." + ) + print( + f"{RAW_STREAM} missing -- detecting AprilTags over {CAMERA} " + f"(tag_size={MARKER_LENGTH_M} m, dict={DICTIONARY})...", + flush=True, + ) + camera_matrix = np.array(intrinsics["intrinsics"], float).reshape(3, 3) + distortion = np.array(intrinsics.get("distortion", []), float) + raw_detections, _, n_images = detect_raw_detections( + store, + camera_matrix, + distortion, + image_stream=CAMERA, + marker_length=MARKER_LENGTH_M, + dictionary=DICTIONARY, + ) + _write_tag_stream(store, RAW_STREAM, raw_detections, diagnostics=True) + print( + f"wrote {RAW_STREAM}: {len(raw_detections)} raw detections over {n_images} frames", + flush=True, ) @@ -194,6 +227,20 @@ def passes(detection): for observation in store.stream(RAW_STREAM): pose = observation.data tags = observation.tags + # distance_m / view_angle_deg are derived from the tag pose (not required in the + # tag stream), so any detector's raw stream works as long as it carries the pose + # + the image-only diagnostics (sharpness/reproj_px/tag_px). speeds default to -1 + # ("unknown", always passes) when a stream doesn't record them. + tag_pose = [ + pose.x, + pose.y, + pose.z, + pose.orientation.x, + pose.orientation.y, + pose.orientation.z, + pose.orientation.w, + ] + distance_m, view_angle_deg = view_quality(tag_pose) raw_detections.append( dict( ts=float(observation.ts), @@ -205,17 +252,12 @@ def passes(detection): Point3(pose.x, pose.y, pose.z), ), reproj_px=float(tags["reproj_px"]), - **{ - diagnostic: float(tags[diagnostic]) - for diagnostic in ( - "sharpness", - "tag_px", - "distance_m", - "view_angle_deg", - "lin_speed", - "ang_speed", - ) - }, + sharpness=float(tags["sharpness"]), + tag_px=float(tags["tag_px"]), + distance_m=float(distance_m), + view_angle_deg=float(view_angle_deg), + lin_speed=float(tags.get("lin_speed", -1.0)), + ang_speed=float(tags.get("ang_speed", -1.0)), ) ) gated_detections = [ @@ -539,6 +581,61 @@ def pose_tuple(pose): ) +# Persist the PGO's internal artifacts as REAL streams (true payload types): +# gt_tf_deformation_nodes (DeformationNode) -- per keyframe, the raw pose (original) +# then the optimized pose (current); a deformation-aware tf.get can replay the +# loop-closure correction from these exactly like the online gsc_pgo stream. +# pose_graph (Graph3D) -- the optimized keyframe nodes + sequential odom edges. +_tf_edge_id = tf_id_for("map", "odom") +_deform_name = f"gt_tf_deformation_nodes{SUFFIX}" +if _deform_name in store.list_streams(): + store.delete_stream(_deform_name) +_deform_stream = store.stream(_deform_name, DeformationNode) +for index in range(num_keyframes): + node_ts = float(keyframe_times[index]) + for keyframe_pose in (raw_keyframe_poses[index], estimate.atPose3(index)): # original, current + px, py, pz, qx, qy, qz, qw = pose_tuple(keyframe_pose) + _deform_stream.append( + DeformationNode( + id=index, + tf_id=_tf_edge_id, + pose=PoseStamped( + ts=node_ts, frame_id="map", position=[px, py, pz], orientation=[qx, qy, qz, qw] + ), + ), + ts=node_ts, + pose=None, + tags={"tf_id": str(_tf_edge_id), "id": str(index)}, + ) +print(f"wrote {_deform_name}: {num_keyframes} keyframes (raw+optimized)", flush=True) + +_graph_name = f"pose_graph{SUFFIX}" +if _graph_name in store.list_streams(): + store.delete_stream(_graph_name) +_graph_nodes = [] +for index in range(num_keyframes): + px, py, pz, qx, qy, qz, qw = pose_tuple(estimate.atPose3(index)) + _graph_nodes.append( + Graph3D.Node3D( + pose=PoseStamped( + ts=float(keyframe_times[index]), + frame_id="map", + position=[px, py, pz], + orientation=[qx, qy, qz, qw], + ), + id=index, + ) + ) +_graph_edges = [ + Graph3D.Edge(index, index + 1, float(keyframe_times[index + 1])) + for index in range(num_keyframes - 1) +] +_graph_ts = float(keyframe_times[-1]) +store.stream(_graph_name, Graph3D).append( + Graph3D(ts=_graph_ts, nodes=_graph_nodes, edges=_graph_edges), ts=_graph_ts, pose=None +) +print(f"wrote {_graph_name}: {num_keyframes} nodes, {len(_graph_edges)} edges", flush=True) + if WHAT in ("odom", "both"): out_name = f"{OUT_PREFIX}_odometry{SUFFIX}" if out_name in store.list_streams(): From e87fca4287f2102a3afea527673d2bba7163254e Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Jul 2026 07:52:59 -0500 Subject: [PATCH 51/69] fix(eval): replace removed odometry_pose7_lookup with tf-based lookup eval.py's odom-pose baseline was replaced by tf_pose_samples, but eval_ground_truth_tag.py still imported the deleted odometry_pose7_lookup. Build the raw pose7 lookup from tf_pose_samples + pose7_lookup, mirroring evaluate(). Fixes the broken import (full-repo mypy). --- .../components/loop_closure/eval_ground_truth_tag.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py b/dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py index 3cb34b9018..766aeefe2d 100644 --- a/dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py +++ b/dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py @@ -51,16 +51,17 @@ from scipy.spatial.transform import Rotation from dimos.navigation.jnav.components.loop_closure.eval import ( - odometry_pose7_lookup, run_module_graph, + tf_pose_samples, ) from dimos.navigation.jnav.utils.module_loading import ( filter_config_for_module, load_module_class, ) -from dimos.navigation.jnav.utils.recording_db import store +from dimos.navigation.jnav.utils.recording_db import ODOM_MATCH_TOLERANCE_S, store from dimos.navigation.jnav.utils.trajectory_metrics import ( drift_delta_lookup, + pose7_lookup, rigid_align_rmse, ) @@ -141,6 +142,8 @@ def tag_constellation( odom_stream: str, optical_in_base: list[float], ignore_tags: set[int], + world_frame: str = "world", + body_frame: str = "base_link", ) -> tuple[dict[int, np.ndarray], int]: """Run the PGO, then place each tag sighting in the world via the corrected trajectory + extrinsic, averaging per tag. Returns {marker_id: world centroid} @@ -155,7 +158,10 @@ def tag_constellation( odom_stream=odom_stream, lockstep=True, ) - raw_pose7_lookup = odometry_pose7_lookup(db_path, odom_stream) + raw_times, raw_poses7 = tf_pose_samples( + db_path, odom_stream, world_frame=world_frame, body_frame=body_frame + ) + raw_pose7_lookup = pose7_lookup(raw_times, raw_poses7, ODOM_MATCH_TOLERANCE_S) delta_lookup = drift_delta_lookup(graph, raw_pose7_lookup) by_tag: dict[int, list[np.ndarray]] = defaultdict(list) for timestamp, marker_id, tag_pose7 in tag_in_camera_sightings(db_path): From bf5469eb5926a7ef22945b7d95af9182d2bdf1bb Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Jul 2026 08:05:48 -0500 Subject: [PATCH 52/69] test(memory2): skip DbTfSql tests when sqlite-vec can't load test_db_tf.py instantiates SqliteStore directly (not via the sqlite_store fixture), so it bypassed conftest's _SKIP_SQLITE_VEC guard and failed on the ubuntu-24.04-arm runner (vec0.so wrong ELF class). Add a module-level skipif mirroring memory2/conftest.py so these 6 tests skip instead of fail there. --- dimos/memory2/test_db_tf.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/dimos/memory2/test_db_tf.py b/dimos/memory2/test_db_tf.py index 96d33fc2f7..420985c953 100644 --- a/dimos/memory2/test_db_tf.py +++ b/dimos/memory2/test_db_tf.py @@ -22,6 +22,9 @@ import math from pathlib import Path +import platform + +import pytest from dimos.memory2.db_tf_sql import DbTfSql from dimos.memory2.store.sqlite import SqliteStore @@ -33,6 +36,11 @@ from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.protocol.tf.tf import MultiTBuffer +# sqlite-vec fails to load on Linux ARM (32-bit binary in the aarch64 wheel) and +# on macOS in CI; SqliteStore loads it on init, so skip like memory2/conftest.py. +_SKIP_SQLITE_VEC = platform.machine() == "aarch64" or platform.system() == "Darwin" +pytestmark = pytest.mark.skipif(_SKIP_SQLITE_VEC, reason="sqlite-vec extension not loadable here") + _T0 = 1000.0 _DYN_RATE = 30.0 _DURATION = 10.0 From c214f177897c791847164f8168497c64987e6e41 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Thu, 2 Jul 2026 11:11:10 -0500 Subject: [PATCH 53/69] fix(post_process): honor lidar scan frame_id instead of hardcoded frame world_points() always ran tf.get(world, LIDAR_FRAME=mid360_link) regardless of the scan's actual frame. Scans already in the world frame (e.g. pointlio_lidar, header.frame_id='world') got the full robot pose applied on top, double- registering them into a smear that grows with distance from origin. Read the scan's own header.frame_id: leave already-world scans untouched, else tf- register world<-frame_id. LIDAR_FRAME is now only a fallback for a missing frame. --- .../gsc_pgo/scripts/post_process.py | 22 ++++++++++++++----- 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py index e6458231ef..ca723677c3 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py @@ -181,15 +181,21 @@ def transform_matrix(transform): def world_points(observation): """Nx3 world-registered points for a lidar observation. - Primary: the recording's tf (``world <- LIDAR_FRAME``) at the scan time. - Fallback: the observation's stored pose as the transform. Otherwise the - scan is assumed already world-registered (legacy recordings). + The scan's own ``header.frame_id`` decides what to do: a scan already in + WORLD_FRAME is returned untouched (transforming it again double-registers + it), otherwise it's brought into world via tf (``world <- frame_id``) at the + scan time. Falls back to LIDAR_FRAME when the header carries no frame, then + to the observation's stored pose, then to assuming it's already world. """ points = np.asarray(observation.data.points_f32()) if not len(points): return points + header = getattr(observation.data, "header", None) + scan_frame = getattr(header, "frame_id", "") or LIDAR_FRAME + if scan_frame == WORLD_FRAME: + return points # already world-registered per its own header if _TF_AVAILABLE: - transform = store.tf.get(WORLD_FRAME, LIDAR_FRAME, float(observation.ts), TF_TOL) + transform = store.tf.get(WORLD_FRAME, scan_frame, float(observation.ts), TF_TOL) if transform is not None: rotation, translation = transform_matrix(transform) return points @ rotation.T + translation @@ -197,12 +203,16 @@ def world_points(observation): if isinstance(pose, (tuple, list)) and len(pose) >= 7: rotation = Rot3.Quaternion(pose[6], pose[3], pose[4], pose[5]).matrix() return points @ rotation.T + np.array(pose[:3], float) - return points # already world-registered + return points # unknown frame, assume already world-registered print(f"recording: {REC}", flush=True) if _TF_AVAILABLE: - print(f"world-registering {LIDAR_STREAM} via tf ({WORLD_FRAME} <- {LIDAR_FRAME})", flush=True) + print( + f"world-registering {LIDAR_STREAM} via tf into {WORLD_FRAME} " + f"(per-scan header.frame_id; already-{WORLD_FRAME} scans left as-is)", + flush=True, + ) print( f"streams: tags={RAW_STREAM} odom={ODOM_STREAM} lidar={LIDAR_STREAM} -> out={OUT_PREFIX}{SUFFIX}", flush=True, From 4292948d880cb35e795a4a8f835ee88b5b4f8c9e Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sun, 5 Jul 2026 13:29:25 -0700 Subject: [PATCH 54/69] - --- .../components/loop_closure/gsc_pgo/scripts/make_rrd.py | 5 +++-- .../loop_closure/gsc_pgo/scripts/post_process.py | 9 ++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py index 561cc64430..0185ffb0f1 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py @@ -121,8 +121,9 @@ def landmarks(gt_odom): float(tag_metrics["sharpness"]) >= GATE["s"] and float(tag_metrics["reproj_px"]) <= GATE["r"] and float(tag_metrics["tag_px"]) >= GATE["px"] - and float(tag_metrics["distance_m"]) <= GATE["d"] - and float(tag_metrics["view_angle_deg"]) <= GATE["a"] + # older tag streams lack distance/view-angle; unknown passes the gate + and float(tag_metrics.get("distance_m", 0.0)) <= GATE["d"] + and float(tag_metrics.get("view_angle_deg", 0.0)) <= GATE["a"] and ( float(tag_metrics["lin_speed"]) < 0 or float(tag_metrics["lin_speed"]) <= GATE["lv"] diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py index ca723677c3..6c41d01713 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py @@ -181,17 +181,16 @@ def transform_matrix(transform): def world_points(observation): """Nx3 world-registered points for a lidar observation. - The scan's own ``header.frame_id`` decides what to do: a scan already in + The scan's own ``frame_id`` decides what to do: a scan already in WORLD_FRAME is returned untouched (transforming it again double-registers it), otherwise it's brought into world via tf (``world <- frame_id``) at the - scan time. Falls back to LIDAR_FRAME when the header carries no frame, then + scan time. Falls back to LIDAR_FRAME when the scan carries no frame, then to the observation's stored pose, then to assuming it's already world. """ points = np.asarray(observation.data.points_f32()) if not len(points): return points - header = getattr(observation.data, "header", None) - scan_frame = getattr(header, "frame_id", "") or LIDAR_FRAME + scan_frame = getattr(observation.data, "frame_id", "") or LIDAR_FRAME if scan_frame == WORLD_FRAME: return points # already world-registered per its own header if _TF_AVAILABLE: @@ -210,7 +209,7 @@ def world_points(observation): if _TF_AVAILABLE: print( f"world-registering {LIDAR_STREAM} via tf into {WORLD_FRAME} " - f"(per-scan header.frame_id; already-{WORLD_FRAME} scans left as-is)", + f"(per-scan frame_id; already-{WORLD_FRAME} scans left as-is)", flush=True, ) print( From 24a162e5d81485e16b79f4f2f718bc2459a003f5 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 08:35:21 -0700 Subject: [PATCH 55/69] refactor(jnav): use main's StreamTF, drop branch-local db_tf from memory2 Replace the branch's memory2 tf-tree/DbTf layer with a jnav-local RecordingTF(StreamTF) built from the recording store, leaving memory2 net-zero vs origin/main (no changes to Ivan's code). RecordingTF full-loads the recording's tf once so one-shot static frames aren't evicted by StreamTF's windowed cache. post_process.py and eval.py now construct RecordingTF.from_store(store); the live-path deformation feature in protocol/tf is kept and its stale DbTfSql comment refs scrubbed. --- dimos/memory2/db_tf.py | 66 --- dimos/memory2/db_tf_live.py | 71 --- dimos/memory2/db_tf_sql.py | 501 ------------------ dimos/memory2/module.py | 151 +----- dimos/memory2/store/base.py | 15 - dimos/memory2/store/sqlite.py | 16 +- dimos/memory2/test_db_tf.py | 333 ------------ .../jnav/components/loop_closure/eval.py | 10 +- .../gsc_pgo/scripts/post_process.py | 12 +- dimos/navigation/jnav/utils/recording_tf.py | 43 ++ dimos/protocol/tf/deformation.py | 7 +- dimos/protocol/tf/tf.py | 2 +- 12 files changed, 79 insertions(+), 1148 deletions(-) delete mode 100644 dimos/memory2/db_tf.py delete mode 100644 dimos/memory2/db_tf_live.py delete mode 100644 dimos/memory2/db_tf_sql.py delete mode 100644 dimos/memory2/test_db_tf.py create mode 100644 dimos/navigation/jnav/utils/recording_tf.py diff --git a/dimos/memory2/db_tf.py b/dimos/memory2/db_tf.py deleted file mode 100644 index d1006bb68e..0000000000 --- a/dimos/memory2/db_tf.py +++ /dev/null @@ -1,66 +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. - -""" -The read-side transform-lookup interface for memory2, plus the graph-stream payload. - -Kept dependency-light (no sqlite/numpy) so importing :class:`DbTf` or :class:`TfGraph` -costs nothing. The implementations live in sibling modules: - * :class:`dimos.memory2.db_tf_live.DbTfLive` — in-RAM buffer; default for any Store. - * :class:`dimos.memory2.db_tf_sql.DbTfSql` — sqlite graph-stream; sqlite store only. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, Protocol - -if TYPE_CHECKING: - from dimos.msgs.geometry_msgs.Transform import Transform - - -class DbTf(Protocol): - """Transform-lookup interface over a recording's transforms. Implemented by - :class:`DbTfLive` (default) and :class:`DbTfSql` (sqlite).""" - - def get( - self, - target_frame: str, - source_frame: str, - time_point: float | None = None, - time_tolerance: float | None = None, - ) -> Transform | None: ... - - def has_transforms(self) -> bool: ... - - -class TfGraph: - """A tf topology snapshot, recorded one per structure change. - - ``structure`` maps each child frame to ``{"parent": str, "static": bool}`` — - the full tf tree as of this message's timestamp. The stream of these snapshots - (the ``tf_graph`` stream) is the topology change-log that transform lookups - walk to resolve a source->target chain at any past time. Defined here (not under - ``dimos/msgs``) because it is a recording-internal payload, not a wire message; - it is stored via the pickle codec.""" - - structure: dict[str, dict[str, Any]] - msg_name = "tf2_msgs.TfGraph" - - def __init__(self, structure: dict[str, dict[str, Any]]) -> None: - # copy so later mutations of the writer's running structure don't alter an - # already-recorded snapshot - self.structure = {child: dict(entry) for child, entry in structure.items()} - - def __repr__(self) -> str: - return f"TfGraph({len(self.structure)} frames)" diff --git a/dimos/memory2/db_tf_live.py b/dimos/memory2/db_tf_live.py deleted file mode 100644 index e67ad9dfdc..0000000000 --- a/dimos/memory2/db_tf_live.py +++ /dev/null @@ -1,71 +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. - -"""In-RAM transform lookups — the default for any :class:`Store`.""" - -from __future__ import annotations - -import math -import threading -from typing import TYPE_CHECKING - -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.protocol.tf.tf import MultiTBuffer - -if TYPE_CHECKING: - from dimos.memory2.store.base import Store - from dimos.msgs.geometry_msgs.Transform import Transform - -# The transform streams the buffer loads (dynamic + latched static). -TF_STREAMS = ("tf", "tf_static") - - -class DbTfLive: - """Transform lookups by loading the ``tf``/``tf_static`` streams into an in-RAM - :class:`MultiTBuffer`. The default for any :class:`Store` — backend-agnostic (uses - only the generic stream API). Holds the full history in memory, loaded once.""" - - def __init__(self, store: Store) -> None: - self._store = store - self._lock = threading.Lock() - self._buffer: MultiTBuffer | None = None - - def _ensure_loaded(self) -> MultiTBuffer: - if self._buffer is not None: - return self._buffer - with self._lock: - if self._buffer is not None: - return self._buffer - buffer = MultiTBuffer(buffer_size=math.inf) - available = set(self._store.list_streams()) - for name in TF_STREAMS: - if name not in available: - continue - for observation in self._store.stream(name, TFMessage): - transforms = getattr(observation.data, "transforms", None) or [observation.data] - buffer.receive_transform(*transforms) - self._buffer = buffer - return buffer - - def has_transforms(self) -> bool: - return bool(self._ensure_loaded().buffers) - - def get( - self, - target_frame: str, - source_frame: str, - time_point: float | None = None, - time_tolerance: float | None = None, - ) -> Transform | None: - return self._ensure_loaded().lookup(target_frame, source_frame, time_point, time_tolerance) diff --git a/dimos/memory2/db_tf_sql.py b/dimos/memory2/db_tf_sql.py deleted file mode 100644 index 888ea50d9c..0000000000 --- a/dimos/memory2/db_tf_sql.py +++ /dev/null @@ -1,501 +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. - -"""SQLite transform lookups — the sqlite graph-stream path (indexed per-lookup queries -plus loop-closure deformation correction). Used only by the sqlite store.""" - -from __future__ import annotations - -import bisect -import math -import sqlite3 -from typing import TYPE_CHECKING, Any, cast - -import numpy as np - -from dimos.memory2.db_tf import TfGraph -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.nav_msgs.DeformationNode import DeformationNode, tf_id_for -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.protocol.tf.deformation import apply_delta, blend_se3 -from dimos.protocol.tf.tf import MultiTBuffer -from dimos.utils.logging_config import setup_logger - -if TYPE_CHECKING: - from dimos.memory2.store.base import Store - from dimos.memory2.store.sqlite import SqliteStoreConfig - from dimos.memory2.stream import Stream - -logger = setup_logger() - -GRAPH_STREAM = "tf_graph" # changes to the tf tree (with room for non-tree structures) -DEFORMATION_STREAM = "tf_deformation_nodes" # loop closure and other deformations -# 99% of the time there are less than 10 tree changes -# this cache allows us to avoid a DB query in the offline/map-load case -DEFAULT_MAX_GRAPH_CHANGES_IN_RAM = 64 -# A frame is "static" if its pose never changes; poses are compared rounded to this -# many decimals (~nanometre / nanoradian) so float noise doesn't read as motion. -POSE_EQUALITY_DECIMALS = 9 - - -class DbTfSql: - """Transform lookups over the sqlite graph-stream: a topology change-log plus - per-lookup indexed queries for the bracketing samples, with loop-closure - deformation correction. Used only by the sqlite store.""" - - def __init__( - self, - store: Store, - max_graph_changes_in_ram: int = DEFAULT_MAX_GRAPH_CHANGES_IN_RAM, - ) -> None: - self._store = store - self._max_in_ram = max_graph_changes_in_ram - self._conn: sqlite3.Connection | None = None - self._built = False - # graph cache: either the whole change-log in RAM, or None (query per lookup) - self._graph_in_ram: list[tuple[float, dict[str, Any]]] | None = None - self._graph_loaded = False - self._static_cache: dict[str, Transform] = {} - # The set of tf_ids present on the deformation stream (the edges loop closure - # corrects). Cached once so a recording with no deformation nodes pays nothing. - self._deformation_tf_ids: set[int] | None = None - self.rows_fetched = 0 - self.graph_queries = 0 - - def _connection(self) -> sqlite3.Connection: - conn = self._conn - if conn is None: - conn = _connect(cast("SqliteStoreConfig", self._store.config).path) - self._conn = conn - return conn - - def has_transforms(self) -> bool: - if "tf" not in set(self._store.list_streams()): - return False - (n_rows,) = self._connection().execute('SELECT count(*) FROM "tf"').fetchone() - return bool(n_rows) - - def _graph_stream(self) -> Stream[TfGraph]: - return self._store.stream(GRAPH_STREAM, TfGraph) - - def _graph_count(self) -> int: - if GRAPH_STREAM not in set(self._store.list_streams()): - return 0 - return self._graph_stream().count() - - def _ensure_built(self) -> None: - """First use: if the recording has tf rows but no graph stream (a recording - that predates it / wasn't written by the recorder), build the graph stream once - by replaying the tf rows, then make sure the seek index exists.""" - if self._built: - return - conn = self._connection() - (n_rows,) = conn.execute('SELECT count(*) FROM "tf"').fetchone() - if n_rows and self._graph_count() == 0: - logger.warning( - "\n========================================================================\n" - " tf graph stream MISSING for %r. Building it (one-time): tagging tf rows\n" - " with child_frame and writing the topology change-log.\n" - "========================================================================", - "tf", - ) - self._build_graph_stream() - if n_rows: - _ensure_child_index(conn) # tf table exists now - self._built = True - - def _build_graph_stream(self) -> None: - """One-time migration: decode every tf row, tag it with its child_frame, and - append a ``TfGraph`` snapshot whenever the topology changes. A frame counts as - static if its pose never varies across the whole recording.""" - # one decode pass: collect (id, ts, child, parent, pose-key) + per-child poses - rows: list[tuple[int, float, str, str]] = [] - poses_per_child: dict[str, set[tuple[float, ...]]] = {} - for obs in self._store.stream("tf", TFMessage).order_by("ts"): - for transform in getattr(obs.data, "transforms", None) or [obs.data]: - pose_key = tuple( - round(value, POSE_EQUALITY_DECIMALS) - for value in ( - transform.translation.x, - transform.translation.y, - transform.translation.z, - transform.rotation.x, - transform.rotation.y, - transform.rotation.z, - transform.rotation.w, - ) - ) - rows.append((obs.id, obs.ts, transform.child_frame_id, transform.frame_id)) - poses_per_child.setdefault(transform.child_frame_id, set()).add(pose_key) - static_frames = {child for child, poses in poses_per_child.items() if len(poses) == 1} - - # tag each tf row with its child_frame (raw UPDATE on the tf table) - conn = self._connection() - for row_id, _ts, child, _parent in rows: - conn.execute( - "UPDATE \"tf\" SET tags = json_set(tags, '$.child_frame', ?) WHERE id = ?", - (child, row_id), - ) - conn.commit() - - # build the change-log as a first-class stream: one snapshot per change - graph_stream = self._store.stream(GRAPH_STREAM, TfGraph) - structure: dict[str, dict[str, Any]] = {} - written = 0 - for _row_id, ts, child, parent in rows: - entry = {"parent": parent, "static": child in static_frames} - if structure.get(child) == entry: - continue - structure[child] = entry - graph_stream.append(TfGraph(structure), ts=ts) - written += 1 - logger.warning("tf graph built: %d topology changes for %r.", written, "tf") - - def _load_graph_if_small(self) -> None: - if self._graph_loaded: - return - if self._graph_count() < self._max_in_ram: - # Sort by (ts, id): several topology changes can share a timestamp (e.g. - # every static frame latched at t0), and the LAST-inserted of those is the - # complete snapshot — a plain ts sort leaves same-ts order undefined. - snapshots = sorted( - ((obs.ts, obs.id, obs.data.structure) for obs in self._graph_stream()), - key=lambda row: (row[0], row[1]), - ) - self._graph_in_ram = [(ts, structure) for ts, _id, structure in snapshots] - else: - self._graph_in_ram = None # too many -> query per lookup - self._graph_loaded = True - - def _graph_at(self, query_time: float) -> dict[str, Any] | None: - if self._graph_in_ram is not None: - # in-RAM: binary search the latest change at-or-before query_time - stamps = [ts for ts, _ in self._graph_in_ram] - index = bisect.bisect_right(stamps, query_time) - 1 - if index < 0: - return self._graph_in_ram[0][1] # before first -> earliest - return self._graph_in_ram[index][1] - # fallback: one indexed query for the latest snapshot at or before query_time. - # Tie-break by id (DESC) so same-timestamp changes resolve to the complete one. - self.graph_queries += 1 - conn = self._connection() - graph, blob = f'"{GRAPH_STREAM}"', f'"{GRAPH_STREAM}_blob"' - row = conn.execute( - f"SELECT x.data FROM {graph} g JOIN {blob} x ON x.id = g.id " - "WHERE g.ts <= ? ORDER BY g.ts DESC, g.id DESC LIMIT 1", - (query_time,), - ).fetchone() - if row is None: # before the first snapshot -> earliest - row = conn.execute( - f"SELECT x.data FROM {graph} g JOIN {blob} x ON x.id = g.id " - "ORDER BY g.ts ASC, g.id ASC LIMIT 1" - ).fetchone() - if row is None: - return None - codec = cast("Any", self._store.stream(GRAPH_STREAM, TfGraph)._source).codec - return cast("TfGraph", codec.decode(row[0])).structure - - def _chain_frames(self, graph: dict[str, Any], source: str, target: str) -> list[str] | None: - def to_root(frame: str) -> list[str]: - path = [frame] - seen = {frame} - while ( - frame in graph - and graph[frame].get("parent") in graph - and graph[frame]["parent"] not in seen - ): - frame = graph[frame]["parent"] - path.append(frame) - seen.add(frame) - # include a final parent that is itself a root (not a key in graph) - if frame in graph and graph[frame].get("parent") and graph[frame]["parent"] not in seen: - path.append(graph[frame]["parent"]) - return path - - source_path = to_root(source) - target_path = to_root(target) - common = next((f for f in source_path if f in set(target_path)), None) - if common is None: - return None # disjoint graph: no transform between them - frames = source_path[: source_path.index(common) + 1] - frames += target_path[: target_path.index(common)] - return frames - - def _decode_blob(self, data: bytes, frame: str) -> Transform: - # The blob is the codec-encoded message; pick the transform for `frame` - # (rows normally hold one; legacy rows may pack several). - codec = cast("Any", self._store.stream("tf", TFMessage)._source).codec - message = codec.decode(data) - transforms = getattr(message, "transforms", None) or [message] - for transform in transforms: - if transform.child_frame_id == frame: - return cast("Transform", transform) - return cast("Transform", transforms[0]) - - def _fetch_rows( - self, dynamic: list[str], static: list[str], query_time: float - ) -> dict[tuple[str, str], bytes]: - """ONE query: for each dynamic frame the bracketing rows ('lo' = latest at - or before query_time, 'hi' = earliest at or after), and for each (uncached) - static frame its latest row ('st') — all joined to the blob data. Keyed by - (frame, kind) -> blob bytes.""" - cf = "json_extract(tags, '$.child_frame')" - # One UNION of per-frame, index-served LIMIT-1 subqueries: each is a direct - # (child_frame, ts) range seek — far cheaper than a window-function scan, and - # still a single round-trip. - parts: list[str] = [] - params: list[Any] = [] - - def pick(frame: str, kind: str, where_ts: str, order: str) -> None: - parts.append( - f"SELECT ? AS cf, ? AS kind, " - f'(SELECT id FROM "tf" WHERE {cf} = ?{where_ts} ORDER BY ts {order} LIMIT 1) AS id' - ) - params.extend([frame, kind, frame]) - - for frame in dynamic: - pick(frame, "lo", " AND ts <= ?", "DESC") - params.append(query_time) - pick(frame, "hi", " AND ts >= ?", "ASC") - params.append(query_time) - for frame in static: - pick(frame, "st", "", "DESC") - if not parts: - return {} - union = " UNION ALL ".join(parts) - sql = f'SELECT t.cf, t.kind, b.data FROM ({union}) t JOIN "tf_blob" b ON b.id = t.id' - rows: dict[tuple[str, str], bytes] = {} - for cf_val, kind, data in self._connection().execute(sql, params): - rows[(cf_val, kind)] = data - self.rows_fetched += 1 - return rows - - def _load_deformation_ids(self) -> set[int]: - """The distinct tf_ids present on the deformation stream, cached once. Empty - when there's no such stream — so the correction path is skipped at zero cost - for recordings without loop closure.""" - if self._deformation_tf_ids is not None: - return self._deformation_tf_ids - ids: set[int] = set() - if DEFORMATION_STREAM in set(self._store.list_streams()): - for (tf_id,) in self._connection().execute( - f"SELECT DISTINCT json_extract(tags, '$.tf_id') FROM \"{DEFORMATION_STREAM}\"" - ): - if tf_id is not None: - ids.add(int(tf_id)) - self._deformation_tf_ids = ids - return ids - - def _node_pose_matrix(self, order: str, tf_id: int, node_id: str) -> np.ndarray | None: - """The 4x4 pose of one keyframe's first (``ASC``) or latest (``DESC``) recorded - version. Versions of a node share the keyframe ts, so they're ordered by row id - (insertion order), not ts.""" - stream, blob = f'"{DEFORMATION_STREAM}"', f'"{DEFORMATION_STREAM}_blob"' - row = ( - self._connection() - .execute( - f"SELECT b.data FROM {stream} s JOIN {blob} b ON b.id = s.id " - f"WHERE json_extract(s.tags, '$.tf_id') = ? AND json_extract(s.tags, '$.id') = ? " - f"ORDER BY s.id {order} LIMIT 1", - (str(tf_id), node_id), - ) - .fetchone() - ) - if row is None: - return None - codec = cast("Any", self._store.stream(DEFORMATION_STREAM, DeformationNode)._source).codec - node = cast("DeformationNode", codec.decode(row[0])) - return Transform.from_pose(node.pose.frame_id, node.pose).to_matrix() - - def _node_delta(self, tf_id: int, node_id: str) -> np.ndarray | None: - """How far the optimizer has moved a keyframe since it was first recorded: - ``current ∘ inv(original)`` — the SE(3) correction this node contributes.""" - original = self._node_pose_matrix("ASC", tf_id, node_id) - current = self._node_pose_matrix("DESC", tf_id, node_id) - if original is None or current is None: - return None - return current @ np.linalg.inv(original) - - def _edge_delta(self, tf_id: int, query_time: float) -> np.ndarray | None: - """The blended correction for an edge at ``query_time``: take the keyframes - bracketing the time (latest at-or-before, earliest at-or-after), each node's - ``current ∘ inv(original)`` delta, and linear-blend-skin between them.""" - stream = f'"{DEFORMATION_STREAM}"' - id_tag = "json_extract(tags, '$.id')" - tf_tag = "json_extract(tags, '$.tf_id')" - conn = self._connection() - lo = conn.execute( - f"SELECT {id_tag}, ts FROM {stream} WHERE {tf_tag} = ? AND ts <= ? " - "ORDER BY ts DESC, id DESC LIMIT 1", - (str(tf_id), query_time), - ).fetchone() - hi = conn.execute( - f"SELECT {id_tag}, ts FROM {stream} WHERE {tf_tag} = ? AND ts >= ? " - "ORDER BY ts ASC, id ASC LIMIT 1", - (str(tf_id), query_time), - ).fetchone() - samples: list[tuple[float, np.ndarray]] = [] - seen_ids: set[str] = set() - for node_id, ts in (row for row in (lo, hi) if row is not None): - if node_id in seen_ids: - continue - seen_ids.add(node_id) - delta = self._node_delta(tf_id, node_id) - if delta is not None: - samples.append((ts, delta)) - if not samples: - return None - if len(samples) == 1: - return samples[0][1] - (ts_lo, mat_lo), (ts_hi, mat_hi) = sorted(samples, key=lambda item: item[0]) - weight = 0.0 if ts_hi == ts_lo else (query_time - ts_lo) / (ts_hi - ts_lo) - return blend_se3(mat_lo, mat_hi, weight) - - def _edge_corrections( - self, graph: dict[str, Any], edges: list[str], query_time: float - ) -> dict[str, np.ndarray]: - """For each edge whose tf_id matches the deformation stream, its blended SE(3) - correction (applied on the parent side of the edge). Empty when nothing matches.""" - tf_ids = self._load_deformation_ids() - if not tf_ids: - return {} - corrections: dict[str, np.ndarray] = {} - for frame in edges: - tf_id = tf_id_for(graph[frame]["parent"], frame) - if tf_id not in tf_ids: - continue - delta = self._edge_delta(tf_id, query_time) - if delta is not None: - corrections[frame] = delta - return corrections - - def get( - self, - target_frame: str, - source_frame: str, - time_point: float | None = None, - time_tolerance: float | None = None, - ) -> Transform | None: - """Transform that maps a point in ``source_frame`` into ``target_frame``, - or ``None`` if no chain connects them at the requested time.""" - self._ensure_built() - self._load_graph_if_small() - query_time = time_point if time_point is not None else 0.0 - graph = self._graph_at(query_time) # 0 queries when the graph is in RAM - if graph is None: - return None - frames = self._chain_frames(graph, source_frame, target_frame) - if frames is None: - return None - - edges = [f for f in frames if f in graph] # roots have no incoming edge - dynamic = [f for f in edges if not graph[f].get("static")] - static = [f for f in edges if graph[f].get("static")] - uncached_static = [f for f in static if f not in self._static_cache] - - rows = self._fetch_rows(dynamic, uncached_static, query_time) # ONE detail query - # Per-edge loop-closure corrections (empty unless a deformation stream matches). - corrections = self._edge_corrections(graph, edges, query_time) - - buffer = MultiTBuffer(buffer_size=math.inf) - for frame in static: - transform = self._static_cache.get(frame) - if transform is None: - data = rows.get((frame, "st")) - if data is None: - return None - transform = self._decode_blob(data, frame) - self._static_cache[frame] = transform - # restamp the constant to query_time so the buffer's tolerance never - # rejects a static that was recorded long ago (latched once). - transform = _restamp(transform, query_time) - if frame in corrections: - transform = apply_delta(corrections[frame], transform, query_time) - buffer.receive_transform(transform) - for frame in dynamic: - lo = rows.get((frame, "lo")) - hi = rows.get((frame, "hi")) - chosen = lo if lo is not None else hi - if chosen is None: - return None - if frame in corrections: - # Resolve the raw edge at query_time, then deform it. The correction - # already carries the time blend, so a single corrected sample suffices. - raw = self._interpolate_dynamic( - frame, graph[frame]["parent"], lo, hi, query_time, time_tolerance - ) - if raw is None: - return None - buffer.receive_transform(apply_delta(corrections[frame], raw, query_time)) - continue - buffer.receive_transform(self._decode_blob(chosen, frame)) - other = hi if hi is not None else lo - if other is not None and other is not chosen: - buffer.receive_transform(self._decode_blob(other, frame)) - return buffer.lookup(target_frame, source_frame, time_point, time_tolerance) - - def _interpolate_dynamic( - self, - frame: str, - parent: str, - lo: bytes | None, - hi: bytes | None, - query_time: float, - time_tolerance: float | None, - ) -> Transform | None: - """The raw ``parent <- frame`` transform at ``query_time``, interpolated from - the bracketing rows (its own small buffer so the chain buffer only ever sees - the corrected result).""" - edge_buffer = MultiTBuffer(buffer_size=math.inf) - chosen = lo if lo is not None else hi - if chosen is None: - return None - edge_buffer.receive_transform(self._decode_blob(chosen, frame)) - other = hi if hi is not None else lo - if other is not None and other is not chosen: - edge_buffer.receive_transform(self._decode_blob(other, frame)) - return edge_buffer.lookup(parent, frame, query_time, time_tolerance) - - -def _restamp(transform: Transform, ts: float) -> Transform: - return Transform( - translation=transform.translation, - rotation=transform.rotation, - frame_id=transform.frame_id, - child_frame_id=transform.child_frame_id, - ts=ts, - ) - - -def _connect(db_path: str) -> sqlite3.Connection: - conn = sqlite3.connect(db_path, check_same_thread=False) - conn.execute("PRAGMA busy_timeout=5000") - return conn - - -def _ensure_child_index(conn: sqlite3.Connection) -> None: - """Index the child_frame json tag on the tf rows so per-frame time queries - seek. The live recorder gets this for free (the store auto-indexes tag keys on - tagged appends); this is for migrated recordings and the read side. Requires - the tf table to exist.""" - # Composite (child_frame, ts) so a per-frame "latest at/before T" is a direct - # index range seek, not a scan+sort. Index names share SQLite's global namespace - # with tables, so the name is double-underscore-namespaced to keep it clear of any - # real stream/table name (no stream would contain "__dbtf_"). - index_name = "tf__dbtf_child_ts_idx" - conn.execute( - f'CREATE INDEX IF NOT EXISTS "{index_name}" ' - "ON \"tf\"(json_extract(tags, '$.child_frame'), ts)" - ) - conn.commit() diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index c2767b225a..e16c1527e7 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -30,8 +30,6 @@ from dimos.constants import DIMOS_PROJECT_ROOT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig -from dimos.core.stream import In -from dimos.memory2.db_tf import TfGraph from dimos.memory2.embed import EmbedImages from dimos.memory2.store.null import NullStore from dimos.memory2.store.sqlite import SqliteStore @@ -40,7 +38,6 @@ from dimos.memory2.type.observation import EmbeddedObservation, Observation from dimos.models.embedding.base import EmbeddingModel from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.DeformationNode import DeformationNode from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.utils.data import backup_file @@ -49,7 +46,7 @@ if TYPE_CHECKING: from reactivex.abc import DisposableBase - from dimos.core.stream import Out + from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs.Pose import Pose logger = setup_logger() @@ -199,10 +196,6 @@ def store(self) -> SqliteStore: class SemanticSearchConfig(MemoryModuleConfig): embedding_model: type[EmbeddingModel] | None = None - # The pose of a match is resolved at QUERY time via store.tf (so it reflects loop - # closure), looking up root_frame <- the observation's frame_id at its timestamp. - root_frame: str = "world" - tf_tolerance: float = 0.5 class SemanticSearch(MemoryModule): @@ -253,21 +246,8 @@ def _similarity(obs: Observation[Any]) -> float: return cast("EmbeddedObservation[Any]", obs).similarity or 0.0 best = results.transform(peaks(key=_similarity, distance=1.0)).last() - - # Resolve the match's world pose at QUERY time via tf, so it reflects any loop - # closure since the observation was recorded (the baked pose_stamped is stale). - frame_id = getattr(best.data, "frame_id", None) or self.config.root_frame - transform = self.store.tf.get( - self.config.root_frame, frame_id, best.ts, self.config.tf_tolerance - ) - if transform is not None: - return transform.to_pose(ts=best.ts) - # No tf chain at that time — fall back to the pose baked at record time. if best.pose_stamped is None: - raise LookupError("No tf and no baked pose on best search result") - logger.warning( - "SemanticSearch: no tf for %s @ %s; using stale baked pose", frame_id, best.ts - ) + raise LookupError("No pose on best search result") return best.pose_stamped @@ -337,13 +317,7 @@ async def _lidar_pose(self, msg): config: RecorderConfig - tf_static: In[TFMessage] - tf_deformation_nodes: In[DeformationNode] - _pose_setters: dict[str, Any] = {} - # Per-stream count of frames lost to the dispatcher's LATEST coalescing - # (sink slower than input). Populated lazily as drops happen. - _dropped_frames: dict[str, int] = {} @rpc def start(self) -> None: @@ -356,7 +330,6 @@ def start(self) -> None: return self._pose_setters = self._collect_pose_setters() - self._dropped_frames = {} # TODO: store reset API/logic is not implemented yet. This module # shouldn't need to know about files (SqliteStore specific), and @@ -385,10 +358,6 @@ def start(self) -> None: return for name, port in self.inputs.items(): - if name == "tf_static": - continue # folded into the "tf" stream + graph by _record_tf - if name == "tf_deformation_nodes": - continue # recorded by _record_tf into its own stream (carries its own pose) stream_name = self.config.stream_remapping.get(name, name) stream: Stream[Any] = self.store.stream(stream_name, port.type) self._port_to_stream(name, port, stream) @@ -422,46 +391,17 @@ async def on_msg(msg: Any) -> None: ) stream.append(msg, ts=ts, pose=pose) - self.process_observable( - input_topic.pure_observable(), on_msg, on_drop=lambda: self._on_frame_dropped(name) - ) - - def _on_frame_dropped(self, name: str) -> None: - """A frame for *name* was dropped because the sink couldn't keep up with - the input rate (dispatcher LATEST coalescing). Count it and warn — once, - then on each power-of-ten — so silent data loss is visible without - flooding the log.""" - count = self._dropped_frames.get(name, 0) + 1 - self._dropped_frames[name] = count - if count == 1 or count % 1000 == 0: - logger.warning( - "[%s] Recorder dropped %d frame(s) — sink slower than input; recording is lossy", - name, - count, - ) + self.process_observable(input_topic.pure_observable(), on_msg) def _prepare_streams(self) -> None: - """On APPEND, clear only the streams this recorder will actually (re)write — - i.e. those with a connected source — so a re-run overwrites them cleanly. - Streams whose source is unconnected are left untouched: recording into an - existing db never drops data it won't repopulate (e.g. pcap_to_db, which has - no tf source, must not wipe an existing tf / tf_deformation_nodes).""" + """On APPEND, drop the streams this recorder is about to (re)write — the + remapped In-port streams plus ``tf`` — so a re-run replaces them instead + of duplicating, while leaving any other streams in the db untouched.""" if self.config.on_existing is not OnExisting.APPEND: return - targets: set[str] = set() - for name, port in self.inputs.items(): - if name in ("tf_static", "tf_deformation_nodes"): - continue # tf-family, handled below only when their source is connected - if port.transport is None: - continue # unconnected -> nothing to write, leave any existing data - targets.add(self.config.stream_remapping.get(name, name)) + targets = {self.config.stream_remapping.get(name, name) for name in self.inputs} if self.config.record_tf: - topic = getattr(self.tf.config, "topic", None) - pubsub = getattr(self.tf, "pubsub", None) - if (topic and pubsub is not None) or self.tf_static.transport is not None: - targets.update(("tf", "tf_graph")) - if self.tf_deformation_nodes.transport is not None: - targets.add("tf_deformation_nodes") + targets.add("tf") for stream in targets.intersection(self.store.list_streams()): self.store.delete_stream(stream) @@ -492,77 +432,20 @@ def _collect_pose_setters(self) -> dict[str, PoseSetter]: return setters def _record_tf(self) -> None: - """Record tf into the "tf" stream + the topology change-log ("tf_graph"). - - Two inputs, both folded into one tf stream: the live (dynamic) tf via the - module's tf interface, and the optional ``tf_static`` In-port stream (a - future system publishes latched mount/extrinsic transforms there). Frames - from tf_static are flagged static in the graph; latched statics resolve - for all time, dynamic frames bracket+interpolate.""" - tf_stream = self.store.stream("tf", TFMessage) - graph_stream = self.store.stream("tf_graph", TfGraph) - # Running tf topology; a TfGraph snapshot is appended whenever it changes, so - # the "tf_graph" stream is the topology change-log transform lookups walk. - structure: dict[str, dict[str, Any]] = {} - - def record(msg: TFMessage, is_static: bool) -> None: - try: - for transform in msg.transforms: - tf_stream.append( - TFMessage(transform), - ts=transform.ts, - pose=None, - tags={"child_frame": transform.child_frame_id}, - ) - entry = {"parent": transform.frame_id, "static": is_static} - if structure.get(transform.child_frame_id) != entry: - structure[transform.child_frame_id] = entry - graph_stream.append(TfGraph(structure), ts=transform.ts) - except sqlite3.ProgrammingError: - # A late callback raced teardown and hit the closed store. - pass - - # static tf: the tf_static In-port stream. Only subscribe when something is - # wired to it (its transport is set on connect) — unconnected today. - if self.tf_static.transport is not None: - - async def on_static(msg: TFMessage) -> None: - record(msg, is_static=True) - - self.process_observable(self.tf_static.pure_observable(), on_static) - - self._record_deformation_nodes() - - # dynamic tf: the module's live tf interface + """Record the live tf stream under "tf" (no-op without a pubsub tf).""" topic = getattr(self.tf.config, "topic", None) pubsub = getattr(self.tf, "pubsub", None) if not topic or pubsub is None: - logger.warning("Recorder: no pubsub tf available — recording static tf only") + logger.warning("Recorder: no pubsub tf available — not recording tf") return - self.register_disposable( - Disposable(pubsub.subscribe(topic, lambda msg, _t: record(msg, False))) - ) - - def _record_deformation_nodes(self) -> None: - """Record the optional ``tf_deformation_nodes`` In-port into its own stream. - - Each DeformationNode is one pose-graph keyframe; the backend re-publishes a - node (same ``id``) when the optimizer moves it, so rows accumulate and a query - takes the latest per ``id``. Tagged by ``tf_id`` (which transform edge) and - ``id`` so lookups can filter by edge and dedup by node. Unconnected = no-op.""" - if self.tf_deformation_nodes.transport is None: - return - stream = self.store.stream("tf_deformation_nodes", DeformationNode) + tf_stream = self.store.stream("tf", TFMessage) - async def on_node(node: DeformationNode) -> None: + def on_tf(msg: TFMessage, _topic: Any) -> None: try: - stream.append( - node, - ts=node.pose.ts, - pose=None, - tags={"tf_id": str(node.tf_id), "id": str(node.id)}, - ) + for transform in msg.transforms: + tf_stream.append(TFMessage(transform), ts=transform.ts, pose=None) except sqlite3.ProgrammingError: - pass # late callback raced teardown and hit the closed store + # A late LCM callback raced teardown and hit the closed store. + pass - self.process_observable(self.tf_deformation_nodes.pure_observable(), on_node) + self.register_disposable(Disposable(pubsub.subscribe(topic, on_tf))) diff --git a/dimos/memory2/store/base.py b/dimos/memory2/store/base.py index 7818950f7a..7a7162a6d1 100644 --- a/dimos/memory2/store/base.py +++ b/dimos/memory2/store/base.py @@ -29,7 +29,6 @@ from dimos.protocol.service.spec import BaseConfig, Configurable if TYPE_CHECKING: - from dimos.memory2.db_tf import DbTf from dimos.memory2.replay import Replay T = TypeVar("T") @@ -106,26 +105,12 @@ def __init__(self, **kwargs: Any) -> None: Configurable.__init__(self, **kwargs) CompositeResource.__init__(self) self._streams: dict[str, Stream[Any]] = {} - self._tf: DbTf | None = None @property def streams(self) -> StreamAccessor[Stream[Any]]: """Attribute-style access to streams: ``store.streams.name``.""" return StreamAccessor(self) - @property - def tf(self) -> DbTf: - """Transform lookups over the recording's ``tf``/``tf_static`` streams. - - ``store.tf.get(target_frame, source_frame, ts)`` composes multi-hop chains - (e.g. ``world -> ... -> mid360_link``) from the recorded transforms. - """ - if self._tf is None: - from dimos.memory2.db_tf_live import DbTfLive - - self._tf = DbTfLive(self) - return self._tf - def replay( self, *, diff --git a/dimos/memory2/store/sqlite.py b/dimos/memory2/store/sqlite.py index fe7a6ff393..229961a126 100644 --- a/dimos/memory2/store/sqlite.py +++ b/dimos/memory2/store/sqlite.py @@ -16,7 +16,7 @@ import os import sqlite3 -from typing import TYPE_CHECKING, Annotated, Any +from typing import Annotated, Any from pydantic import BeforeValidator @@ -32,9 +32,6 @@ from dimos.memory2.vectorstore.base import VectorStore from dimos.memory2.vectorstore.sqlite import SqliteVectorStore -if TYPE_CHECKING: - from dimos.memory2.db_tf import DbTf - class SqliteStoreConfig(StoreConfig): """Config for SQLite-backed store.""" @@ -64,17 +61,6 @@ def __init__(self, **kwargs: Any) -> None: self._registry_conn = self._open_connection() self._registry = RegistryStore(conn=self._registry_conn) - @property - def tf(self) -> DbTf: - """Transform lookups over the sqlite graph stream (see :class:`DbTfSql`): - indexed per-lookup queries + loop-closure deformation correction, rather than - the base :class:`DbTfLive` in-RAM buffer.""" - if self._tf is None: - from dimos.memory2.db_tf_sql import DbTfSql - - self._tf = DbTfSql(self) - return self._tf - def _open_connection(self) -> sqlite3.Connection: """Open a new WAL-mode connection with sqlite-vec loaded.""" disposable, connection = open_disposable_sqlite_connection(self.config.path) diff --git a/dimos/memory2/test_db_tf.py b/dimos/memory2/test_db_tf.py deleted file mode 100644 index 420985c953..0000000000 --- a/dimos/memory2/test_db_tf.py +++ /dev/null @@ -1,333 +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. - -"""Integration tests for DbTf (graph-stream transform lookups). Data is written -one-transform-per-row + child_frame tag + topology change-log, exactly as the live -recorder does. These cover the behaviours that matter end-to-end: interpolation -against a full-load buffer, a latched static, time-varying topology (relocalization -re-parents a frame), and a disjoint multi-robot graph.""" - -from __future__ import annotations - -import math -from pathlib import Path -import platform - -import pytest - -from dimos.memory2.db_tf_sql import DbTfSql -from dimos.memory2.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -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.DeformationNode import DeformationNode, tf_id_for -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.protocol.tf.tf import MultiTBuffer - -# sqlite-vec fails to load on Linux ARM (32-bit binary in the aarch64 wheel) and -# on macOS in CI; SqliteStore loads it on init, so skip like memory2/conftest.py. -_SKIP_SQLITE_VEC = platform.machine() == "aarch64" or platform.system() == "Darwin" -pytestmark = pytest.mark.skipif(_SKIP_SQLITE_VEC, reason="sqlite-vec extension not loadable here") - -_T0 = 1000.0 -_DYN_RATE = 30.0 -_DURATION = 10.0 - - -def _yaw(theta: float) -> Quaternion: - return Quaternion(0.0, 0.0, math.sin(theta / 2.0), math.cos(theta / 2.0)) - - -def _static(parent: str, child: str, xyz: tuple[float, float, float], ts: float) -> Transform: - return Transform( - translation=Vector3(*xyz), - rotation=Quaternion(0, 0, 0, 1), - frame_id=parent, - child_frame_id=child, - ts=ts, - ) - - -def _append(store: SqliteStore, transform: Transform) -> None: - """Record one transform as the recorder does: one row tagged with its child_frame. - The graph stream is built lazily by DbTf on first lookup (static-ness inferred - from whether a frame's pose ever changes), so tests exercise that migration.""" - store.stream("tf", TFMessage).append( - TFMessage(transform), - ts=transform.ts, - pose=None, - tags={"child_frame": transform.child_frame_id}, - ) - - -def _ref(transforms: list[Transform]) -> MultiTBuffer: - buffer = MultiTBuffer(buffer_size=math.inf) - buffer.receive_transform(*transforms) - return buffer - - -def _diff(a: Transform, b: Transform) -> float: - return ( - abs(a.translation.x - b.translation.x) - + abs(a.translation.y - b.translation.y) - + abs(a.translation.z - b.translation.z) - + abs(a.rotation.x - b.rotation.x) - + abs(a.rotation.y - b.rotation.y) - + abs(a.rotation.z - b.rotation.z) - + abs(a.rotation.w - b.rotation.w) - ) - - -def _record_single_robot(path: Path, *, static_repeat: bool) -> list[Transform]: - """world->map->odom->base_link->sensor. Statics emitted once (latched) unless - static_repeat, in which case they're re-emitted each second.""" - store = SqliteStore(path=str(path)) - written: list[Transform] = [] - - statics = [ - ("world", "map", (0.0, 0.0, 0.0)), - ("map", "odom", (0.0, 0.0, 0.0)), - ("base_link", "sensor", (0.0, 0.0, 0.3)), - ] - static_times = [_T0 + j for j in range(int(_DURATION))] if static_repeat else [_T0] - for ts in static_times: - for parent, child, xyz in statics: - transform = _static(parent, child, xyz, ts) - _append(store, transform) - written.append(transform) - - for i in range(int(_DURATION * _DYN_RATE)): - ts = _T0 + i / _DYN_RATE - transform = Transform( - translation=Vector3(0.5 * i / _DYN_RATE, 0.1 * i / _DYN_RATE, 0.0), - rotation=_yaw(0.02 * i), - frame_id="odom", - child_frame_id="base_link", - ts=ts, - ) - _append(store, transform) - written.append(transform) - - store.stop() - return written - - -def test_interpolates_and_matches_full_load(tmp_path: Path) -> None: - """At off-sample (interpolated) query times, DbTf matches a naive full-load - buffer within tolerance — proving the chain compose + per-edge interpolation.""" - transforms = _record_single_robot(tmp_path / "r.db", static_repeat=True) - reference = _ref(transforms) - store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) - db = DbTfSql(store) - compared = 0 - for k in range(25): - q = _T0 + 0.013 + k * 0.317 - want = reference.lookup("world", "sensor", q, 0.5) - got = db.get("world", "sensor", q, 0.5) - assert (want is None) == (got is None), f"None mismatch at {q}" - if want is not None and got is not None: - assert _diff(want, got) < 1e-6, f"diff at {q}: {_diff(want, got)}" - compared += 1 - assert compared >= 20 - store.stop() - - -def test_latched_static_resolves(tmp_path: Path) -> None: - """A static recorded once at the very start still resolves at a much later time - (no bracket, no tolerance) — the case a plain time-bracket would drop.""" - _record_single_robot(tmp_path / "r.db", static_repeat=False) - store = SqliteStore(path=str(tmp_path / "r.db"), must_exist=True) - db = DbTfSql(store) - assert db.get("world", "sensor", _T0 + 9.5, 0.5) is not None # ~9.5s after statics - store.stop() - - -def test_reparent_midrun_uses_graph_as_of_query_time(tmp_path: Path) -> None: - """world->map(+10) and map->odom(+100) are non-identity statics; at t0+5 a - relocalization re-parents base_link from odom to map. An early lookup composes - the odom branch (+100), a late one does not (+10) — wrong topology is off by - ~100. Exercises time-varying / multi-robot topology.""" - path = tmp_path / "reparent.db" - store = SqliteStore(path=str(path)) - _append(store, _static("world", "map", (10.0, 0, 0), _T0)) - _append(store, _static("map", "odom", (100.0, 0, 0), _T0)) - - era1: list[Transform] = [] - for i in range(150): # [t0, t0+5) - ts = _T0 + i / _DYN_RATE - transform = Transform( - translation=Vector3(0.5 * i / _DYN_RATE, 0, 0), - rotation=_yaw(0.01 * i), - frame_id="odom", - child_frame_id="base_link", - ts=ts, - ) - _append(store, transform) - era1.append(transform) - switch = _T0 + 5.0 - era2: list[Transform] = [] - for i in range(150): # [t0+5, t0+10): base_link now hangs off map directly - ts = switch + i / _DYN_RATE - transform = Transform( - translation=Vector3(7.0 + 0.3 * i / _DYN_RATE, 0, 0), - rotation=_yaw(0.01 * i), - frame_id="map", - child_frame_id="base_link", - ts=ts, - ) - _append(store, transform) - era2.append(transform) - store.stop() - - store = SqliteStore(path=str(path), must_exist=True) - db = DbTfSql(store) - - q1 = _T0 + 2.013 - want1 = _ref( - [ - _static("world", "map", (10.0, 0, 0), q1), - _static("map", "odom", (100.0, 0, 0), q1), - *era1, - ] - ).lookup("world", "base_link", q1, 0.5) - got1 = db.get("world", "base_link", q1, 0.5) - assert want1 is not None and got1 is not None - assert _diff(want1, got1) < 1e-6 - assert got1.translation.x > 100.0 # odom branch present - - q2 = switch + 2.013 - want2 = _ref([_static("world", "map", (10.0, 0, 0), q2), *era2]).lookup( - "world", "base_link", q2, 0.5 - ) - got2 = db.get("world", "base_link", q2, 0.5) - assert want2 is not None and got2 is not None - assert _diff(want2, got2) < 1e-6 - assert got2.translation.x < 20.0 # odom branch gone - store.stop() - - -def _append_deformation( - store: SqliteStore, tf_id: int, node_id: int, ts: float, xyz: tuple[float, float, float] -) -> None: - """Record one DeformationNode version, exactly as the recorder does: tagged by - tf_id (the edge) and id (the keyframe). Re-call with the same id to add a moved - version (the optimizer relocating that keyframe).""" - node = DeformationNode( - id=node_id, - tf_id=tf_id, - pose=PoseStamped(ts=ts, frame_id="map", position=list(xyz), orientation=[0, 0, 0, 1]), - ) - store.stream("tf_deformation_nodes", DeformationNode).append( - node, ts=ts, pose=None, tags={"tf_id": str(tf_id), "id": str(node_id)} - ) - - -def _record_corrected_edge(path: Path) -> None: - """A map->odom edge (raw = identity) with base_link hanging off odom.""" - store = SqliteStore(path=str(path)) - for i in range(20): - ts = _T0 + i / _DYN_RATE - _append(store, _static("map", "odom", (0.0, 0.0, 0.0), ts)) - _append( - store, - Transform( - translation=Vector3(i * 0.1, 0, 0), - rotation=_yaw(0), - frame_id="odom", - child_frame_id="base_link", - ts=ts, - ), - ) - store.stop() - - -def test_loop_closure_deformation_corrects_matched_edge(tmp_path: Path) -> None: - """A deformation node whose tf_id = hash(map|odom) deforms the raw map<-odom edge - by current ∘ inv(original); an edge with no matching tf_id is untouched.""" - path = tmp_path / "lc.db" - _record_corrected_edge(path) - store = SqliteStore(path=str(path)) - tf_id = tf_id_for("map", "odom") - # one keyframe at t0+0.3: the optimizer later moved it from origin to (1, 0, 0) - _append_deformation(store, tf_id, 11, _T0 + 0.3, (0.0, 0.0, 0.0)) # original - _append_deformation(store, tf_id, 11, _T0 + 0.3, (1.0, 0.0, 0.0)) # current - store.stop() - - store = SqliteStore(path=str(path), must_exist=True) - db = DbTfSql(store) - got = db.get("map", "odom", _T0 + 0.3, 0.5) - assert got is not None and abs(got.translation.x - 1.0) < 1e-6 # raw identity + delta - base = db.get("odom", "base_link", _T0 + 0.3, 0.5) # unmatched edge, unchanged - assert base is not None and abs(base.translation.x - 0.9) < 0.2 - store.stop() - - -def test_loop_closure_deformation_blends_between_keyframes(tmp_path: Path) -> None: - """Two bracketing keyframes (delta 0 at t0, delta 2 at t0+? ) blend at the midpoint - to delta 1 — linear blend skinning across the trajectory.""" - path = tmp_path / "lc_blend.db" - store = SqliteStore(path=str(path)) - for i in range(int(_DURATION * _DYN_RATE)): # map->odom identity across [t0, t0+10) - ts = _T0 + i / _DYN_RATE - _append(store, _static("map", "odom", (0.0, 0.0, 0.0), ts)) - tf_id = tf_id_for("map", "odom") - _append_deformation(store, tf_id, 1, _T0, (0.0, 0.0, 0.0)) # kf A: original - _append_deformation(store, tf_id, 1, _T0, (0.0, 0.0, 0.0)) # kf A: unmoved - _append_deformation(store, tf_id, 2, _T0 + 8.0, (0.0, 0.0, 0.0)) # kf B: original - _append_deformation(store, tf_id, 2, _T0 + 8.0, (2.0, 0.0, 0.0)) # kf B: moved +2 - store.stop() - - store = SqliteStore(path=str(path), must_exist=True) - db = DbTfSql(store) - got = db.get("map", "odom", _T0 + 4.0, 0.5) # midpoint of [t0, t0+8] - assert got is not None and abs(got.translation.x - 1.0) < 1e-6 - store.stop() - - -def test_disjoint_multirobot_returns_none(tmp_path: Path) -> None: - """Two unconnected components (two robots, no shared frame) → a cross-component - query is None, an in-component query resolves.""" - path = tmp_path / "two.db" - store = SqliteStore(path=str(path)) - for i in range(20): - ts = _T0 + i / _DYN_RATE - _append( - store, - Transform( - translation=Vector3(i * 0.1, 0, 0), - rotation=_yaw(0), - frame_id="worldA", - child_frame_id="baseA", - ts=ts, - ), - ) - _append( - store, - Transform( - translation=Vector3(0, i * 0.1, 0), - rotation=_yaw(0), - frame_id="worldB", - child_frame_id="baseB", - ts=ts, - ), - ) - store.stop() - - store = SqliteStore(path=str(path), must_exist=True) - db = DbTfSql(store) - q = _T0 + 0.3 - assert db.get("baseA", "worldA", q, 0.5) is not None # same component - assert db.get("baseB", "baseA", q, 0.5) is None # different components - store.stop() diff --git a/dimos/navigation/jnav/components/loop_closure/eval.py b/dimos/navigation/jnav/components/loop_closure/eval.py index f694378629..d3ea0e7077 100644 --- a/dimos/navigation/jnav/components/loop_closure/eval.py +++ b/dimos/navigation/jnav/components/loop_closure/eval.py @@ -101,6 +101,7 @@ store, stream_count, ) +from dimos.navigation.jnav.utils.recording_tf import RecordingTF from dimos.navigation.jnav.utils.trajectory_metrics import ( GraphPose, drifted_lookup, @@ -544,8 +545,6 @@ def run_module_graph( return graph, int(data["closures"]), replay_stats # type: ignore[return-value] -# Tolerance for composing world->body from the recording's tf at an odom time. -TF_LOOKUP_TOLERANCE_S = 0.1 # tf may come online slightly after the odom stream starts (sensor/static warmup); # odom samples before tf is available are skipped only within this opening window. TF_STARTUP_TOLERANCE_S = 2.0 @@ -567,7 +566,8 @@ def tf_pose_samples( odom start. Missing samples after tf is online (sporadic lookup gaps) are skipped.""" db_store = store(db_path) - if "tf" not in db_store.list_streams() or not db_store.tf.has_transforms(): + tf = RecordingTF.from_store(db_store) + if tf is None: raise SystemExit( f"{db_path}: no tf tree — eval requires tf (run add_tf first); odom fallback removed" ) @@ -578,7 +578,9 @@ def tf_pose_samples( for timestamp, _payload in iterate_stream(db_path, odom_stream): if odom_t0 is None: odom_t0 = timestamp - transform = db_store.tf.get(world_frame, body_frame, timestamp, TF_LOOKUP_TOLERANCE_S) + # tolerance=None -> nearest recorded sample per edge; RecordingTF keeps the + # whole recording buffered so one-shot static frames stay resolvable. + transform = tf.get(world_frame, body_frame, timestamp, None) if transform is None: if tf_online: continue # sporadic gap once tf is up — skip this sample diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py index 6c41d01713..4aa6ec33cd 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py @@ -73,6 +73,7 @@ detect_raw_detections, view_quality, ) +from dimos.navigation.jnav.utils.recording_tf import RecordingTF VISIT_GAP_S = 30.0 WHAT = sys.argv[1] if len(sys.argv) > 1 and not sys.argv[1].startswith("-") else "both" @@ -110,7 +111,6 @@ def arg(flag, default=""): LIDAR_FRAME = arg("--lidar-frame", "mid360_link") # frame the raw lidar scans live in WORLD_FRAME = arg("--world-frame", "world") # frame to register scans into USE_TF = "--no-tf" not in sys.argv # world-register via recording tf (fallback: obs.pose) -TF_TOL = float(arg("--tf-tol", "0.5")) # max seconds to the nearest recorded transform # Per-glimpse gates (speed == -1 means "unknown" and always passes). GATE = dict( @@ -175,7 +175,8 @@ def transform_matrix(transform): return rotation, translation -_TF_AVAILABLE = USE_TF and store.tf.has_transforms() +_STORE_TF = RecordingTF.from_store(store) if USE_TF else None +_TF_AVAILABLE = _STORE_TF is not None def world_points(observation): @@ -193,8 +194,11 @@ def world_points(observation): scan_frame = getattr(observation.data, "frame_id", "") or LIDAR_FRAME if scan_frame == WORLD_FRAME: return points # already world-registered per its own header - if _TF_AVAILABLE: - transform = store.tf.get(WORLD_FRAME, scan_frame, float(observation.ts), TF_TOL) + if _STORE_TF is not None: + # tolerance=None -> nearest recorded sample per edge; RecordingTF keeps the + # whole recording buffered, so one-shot static frames stay resolvable and the + # densely-sampled odom->base_link edge lands within a few ms of the scan. + transform = _STORE_TF.get(WORLD_FRAME, scan_frame, float(observation.ts), None) if transform is not None: rotation, translation = transform_matrix(transform) return points @ rotation.T + translation diff --git a/dimos/navigation/jnav/utils/recording_tf.py b/dimos/navigation/jnav/utils/recording_tf.py new file mode 100644 index 0000000000..e279d7d021 --- /dev/null +++ b/dimos/navigation/jnav/utils/recording_tf.py @@ -0,0 +1,43 @@ +# 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. + +"""Batch-friendly ``StreamTF`` for offline post-processing of a recording.""" + +from __future__ import annotations + +import math + +from dimos.memory2.tf import StreamTF + + +class RecordingTF(StreamTF): + """``StreamTF`` that caches the whole recording's tf in one pass. + + ``StreamTF`` is built for live/windowed use: each lookup keeps only a window + around the query time and evicts everything else. Recordings often publish + near-static frames (sensor mounts, ``world->map``) exactly once at the start + of the ``tf`` stream rather than into ``tf_static``, so a windowed lookup at + any later time drops those edges and the transform chain breaks. Loading the + full stream once keeps them buffered for the entire run; the only time-varying + edge (``odom->base_link``) is densely sampled, so a nearest lookup + (``time_tolerance=None``) reproduces the pose at each query time. + """ + + def _ensure(self, lo: float, hi: float) -> None: + with self._cv: + if self._covered is not None: + return + for observation in self.stream: + self.receive_transform(*observation.data.transforms) + self._covered = (-math.inf, math.inf) diff --git a/dimos/protocol/tf/deformation.py b/dimos/protocol/tf/deformation.py index a56c2a2698..2cd8135b06 100644 --- a/dimos/protocol/tf/deformation.py +++ b/dimos/protocol/tf/deformation.py @@ -23,8 +23,8 @@ delta(t) = slerp/lerp blend of the two keyframes bracketing t corrected_edge = delta(t) ∘ raw_edge -Used by both the recorded path (:class:`DbTfSql`) and the live path -(:class:`PubSubTF`); the SE(3) helpers are shared so the two stay identical. +Used by the live tf path (:class:`PubSubTF`) so that a running consumer's +tf lookups reflect loop-closure corrections as the optimizer republishes nodes. """ from __future__ import annotations @@ -88,8 +88,7 @@ class DeformationBuffer: Feed it :class:`DeformationNode` messages via :meth:`receive`; it keeps, per node, the original (first-seen) and current (latest) pose. :meth:`correct` applies the - blended loop-closure delta to a raw edge transform at a query time. This is the - live-path analogue of the DB-backed correction in :class:`DbTfSql`.""" + blended loop-closure delta to a raw edge transform at a query time.""" def __init__(self) -> None: # tf_id -> node_id -> [original_4x4, current_4x4, keyframe_ts] diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 2d9ed990b5..c9fe822cf3 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -373,7 +373,7 @@ class PubSubTFConfig(TFConfig): topic: Topic | None = None # Required field but needs default for dataclass inheritance # Optional per-keyframe loop-closure deformation stream (e.g. gsc_pgo's # tf_deformation_nodes). When wired, get()/lookup() return loop-closure-corrected - # transforms live, matching the recorded DbTfSql behaviour. Unset = raw tf only. + # transforms live. Unset = raw tf only. deformation_topic: Topic | None = None pubsub: type[PubSub] | PubSub | None = None # type: ignore[type-arg] autostart: bool = True From 2cb0274fa9dd0e9eb3825b517c05e500d75f047f Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 08:57:46 -0700 Subject: [PATCH 56/69] refactor(jnav): strip non-load-bearing changes outside jnav Keep the PR focused on dimos/navigation/jnav. Revert protocol/tf, core, and lidar recorders to origin/main; drop the unused deformation buffer and DeformationNode consumer wiring; move DeformationNode into jnav (its only consumer is the gsc_pgo module). - revert dimos/protocol/tf/tf.py to main (lookup() had no callers) - delete dimos/protocol/tf/deformation.py + test (unused) - revert core/module.py on_drop plumbing (dead code) - revert fastlio2/pointlio recorders and go2 static transforms to main - move DeformationNode.py into dimos/navigation/jnav/msgs --- dimos/core/module.py | 18 +-- .../sensors/lidar/fastlio2/recorder.py | 4 +- .../sensors/lidar/pointlio/recorder.py | 2 - .../components/loop_closure/gsc_pgo/module.py | 2 +- .../gsc_pgo/scripts/post_process.py | 2 +- .../jnav/msgs}/DeformationNode.py | 0 dimos/protocol/tf/deformation.py | 142 ------------------ dimos/protocol/tf/test_deformation.py | 74 --------- dimos/protocol/tf/tf.py | 64 -------- .../go2/go2_mid360_static_transforms.py | 12 +- 10 files changed, 8 insertions(+), 312 deletions(-) rename dimos/{msgs/nav_msgs => navigation/jnav/msgs}/DeformationNode.py (100%) delete mode 100644 dimos/protocol/tf/deformation.py delete mode 100644 dimos/protocol/tf/test_deformation.py diff --git a/dimos/core/module.py b/dimos/core/module.py index 23b2feb2f9..4f7aa9f964 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -478,16 +478,14 @@ def process_observable( self, observable: "Observable[Any]", async_cb: Callable[[Any], Any], - on_drop: Callable[[], None] | None = None, ) -> "DisposableBase": """Subscribe `async_cb` (an async function) to `observable`, dispatching each emitted value onto self._loop. Invocations are serialized through a - per-subscription dispatcher task with LATEST coalescing. `on_drop`, if - given, fires once per message dropped by that coalescing. The subscription + per-subscription dispatcher task with LATEST coalescing. The subscription is registered for cleanup on stop().""" if not inspect.iscoroutinefunction(async_cb): raise TypeError("process_observable requires an `async def` callback") - on_msg, dispatcher_disp = self._make_async_dispatch(async_cb, on_drop) + on_msg, dispatcher_disp = self._make_async_dispatch(async_cb) sub = observable.subscribe(on_msg) return self.register_disposable(CompositeDisposable(sub, dispatcher_disp)) @@ -638,9 +636,7 @@ def _auto_bind_handlers(self) -> None: self.process_observable(in_stream.pure_observable(), handler) def _make_async_dispatch( - self, - async_handler: Callable[[Any], Any], - on_drop: Callable[[], None] | None = None, + self, async_handler: Callable[[Any], Any] ) -> tuple[Callable[[Any], None], "DisposableBase"]: """Build a sync callback that delivers `msg` into a single-slot LATEST mailbox drained by a dedicated dispatcher task on `self._loop`. @@ -650,9 +646,7 @@ def _make_async_dispatch( awaits). - If messages arrive faster than the handler can process them, intermediate messages are dropped and only the most recent unprocessed - message is kept (LATEST policy). `on_drop`, if given, is called once - per dropped message (on the loop thread) so callers that need every - message can surface the loss. + message is kept (LATEST policy). - The returned Disposable cancels the dispatcher task. """ loop = self._loop @@ -692,10 +686,6 @@ def on_msg(msg: Any) -> None: return def _set() -> None: - # A slot that still holds an unconsumed value is about to be - # overwritten — that queued message is being dropped (LATEST). - if slot["has_value"] and on_drop is not None: - on_drop() slot["value"] = msg slot["has_value"] = True event.set() diff --git a/dimos/hardware/sensors/lidar/fastlio2/recorder.py b/dimos/hardware/sensors/lidar/fastlio2/recorder.py index 5f081d13e9..3ac40ae4bf 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/recorder.py +++ b/dimos/hardware/sensors/lidar/fastlio2/recorder.py @@ -52,6 +52,4 @@ async def _odom_pose(self, msg: Odometry) -> Pose | None: @pose_setter_for("fastlio_lidar") async def _lidar_pose(self, msg: PointCloud2) -> Pose | None: - # Most-recent odometry pose, stamped directly (no tf). None before the - # first odometry -> frame stored unposed, map-skipped. - return self._last_odom_pose + return getattr(self, "_last_odom_pose", None) diff --git a/dimos/hardware/sensors/lidar/pointlio/recorder.py b/dimos/hardware/sensors/lidar/pointlio/recorder.py index f7eff72c78..e859190351 100644 --- a/dimos/hardware/sensors/lidar/pointlio/recorder.py +++ b/dimos/hardware/sensors/lidar/pointlio/recorder.py @@ -34,8 +34,6 @@ class PointlioRecorderConfig(RecorderConfig): # Append into a populated db (keep other streams); replace only our own. on_existing: OnExisting = OnExisting.APPEND - # pcap_to_db has no tf source of its own — don't record (and thus don't touch) tf. - record_tf: bool = False class PointlioRecorder(Recorder): diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py index 6f61728bb8..1487db611a 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py @@ -27,9 +27,9 @@ from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import In, Out from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.nav_msgs.DeformationNode import DeformationNode from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.navigation.jnav.msgs.DeformationNode import DeformationNode from dimos.navigation.jnav.msgs.Graph3D import Graph3D from dimos.navigation.jnav.msgs.GraphDelta3D import GraphDelta3D from dimos.navigation.jnav.msgs.LocationConstraint import LocationConstraint diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py index 4aa6ec33cd..2197f4ab5a 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py @@ -56,9 +56,9 @@ from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.DeformationNode import DeformationNode, tf_id_for from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.navigation.jnav.msgs.DeformationNode import DeformationNode, tf_id_for from dimos.navigation.jnav.msgs.Graph3D import Graph3D from dimos.navigation.jnav.utils import recording_db as rdb from dimos.navigation.jnav.utils.apriltags import ( diff --git a/dimos/msgs/nav_msgs/DeformationNode.py b/dimos/navigation/jnav/msgs/DeformationNode.py similarity index 100% rename from dimos/msgs/nav_msgs/DeformationNode.py rename to dimos/navigation/jnav/msgs/DeformationNode.py diff --git a/dimos/protocol/tf/deformation.py b/dimos/protocol/tf/deformation.py deleted file mode 100644 index 2cd8135b06..0000000000 --- a/dimos/protocol/tf/deformation.py +++ /dev/null @@ -1,142 +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. - -"""Loop-closure deformation correction for transform lookups. - -A loop-closure backend (e.g. gsc_pgo) publishes a stream of :class:`DeformationNode` -keyframes — one per pose-graph node, re-published (same id) when the optimizer moves -it, tagged with the tf_id of the edge it corrects. This module turns that stream into -a query-time correction on a raw edge transform: - - delta[k] = current[k] ∘ inv(original[k]) # how far keyframe k moved since first seen - delta(t) = slerp/lerp blend of the two keyframes bracketing t - corrected_edge = delta(t) ∘ raw_edge - -Used by the live tf path (:class:`PubSubTF`) so that a running consumer's -tf lookups reflect loop-closure corrections as the optimizer republishes nodes. -""" - -from __future__ import annotations - -import bisect -from typing import Any, cast - -import numpy as np - -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.nav_msgs.DeformationNode import DeformationNode, tf_id_for - - -def apply_delta(delta: np.ndarray, transform: Transform, ts: float) -> Transform: - """Deform ``transform`` by the SE(3) correction ``delta`` on the parent (frame_id) - side: ``corrected = delta @ raw``. Keeps the edge's frame names.""" - corrected = delta @ transform.to_matrix() - return Transform.from_matrix( - corrected, - ts=ts, - frame_id=transform.frame_id, - child_frame_id=transform.child_frame_id, - ) - - -def quat_slerp(q_lo: np.ndarray, q_hi: np.ndarray, weight: float) -> np.ndarray: - """Spherical-linear interpolation between two quaternions ``[x, y, z, w]``.""" - dot = float(np.dot(q_lo, q_hi)) - if dot < 0.0: # take the shorter arc - q_hi = -q_hi - dot = -dot - if dot > 0.9995: # nearly parallel: lerp + renormalize avoids a divide-by-~0 - result = q_lo + weight * (q_hi - q_lo) - return cast("np.ndarray", result / np.linalg.norm(result)) - theta_0 = np.arccos(np.clip(dot, -1.0, 1.0)) - sin_0 = np.sin(theta_0) - scale_lo = np.sin((1.0 - weight) * theta_0) / sin_0 - scale_hi = np.sin(weight * theta_0) / sin_0 - return cast("np.ndarray", scale_lo * q_lo + scale_hi * q_hi) - - -def blend_se3(mat_lo: np.ndarray, mat_hi: np.ndarray, weight: float) -> np.ndarray: - """Linear-blend-skin two SE(3) deltas by ``weight`` in [0, 1] (0 -> lo, 1 -> hi): - lerp the translation, slerp the rotation.""" - quat_lo = Quaternion.from_rotation_matrix(mat_lo[:3, :3]) - quat_hi = Quaternion.from_rotation_matrix(mat_hi[:3, :3]) - blended = quat_slerp( - np.array([quat_lo.x, quat_lo.y, quat_lo.z, quat_lo.w]), - np.array([quat_hi.x, quat_hi.y, quat_hi.z, quat_hi.w]), - weight, - ) - out = np.eye(4) - out[:3, :3] = Quaternion(blended).to_rotation_matrix() - out[:3, 3] = (1.0 - weight) * mat_lo[:3, 3] + weight * mat_hi[:3, 3] - return out - - -class DeformationBuffer: - """In-RAM rolling store of pose-graph keyframes, keyed by tf_id (the corrected edge). - - Feed it :class:`DeformationNode` messages via :meth:`receive`; it keeps, per node, - the original (first-seen) and current (latest) pose. :meth:`correct` applies the - blended loop-closure delta to a raw edge transform at a query time.""" - - def __init__(self) -> None: - # tf_id -> node_id -> [original_4x4, current_4x4, keyframe_ts] - self._edges: dict[int, dict[int, list[Any]]] = {} - - def receive(self, node: DeformationNode) -> None: - matrix = Transform.from_pose(node.pose.frame_id, node.pose).to_matrix() - edge = self._edges.setdefault(node.tf_id, {}) - state = edge.get(node.id) - if state is None: - edge[node.id] = [matrix, matrix, node.pose.ts] # original, current, ts - else: - state[1] = matrix # optimizer moved it -> update current only - - @property - def has_corrections(self) -> bool: - return bool(self._edges) - - def _edge_delta(self, tf_id: int, query_time: float) -> np.ndarray | None: - edge = self._edges.get(tf_id) - if not edge: - return None - # keyframes sorted by ts; bracket the query time (latest<=t, earliest>=t) - ordered = sorted(edge.values(), key=lambda state: state[2]) - stamps = [state[2] for state in ordered] - lo_index = bisect.bisect_right(stamps, query_time) - 1 # last at or before - hi_index = bisect.bisect_left(stamps, query_time) # first at or after - picks = [] - if lo_index >= 0: - picks.append(ordered[lo_index]) - if hi_index < len(ordered) and (not picks or ordered[hi_index] is not picks[0]): - picks.append(ordered[hi_index]) - if not picks: - return None - # per keyframe: delta = current ∘ inv(original) - samples = [(state[2], state[1] @ np.linalg.inv(state[0])) for state in picks] - if len(samples) == 1 or samples[0][0] == samples[1][0]: - return cast("np.ndarray", samples[0][1]) - (ts_lo, mat_lo), (ts_hi, mat_hi) = samples - weight = (query_time - ts_lo) / (ts_hi - ts_lo) - return blend_se3(mat_lo, mat_hi, weight) - - def correct( - self, frame_from: str, frame_to: str, raw: Transform, query_time: float - ) -> Transform: - """Return the loop-closure-corrected ``frame_from <- frame_to`` transform, or - ``raw`` unchanged if no deformation nodes match that edge.""" - if not self._edges: - return raw - delta = self._edge_delta(tf_id_for(frame_from, frame_to), query_time) - return apply_delta(delta, raw, raw.ts) if delta is not None else raw diff --git a/dimos/protocol/tf/test_deformation.py b/dimos/protocol/tf/test_deformation.py deleted file mode 100644 index 2389370186..0000000000 --- a/dimos/protocol/tf/test_deformation.py +++ /dev/null @@ -1,74 +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. - -"""Live-path loop-closure correction (DeformationBuffer).""" - -from __future__ import annotations - -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -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.DeformationNode import DeformationNode, tf_id_for -from dimos.protocol.tf.deformation import DeformationBuffer - - -def _node(tf_id: int, node_id: int, ts: float, xyz: tuple[float, float, float]) -> DeformationNode: - return DeformationNode( - id=node_id, - tf_id=tf_id, - pose=PoseStamped(ts=ts, frame_id="map", position=list(xyz), orientation=[0, 0, 0, 1]), - ) - - -def _identity(parent: str, child: str, ts: float) -> Transform: - return Transform( - translation=Vector3(0, 0, 0), - rotation=Quaternion(0, 0, 0, 1), - frame_id=parent, - child_frame_id=child, - ts=ts, - ) - - -def test_single_keyframe_corrects_matched_edge_only() -> None: - buffer = DeformationBuffer() - tf_id = tf_id_for("map", "odom") - buffer.receive(_node(tf_id, 1, 100.0, (0.0, 0.0, 0.0))) # original - buffer.receive(_node(tf_id, 1, 100.0, (1.0, 0.0, 0.0))) # optimizer moved it +1 in x - - corrected = buffer.correct("map", "odom", _identity("map", "odom", 100.0), 100.0) - assert abs(corrected.translation.x - 1.0) < 1e-9 # delta = current ∘ inv(original) applied - - # an edge with no matching tf_id passes through untouched - other = _identity("odom", "base_link", 100.0) - assert buffer.correct("odom", "base_link", other, 100.0).translation.x == 0.0 - - -def test_blends_between_bracketing_keyframes() -> None: - buffer = DeformationBuffer() - tf_id = tf_id_for("map", "odom") - buffer.receive(_node(tf_id, 1, 100.0, (0.0, 0.0, 0.0))) # kf A: unmoved - buffer.receive(_node(tf_id, 1, 100.0, (0.0, 0.0, 0.0))) - buffer.receive(_node(tf_id, 2, 110.0, (0.0, 0.0, 0.0))) # kf B: moved +2 - buffer.receive(_node(tf_id, 2, 110.0, (2.0, 0.0, 0.0))) - - corrected = buffer.correct("map", "odom", _identity("map", "odom", 105.0), 105.0) - assert abs(corrected.translation.x - 1.0) < 1e-9 # midpoint of 0->2 - - -def test_no_nodes_passes_through_identity() -> None: - buffer = DeformationBuffer() - raw = _identity("map", "odom", 1.0) - assert buffer.correct("map", "odom", raw, 1.0) is raw # no copy, no cost diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index c9fe822cf3..464498b718 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -24,13 +24,11 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.nav_msgs.DeformationNode import DeformationNode from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.protocol.pubsub.impl.lcmpubsub import LCM, Topic from dimos.protocol.pubsub.impl.zenohpubsub import Zenoh from dimos.protocol.pubsub.spec import PubSub from dimos.protocol.service.spec import BaseConfig, Service -from dimos.protocol.tf.deformation import DeformationBuffer from dimos.types.timestamped import to_human_readable from dimos.utils.logging_config import setup_logger from dimos.utils.timeseries.inmemory import InMemoryStore @@ -260,19 +258,6 @@ def _wait_get( return None self._cv.wait(timeout=remaining) - def lookup( - self, - parent_frame: str, - child_frame: str, - time_point: float | None = None, - time_tolerance: float | None = None, - ) -> Transform | None: - """Composed transform lookup that does NOT log on a miss (unlike `get`). - - For high-frequency / best-effort lookups where misses are expected and a - per-call warning would spam the log (e.g. per-scan world-registration).""" - return self._get(parent_frame, child_frame, time_point, time_tolerance) - def get( self, parent_frame: str, @@ -371,10 +356,6 @@ def __str__(self) -> str: class PubSubTFConfig(TFConfig): topic: Topic | None = None # Required field but needs default for dataclass inheritance - # Optional per-keyframe loop-closure deformation stream (e.g. gsc_pgo's - # tf_deformation_nodes). When wired, get()/lookup() return loop-closure-corrected - # transforms live. Unset = raw tf only. - deformation_topic: Topic | None = None pubsub: type[PubSub] | PubSub | None = None # type: ignore[type-arg] autostart: bool = True @@ -385,7 +366,6 @@ class PubSubTF(MultiTBuffer, TFSpec): def __init__(self, **kwargs) -> None: # type: ignore[no-untyped-def] TFSpec.__init__(self, **kwargs) MultiTBuffer.__init__(self, self.config.buffer_size) - self._deformation = DeformationBuffer() pubsub_config = getattr(self.config, "pubsub", None) if pubsub_config is not None: @@ -405,44 +385,6 @@ def start(self, sub: bool = True) -> None: topic = getattr(self.config, "topic", None) if topic: self.pubsub.subscribe(topic, self.receive_msg) - deformation_topic = getattr(self.config, "deformation_topic", None) - if deformation_topic: - self.pubsub.subscribe(deformation_topic, self.receive_deformation) - - def receive_deformation(self, msg: DeformationNode, topic: Topic) -> None: - self._deformation.receive(msg) - - def get_transform( - self, - parent_frame: str, - child_frame: str, - time_point: float | None = None, - time_tolerance: float | None = None, - ) -> Transform | None: - """Single-edge lookup, loop-closure-corrected. Correction is applied on the - edge's stored direction, so both direct lookups and BFS chain hops (forward or - reversed) resolve to the corrected transform. Falls back to the raw lookup when - no deformation nodes have arrived.""" - if not self._deformation.has_corrections: - return super().get_transform(parent_frame, child_frame, time_point, time_tolerance) - with self._cv: - forward = (parent_frame, child_frame) - if forward in self.buffers: - raw = self.buffers[forward].get(time_point, time_tolerance) # type: ignore[arg-type] - if raw is None: - return None - ts = time_point if time_point is not None else raw.ts - return self._deformation.correct(parent_frame, child_frame, raw, ts) - reverse = (child_frame, parent_frame) - if reverse in self.buffers: - # child->parent (stored direction) - raw = self.buffers[reverse].get(time_point, time_tolerance) # type: ignore[arg-type] - if raw is None: - return None - ts = time_point if time_point is not None else raw.ts - corrected = self._deformation.correct(child_frame, parent_frame, raw, ts) - return corrected.inverse() - return None def stop(self) -> None: self.pubsub.stop() @@ -496,12 +438,6 @@ def receive_msg(self, msg: TFMessage, topic: Topic) -> None: class LCMPubsubConfig(PubSubTFConfig): topic: Topic = field(default_factory=lambda: Topic("/tf", TFMessage)) - # On by default: subscribe the loop-closure deformation stream so live tf.get is - # corrected whenever a PGO publishes it. Harmless when nothing does — the buffer - # stays empty and correct() is a zero-cost pass-through. - deformation_topic: Topic | None = field( - default_factory=lambda: Topic("/tf_deformation_nodes", DeformationNode) - ) pubsub: type[PubSub] | PubSub | None = LCM # type: ignore[type-arg] autostart: bool = True diff --git a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py index 2f28168961..b430a11429 100644 --- a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py +++ b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py @@ -21,14 +21,7 @@ Mount geometry (measured on the physical rig) --------------------------------------------- - base_link -> front_camera: 32.7cm forward, ~4.3cm up (URDF front_camera mount). -- front_camera -> mid360_link: lidar is 3.2cm back, 12cm up, pitched 44 deg down — - the physical sensor pose (position + tilt). Named ``mid360_link`` to match the LIO - sensor frame_id. -- mid360_link -> mid360_gravity: a counter-rotation (inverse of the mount tilt) that - leaves the position untouched but un-tilts the orientation, so ``mid360_gravity`` - is gravity-aligned. The LIO odometry is gravity-anchored, so the recorded scans - live in this gravity-aligned frame at the physical sensor position — keeping the - physical mount in ``mid360_link`` preserves the geometry. +- front_camera -> mid360_link: lidar is 3.2cm back, 12cm up, pitched 44 deg down. - front_camera -> camera_optical: the standard ROS optical rotation (x-right, y-down, z-forward). """ @@ -52,10 +45,7 @@ FRAMES: list[FrameSpec] = [ ("base_link", None, (0.0, 0.0, 0.0), (0.0, 0.0, 0.0)), ("front_camera", "base_link", (0.32715, -0.00003, 0.04297), (0.0, 0.0, 0.0)), - # physical sensor pose: position + the 44 deg downward tilt (matches LIO frame_id) ("mid360_link", "front_camera", (-0.032, 0.0, 0.12), (0.0, MID360_PITCH_DOWN, 0.0)), - # gravity-aligned derived frame: same position, tilt counter-rotated out - ("mid360_gravity", "mid360_link", (0.0, 0.0, 0.0), (0.0, -MID360_PITCH_DOWN, 0.0)), ("camera_optical", "front_camera", (0.0, 0.0, 0.0), OPTICAL_RPY), ] From a42693ee3004c6209df037f3b29caf7cab65e70b Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 10:03:47 -0700 Subject: [PATCH 57/69] docs(jnav): remove pgo_migration_plan.md planning artifact --- experimental/docs/jnav/pgo_migration_plan.md | 144 ------------------- 1 file changed, 144 deletions(-) delete mode 100644 experimental/docs/jnav/pgo_migration_plan.md diff --git a/experimental/docs/jnav/pgo_migration_plan.md b/experimental/docs/jnav/pgo_migration_plan.md deleted file mode 100644 index 713aa22444..0000000000 --- a/experimental/docs/jnav/pgo_migration_plan.md +++ /dev/null @@ -1,144 +0,0 @@ -# PGO → jnav Loop-Closure Migration Plan - -Goal: land the pose-graph-optimization (PGO) / loop-closure work on -`jeff/feat/jnav_pgo` in the new `jnav` layout, extract the C++ + nix flake into a -standalone `github:jeff-hykin/gsc_pgo` repo, and merge in the offline AprilTag -map-postprocessing tooling. - -## Source material - -Two branches hold the pieces; neither alone is the target. - -| Source | Has | Layout | -|---|---|---| -| `jeff/feat/jnav` | `jnav/{msgs,utils}/`, `loop_closure/{gsc_pgo,ivan_pgo,ivan_pgo_transformer,unrefined_pgo}`, `eval.py`, `eval_all.py`, Scan-Context + Landmark C++ | `dimos/navigation/jnav/modules/...` (inline C++) | -| `jeff/feat/better_pgo` | postprocessing scripts (`add_april.py`, `detect_tags.py`, `post_process.py`, `make_rrd.py`), `map_postprocessing.md` doc | `dimos/navigation/nav_stack/modules/pgo/scripts/...` | -| `jeff/feat/jnav_pgo` (current) | nothing yet — only `nav_stack→cmu_nav` rename on top of `main` | — | - -`jeff/feat/jnav` is the more evolved navigation reorg; `better_pgo` is the source -for the AprilTag postprocessing scripts/doc. The migration is the union of both. - -## Target layout (on `jeff/feat/jnav_pgo`) - -``` -dimos/navigation/jnav/ - msgs/ # Graph3D/GraphDelta3D/Landmark/Marker (.py + .hpp + tests) - utils/ # ALL pgo utils (apriltags, trajectory_metrics, recording_db, ...) - components/ # renamed from jnav's "modules/" - loop_closure/ - eval.py - eval_all.py - spec.py # LoopClosure Protocol - gsc_pgo/ - module.py # NativeModule -> builds external gsc_pgo flake - scripts/ - post_process.py # ported from better_pgo scripts/post_process.py - (add_april.py, detect_tags.py, make_rrd.py) # the other postprocess scripts -experimental/docs/jnav/ - map_postprocessing.md # adapted from docs/capabilities/navigation/map_postprocessing.md -``` - -External repo: -``` -github:jeff-hykin/gsc_pgo # all C++ + flake.nix/flake.lock - flake.nix, flake.lock, CMakeLists.txt - main.cpp, simple_pgo.{cpp,h}, scan_context.{cpp,h}, - commons.{cpp,h}, point_cloud_utils.hpp, dimos_native_module.hpp, - pgo_landmark_test.cpp - msgs/ (Graph3D.hpp, GraphDelta3D.hpp, Landmark.hpp) # C++ wire helpers -``` - -## Decisions (confirmed by Jeff, 2026-06-22) - -1. **Base branch.** ✅ Base `jnav_pgo` on `jeff/feat/jnav`. BUT `post_process.py` - must come from `better_pgo` (not jnav, which lacks it). -2. **`components/` scope.** ⏳ pending (whole `jnav/modules/`→`components/` vs only - `loop_closure`). -3. **`gsc_pgo` repo.** ✅ **public**, created at - `github.com/jeff-hykin/gsc_pgo`, initial rev `494e7a1d657c3702ec805c9e3d251a2fe8bc9529`. - Flake input will pin to that rev: `github:jeff-hykin/gsc_pgo/494e7a1...#default`. - -## Work breakdown - -### Phase 0 — branch setup -- Confirm `jnav_pgo` base (decision 1). If basing on jnav: merge/cherry-pick the - `dimos/navigation/jnav/` tree onto current `jnav_pgo` (which already has the - `cmu_nav` rename). Resolve any overlap with cmu_nav's own `pgo` module. - -### Phase 1 — external `gsc_pgo` repo -- `gh repo create jeff-hykin/gsc_pgo` (visibility per decision 3). *(outward-facing — confirm first)* -- Move `gsc_pgo/cpp/*` (C++, CMakeLists, flake.nix, flake.lock, `msgs/*.hpp`, - `pgo_landmark_test.cpp`) into the new repo. Keep the flake's - `lcm-extended` / `dimos-lcm` / `gtsam-extended` inputs as-is. -- `flake.nix`: `src = ./.;` already self-contained — switching to a tracked git - repo *fixes* the untracked-files gotcha that forced `path:$PWD#default` - (see memory: nix untracked-files gotcha). Verify `nix build .#default` from a - clean checkout. -- Push, tag/record the rev for the flake pin. - -### Phase 2 — `gsc_pgo` module.py -- Update `PGOConfig`: - - `build_command = 'nix build "github:jeff-hykin/gsc_pgo/#default" --no-write-lock-file'` - - drop the `path:$PWD` workaround comment; `cwd` no longer needs a local `cpp/`. - - `executable` resolves from the nix `result/bin/pgo` (confirm NativeModule - out-of-tree build path handling). -- Keep `In/Out` stream wiring and all loop-closure hyperparameters unchanged. -- Confirm C++ `msgs/*.hpp` wire format still matches `jnav/msgs/*.py` decoders - (they're hand-synced; the `.hpp` headers travel with the C++ repo). - -### Phase 3 — msgs/ + utils/ -- `jnav/msgs/`: already present on jnav (Graph3D, GraphDelta3D, Landmark, Marker + - `.hpp` + tests). Bring over as-is. The `.hpp` files are duplicated into the - external repo (C++ side); decide whether the canonical `.hpp` lives in - `gsc_pgo` and `jnav/msgs/*.hpp` is a mirror, or vice-versa. Recommend: canonical - in `gsc_pgo`, keep a note in `jnav/msgs` pointing at it. -- `jnav/utils/`: bring all utils. `better_pgo`'s `eval_utils/` (apriltags, - apriltag_agreement, trajectory_metrics, recording_db, voxel_map, module_loading) - are already absorbed into jnav `utils/` — verify no newer logic in - `better_pgo` was lost (diff the 6 overlapping files). - -### Phase 4 — eval.py / eval_all.py -- Copy from jnav `loop_closure/`. Fix imports to `dimos.navigation.jnav.utils.*` - and `...jnav.msgs.*` (already correct on jnav). -- `eval_all.py` enumerates the comparison modules (`gsc_pgo`, `ivan_pgo`, - `ivan_pgo_transformer`, `unrefined_pgo`) → port all four component dirs so eval - doesn't break. (If baselines are unwanted, prune eval_all's list instead.) - -### Phase 5 — postprocessing scripts (from better_pgo) -- Port `scripts/post_process.py` → `components/loop_closure/gsc_pgo/scripts/post_process.py`. -- Port `add_april.py`, `detect_tags.py`, `make_rrd.py` alongside it (the doc's - 3-step flow references all of them). -- Rewrite their internal imports: `eval_utils.apriltags` → `dimos.navigation.jnav.utils.apriltags`, etc. -- Tag quality gates are single-sourced in `utils/apriltags.py` (`DEFAULT_*`); keep - post_process importing from there (don't duplicate constants). - -### Phase 6 — docs -- Adapt `docs/capabilities/navigation/map_postprocessing.md` → - `experimental/docs/jnav/map_postprocessing.md`. -- Rewrite every script path in the doc: - `dimos/navigation/nav_stack/modules/pgo/scripts/X.py` - → `dimos/navigation/jnav/components/loop_closure/gsc_pgo/X.py`. -- Update the `eval_utils/apriltags.py` reference → `jnav/utils/apriltags.py`. -- Add a back-link from the navigation readme if appropriate. - -### Phase 7 — wiring + tests -- `spec.py` `LoopClosure` Protocol: bring over, update any module paths. -- Blueprints: if any blueprint references the old `nav_stack`/`pgo` module path, - repoint to `jnav.components.loop_closure.gsc_pgo`. Regenerate - `all_blueprints.py` via `pytest dimos/robot/test_all_blueprints_generation.py`. -- Port tests (`test_pgo_synthetic_drift.py`, msgs tests, utils tests). Run - `uv run pytest` for the loop_closure + msgs + utils dirs. -- `ruff check --fix && ruff format`. - -## Risks / cut-corners to flag -- **C++ wire-format drift**: `.hpp` (C++) and `.py` (decode) are hand-synced. Once - the `.hpp` lives in an external repo, a change there can silently desync the - Python decoder. Mitigate: keep a `test_Graph3D` round-trip test in `jnav/msgs` - that builds against the pinned flake rev, or at least a schema-comment check. -- **Private flake in CI**: if `gsc_pgo` is private, dimos CI nix builds will fail - without a token — do not add a token to CI without asking (secrets rule). -- **eval_results/**: jnav carries committed `eval_results/*/summary.json` snapshots. - Decide whether to bring those (history/benchmarks) or regenerate. -- **Baseline PGO impls**: `ivan_pgo*` / `unrefined_pgo` carry their own inline C++ - (`unrefined_pgo/cpp`). If only `gsc_pgo` goes external, the layout is asymmetric - — acceptable (they're experimental baselines) but worth noting. From c5d112860c0fb4432c9a3c94a82a680f3b339ff1 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 11:05:03 -0700 Subject: [PATCH 58/69] refactor(jnav): fold detect_tags into add_april detect_tags.py duplicated ensure_april_streams, which add_april.py already calls to build raw_april_tags (same gate diagnostics) plus the gated stream and summary. Repoint the two docstring/doc references to add_april.py. --- .../gsc_pgo/scripts/detect_tags.py | 189 ------------------ .../gsc_pgo/scripts/post_process.py | 2 +- experimental/docs/jnav/map_postprocessing.md | 2 +- 3 files changed, 2 insertions(+), 191 deletions(-) delete mode 100644 dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py deleted file mode 100644 index 4e6cf00527..0000000000 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py +++ /dev/null @@ -1,189 +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. - -# Untyped analysis script: gtsam/open3d/cv2 lack type stubs. -# mypy: ignore-errors -"""Build the raw AprilTag stream: EVERY detection over the camera image, NO filtering whatsoever -(no blur/reproj/distance/angle/motion/size gate, no time-clustering). One row per per-frame -detection that yields a valid PnP pose. Each row carries its gate diagnostics in tags -(sharpness, reproj_px, tag_px, distance_m, view_angle_deg, lin_speed, ang_speed) so downstream -gate tuning in post_process.py needs no re-detection. - -Prints the raw per-marker histogram + visit structure (visit = sightings >30s apart). - -Usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py --rec=PATH - [--camera=color_image] [--tag-size=0.10] - [--dict=DICT_APRILTAG_36h11] [--intrinsics=PATH] [--out=raw_april_tags] -""" - -import json -from pathlib import Path -import sys - -import cv2 -import numpy as np -from scipy.spatial.transform import Rotation - -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.Image import Image -from dimos.navigation.jnav.utils import recording_db as rdb -from dimos.navigation.jnav.utils.apriltags import ( - _camera_speeds, - estimate_marker_pose, - make_detector, - reprojection_error_px, - tag_pixel_size, - tag_sharpness, - view_quality, -) - - -def arg(flag, default=None): - return next( - (argument.split("=", 1)[1] for argument in sys.argv if argument.startswith(flag + "=")), - default, - ) - - -REC_ARG = arg("--rec") -if not REC_ARG: - sys.exit( - "usage: python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/detect_tags.py --rec=PATH [--camera=...] [--tag-size=...] " - "[--dict=...] [--intrinsics=PATH] [--out=...] (--rec is required)" - ) -REC = Path(REC_ARG).expanduser() -CAMERA = arg("--camera", "color_image") # camera image stream the tags are detected on -MARKER_LENGTH_M = float(arg("--tag-size", "0.10")) -DICTIONARY = arg("--dict", "DICT_APRILTAG_36h11") -STREAM = arg("--out", "raw_april_tags") -VISIT_GAP_S = 30.0 - -DB = REC / "mem2.db" -intrinsics_path = Path(arg("--intrinsics", str(REC / "camera_intrinsics.json"))).expanduser() -intrinsics = json.loads(intrinsics_path.read_text()) -camera_matrix = np.array(intrinsics["intrinsics"], float).reshape(3, 3) -distortion = np.array(intrinsics.get("distortion", []), float) -store = rdb.store(DB) -if CAMERA not in store.list_streams(): - sys.exit(f"!! camera stream '{CAMERA}' not in db — available: {store.list_streams()}") - -detector = make_detector(DICTIONARY) -print(f"loading '{CAMERA}' frames...") -images = store.stream(CAMERA, Image).to_list() -speed_by_ts, speed_available = _camera_speeds(images) -print( - f"detecting over {len(images)} frames (unfiltered), tag_size={MARKER_LENGTH_M} m, dict={DICTIONARY}..." -) - -rows = [] -for image_obs in images: - image = image_obs.data - bgr = image.numpy() if hasattr(image, "numpy") else np.asarray(image.data) - gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) if bgr.ndim == 3 else bgr - all_corners, marker_ids, _ = detector.detectMarkers(bgr) - if marker_ids is None: - continue - for corners, marker_id in zip(all_corners, marker_ids.flatten(), strict=False): - pose = estimate_marker_pose(corners, MARKER_LENGTH_M, camera_matrix, distortion) - if pose is None: - continue - rotation_vector, translation_vector = pose - quaternion = Rotation.from_rotvec(rotation_vector.reshape(3)).as_quat() # x,y,z,w - translation = translation_vector.reshape(3) - tag_pose = [ - float(translation[0]), - float(translation[1]), - float(translation[2]), - float(quaternion[0]), - float(quaternion[1]), - float(quaternion[2]), - float(quaternion[3]), - ] - distance, view_angle = view_quality(tag_pose) - speed = speed_by_ts.get(float(image_obs.ts)) - rows.append( - { - "ts": float(image_obs.ts), - "marker_id": int(marker_id), - "tcm": tag_pose, - "sharpness": float(tag_sharpness(gray, corners)), - "reproj_px": float( - reprojection_error_px( - corners, - rotation_vector, - translation_vector, - MARKER_LENGTH_M, - camera_matrix, - distortion, - ) - ), - "tag_px": float(tag_pixel_size(corners)), - "distance_m": float(distance), - "view_angle_deg": float(view_angle), - "lin_speed": float(speed[0]) if speed else -1.0, - "ang_speed": float(speed[1]) if speed else -1.0, - } - ) -rows.sort(key=lambda row: row["ts"]) - -if STREAM in store.list_streams(): - store.delete_stream(STREAM) -out_stream = store.stream(STREAM, PoseStamped) -for row in rows: - tag_pose = row["tcm"] - out_stream.append( - PoseStamped(ts=row["ts"], position=tag_pose[:3], orientation=tag_pose[3:]), - ts=row["ts"], - pose=tuple(tag_pose), - tags={ - tag_key: row[tag_key] - for tag_key in ( - "marker_id", - "sharpness", - "reproj_px", - "tag_px", - "distance_m", - "view_angle_deg", - "lin_speed", - "ang_speed", - ) - }, - ) -print(f"\nwrote {STREAM}: {len(rows)} unfiltered detections") - -rows_by_marker = {} -for row in rows: - rows_by_marker.setdefault(row["marker_id"], []).append(row) -print(f"\n=== RAW per-marker (visit = >{VISIT_GAP_S:.0f}s apart) ===") -print( - f"{'tag':>4} {'det':>4} {'visits':>6} {'dist_m':>12} {'sharp>=60%':>10} {'reproj<=2%':>10} {'span_s':>7}" -) -for marker_id in sorted(rows_by_marker): - marker_rows = sorted(rows_by_marker[marker_id], key=lambda row: row["ts"]) - times = [row["ts"] for row in marker_rows] - visits = [[times[0]]] - for timestamp in times[1:]: - ( - visits[-1].append(timestamp) - if timestamp - visits[-1][-1] <= VISIT_GAP_S - else visits.append([timestamp]) - ) - distances = [row["distance_m"] for row in marker_rows] - sharp_ok = 100 * np.mean([row["sharpness"] >= 60 for row in marker_rows]) - reproj_ok = 100 * np.mean([row["reproj_px"] <= 2 for row in marker_rows]) - print( - f"{marker_id:>4} {len(marker_rows):>4} {len(visits):>6} {min(distances):5.2f}-{max(distances):5.2f} " - f"{sharp_ok:9.0f}% {reproj_ok:9.0f}% {times[-1] - times[0]:7.0f}" - ) -print("\nmarkers present:", sorted(rows_by_marker)) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py index 2197f4ab5a..f8e8222fdc 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py @@ -16,7 +16,7 @@ # mypy: ignore-errors """AprilTag-loop-closed + ICP-refined ground-truth post-processing for a go2 recording. -Source of tags = the UNFILTERED tag stream (build it first with detect_tags.py). Each raw +Source of tags = the UNFILTERED tag stream (build it first with add_april.py). Each raw detection carries its gate diagnostics, so gates are applied here post-hoc (no re-detection) and are easy to relax. Factors: one robust (best-reproj) observation per keyframe x marker -> denser, balanced loop closure than one-medoid-per-visit. diff --git a/experimental/docs/jnav/map_postprocessing.md b/experimental/docs/jnav/map_postprocessing.md index e579e06b79..8fe2e18ef9 100644 --- a/experimental/docs/jnav/map_postprocessing.md +++ b/experimental/docs/jnav/map_postprocessing.md @@ -27,7 +27,7 @@ The scripts live in `dimos/navigation/jnav/components/loop_closure/gsc_pgo/scrip python dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py --rec=PATH ``` -`add_april.py` writes `raw_april_tags` (every detection, unfiltered, with its quality numbers attached — this is what postprocessing reads), the gated/clustered `april_tags` stream, and an `april_tags` section in the recording's `summary.json` (see below). Leaving `raw_april_tags` unfiltered matters: you tune the quality gates later without re-running detection, which is the slow part. (`detect_tags.py` writes just the raw stream if that's all you want; `--dynamic 17` keeps a moving tag in raw but drops it from the gated stream.) +`add_april.py` writes `raw_april_tags` (every detection, unfiltered, with its quality numbers attached — this is what postprocessing reads), the gated/clustered `april_tags` stream, and an `april_tags` section in the recording's `summary.json` (see below). Leaving `raw_april_tags` unfiltered matters: you tune the quality gates later without re-running detection, which is the slow part. (`--dynamic 17` keeps a moving tag in raw but drops it from the gated stream.) Inspect what was found without rebuilding anything — per tag, raw count and revisit count, flagging any tag never revisited (your fast check for whether a recording even has loop constraints): From 70790ac4e9d34b4ebc02dbbaa47a354d191fc988 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 11:05:11 -0700 Subject: [PATCH 59/69] chore(jnav): pin gsc_pgo to v1.0.0 branch instead of commit hash Point the nix flake ref at github:jeff-hykin/gsc_pgo/v1.0.0 (a version branch at the previously-pinned rev) so the pin reads as a version and can be bumped on the gsc_pgo side without editing a hash here. --- .../jnav/components/loop_closure/gsc_pgo/module.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py index 1487db611a..47fe9759f6 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py @@ -40,13 +40,12 @@ class PGOConfig(NativeModuleConfig): # C++ + nix flake live in the standalone repo github.com/jeff-hykin/gsc_pgo. - # Pinned to a rev for reproducibility; bump when the C++ changes. The build - # runs in this module dir and drops a `result` symlink here (gitignored). + # Pinned to a version branch for reproducibility; bump when the C++ changes. + # The build runs in this module dir and drops a `result` symlink here (gitignored). cwd: str | None = str(Path(__file__).resolve().parent) executable: str = "result/bin/pgo" build_command: str | None = ( - 'nix build "github:jeff-hykin/gsc_pgo/e8d910f305719b81efabc6b8e1035ee7020369e9#default"' - " --no-write-lock-file" + 'nix build "github:jeff-hykin/gsc_pgo/v1.0.0#default" --no-write-lock-file' ) frame_id: str = "map" From 49ba7ee618f3281d5825faea7a55cc3afe1ab82d Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 11:09:07 -0700 Subject: [PATCH 60/69] feat(jnav): add map field to LocationConstraint Scopes the "to" location variable to a named map so a constraint can close a loop across maps (cross-map closure). Empty = current/default map. Appended to the wire string group; roundtrip + defaults tests cover it. --- dimos/navigation/jnav/msgs/LocationConstraint.py | 13 ++++++++++--- .../navigation/jnav/msgs/test_LocationConstraint.py | 3 +++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/dimos/navigation/jnav/msgs/LocationConstraint.py b/dimos/navigation/jnav/msgs/LocationConstraint.py index dc6a9c1e7a..68162fad13 100644 --- a/dimos/navigation/jnav/msgs/LocationConstraint.py +++ b/dimos/navigation/jnav/msgs/LocationConstraint.py @@ -36,6 +36,9 @@ constraint reusing the same ``constraint_instance_id`` removes the committed factors carrying it, letting an external estimator do rolling outlier removal / revision (e.g. as a tag/GPS lock improves). +- ``map`` names the map the ``to`` location belongs to. Empty means the current / + default map; a non-empty value scopes the location variable to another map so a + constraint can close a loop across maps (cross-map closure). """ from __future__ import annotations @@ -68,6 +71,7 @@ class LocationConstraint(Timestamped): to_id: str # the BetweenFactor "to" (location variable id), URL-like frame_id: str # the BetweenFactor "from" frame (== body frame for now) constraint_instance_id: str # external instance id, for revision/removal + map: str # map the "to" location belongs to ("" = current/default map) pose: Pose # relative transform frame_id -> location covariance: list[float] # 6x6 row-major, tangent order [rot(3), trans(3)] @@ -78,12 +82,14 @@ def __init__( pose: Pose | None = None, covariance: list[float] | None = None, constraint_instance_id: str = "", + map: str = "", ts: float = 0.0, ) -> None: self.ts = ts if ts != 0 else time.time() self.to_id = to_id self.frame_id = frame_id self.constraint_instance_id = constraint_instance_id + self.map = map self.pose = pose if pose is not None else Pose() if covariance is None: self.covariance = _identity_covariance() @@ -97,7 +103,7 @@ def __init__( def lcm_encode(self) -> bytes: parts: list[bytes] = [struct.pack(">d", self.ts)] - for text in (self.to_id, self.frame_id, self.constraint_instance_id): + for text in (self.to_id, self.frame_id, self.constraint_instance_id, self.map): encoded = text.encode("utf-8") parts.append(struct.pack(">I", len(encoded))) parts.append(encoded) @@ -124,12 +130,12 @@ def lcm_decode(cls, data: bytes | BinaryIO) -> LocationConstraint: (ts,) = struct.unpack_from(">d", buf, offset) offset += 8 texts: list[str] = [] - for _ in range(3): + for _ in range(4): (length,) = struct.unpack_from(">I", buf, offset) offset += 4 texts.append(buf[offset : offset + length].decode("utf-8")) offset += length - to_id, frame_id, constraint_instance_id = texts + to_id, frame_id, constraint_instance_id, map_name = texts px, py, pz, qx, qy, qz, qw = struct.unpack_from(">7d", buf, offset) offset += 56 pose = Pose() @@ -143,5 +149,6 @@ def lcm_decode(cls, data: bytes | BinaryIO) -> LocationConstraint: pose=pose, covariance=covariance, constraint_instance_id=constraint_instance_id, + map=map_name, ts=ts, ) diff --git a/dimos/navigation/jnav/msgs/test_LocationConstraint.py b/dimos/navigation/jnav/msgs/test_LocationConstraint.py index 08e7b714c6..5b92312176 100644 --- a/dimos/navigation/jnav/msgs/test_LocationConstraint.py +++ b/dimos/navigation/jnav/msgs/test_LocationConstraint.py @@ -45,6 +45,7 @@ def test_roundtrip_preserves_all_fields() -> None: pose=_pose(1.5, -2.0, 0.3), covariance=cov, constraint_instance_id="tag5#42", + map="hk_village", ts=1781565207.5, ) decoded = LocationConstraint.lcm_decode(constraint.lcm_encode()) @@ -52,6 +53,7 @@ def test_roundtrip_preserves_all_fields() -> None: assert decoded.to_id == "apriltag://36h11/40cm/5" assert decoded.frame_id == "base_link" assert decoded.constraint_instance_id == "tag5#42" + assert decoded.map == "hk_village" assert decoded.ts == constraint.ts assert decoded.pose.position.x == 1.5 assert decoded.pose.position.y == -2.0 @@ -64,6 +66,7 @@ def test_defaults() -> None: assert constraint.to_id == "" assert constraint.frame_id == "" assert constraint.constraint_instance_id == "" + assert constraint.map == "" assert constraint.ts > 0 # auto-stamped # Default covariance is a non-degenerate identity (unit variance per DOF). assert constraint.covariance[0] == 1.0 and constraint.covariance[35] == 1.0 From 36798c2cbf5a0184103fc0d32f81f7b250fd3ed2 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 11:42:56 -0700 Subject: [PATCH 61/69] chore(jnav): bump gsc_pgo pin to v1.1.0 (LocationConstraint map field) --- dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py index 47fe9759f6..5636167d86 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py @@ -45,7 +45,7 @@ class PGOConfig(NativeModuleConfig): cwd: str | None = str(Path(__file__).resolve().parent) executable: str = "result/bin/pgo" build_command: str | None = ( - 'nix build "github:jeff-hykin/gsc_pgo/v1.0.0#default" --no-write-lock-file' + 'nix build "github:jeff-hykin/gsc_pgo/v1.1.0#default" --no-write-lock-file' ) frame_id: str = "map" From 1dbdd62b5ee0365f94beede45500368e0f345f2b Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 12:33:23 -0700 Subject: [PATCH 62/69] refactor(jnav): consolidate LocationConstraint map field to map_id + kind Supersedes the earlier single `map` string (added as a 4th field inside the existing string group) with the consolidated design from jeff/feat/repulsive-field-local-planner: map_id (multi-map frame system) and kind (coarse category, defaults to the to_id URL scheme) ride as tolerant TAIL fields after the covariance, so pre-merge payloads and the native gsc_pgo fixed-sequence decoder both stay compatible. Also adds from_confidence()/confidence and the (1-c)/c variance convention. Mirrors gsc_pgo v1.1.0 C++ decode (map_id + kind tolerant tail). --- .../jnav/msgs/LocationConstraint.py | 177 ++++++++++++++---- .../jnav/msgs/test_LocationConstraint.py | 34 +++- 2 files changed, 171 insertions(+), 40 deletions(-) diff --git a/dimos/navigation/jnav/msgs/LocationConstraint.py b/dimos/navigation/jnav/msgs/LocationConstraint.py index 68162fad13..9366a3853d 100644 --- a/dimos/navigation/jnav/msgs/LocationConstraint.py +++ b/dimos/navigation/jnav/msgs/LocationConstraint.py @@ -12,33 +12,50 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""LocationConstraint: the ready-made args for a GTSAM ``BetweenFactor``. - -A LocationConstraint is a relative-pose measurement from a ``from`` graph node to -a ``to`` location variable, plus the 6x6 covariance that becomes the factor's -noise model directly. The PGO turns each one into its own pose node (placed from -interpolated odometry at ``ts``) and a ``BetweenFactor(node, location)``. - -Field meanings (mapping onto BetweenFactor): -- ``to_id`` is the BetweenFactor "to": the location variable's identity, URL-like - (e.g. ``apriltag://36h11/40cm/5`` or ``gps://fix`` or ``ui_click://...``). Two - constraints sharing a ``to_id`` observe the same graph variable, which closes - the loop. -- ``frame_id`` is the BetweenFactor "from": the frame the ``pose`` is expressed - in. For now the PGO enforces ``frame_id == body_frame`` (no full C++ tf yet), - and re-bases the measurement onto the node it creates. -- ``pose`` is the relative transform ``frame_id -> location``. -- ``covariance`` is the 6x6 measurement covariance in GTSAM Pose3 tangent order - ``[rot(3), trans(3)]`` (row-major, 36 values). It is used as the factor's noise - model directly — degenerate DOFs (e.g. a position-only fix) get a huge variance - on the rotation block. -- ``constraint_instance_id`` identifies this specific external instance. A later - constraint reusing the same ``constraint_instance_id`` removes the committed - factors carrying it, letting an external estimator do rolling outlier removal / - revision (e.g. as a tag/GPS lock improves). -- ``map`` names the map the ``to`` location belongs to. Empty means the current / - default map; a non-empty value scopes the location variable to another map so a - constraint can close a loop across maps (cross-map closure). +"""LocationConstraint: THE constraint message — factor injection AND map bridging. + +The consolidation of the old ``MapConstraint`` (multi-map bridging) into the +GTSAM-shaped constraint (Jeff): one message serves both consumers. + +- The PGO turns one into its own pose node (placed from interpolated odometry at + ``ts``) and a ``BetweenFactor(node, location)`` whose noise model is + ``covariance`` directly. +- MultiMap bridges maps through them: two maps observing the same ``to_id`` are + bridgeable (the shared location IS the bridge anchor), with ``map_id`` naming + which map's frame system each observation lives in. + +Field meanings: +- ``to_id``: the location variable's identity — URL-like or a random UUID + (e.g. ``apriltag://36h11/40cm/5``, ``reloc://map0/dim_city``, ``gps://fix``). + The URL form encodes the exact source so identical tag numbers from different + families/sizes can't false-bridge. Constraints sharing a ``to_id`` observe the + same variable — which is what closes loops and bridges maps. +- ``map_id``: which map's frame system this observation is in (multi-map + bridging). Single-graph consumers (the PGO) ignore it; ``""`` is fine there. +- ``frame_id``: the "from" — the OBSERVATION frame the ``pose`` is relative to + (camera frame for a tag, odom frame for a reloc fix), NOT the map root, so the + uncertainty stays honest. MultiMap resolves it to the map root via the map's + recorded tf at ``ts``; the PGO currently enforces ``frame_id == body_frame``. +- ``pose``: the relative transform ``frame_id -> location``. +- ``covariance``: 6x6 measurement covariance in GTSAM Pose3 tangent order + ``[rot(3), trans(3)]`` (row-major, 36 values). Degenerate DOFs (e.g. a + position-only fix) get a huge variance on the rotation block. For consumers + that want a coarse scalar, ``.confidence`` derives one (see below); producers + with only a scalar use :meth:`from_confidence`. +- ``constraint_instance_id``: identifies this specific external instance. A + later constraint reusing the same instance id REPLACES the earlier one + (rolling revision — e.g. a perceiver refining a tag lock); a fresh id per + event means additive observations (e.g. successive reloc fixes, which + averaging consumers want to keep). This subsumes the old ``replacement`` + time-window mechanism. +- ``kind``: optional coarse category ("apriltag"/"reloc"/"ui_click"/...). When + left empty it defaults to the ``to_id`` URL scheme (``reloc://...`` -> + ``"reloc"``), so URL-style producers get it for free. + +Confidence <-> covariance convention (shared by every producer/consumer that +thinks in ``[0, 1]`` scalars): per-axis variance ``= (1 - c) / c`` — c=1 is a +perfect measurement (variance 0), c→0 is worthless (variance → inf) — and the +derived scalar is ``1 / (1 + mean(diagonal))``, which round-trips a uniform c. """ from __future__ import annotations @@ -54,6 +71,9 @@ # 6x6 covariance, row-major. _COVARIANCE_LENGTH = 36 +# Diagonal tangent order [rot(3), trans(3)]: per-axis confidence names -> index. +_AXIS_ORDER = ("roll", "pitch", "yaw", "x", "y", "z") +_MIN_CONFIDENCE = 1e-6 def _identity_covariance() -> list[float]: @@ -64,14 +84,26 @@ def _identity_covariance() -> list[float]: return cov +def confidence_to_variance(confidence: float) -> float: + """The shared [0,1]-confidence -> variance convention (see module docstring).""" + clamped = min(max(float(confidence), _MIN_CONFIDENCE), 1.0) + return (1.0 - clamped) / clamped + + +def variance_to_confidence(variance: float) -> float: + """Inverse of :func:`confidence_to_variance` (for the derived scalar).""" + return 1.0 / (1.0 + max(float(variance), 0.0)) + + class LocationConstraint(Timestamped): msg_name = "jnav.LocationConstraint" ts: float - to_id: str # the BetweenFactor "to" (location variable id), URL-like - frame_id: str # the BetweenFactor "from" frame (== body frame for now) - constraint_instance_id: str # external instance id, for revision/removal - map: str # map the "to" location belongs to ("" = current/default map) + to_id: str # location variable id — URL-like or UUID; the bridge/loop key + map_id: str # whose frame system (multi-map); "" for single-graph consumers + frame_id: str # the "from": the observation frame the pose is relative to + kind: str # optional category; defaults to the to_id URL scheme + constraint_instance_id: str # same id -> replaces the earlier instance pose: Pose # relative transform frame_id -> location covariance: list[float] # 6x6 row-major, tangent order [rot(3), trans(3)] @@ -82,14 +114,17 @@ def __init__( pose: Pose | None = None, covariance: list[float] | None = None, constraint_instance_id: str = "", - map: str = "", + map_id: str = "", + kind: str = "", ts: float = 0.0, ) -> None: self.ts = ts if ts != 0 else time.time() self.to_id = to_id + self.map_id = map_id + scheme, separator, _ = to_id.partition("://") + self.kind = kind or (scheme if separator else "") self.frame_id = frame_id self.constraint_instance_id = constraint_instance_id - self.map = map self.pose = pose if pose is not None else Pose() if covariance is None: self.covariance = _identity_covariance() @@ -101,9 +136,59 @@ def __init__( ) self.covariance = list(covariance) + # ---- scalar-confidence bridge (MultiMap gating, simple producers) --------- + + @classmethod + def from_confidence( + cls, + to_id: str = "", + frame_id: str = "", + pose: Pose | None = None, + confidence: float = 1.0, + map_id: str = "", + kind: str = "", + constraint_instance_id: str = "", + ts: float = 0.0, + **per_axis: float, + ) -> LocationConstraint: + """Build with a diagonal covariance from [0,1] confidence(s). + + ``confidence`` is the coarse overall scalar; per-axis keywords + (``roll``/``pitch``/``yaw``/``x``/``y``/``z``) override individual DOFs + (e.g. ``z=0.01`` for a fix that barely constrains height). + """ + unknown = set(per_axis) - set(_AXIS_ORDER) + if unknown: + raise ValueError(f"unknown per-axis confidence(s): {sorted(unknown)}") + cov = [0.0] * _COVARIANCE_LENGTH + for index, axis in enumerate(_AXIS_ORDER): + cov[index * 6 + index] = confidence_to_variance(per_axis.get(axis, confidence)) + return cls( + to_id=to_id, + frame_id=frame_id, + pose=pose, + covariance=cov, + constraint_instance_id=constraint_instance_id, + map_id=map_id, + kind=kind, + ts=ts, + ) + + @property + def confidence(self) -> float: + """Coarse [0,1] scalar derived from the covariance diagonal (mean variance). + + Round-trips a uniform :meth:`from_confidence`; consumers with [0,1] + thresholds (MultiMap's ``marker_min_confidence``) gate on this. + """ + diagonal = [self.covariance[axis * 6 + axis] for axis in range(6)] + return variance_to_confidence(sum(diagonal) / len(diagonal)) + + # ---- wire format ----------------------------------------------------------- + def lcm_encode(self) -> bytes: parts: list[bytes] = [struct.pack(">d", self.ts)] - for text in (self.to_id, self.frame_id, self.constraint_instance_id, self.map): + for text in (self.to_id, self.frame_id, self.constraint_instance_id): encoded = text.encode("utf-8") parts.append(struct.pack(">I", len(encoded))) parts.append(encoded) @@ -121,6 +206,13 @@ def lcm_encode(self) -> bytes: ) ) parts.append(struct.pack(">36d", *self.covariance)) + # map_id + kind ride at the TAIL: pre-merge payloads decode fine + # (missing -> ""), and a fixed-sequence decoder that stops after the + # covariance (the native gsc_pgo struct) tolerates the trailing bytes. + for text in (self.map_id, self.kind): + encoded = text.encode("utf-8") + parts.append(struct.pack(">I", len(encoded))) + parts.append(encoded) return b"".join(parts) @classmethod @@ -130,12 +222,12 @@ def lcm_decode(cls, data: bytes | BinaryIO) -> LocationConstraint: (ts,) = struct.unpack_from(">d", buf, offset) offset += 8 texts: list[str] = [] - for _ in range(4): + for _ in range(3): (length,) = struct.unpack_from(">I", buf, offset) offset += 4 texts.append(buf[offset : offset + length].decode("utf-8")) offset += length - to_id, frame_id, constraint_instance_id, map_name = texts + to_id, frame_id, constraint_instance_id = texts px, py, pz, qx, qy, qz, qw = struct.unpack_from(">7d", buf, offset) offset += 56 pose = Pose() @@ -143,12 +235,23 @@ def lcm_decode(cls, data: bytes | BinaryIO) -> LocationConstraint: pose.orientation = Quaternion(qx, qy, qz, qw) covariance = list(struct.unpack_from(">36d", buf, offset)) offset += _COVARIANCE_LENGTH * 8 + tail: list[str] = [] + for _ in range(2): # map_id, kind — absent on pre-merge payloads + if offset + 4 > len(buf): + tail.append("") + continue + (length,) = struct.unpack_from(">I", buf, offset) + offset += 4 + tail.append(buf[offset : offset + length].decode("utf-8")) + offset += length + map_id, kind = tail return cls( to_id=to_id, frame_id=frame_id, pose=pose, covariance=covariance, constraint_instance_id=constraint_instance_id, - map=map_name, + map_id=map_id, + kind=kind, ts=ts, ) diff --git a/dimos/navigation/jnav/msgs/test_LocationConstraint.py b/dimos/navigation/jnav/msgs/test_LocationConstraint.py index 5b92312176..4be61f82e6 100644 --- a/dimos/navigation/jnav/msgs/test_LocationConstraint.py +++ b/dimos/navigation/jnav/msgs/test_LocationConstraint.py @@ -16,6 +16,8 @@ from __future__ import annotations +import struct + import pytest from dimos.memory2.codecs.base import codec_for @@ -45,7 +47,8 @@ def test_roundtrip_preserves_all_fields() -> None: pose=_pose(1.5, -2.0, 0.3), covariance=cov, constraint_instance_id="tag5#42", - map="hk_village", + map_id="hk_village", + kind="apriltag", ts=1781565207.5, ) decoded = LocationConstraint.lcm_decode(constraint.lcm_encode()) @@ -53,7 +56,8 @@ def test_roundtrip_preserves_all_fields() -> None: assert decoded.to_id == "apriltag://36h11/40cm/5" assert decoded.frame_id == "base_link" assert decoded.constraint_instance_id == "tag5#42" - assert decoded.map == "hk_village" + assert decoded.map_id == "hk_village" + assert decoded.kind == "apriltag" assert decoded.ts == constraint.ts assert decoded.pose.position.x == 1.5 assert decoded.pose.position.y == -2.0 @@ -66,13 +70,37 @@ def test_defaults() -> None: assert constraint.to_id == "" assert constraint.frame_id == "" assert constraint.constraint_instance_id == "" - assert constraint.map == "" + assert constraint.map_id == "" + assert constraint.kind == "" assert constraint.ts > 0 # auto-stamped # Default covariance is a non-degenerate identity (unit variance per DOF). assert constraint.covariance[0] == 1.0 and constraint.covariance[35] == 1.0 assert sum(constraint.covariance) == 6.0 +def test_kind_defaults_to_to_id_scheme() -> None: + assert LocationConstraint(to_id="reloc://map0/dim_city").kind == "reloc" + assert LocationConstraint(to_id="apriltag://36h11/40cm/5").kind == "apriltag" + # An explicit kind wins; a to_id without a URL scheme leaves kind empty. + assert LocationConstraint(to_id="gps://fix", kind="override").kind == "override" + assert LocationConstraint(to_id="bare-uuid").kind == "" + + +def test_pre_merge_payload_decodes_tail_as_empty() -> None: + """A payload written before map_id/kind existed (stops after covariance).""" + parts: list[bytes] = [struct.pack(">d", 123.0)] + for text in ("to", "frame", "instance"): + encoded = text.encode("utf-8") + parts.append(struct.pack(">I", len(encoded))) + parts.append(encoded) + parts.append(struct.pack(">7d", 0, 0, 0, 0, 0, 0, 1)) + parts.append(struct.pack(">36d", *([0.0] * 36))) + decoded = LocationConstraint.lcm_decode(b"".join(parts)) + assert decoded.to_id == "to" + assert decoded.map_id == "" + assert decoded.kind == "" + + def test_full_6x6_covariance_roundtrips_offdiagonals() -> None: cov = [float(i) for i in range(36)] # all entries distinct, incl. off-diagonal constraint = LocationConstraint(to_id="x", frame_id="base_link", covariance=cov) From 33bf602bcadbca26d51480bbdd4d6fa3ea92679f Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 12:56:52 -0700 Subject: [PATCH 63/69] style(jnav): drop section-marker comments from LocationConstraint The repo's test_no_section_markers convention check rejects the '# ----' banner comments; remove them (carried over from the source file). --- dimos/navigation/jnav/msgs/LocationConstraint.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/dimos/navigation/jnav/msgs/LocationConstraint.py b/dimos/navigation/jnav/msgs/LocationConstraint.py index 9366a3853d..36843a4109 100644 --- a/dimos/navigation/jnav/msgs/LocationConstraint.py +++ b/dimos/navigation/jnav/msgs/LocationConstraint.py @@ -136,8 +136,6 @@ def __init__( ) self.covariance = list(covariance) - # ---- scalar-confidence bridge (MultiMap gating, simple producers) --------- - @classmethod def from_confidence( cls, @@ -184,8 +182,6 @@ def confidence(self) -> float: diagonal = [self.covariance[axis * 6 + axis] for axis in range(6)] return variance_to_confidence(sum(diagonal) / len(diagonal)) - # ---- wire format ----------------------------------------------------------- - def lcm_encode(self) -> bytes: parts: list[bytes] = [struct.pack(">d", self.ts)] for text in (self.to_id, self.frame_id, self.constraint_instance_id): From 857faf5ac33ce41b695109a784d88defb9193bea Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 14:20:48 -0700 Subject: [PATCH 64/69] refactor(jnav): delegate apriltag pose math to marker_pose Drop the local make_detector/_object_points/estimate_marker_pose/ reprojection_error_px in apriltags.py in favor of the shared create_aruco_detector/estimate_marker_pose/marker_reprojection_error in perception/fiducial/marker_pose.py. --- dimos/navigation/jnav/utils/apriltags.py | 76 ++++++------------------ 1 file changed, 17 insertions(+), 59 deletions(-) diff --git a/dimos/navigation/jnav/utils/apriltags.py b/dimos/navigation/jnav/utils/apriltags.py index 347065af38..eefa55b19f 100644 --- a/dimos/navigation/jnav/utils/apriltags.py +++ b/dimos/navigation/jnav/utils/apriltags.py @@ -43,6 +43,11 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.Image import Image +from dimos.perception.fiducial.marker_pose import ( + create_aruco_detector, + estimate_marker_pose, + marker_reprojection_error, +) DEFAULT_MAX_DISTANCE_M = 1.0 DEFAULT_MAX_VIEW_ANGLE_DEG = 45.0 @@ -63,35 +68,6 @@ Detection = dict[str, Any] -def make_detector(dictionary_name: str) -> cv2.aruco.ArucoDetector: - d = cv2.aruco.getPredefinedDictionary(getattr(cv2.aruco, dictionary_name)) - return cv2.aruco.ArucoDetector(d, cv2.aruco.DetectorParameters()) - - -def _object_points(marker_length_m: float) -> np.ndarray: - h = marker_length_m / 2.0 - return np.array([[-h, h, 0.0], [h, h, 0.0], [h, -h, 0.0], [-h, -h, 0.0]], dtype=np.float32) - - -def estimate_marker_pose( - corners_pixels: np.ndarray, - marker_length_m: float, - intrinsics: np.ndarray, - distortion: np.ndarray, -) -> tuple[np.ndarray, np.ndarray] | None: - """solvePnP a single tag -> (rotation_vector, translation_vector) in the - camera_optical frame, or None if it failed.""" - image_corners = corners_pixels.reshape(4, 1, 2).astype(np.float32) - found, rotation_vector, translation_vector = cv2.solvePnP( - _object_points(marker_length_m), - image_corners, - intrinsics, - distortion, - flags=cv2.SOLVEPNP_IPPE_SQUARE, - ) - return (rotation_vector, translation_vector) if found else None - - def view_quality(t_cam_marker: list[float]) -> tuple[float, float]: """(distance_m, view_angle_deg) for a tag pose in the camera optical frame. @@ -148,24 +124,6 @@ def cluster_medoid(cluster: list[Detection], rotation_weight_m_per_rad: float) - return cluster[best_index] -def reprojection_error_px( - corners_pixels: np.ndarray, - rotation_vector: np.ndarray, - translation_vector: np.ndarray, - marker_length_m: float, - intrinsics: np.ndarray, - distortion: np.ndarray, -) -> float: - """RMS pixel distance between detected corners and the solvePnP pose reprojected - back onto the image — a direct measure of how well the pose explains the tag.""" - projected, _ = cv2.projectPoints( - _object_points(marker_length_m), rotation_vector, translation_vector, intrinsics, distortion - ) - measured = corners_pixels.reshape(4, 2).astype(np.float64) - diff = projected.reshape(4, 2) - measured - return float(np.sqrt(np.mean(np.sum(diff * diff, axis=1)))) - - def tag_pixel_size(corners_pixels: np.ndarray) -> float: """Tag side length in pixels (sqrt of the quad's image area); small = unreliable.""" quad = corners_pixels.reshape(4, 2).astype(np.float32) @@ -303,7 +261,7 @@ def detect_apriltags( far/oblique views, fast motion), cluster same-id detections by time, drop thin clusters, and (re)write the `april_tags` stream from one Huber-refined medoid representative per cluster. Returns that list of representatives.""" - detector = make_detector(dictionary) + detector = create_aruco_detector(dictionary) raw_detections: list[Detection] = [] images = store.stream(image_stream, Image).to_list() speed_by_ts, speed_available = _camera_speeds(images) @@ -336,13 +294,13 @@ def detect_apriltags( "marker_id": int(marker_id), "t_cam_marker": tag_in_camera, "sharpness": tag_sharpness(gray, corners), - "reproj_px": reprojection_error_px( + "reproj_px": marker_reprojection_error( corners, - rotation_vector, - translation_vector, marker_length, intrinsics, distortion, + rotation_vector, + translation_vector, ), "tag_px": tag_pixel_size(corners), "speed": speed_by_ts.get(float(image_obs.ts)), @@ -508,7 +466,7 @@ def pure_get_tags(img: Any, camera_intrinsics: Any) -> list[TagInfo]: if bgr.ndim == 2: bgr = cv2.cvtColor(bgr, cv2.COLOR_GRAY2BGR) - detector = make_detector(dictionary) + detector = create_aruco_detector(dictionary) all_corners, marker_ids, _ = detector.detectMarkers(bgr) if marker_ids is None: return [] @@ -530,13 +488,13 @@ def pure_get_tags(img: Any, camera_intrinsics: Any) -> list[TagInfo]: float(quaternion[2]), float(quaternion[3]), ] - reproj = reprojection_error_px( + reproj = marker_reprojection_error( corners, - rotation_vector, - translation_vector, marker_length, intrinsics, distortion, + rotation_vector, + translation_vector, ) confidence = max(0.0, 1.0 - reproj / DEFAULT_MAX_REPROJ_PX) tags.append( @@ -673,7 +631,7 @@ def detect_raw_detections( dictionary: str = "DICT_APRILTAG_36h11", ) -> tuple[list[Detection], bool, int]: """Every valid-PnP tag detection over `image_stream`, unfiltered, with gate diagnostics.""" - detector = make_detector(dictionary) + detector = create_aruco_detector(dictionary) raw_detections: list[Detection] = [] images = store.stream(image_stream, Image).to_list() speed_by_ts, speed_available = _camera_speeds(images) @@ -706,13 +664,13 @@ def detect_raw_detections( "marker_id": int(marker_id), "t_cam_marker": tag_in_camera, "sharpness": tag_sharpness(gray, corners), - "reproj_px": reprojection_error_px( + "reproj_px": marker_reprojection_error( corners, - rotation_vector, - translation_vector, marker_length, intrinsics, distortion, + rotation_vector, + translation_vector, ), "tag_px": tag_pixel_size(corners), "speed": speed_by_ts.get(float(image_obs.ts)), From cdf22c5079d8184ef46d7ebaa2c20b352782129a Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 16:38:33 -0700 Subject: [PATCH 65/69] refactor(pointlio): remove dead config values not read by the binary cube_side_length, det_range, odom_only, and timestamp_unit are parsed and forwarded as CLI args but never consumed by the compiled Point-LIO binary: cube_len/DET_RANGE/odom_only are write-only globals (the iVox port dropped the fov-segment logic that used them) and preprocess.cpp hardcodes the ms timestamp scale, ignoring time_unit. Drop them from PointLioConfig, main.cpp arg parsing, and the pcap_to_db CLI. --- dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp | 6 ------ dimos/hardware/sensors/lidar/pointlio/module.py | 5 ----- .../hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py | 8 -------- 3 files changed, 19 deletions(-) diff --git a/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp b/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp index 91c5c53e0c..140a761dea 100644 --- a/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp +++ b/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp @@ -325,9 +325,6 @@ int main(int argc, char** argv) { std::string lidar_type = mod.arg("lidar_type", "avia"); params.lidar_type = lidar_type == "velodyne" ? 2 : lidar_type == "ouster" ? 3 : lidar_type == "hesai" ? 4 : lidar_type == "unilidar" ? 5 : 1; - std::string ts_unit = mod.arg("timestamp_unit", "nanosecond"); - params.timestamp_unit = ts_unit == "second" ? 0 : ts_unit == "millisecond" ? 1 : - ts_unit == "microsecond" ? 2 : 3; // mapping params.use_imu_as_input = mod.arg_bool("use_imu_as_input", params.use_imu_as_input); params.prop_at_freq_of_imu = mod.arg_bool("prop_at_freq_of_imu", params.prop_at_freq_of_imu); @@ -344,8 +341,6 @@ int main(int argc, char** argv) { std::string ivox_nearby = mod.arg("ivox_nearby_type", "nearby6"); params.ivox_nearby_type = ivox_nearby == "center" ? 0 : ivox_nearby == "nearby18" ? 18 : ivox_nearby == "nearby26" ? 26 : 6; - params.cube_side_length = mod.arg_float("cube_side_length", params.cube_side_length); - params.det_range = mod.arg_float("det_range", params.det_range); params.fov_degree = mod.arg_float("fov_degree", params.fov_degree); params.imu_en = mod.arg_bool("imu_en", params.imu_en); params.start_in_aggressive_motion = mod.arg_bool("start_in_aggressive_motion", params.start_in_aggressive_motion); @@ -370,7 +365,6 @@ int main(int argc, char** argv) { // odometry params.publish_odometry_without_downsample = mod.arg_bool("publish_odometry_without_downsample", params.publish_odometry_without_downsample); - params.odom_only = mod.arg_bool("odom_only", params.odom_only); // Point-LIO internal processing rates double msr_freq = mod.arg_float("msr_freq", 50.0f); diff --git a/dimos/hardware/sensors/lidar/pointlio/module.py b/dimos/hardware/sensors/lidar/pointlio/module.py index 352768254b..67aff4ddce 100644 --- a/dimos/hardware/sensors/lidar/pointlio/module.py +++ b/dimos/hardware/sensors/lidar/pointlio/module.py @@ -66,7 +66,6 @@ # LID_TYPE enum (Point-LIO src/preprocess.h). avia = 1 selects the Livox branch # the Mid-360 emits. LidarType = Literal["avia", "velodyne", "ouster", "hesai", "unilidar"] -TimestampUnit = Literal["second", "millisecond", "microsecond", "nanosecond"] # iVox local-map neighbour stencil. IvoxNearbyType = Literal["center", "nearby6", "nearby18", "nearby26"] @@ -106,7 +105,6 @@ class PointLioConfig(NativeModuleConfig): lidar_type: LidarType = "avia" # 1 = AVIA (Livox) branch the Mid-360 emits scan_line: int = 4 scan_rate: int = 10 - timestamp_unit: TimestampUnit = "nanosecond" blind: float = 0.5 # spherical min range (m) point_filter_num: int = 3 # pre-KF decimation: keep every Nth raw point (1 = all) # mapping @@ -123,8 +121,6 @@ class PointLioConfig(NativeModuleConfig): filter_size_map: float = 0.5 ivox_grid_resolution: float = 2.0 # iVox local-map grid (m) ivox_nearby_type: IvoxNearbyType = "nearby6" - cube_side_length: float = 1000.0 - det_range: float = 100.0 fov_degree: float = 360.0 imu_en: bool = True start_in_aggressive_motion: bool = False @@ -150,7 +146,6 @@ class PointLioConfig(NativeModuleConfig): ) # odometry publish_odometry_without_downsample: bool = False - odom_only: bool = False # SDK port configuration (see livox/ports.py for defaults) cmd_data_port: int = SDK_CMD_DATA_PORT diff --git a/dimos/hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py b/dimos/hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py index e362a34461..d5020aa4ab 100644 --- a/dimos/hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py +++ b/dimos/hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py @@ -86,11 +86,6 @@ ), ("scan_line", "int", "number of scan lines"), ("scan_rate", "int", "scan rate (Hz)"), - ( - "timestamp_unit", - ("second", "millisecond", "microsecond", "nanosecond"), - "per-point timestamp unit", - ), ("blind", "float", "spherical min range (m); nearer points dropped"), ("point_filter_num", "int", "keep every Nth raw point (1 = all)"), # mapping @@ -111,8 +106,6 @@ ("center", "nearby6", "nearby18", "nearby26"), "iVox neighbour stencil", ), - ("cube_side_length", "float", "map cube side length (m)"), - ("det_range", "float", "max detection range (m)"), ("fov_degree", "float", "horizontal FOV (deg)"), ("imu_en", "bool", "use the IMU"), ("start_in_aggressive_motion", "bool", "skip the static IMU-init assumption"), @@ -136,7 +129,6 @@ ("extrinsic_r", "vec", "IMU->lidar rotation: 9 values row-major"), # odometry ("publish_odometry_without_downsample", "bool", "publish odom per scan, no downsample"), - ("odom_only", "bool", "odometry only, skip map publishing"), ) From 51f5d567d9bc3f3a3577fc788e19fe7997c52cb5 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 16:44:46 -0700 Subject: [PATCH 66/69] jetson recording fix --- pyproject.toml | 4 +++- uv.lock | 14 +++++++------- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c7a043562d..6e324c79bd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -139,7 +139,9 @@ dependencies = [ "toolz>=1.1.0", "protobuf>=6.33.5,<7", "psutil>=7.0.0", - "sqlite-vec>=0.1.6", + # >=0.1.7: the 0.1.6 linux-aarch64 wheel shipped a 32-bit armv7 vec0.so + # (wrong ELF class), which fails to load on aarch64. Fixed from 0.1.7 on. + "sqlite-vec>=0.1.7", "lz4>=4.4.5", "bleak>=3.0.2", "cryptography>=46.0.5", diff --git a/uv.lock b/uv.lock index 493e30ce45..b638be5215 100644 --- a/uv.lock +++ b/uv.lock @@ -2044,7 +2044,7 @@ requires-dist = [ { name = "sortedcontainers", specifier = "==2.4.0" }, { name = "sounddevice", marker = "extra == 'agents'" }, { name = "soundfile", marker = "extra == 'web'" }, - { name = "sqlite-vec", specifier = ">=0.1.6" }, + { name = "sqlite-vec", specifier = ">=0.1.7" }, { name = "sse-starlette", marker = "extra == 'web'", specifier = ">=2.2.1" }, { name = "structlog", specifier = ">=25.5.0,<26" }, { name = "tensorboard", marker = "extra == 'misc'", specifier = "==2.20.0" }, @@ -8154,14 +8154,14 @@ wheels = [ [[package]] name = "sqlite-vec" -version = "0.1.6" +version = "0.1.9" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/88/ed/aabc328f29ee6814033d008ec43e44f2c595447d9cccd5f2aabe60df2933/sqlite_vec-0.1.6-py3-none-macosx_10_6_x86_64.whl", hash = "sha256:77491bcaa6d496f2acb5cc0d0ff0b8964434f141523c121e313f9a7d8088dee3", size = 164075, upload-time = "2024-11-20T16:40:29.847Z" }, - { url = "https://files.pythonhosted.org/packages/a7/57/05604e509a129b22e303758bfa062c19afb020557d5e19b008c64016704e/sqlite_vec-0.1.6-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fdca35f7ee3243668a055255d4dee4dea7eed5a06da8cad409f89facf4595361", size = 165242, upload-time = "2024-11-20T16:40:31.206Z" }, - { url = "https://files.pythonhosted.org/packages/f2/48/dbb2cc4e5bad88c89c7bb296e2d0a8df58aab9edc75853728c361eefc24f/sqlite_vec-0.1.6-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7b0519d9cd96164cd2e08e8eed225197f9cd2f0be82cb04567692a0a4be02da3", size = 103704, upload-time = "2024-11-20T16:40:33.729Z" }, - { url = "https://files.pythonhosted.org/packages/80/76/97f33b1a2446f6ae55e59b33869bed4eafaf59b7f4c662c8d9491b6a714a/sqlite_vec-0.1.6-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux1_x86_64.whl", hash = "sha256:823b0493add80d7fe82ab0fe25df7c0703f4752941aee1c7b2b02cec9656cb24", size = 151556, upload-time = "2024-11-20T16:40:35.387Z" }, - { url = "https://files.pythonhosted.org/packages/6a/98/e8bc58b178266eae2fcf4c9c7a8303a8d41164d781b32d71097924a6bebe/sqlite_vec-0.1.6-py3-none-win_amd64.whl", hash = "sha256:c65bcfd90fa2f41f9000052bcb8bb75d38240b2dae49225389eca6c3136d3f0c", size = 281540, upload-time = "2024-11-20T16:40:37.296Z" }, + { url = "https://files.pythonhosted.org/packages/68/85/9fad0045d8e7c8df3e0fa5a56c630e8e15ad6e5ca2e6106fceb666aa6638/sqlite_vec-0.1.9-py3-none-macosx_10_6_x86_64.whl", hash = "sha256:1b62a7f0a060d9475575d4e599bbf94a13d85af896bc1ce86ee80d1b5b48e5fb", size = 131171, upload-time = "2026-03-31T08:02:31.717Z" }, + { url = "https://files.pythonhosted.org/packages/a4/3d/3677e0cd2f92e5ebc43cd29fbf565b75582bff1ccfa0b8327c7508e1084f/sqlite_vec-0.1.9-py3-none-macosx_11_0_arm64.whl", hash = "sha256:1d52e30513bae4cc9778ddbf6145610434081be4c3afe57cd877893bad9f6b6c", size = 165434, upload-time = "2026-03-31T08:02:32.712Z" }, + { url = "https://files.pythonhosted.org/packages/00/d4/f2b936d3bdc38eadcbd2a87875815db36430fab0363182ba5d12cd8e0b51/sqlite_vec-0.1.9-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4e921e592f24a5f9a18f590b6ddd530eb637e2d474e3b1972f9bbeb773aa3cb9", size = 160076, upload-time = "2026-03-31T08:02:33.796Z" }, + { url = "https://files.pythonhosted.org/packages/6f/ad/6afd073b0f817b3e03f9e37ad626ae341805891f23c74b5292818f49ac63/sqlite_vec-0.1.9-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux1_x86_64.whl", hash = "sha256:1515727990b49e79bcaf75fdee2ffc7d461f8b66905013231251f1c8938e7786", size = 163388, upload-time = "2026-03-31T08:02:34.888Z" }, + { url = "https://files.pythonhosted.org/packages/42/89/81b2907cda14e566b9bf215e2ad82fc9b349edf07d2010756ffdb902f328/sqlite_vec-0.1.9-py3-none-win_amd64.whl", hash = "sha256:4a28dc12fa4b53d7b1dced22da2488fade444e96b5d16fd2d698cd670675cf32", size = 292804, upload-time = "2026-03-31T08:02:36.035Z" }, ] [[package]] From f263f6167158a206e4bd2fcab09e44472fc3178d Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 16:59:27 -0700 Subject: [PATCH 67/69] Revert "refactor(pointlio): remove dead config values not read by the binary" This reverts commit cdf22c5079d8184ef46d7ebaa2c20b352782129a. --- dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp | 6 ++++++ dimos/hardware/sensors/lidar/pointlio/module.py | 5 +++++ .../hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py | 8 ++++++++ 3 files changed, 19 insertions(+) diff --git a/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp b/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp index 140a761dea..91c5c53e0c 100644 --- a/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp +++ b/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp @@ -325,6 +325,9 @@ int main(int argc, char** argv) { std::string lidar_type = mod.arg("lidar_type", "avia"); params.lidar_type = lidar_type == "velodyne" ? 2 : lidar_type == "ouster" ? 3 : lidar_type == "hesai" ? 4 : lidar_type == "unilidar" ? 5 : 1; + std::string ts_unit = mod.arg("timestamp_unit", "nanosecond"); + params.timestamp_unit = ts_unit == "second" ? 0 : ts_unit == "millisecond" ? 1 : + ts_unit == "microsecond" ? 2 : 3; // mapping params.use_imu_as_input = mod.arg_bool("use_imu_as_input", params.use_imu_as_input); params.prop_at_freq_of_imu = mod.arg_bool("prop_at_freq_of_imu", params.prop_at_freq_of_imu); @@ -341,6 +344,8 @@ int main(int argc, char** argv) { std::string ivox_nearby = mod.arg("ivox_nearby_type", "nearby6"); params.ivox_nearby_type = ivox_nearby == "center" ? 0 : ivox_nearby == "nearby18" ? 18 : ivox_nearby == "nearby26" ? 26 : 6; + params.cube_side_length = mod.arg_float("cube_side_length", params.cube_side_length); + params.det_range = mod.arg_float("det_range", params.det_range); params.fov_degree = mod.arg_float("fov_degree", params.fov_degree); params.imu_en = mod.arg_bool("imu_en", params.imu_en); params.start_in_aggressive_motion = mod.arg_bool("start_in_aggressive_motion", params.start_in_aggressive_motion); @@ -365,6 +370,7 @@ int main(int argc, char** argv) { // odometry params.publish_odometry_without_downsample = mod.arg_bool("publish_odometry_without_downsample", params.publish_odometry_without_downsample); + params.odom_only = mod.arg_bool("odom_only", params.odom_only); // Point-LIO internal processing rates double msr_freq = mod.arg_float("msr_freq", 50.0f); diff --git a/dimos/hardware/sensors/lidar/pointlio/module.py b/dimos/hardware/sensors/lidar/pointlio/module.py index 67aff4ddce..352768254b 100644 --- a/dimos/hardware/sensors/lidar/pointlio/module.py +++ b/dimos/hardware/sensors/lidar/pointlio/module.py @@ -66,6 +66,7 @@ # LID_TYPE enum (Point-LIO src/preprocess.h). avia = 1 selects the Livox branch # the Mid-360 emits. LidarType = Literal["avia", "velodyne", "ouster", "hesai", "unilidar"] +TimestampUnit = Literal["second", "millisecond", "microsecond", "nanosecond"] # iVox local-map neighbour stencil. IvoxNearbyType = Literal["center", "nearby6", "nearby18", "nearby26"] @@ -105,6 +106,7 @@ class PointLioConfig(NativeModuleConfig): lidar_type: LidarType = "avia" # 1 = AVIA (Livox) branch the Mid-360 emits scan_line: int = 4 scan_rate: int = 10 + timestamp_unit: TimestampUnit = "nanosecond" blind: float = 0.5 # spherical min range (m) point_filter_num: int = 3 # pre-KF decimation: keep every Nth raw point (1 = all) # mapping @@ -121,6 +123,8 @@ class PointLioConfig(NativeModuleConfig): filter_size_map: float = 0.5 ivox_grid_resolution: float = 2.0 # iVox local-map grid (m) ivox_nearby_type: IvoxNearbyType = "nearby6" + cube_side_length: float = 1000.0 + det_range: float = 100.0 fov_degree: float = 360.0 imu_en: bool = True start_in_aggressive_motion: bool = False @@ -146,6 +150,7 @@ class PointLioConfig(NativeModuleConfig): ) # odometry publish_odometry_without_downsample: bool = False + odom_only: bool = False # SDK port configuration (see livox/ports.py for defaults) cmd_data_port: int = SDK_CMD_DATA_PORT diff --git a/dimos/hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py b/dimos/hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py index d5020aa4ab..e362a34461 100644 --- a/dimos/hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py +++ b/dimos/hardware/sensors/lidar/pointlio/scripts/pcap_to_db.py @@ -86,6 +86,11 @@ ), ("scan_line", "int", "number of scan lines"), ("scan_rate", "int", "scan rate (Hz)"), + ( + "timestamp_unit", + ("second", "millisecond", "microsecond", "nanosecond"), + "per-point timestamp unit", + ), ("blind", "float", "spherical min range (m); nearer points dropped"), ("point_filter_num", "int", "keep every Nth raw point (1 = all)"), # mapping @@ -106,6 +111,8 @@ ("center", "nearby6", "nearby18", "nearby26"), "iVox neighbour stencil", ), + ("cube_side_length", "float", "map cube side length (m)"), + ("det_range", "float", "max detection range (m)"), ("fov_degree", "float", "horizontal FOV (deg)"), ("imu_en", "bool", "use the IMU"), ("start_in_aggressive_motion", "bool", "skip the static IMU-init assumption"), @@ -129,6 +136,7 @@ ("extrinsic_r", "vec", "IMU->lidar rotation: 9 values row-major"), # odometry ("publish_odometry_without_downsample", "bool", "publish odom per scan, no downsample"), + ("odom_only", "bool", "odometry only, skip map publishing"), ) From f27dc6e3a573e5bf2383383c9e867ef1f721b55f Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 18:51:34 -0700 Subject: [PATCH 68/69] refactor(jnav): merge apriltag_agreement into apriltags The agreement metric only needs numpy plus the local trajectory_metrics helper, nothing apriltags.py doesn't already import, so fold it in and drop the separate module. Repoints the three importers and renames the test to test_apriltags.py. --- .../jnav/components/loop_closure/eval.py | 6 +- .../loop_closure/gsc_pgo/scripts/add_april.py | 6 +- .../jnav/utils/apriltag_agreement.py | 156 ------------------ dimos/navigation/jnav/utils/apriltags.py | 133 +++++++++++++++ ...priltag_agreement.py => test_apriltags.py} | 0 5 files changed, 137 insertions(+), 164 deletions(-) delete mode 100644 dimos/navigation/jnav/utils/apriltag_agreement.py rename dimos/navigation/jnav/utils/{test_apriltag_agreement.py => test_apriltags.py} (100%) diff --git a/dimos/navigation/jnav/components/loop_closure/eval.py b/dimos/navigation/jnav/components/loop_closure/eval.py index d3ea0e7077..3df36ad448 100644 --- a/dimos/navigation/jnav/components/loop_closure/eval.py +++ b/dimos/navigation/jnav/components/loop_closure/eval.py @@ -74,17 +74,15 @@ from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.navigation.jnav.msgs.Graph3D import Graph3D from dimos.navigation.jnav.msgs.GraphDelta3D import GraphDelta3D -from dimos.navigation.jnav.utils.apriltag_agreement import ( +from dimos.navigation.jnav.utils.apriltags import ( VISIT_GAP_S, AgreementReport, agreement_improvement, agreement_report, - paired_tag_visit_positions, -) -from dimos.navigation.jnav.utils.apriltags import ( detect_apriltags, load_intrinsics_json, load_or_detect_sightings, + paired_tag_visit_positions, ) from dimos.navigation.jnav.utils.module_loading import ( filter_config_for_module, diff --git a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py index 023c1d033b..c4f664cc03 100644 --- a/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py @@ -27,14 +27,12 @@ from typing import Any from dimos.navigation.jnav.utils import recording_db -from dimos.navigation.jnav.utils.apriltag_agreement import ( - VISIT_GAP_S, - split_visits, -) from dimos.navigation.jnav.utils.apriltags import ( + VISIT_GAP_S, ensure_april_streams, gate_params, load_intrinsics_json, + split_visits, ) RAW_STREAM = "raw_april_tags" diff --git a/dimos/navigation/jnav/utils/apriltag_agreement.py b/dimos/navigation/jnav/utils/apriltag_agreement.py deleted file mode 100644 index bab26e2305..0000000000 --- a/dimos/navigation/jnav/utils/apriltag_agreement.py +++ /dev/null @@ -1,156 +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. - -"""April-tag agreement: a drift-quality metric for a corrected trajectory. - -A fixed April tag observed many times along a trajectory should map to ONE world -position. Odometry drift scatters those estimates (the same tag lands in a -different spot each lap); a good loop-closure / PGO correction pulls them back -together. So the spread of a tag's repeated world-position estimates is a -ground-truth-free measure of trajectory consistency — and the *drop* in spread -from raw odometry to PGO-corrected odometry is how much PGO helped. - -Pure functions over plain arrays; no PGO, sim, or I/O here (the benchmark harness -in gsc_pgo/pgo_apriltag_benchmark.py feeds these from real hk_village data). -""" - -from __future__ import annotations - -from dataclasses import dataclass - -import numpy as np - -from dimos.navigation.jnav.utils.trajectory_metrics import PoseLookup - -# sightings further apart than this are treated as separate visits -VISIT_GAP_S = 20.0 - - -@dataclass(frozen=True) -class TagAgreement: - """Per-tag spread of repeated world-position estimates (metres).""" - - tag_id: int - observations: int - spread: float # RMS distance of estimates to their centroid - - -@dataclass(frozen=True) -class AgreementReport: - """Agreement across all multiply-observed tags.""" - - per_tag: tuple[TagAgreement, ...] - mean_spread: float # mean per-tag spread (metres); lower = better agreement - total_observations: int - - @property - def tag_count(self) -> int: - return len(self.per_tag) - - -def tag_spread(positions: np.ndarray) -> float: - """RMS distance of a tag's world-position estimates to their centroid.""" - if len(positions) < 2: - return 0.0 - centroid = positions.mean(axis=0) - return float(np.sqrt(np.mean(np.sum((positions - centroid) ** 2, axis=1)))) - - -def agreement_report( - tag_positions: dict[int, np.ndarray], *, min_observations: int = 2 -) -> AgreementReport: - """Score per-tag agreement from ``tag_id -> (N, 3) world positions``. - - Tags seen fewer than ``min_observations`` times carry no agreement signal and - are excluded from ``mean_spread``. - """ - per_tag: list[TagAgreement] = [] - total = 0 - for tag_id in sorted(tag_positions): - positions = np.asarray(tag_positions[tag_id], dtype=np.float64).reshape(-1, 3) - total += len(positions) - if len(positions) < min_observations: - continue - per_tag.append(TagAgreement(tag_id, len(positions), tag_spread(positions))) - mean_spread = float(np.mean([t.spread for t in per_tag])) if per_tag else 0.0 - return AgreementReport(tuple(per_tag), mean_spread, total) - - -def agreement_improvement(raw: AgreementReport, corrected: AgreementReport) -> float: - """Fractional drop in mean spread from ``raw`` to ``corrected`` (1.0 = perfect). - - Positive means the correction tightened tag agreement; negative means it made - it worse. Returns 0.0 if there's no raw spread to improve on. - """ - if raw.mean_spread <= 0.0: - return 0.0 - return (raw.mean_spread - corrected.mean_spread) / raw.mean_spread - - -def tag_world_positions( - sightings: dict[int, list[float]], pose_lookup: PoseLookup -) -> dict[int, np.ndarray]: - """Map each tag's sighting times to robot world positions (the proxy estimate).""" - positions: dict[int, np.ndarray] = {} - for tag_id, times in sightings.items(): - located = [p for p in (pose_lookup(t) for t in times) if p is not None] - if located: - positions[tag_id] = np.vstack(located) - return positions - - -def split_visits(times: list[float], *, gap_s: float) -> list[list[float]]: - """Cluster sighting timestamps into visits separated by gaps > ``gap_s``.""" - visits: list[list[float]] = [] - for timestamp in sorted(times): - if visits and timestamp - visits[-1][-1] <= gap_s: - visits[-1].append(timestamp) - else: - visits.append([timestamp]) - return visits - - -def paired_tag_visit_positions( - sightings: dict[int, list[float]], - raw_lookup: PoseLookup, - corrected_lookup: PoseLookup, - *, - gap_s: float, -) -> tuple[dict[int, np.ndarray], dict[int, np.ndarray]]: - """One robot position per tag VISIT under both pose sources, visits paired. - - Outdoors a tag stays visible while the robot walks tens of metres, so - per-sighting spread is dominated by viewing distance, not drift. Visits are - clustered on timestamps only, and a visit is kept only when BOTH pose - sources can place it — so raw and corrected reports always score the exact - same visit set, and a visit outside the pose graph's coverage drops out - instead of skewing one side. - """ - raw_positions: dict[int, np.ndarray] = {} - corrected_positions: dict[int, np.ndarray] = {} - for tag_id, times in sightings.items(): - raw_medians: list[np.ndarray] = [] - corrected_medians: list[np.ndarray] = [] - for visit_times in split_visits(times, gap_s=gap_s): - raw_located = [p for p in (raw_lookup(t) for t in visit_times) if p is not None] - corrected_located = [ - p for p in (corrected_lookup(t) for t in visit_times) if p is not None - ] - if raw_located and corrected_located: - raw_medians.append(np.median(np.vstack(raw_located), axis=0)) - corrected_medians.append(np.median(np.vstack(corrected_located), axis=0)) - if raw_medians: - raw_positions[tag_id] = np.vstack(raw_medians) - corrected_positions[tag_id] = np.vstack(corrected_medians) - return raw_positions, corrected_positions diff --git a/dimos/navigation/jnav/utils/apriltags.py b/dimos/navigation/jnav/utils/apriltags.py index eefa55b19f..bf05c1e56a 100644 --- a/dimos/navigation/jnav/utils/apriltags.py +++ b/dimos/navigation/jnav/utils/apriltags.py @@ -43,6 +43,7 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.sensor_msgs.Image import Image +from dimos.navigation.jnav.utils.trajectory_metrics import PoseLookup from dimos.perception.fiducial.marker_pose import ( create_aruco_detector, estimate_marker_pose, @@ -870,3 +871,135 @@ def ensure_april_streams( n_images, ) return detections + + +# April-tag agreement: a drift-quality metric for a corrected trajectory. +# +# A fixed April tag observed many times along a trajectory should map to ONE +# world position. Odometry drift scatters those estimates (the same tag lands in +# a different spot each lap); a good loop-closure / PGO correction pulls them +# back together. So the spread of a tag's repeated world-position estimates is a +# ground-truth-free measure of trajectory consistency — and the drop in spread +# from raw odometry to PGO-corrected odometry is how much PGO helped. + +# sightings further apart than this are treated as separate visits +VISIT_GAP_S = 20.0 + + +@dataclass(frozen=True) +class TagAgreement: + """Per-tag spread of repeated world-position estimates (metres).""" + + tag_id: int + observations: int + spread: float # RMS distance of estimates to their centroid + + +@dataclass(frozen=True) +class AgreementReport: + """Agreement across all multiply-observed tags.""" + + per_tag: tuple[TagAgreement, ...] + mean_spread: float # mean per-tag spread (metres); lower = better agreement + total_observations: int + + @property + def tag_count(self) -> int: + return len(self.per_tag) + + +def tag_spread(positions: np.ndarray) -> float: + """RMS distance of a tag's world-position estimates to their centroid.""" + if len(positions) < 2: + return 0.0 + centroid = positions.mean(axis=0) + return float(np.sqrt(np.mean(np.sum((positions - centroid) ** 2, axis=1)))) + + +def agreement_report( + tag_positions: dict[int, np.ndarray], *, min_observations: int = 2 +) -> AgreementReport: + """Score per-tag agreement from ``tag_id -> (N, 3) world positions``. + + Tags seen fewer than ``min_observations`` times carry no agreement signal and + are excluded from ``mean_spread``. + """ + per_tag: list[TagAgreement] = [] + total = 0 + for tag_id in sorted(tag_positions): + positions = np.asarray(tag_positions[tag_id], dtype=np.float64).reshape(-1, 3) + total += len(positions) + if len(positions) < min_observations: + continue + per_tag.append(TagAgreement(tag_id, len(positions), tag_spread(positions))) + mean_spread = float(np.mean([t.spread for t in per_tag])) if per_tag else 0.0 + return AgreementReport(tuple(per_tag), mean_spread, total) + + +def agreement_improvement(raw: AgreementReport, corrected: AgreementReport) -> float: + """Fractional drop in mean spread from ``raw`` to ``corrected`` (1.0 = perfect). + + Positive means the correction tightened tag agreement; negative means it made + it worse. Returns 0.0 if there's no raw spread to improve on. + """ + if raw.mean_spread <= 0.0: + return 0.0 + return (raw.mean_spread - corrected.mean_spread) / raw.mean_spread + + +def tag_world_positions( + sightings: dict[int, list[float]], pose_lookup: PoseLookup +) -> dict[int, np.ndarray]: + """Map each tag's sighting times to robot world positions (the proxy estimate).""" + positions: dict[int, np.ndarray] = {} + for tag_id, times in sightings.items(): + located = [p for p in (pose_lookup(t) for t in times) if p is not None] + if located: + positions[tag_id] = np.vstack(located) + return positions + + +def split_visits(times: list[float], *, gap_s: float) -> list[list[float]]: + """Cluster sighting timestamps into visits separated by gaps > ``gap_s``.""" + visits: list[list[float]] = [] + for timestamp in sorted(times): + if visits and timestamp - visits[-1][-1] <= gap_s: + visits[-1].append(timestamp) + else: + visits.append([timestamp]) + return visits + + +def paired_tag_visit_positions( + sightings: dict[int, list[float]], + raw_lookup: PoseLookup, + corrected_lookup: PoseLookup, + *, + gap_s: float, +) -> tuple[dict[int, np.ndarray], dict[int, np.ndarray]]: + """One robot position per tag VISIT under both pose sources, visits paired. + + Outdoors a tag stays visible while the robot walks tens of metres, so + per-sighting spread is dominated by viewing distance, not drift. Visits are + clustered on timestamps only, and a visit is kept only when BOTH pose + sources can place it — so raw and corrected reports always score the exact + same visit set, and a visit outside the pose graph's coverage drops out + instead of skewing one side. + """ + raw_positions: dict[int, np.ndarray] = {} + corrected_positions: dict[int, np.ndarray] = {} + for tag_id, times in sightings.items(): + raw_medians: list[np.ndarray] = [] + corrected_medians: list[np.ndarray] = [] + for visit_times in split_visits(times, gap_s=gap_s): + raw_located = [p for p in (raw_lookup(t) for t in visit_times) if p is not None] + corrected_located = [ + p for p in (corrected_lookup(t) for t in visit_times) if p is not None + ] + if raw_located and corrected_located: + raw_medians.append(np.median(np.vstack(raw_located), axis=0)) + corrected_medians.append(np.median(np.vstack(corrected_located), axis=0)) + if raw_medians: + raw_positions[tag_id] = np.vstack(raw_medians) + corrected_positions[tag_id] = np.vstack(corrected_medians) + return raw_positions, corrected_positions diff --git a/dimos/navigation/jnav/utils/test_apriltag_agreement.py b/dimos/navigation/jnav/utils/test_apriltags.py similarity index 100% rename from dimos/navigation/jnav/utils/test_apriltag_agreement.py rename to dimos/navigation/jnav/utils/test_apriltags.py From 6948cfa1a380e53753df0b1eab89ef18e4ee96f0 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 6 Jul 2026 21:41:20 -0700 Subject: [PATCH 69/69] fix(jnav): repoint test_apriltags import to merged module The rename commit shipped the pre-edit blob (git mv staged the original index content), so test_apriltags.py still imported the deleted apriltag_agreement module and failed collection in CI. --- dimos/navigation/jnav/utils/test_apriltags.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/navigation/jnav/utils/test_apriltags.py b/dimos/navigation/jnav/utils/test_apriltags.py index fc177b3d3d..2486bbc63d 100644 --- a/dimos/navigation/jnav/utils/test_apriltags.py +++ b/dimos/navigation/jnav/utils/test_apriltags.py @@ -18,7 +18,7 @@ import numpy as np -from dimos.navigation.jnav.utils.apriltag_agreement import ( +from dimos.navigation.jnav.utils.apriltags import ( agreement_improvement, agreement_report, tag_spread,