diff --git a/dimos/navigation/jnav/.gitignore b/dimos/navigation/jnav/.gitignore new file mode 100644 index 0000000000..189819aafe --- /dev/null +++ b/dimos/navigation/jnav/.gitignore @@ -0,0 +1,3 @@ +result +result-* +components/loop_closure/eval_results/ 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..3df36ad448 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/eval.py @@ -0,0 +1,997 @@ +# 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. + +The raw-baseline robot trajectory comes from the recording's tf (the +``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 + world position; the spread of its per-visit robot positions measures drift. + 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. + +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.apriltags import ( + VISIT_GAP_S, + AgreementReport, + agreement_improvement, + agreement_report, + 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, + 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, + store, + stream_count, +) +from dimos.navigation.jnav.utils.recording_tf import RecordingTF +from dimos.navigation.jnav.utils.trajectory_metrics import ( + GraphPose, + drifted_lookup, + graph_lookup, + has_drift, + lidar_voxel_agreement, + pose7_lookup, + trajectory_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). +# 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( + 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] + + +# 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 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. + + 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) + 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" + ) + 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 + # 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 + 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( + [ + transform.translation.x, + transform.translation.y, + transform.translation.z, + transform.rotation.x, + transform.rotation.y, + transform.rotation.z, + transform.rotation.w, + ] + ) + if not times: + 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 _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, + 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): + 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: Any = ( + 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") + + # 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, + ) + 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) + + # 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(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: + 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, raw_xyz_base, 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: + write_trajectory_rrd(rrd_path, raw_poses7[:, :3], 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, + "pose_source": "tf", + "world_frame": world_frame, + "body_frame": body_frame, + "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'", + ) + 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)", + ) + 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 = ( + [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, + world_frame=args.world_frame, + body_frame=args.body_frame, + tf_startup_tolerance_s=args.tf_startup_tolerance_s, + ) + + +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..766aeefe2d --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/eval_ground_truth_tag.py @@ -0,0 +1,316 @@ +# 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 ( + 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 ODOM_MATCH_TOLERANCE_S, store +from dimos.navigation.jnav.utils.trajectory_metrics import ( + drift_delta_lookup, + pose7_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 np.asarray(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]]] = [] + 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 + 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], + 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} + 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_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): + 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..b7dfac55a1 --- /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, +) -> 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( + 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/gsc_pgo/module.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py new file mode 100644 index 0000000000..5636167d86 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/module.py @@ -0,0 +1,212 @@ +# 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.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 +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 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/v1.1.0#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 + + # Robust (Huber) kernel on all loop factors (lidar + location). Off = original. + loop_robust_kernel: bool = False + loop_robust_huber_k: float = 1.345 + + # 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) + # 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: 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] + 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] + + @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, correction: Transform) -> None: + self.tf.publish(correction) + + @rpc + def stop(self) -> None: + super().stop() 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 new file mode 100644 index 0000000000..c4f664cc03 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/add_april.py @@ -0,0 +1,222 @@ +# 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/scripts/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 +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" +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_METERS = 0.10 +DEFAULT_DICTIONARY = "DICT_APRILTAG_36h11" + + +def _parse_dynamic(dynamic_arg: str | None) -> list[int]: + if not dynamic_arg: + return [] + return sorted({int(token) for token in dynamic_arg.replace(",", " ").split()}) + + +def _times_by_tag(store: Any, stream_name: str) -> dict[int, list[float]]: + times_by_tag: dict[int, list[float]] = {} + for observation in store.stream(stream_name): + times_by_tag.setdefault(int(observation.tags["marker_id"]), []).append( + float(observation.ts) + ) + return times_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 = _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)" + ) + tags: list[dict[str, Any]] = [] + not_revisited: list[int] = [] + for tag_id in tag_ids: + 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: + 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 "-" + 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" + 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_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, + "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_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 = 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} ===") + + 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() + intrinsics_config = load_intrinsics_json(intrinsics_path) + marker_length = ( + args.tag_size + if args.tag_size is not None + else intrinsics_config.get("marker_length", DEFAULT_MARKER_LENGTH_METERS) + ) + dictionary = args.dictionary or intrinsics_config.get("dictionary", DEFAULT_DICTIONARY) + ensure_april_streams( + store, + intrinsics_config["intrinsics"], + 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/scripts/make_rrd.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py new file mode 100644 index 0000000000..0185ffb0f1 --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/make_rrd.py @@ -0,0 +1,227 @@ +# 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/scripts/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", +): + 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 + 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( + [ + [ + 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): + odom_poses = [ + ( + observation.ts, + Pose3( + Rot3.Quaternion( + 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, + ), + ), + ) + for observation in store.stream(gt_odom) + ] + 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(tag_metrics["sharpness"]) >= GATE["s"] + and float(tag_metrics["reproj_px"]) <= GATE["r"] + and float(tag_metrics["tag_px"]) >= GATE["px"] + # 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"] + ) + and ( + float(tag_metrics["ang_speed"]) < 0 + or float(tag_metrics["ang_speed"]) <= GATE["av"] + ) + ): + continue + 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( + tag_pose.orientation.w, + tag_pose.orientation.x, + tag_pose.orientation.y, + tag_pose.orientation.z, + ), + Point3(tag_pose.x, tag_pose.y, tag_pose.z), + ) + ) + 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_path)) + 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 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( + stream_name + for stream_name in streams + if stream_name.startswith("gt_") and "_odometry" in stream_name + ) + if gt_odoms: + landmark_positions, labels = landmarks(gt_odoms[0]) + if len(landmark_positions): + rr.log( + "landmarks", + rr.Points3D(landmark_positions, colors=[255, 230, 0], radii=0.25, labels=labels), + static=True, + ) + print(f" logged {len(labels)} landmarks") + print("wrote", out_path) + return out_path + + +def _arg(flag, default=None): + return next((arg.split("=", 1)[1] for arg in sys.argv if arg.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/scripts/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/scripts/post_process.py b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py new file mode 100644 index 0000000000..f8e8222fdc --- /dev/null +++ b/dimos/navigation/jnav/components/loop_closure/gsc_pgo/scripts/post_process.py @@ -0,0 +1,825 @@ +# 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 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. + +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/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] +""" + +import json +from pathlib import Path +import re +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.geometry_msgs.PoseStamped import PoseStamped +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 ( + 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, + _write_tag_stream, + 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" + + +def arg(flag, default=""): + return next( + (item.split("=", 1)[1] for item in sys.argv if item.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) +# 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" +) # 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 +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) + +# 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/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)" + ) +REC = Path(REC_ARG).expanduser() +DB = REC / "mem2.db" +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(): + # 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, + ) + + +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 + + +_STORE_TF = RecordingTF.from_store(store) if USE_TF else None +_TF_AVAILABLE = _STORE_TF is not None + + +def world_points(observation): + """Nx3 world-registered points for a lidar observation. + + 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 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 + 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 _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 + pose = getattr(observation, "pose", None) + 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 # unknown frame, assume already world-registered + + +print(f"recording: {REC}", flush=True) +if _TF_AVAILABLE: + print( + f"world-registering {LIDAR_STREAM} via tf into {WORLD_FRAME} " + f"(per-scan 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, +) + + +def passes(detection): + return ( + 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_detections = [] +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), + marker_id=int(tags["marker_id"]), + T_cam_tag=Pose3( + Rot3.Quaternion( + pose.orientation.w, pose.orientation.x, pose.orientation.y, pose.orientation.z + ), + Point3(pose.x, pose.y, pose.z), + ), + reproj_px=float(tags["reproj_px"]), + 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 = [ + detection + for detection in raw_detections + if passes(detection) and detection["marker_id"] not in IGNORE_TAGS +] + +# keyframes from raw odometry +odom_connection = sqlite3.connect(f"file:{DB}?mode=ro", uri=True) +odom_rows = np.array( + list( + 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, +) +odom_connection.close() + + +def row_pose(row): + return Rot3.Quaternion(row[7], row[4], row[5], row[6]), np.array(row[1:4]) + + +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 + ): + keyframe_indices.append(row_index) + prev_rot, prev_pos = rot, pos +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 +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_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 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_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 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_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 +# 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 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 {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 = keyframe_poses[keyframe_index - 1] + graph.add( + BetweenFactorPose3( + keyframe_index - 1, + keyframe_index, + Pose3( + rot_prev.inverse() * rot, + Point3(rot_prev.inverse().rotate(Point3(pos - pos_prev))), + ), + odom_noise, + ) + ) +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( + 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) +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, +# 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 + + 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 + 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(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_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 + 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 + 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])), + ) + 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_index}/{len(candidate_pairs)} pairs, " + f"{accepted_count} accepted, {time.time() - icp_start_time:.0f}s", + flush=True, + ) + if first_index not in submap_clouds or second_index not in submap_clouds: + continue + 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, + initial_guess, + o3d.pipelines.registration.TransformationEstimationPointToPlane(), + ) + 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 {accepted_count}/{len(candidate_pairs)} loop closures " + f"(fit>={ICP_FIT_MIN}, rmse<={ICP_RMSE_MAX}m)", + flush=True, + ) + 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) + 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={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 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(pose): + translation = pose.translation() + quaternion = pose.rotation().toQuaternion() + return ( + translation[0], + translation[1], + translation[2], + quaternion.x(), + quaternion.y(), + quaternion.z(), + quaternion.w(), + ) + + +# 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(): + 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, 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, z, qx, qy, qz, qw), + ), + ts=ts, + pose=(x, y, z, qx, qy, qz, qw), + ) + 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"): + out_name = f"{OUT_PREFIX}_lidar{SUFFIX}" + odom_times = odom_rows[:, 0] + + def base_pose(ts): + 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 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 + # voxel averaging via the color channel. Chunks are collapsed every CHUNK scans to bound memory. + if WRITE_LCM: + import open3d as o3d + + CHUNK = 1000 + 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 + ) + 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: + 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 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_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(downsampled_points):,} pts after voxel, " + f"removing outliers...", + flush=True, + ) + 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) + ) + 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(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 " + 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) 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] diff --git a/dimos/navigation/jnav/msgs/DeformationNode.py b/dimos/navigation/jnav/msgs/DeformationNode.py new file mode 100644 index 0000000000..4238071efc --- /dev/null +++ b/dimos/navigation/jnav/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) 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/LocationConstraint.py b/dimos/navigation/jnav/msgs/LocationConstraint.py new file mode 100644 index 0000000000..36843a4109 --- /dev/null +++ b/dimos/navigation/jnav/msgs/LocationConstraint.py @@ -0,0 +1,253 @@ +# 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 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 + +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 +# 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]: + """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 + + +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 # 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)] + + def __init__( + self, + to_id: str = "", + frame_id: str = "", + pose: Pose | None = None, + covariance: list[float] | None = None, + constraint_instance_id: 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.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) + + @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)) + + 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)) + # 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 + 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 + 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_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 new file mode 100644 index 0000000000..4be61f82e6 --- /dev/null +++ b/dimos/navigation/jnav/msgs/test_LocationConstraint.py @@ -0,0 +1,124 @@ +# 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 struct + +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", + map_id="hk_village", + kind="apriltag", + 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.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 + 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.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) + 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" diff --git a/dimos/navigation/jnav/utils/apriltags.py b/dimos/navigation/jnav/utils/apriltags.py new file mode 100644 index 0000000000..bf05c1e56a --- /dev/null +++ b/dimos/navigation/jnav/utils/apriltags.py @@ -0,0 +1,1005 @@ +# 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 +from dimos.navigation.jnav.utils.trajectory_metrics import PoseLookup +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 +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 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 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 = create_aruco_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": marker_reprojection_error( + corners, + marker_length, + intrinsics, + distortion, + rotation_vector, + translation_vector, + ), + "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 = create_aruco_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 = marker_reprojection_error( + corners, + marker_length, + intrinsics, + distortion, + rotation_vector, + translation_vector, + ) + 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" + + +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 = create_aruco_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": marker_reprojection_error( + corners, + marker_length, + intrinsics, + distortion, + rotation_vector, + translation_vector, + ), + "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 + + +# 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/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..88f1aab647 --- /dev/null +++ b/dimos/navigation/jnav/utils/recording_db.py @@ -0,0 +1,118 @@ +# 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 + +import atexit +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 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() + + +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/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/navigation/jnav/utils/test_apriltags.py b/dimos/navigation/jnav/utils/test_apriltags.py new file mode 100644 index 0000000000..2486bbc63d --- /dev/null +++ b/dimos/navigation/jnav/utils/test_apriltags.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.apriltags 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 diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 61220a0fae..e9561697b1 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -121,6 +121,7 @@ "unitree-go2-memory": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_memory", "unitree-go2-mid360-record": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_mid360_record:unitree_go2_mid360_record", "unitree-go2-nav-3d": "dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_nav_3d:unitree_go2_nav_3d", + "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", @@ -184,6 +185,7 @@ "goal-relay": "dimos.navigation.nav_3d.mls_planner.goal_relay.GoalRelay", "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", @@ -195,6 +197,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", @@ -224,12 +227,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", 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") diff --git a/experimental/docs/jnav/map_postprocessing.md b/experimental/docs/jnav/map_postprocessing.md new file mode 100644 index 0000000000..8fe2e18ef9 --- /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/scripts/`. + +**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/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. (`--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/scripts/add_april.py --rec=PATH --summary +``` + +**2. Solve.** Two stages, one command: + +``` +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. +- **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/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. + +## 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/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]]