diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index 248b7204a3..10bf4395a4 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -12,12 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""NativeModule: blueprint-integrated wrapper for native (C/C++) executables. +"""NativeModule: blueprint-integrated wrapper for native executables. -A NativeModule is a thin Python Module subclass that declares In/Out ports +A NativeModule is a thin Python Module subclass that declares In/Out/IO ports for blueprint wiring but delegates all real work to a managed subprocess. -The native process receives its LCM topic names via CLI args and does -pub/sub directly on the LCM multicast bus. +The native process receives its topic names via CLI args, or as a JSON line on +stdin when ``stdin_config`` is set, and does pub/sub on them directly. Example usage:: @@ -178,15 +178,19 @@ class NativeModule(Module): """ Module that wraps a native executable as a managed subprocess. - Subclass this, declare In/Out ports, and annotate ``config`` with a + Subclass this, declare In/Out/IO ports, and annotate ``config`` with a :class:`NativeModuleConfig` subclass pointing at the executable. On ``start()``, the binary is launched with CLI args:: - -- ... + -- ... -- ... - The native process should parse these args and pub/sub on the given - LCM topics directly. On ``stop()``, the process receives SIGTERM. + Each topic is the wire channel for that port on the transport named by the + ``DIMOS_TRANSPORT`` env var. With ``stdin_config``, those same topics plus + the config and any publisher QoS also arrive as one JSON line on stdin. + + The native process should parse whichever it uses and pub/sub on the given + topics directly. On ``stop()``, the process receives SIGTERM. """ config: NativeModuleConfig @@ -471,7 +475,7 @@ def _maybe_build(self) -> None: def _collect_topics(self) -> dict[str, str]: topics: dict[str, str] = {} - for name in list(self.inputs) + list(self.outputs): + for name in list(self.inputs) + list(self.outputs) + list(self.ios): stream = getattr(self, name, None) if stream is None: continue @@ -484,9 +488,9 @@ def _collect_topics(self) -> dict[str, str]: return topics def _collect_output_qos(self) -> dict[str, dict[str, str]]: - """Publisher QoS per output channel, keyed by channel.""" + """Publisher QoS per published channel, keyed by channel.""" qos_map: dict[str, dict[str, str]] = {} - for name in self.outputs: + for name in list(self.outputs) + list(self.ios): stream = getattr(self, name, None) if stream is None: continue diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index 4a68caedc8..3425d8d7e8 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -18,6 +18,7 @@ The echo script writes received CLI args to a temp file for assertions. """ +import contextlib from io import BytesIO import json from pathlib import Path @@ -33,11 +34,13 @@ from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.native_module import LogFormat, NativeModule, NativeModuleConfig -from dimos.core.stream import In, Out -from dimos.core.transport import LCMTransport +from dimos.core.stream import IO, In, Out +from dimos.core.transport import LCMTransport, ZenohTransport from dimos.msgs.geometry_msgs.Twist import Twist from dimos.msgs.sensor_msgs.Imu import Imu from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.pubsub.impl.zenohpubsub import QOS_NEVER_DROP, Topic as ZenohTopic _ECHO = str(Path(__file__).parent / "demos" / "native_echo.py") @@ -87,6 +90,12 @@ class StubNativeModule(NativeModule): cmd_vel: In[Twist] +class StubIoModule(NativeModule): + config: StubNativeConfig + cmd_vel: In[Twist] + tf: IO[TFMessage] + + class StubConsumer(Module): pointcloud: In[PointCloud2] imu: In[Imu] @@ -160,6 +169,54 @@ def test_manual(dimos_cluster: ModuleCoordinator, args_file: str) -> None: } +def test_io_port_topic_reaches_the_native_process() -> None: + """An IO port is both a subscriber and a publisher, so it needs its topic.""" + module = StubIoModule(executable=_ECHO) + transports = [LCMTransport("/cmd_vel", Twist), LCMTransport("/tf", TFMessage)] + try: + module.set_transport("cmd_vel", transports[0]) + module.set_transport("tf", transports[1]) + + assert module._collect_topics() == { + "cmd_vel": "/cmd_vel#geometry_msgs.Twist", + "tf": "/tf#tf2_msgs.TFMessage", + } + finally: + module.stop() + for transport in transports: + with contextlib.suppress(Exception): + transport.stop() + + +def test_tf_topic_comes_from_the_declared_port_only() -> None: + """No tf port declared means no tf topic, rather than a silently injected one.""" + module = StubNativeModule(executable=_ECHO) + transport = LCMTransport("/cmd_vel", Twist) + try: + module.set_transport("cmd_vel", transport) + + assert module._collect_topics() == {"cmd_vel": "/cmd_vel#geometry_msgs.Twist"} + finally: + module.stop() + with contextlib.suppress(Exception): + transport.stop() + + +def test_io_port_publisher_qos_reaches_the_native_process() -> None: + module = StubIoModule(executable=_ECHO) + transport = ZenohTransport(ZenohTopic("/tf", TFMessage, qos=QOS_NEVER_DROP)) + try: + module.set_transport("tf", transport) + + assert module._collect_output_qos() == { + transport.channel: {"reliability": "reliable", "congestion_control": "block"}, + } + finally: + module.stop() + with contextlib.suppress(Exception): + transport.stop() + + def test_autoconnect(args_file: str) -> None: """autoconnect passes correct topic args to the native subprocess.""" blueprint = autoconnect( diff --git a/dimos/hardware/sensors/lidar/virtual_mid360/Cargo.lock b/dimos/hardware/sensors/lidar/virtual_mid360/Cargo.lock index ef9b748c0b..e71c9027a0 100644 --- a/dimos/hardware/sensors/lidar/virtual_mid360/Cargo.lock +++ b/dimos/hardware/sensors/lidar/virtual_mid360/Cargo.lock @@ -62,6 +62,15 @@ version = "1.0.103" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + [[package]] name = "arc-swap" version = "1.9.2" @@ -187,6 +196,12 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" + [[package]] name = "byteorder" version = "1.5.0" @@ -546,6 +561,8 @@ version = "0.1.0" dependencies = [ "dimos-lcm", "dimos-module-macros", + "lcm-msgs", + "nalgebra", "serde", "serde_json", "tokio", @@ -871,6 +888,30 @@ dependencies = [ "syn", ] +[[package]] +name = "glam" +version = "0.30.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19fc433e8437a212d1b6f1e68c7824af3aed907da60afa994e7f542d18d12aa9" + +[[package]] +name = "glam" +version = "0.31.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556f6b2ea90b8d15a74e0e7bb41671c9bdf38cd9f78c284d750b9ce58a2b5be7" + +[[package]] +name = "glam" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f70749695b063ecbf6b62949ccccde2e733ec3ecbbd71d467dca4e5c6c97cca0" + +[[package]] +name = "glam" +version = "0.33.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f22fb22f065b308be0d8724e3706c7fa3fc2a6c7d6899df4cad7860e7a75436" + [[package]] name = "hashbrown" version = "0.12.3" @@ -1267,6 +1308,14 @@ dependencies = [ "spin 0.9.8", ] +[[package]] +name = "lcm-msgs" +version = "0.1.0" +source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#e7c9428b7201cdfeadecd181c77c9e2d60a14503" +dependencies = [ + "byteorder", +] + [[package]] name = "libc" version = "0.2.186" @@ -1343,6 +1392,16 @@ dependencies = [ "regex-automata", ] +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.2" @@ -1376,6 +1435,37 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "nalgebra" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adc43a60c217b0c6ff46e47f26911015ad8d2e5a8be1af668c67e370d99a4346" +dependencies = [ + "approx", + "glam 0.30.10", + "glam 0.31.1", + "glam 0.32.1", + "glam 0.33.2", + "matrixmultiply", + "nalgebra-macros", + "num-complex", + "num-rational", + "num-traits", + "simba", + "typenum", +] + +[[package]] +name = "nalgebra-macros" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "973e7178a678cfd059ccec50887658d482ce16b0aa9da3888ddeab5cd5eb4889" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "nanorand" version = "0.7.0" @@ -1457,6 +1547,15 @@ dependencies = [ "zeroize", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + [[package]] name = "num-conv" version = "0.2.2" @@ -1483,6 +1582,17 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -1928,6 +2038,12 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "rcgen" version = "0.14.8" @@ -2180,6 +2296,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "safe_arch" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f7caad094bd561859bcd467734a720c3c1f5d1f338995351fefe2190c45efed" +dependencies = [ + "bytemuck", +] + [[package]] name = "same-file" version = "1.0.6" @@ -2460,6 +2585,18 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "simba" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f45c644a9f3a386f9288625d9f0c1e999e1acf07a37df35d0516c7f199d9cb2" +dependencies = [ + "approx", + "num-complex", + "num-traits", + "wide", +] + [[package]] name = "simd-adler32" version = "0.3.9" @@ -3263,6 +3400,16 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "wide" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfdfe6a32973f2d1b268b8895845a8a96cac2f0191e72c27cc929036060dbf89" +dependencies = [ + "bytemuck", + "safe_arch", +] + [[package]] name = "winapi" version = "0.3.9" diff --git a/dimos/mapping/ray_tracing/rust/Cargo.lock b/dimos/mapping/ray_tracing/rust/Cargo.lock index f68f1b28b9..a840b55719 100644 --- a/dimos/mapping/ray_tracing/rust/Cargo.lock +++ b/dimos/mapping/ray_tracing/rust/Cargo.lock @@ -567,6 +567,8 @@ version = "0.1.0" dependencies = [ "dimos-lcm", "dimos-module-macros", + "lcm-msgs", + "nalgebra", "serde", "serde_json", "tokio", diff --git a/dimos/navigation/basic_path_follower/module.py b/dimos/navigation/basic_path_follower/module.py index 5b0b4c491f..ff6ce0db7b 100644 --- a/dimos/navigation/basic_path_follower/module.py +++ b/dimos/navigation/basic_path_follower/module.py @@ -33,6 +33,8 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.nav_msgs.Path import Path +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.navigation.tf_pose import OdomBasePose from dimos.utils.logging_config import setup_logger from dimos.utils.trigonometry import angle_diff @@ -40,6 +42,7 @@ class BasicPathFollowerConfig(ModuleConfig): + base_frame: str = "base_link" speed: float = 0.5 control_frequency: float = 10.0 goal_tolerance: float = 0.3 @@ -67,6 +70,7 @@ class BasicPathFollower(Module): path: In[Path] odometry: In[Odometry] stop_movement: In[Bool] + tf: In[TFMessage] nav_cmd_vel: Out[Twist] goal_reached: Out[Bool] @@ -74,7 +78,8 @@ class BasicPathFollower(Module): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._lock = RLock() - self._current_odom: PoseStamped | None = None + self._base_pose: OdomBasePose | None = None + self._current_pose: PoseStamped | None = None self._waypoints: NDArray[np.float32] | None = None self._stop_event = Event() self._thread: Thread | None = None @@ -98,8 +103,13 @@ def stop(self) -> None: super().stop() def _on_odometry(self, msg: Odometry) -> None: + if self._base_pose is None: + self._base_pose = OdomBasePose(self.tfbuffer, self.config.base_frame) + pose = self._base_pose.resolve(msg) + if pose is None: + return with self._lock: - self._current_odom = msg.to_pose_stamped() + self._current_pose = pose def _on_path(self, path: Path) -> None: # The planner owns path safety: it sends the route as far as it is safe, @@ -124,15 +134,15 @@ def _follow(self) -> None: while not self._stop_event.is_set(): start_time = time.perf_counter() with self._lock: - odom = self._current_odom + pose = self._current_pose waypoints = self._waypoints - if odom is not None and waypoints is not None: - self._step(odom, waypoints) + if pose is not None and waypoints is not None: + self._step(pose, waypoints) elapsed = time.perf_counter() - start_time self._stop_event.wait(max(0.0, period - elapsed)) - def _step(self, odom: PoseStamped, waypoints: NDArray[np.float32]) -> None: - position = np.array([odom.position.x, odom.position.y], dtype=np.float32) + def _step(self, pose: PoseStamped, waypoints: NDArray[np.float32]) -> None: + position = np.array([pose.position.x, pose.position.y], dtype=np.float32) if float(np.linalg.norm(waypoints[-1] - position)) < self.config.goal_tolerance: self.nav_cmd_vel.publish(Twist()) with self._lock: @@ -145,7 +155,7 @@ def _step(self, odom: PoseStamped, waypoints: NDArray[np.float32]) -> None: target = self._lookahead_point(waypoints, position) yaw_error = angle_diff( math.atan2(target[1] - position[1], target[0] - position[0]), - odom.orientation.euler[2], + pose.orientation.euler[2], ) angular = max( diff --git a/dimos/navigation/basic_path_follower/test_module.py b/dimos/navigation/basic_path_follower/test_module.py index 3ee85b7525..c9be69766f 100644 --- a/dimos/navigation/basic_path_follower/test_module.py +++ b/dimos/navigation/basic_path_follower/test_module.py @@ -12,16 +12,69 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.navigation.basic_path_follower.module import lookahead_distance +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.navigation.basic_path_follower.module import BasicPathFollower, lookahead_distance +from dimos.protocol.tf.tf import MultiTBuffer +MOUNT_Z = 0.163 -def test_lookahead_floor_at_low_speed(): + +class FakeTF(MultiTBuffer): + def dispose(self) -> None: + pass + + +def _odom() -> Odometry: + return Odometry( + ts=1.0, + frame_id="odom", + child_frame_id="mid360_link", + pose=Pose(Vector3(1.0, 2.0, 3.0), Quaternion(0.0, 0.0, 0.0, 1.0)), + ) + + +def test_on_odometry_steers_from_the_base_pose() -> None: + tf = FakeTF() + tf.receive_transform( + Transform( + translation=Vector3(0.0, 0.0, MOUNT_Z), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="base_link", + child_frame_id="mid360_link", + ts=1.0, + ) + ) + module = BasicPathFollower() + module._tf = tf + try: + module._on_odometry(_odom()) + assert module._current_pose is not None + assert abs(module._current_pose.position.z - (3.0 - MOUNT_Z)) < 1e-9 + finally: + module.stop() + + +def test_on_odometry_drops_frames_without_the_mount_tf() -> None: + module = BasicPathFollower() + module._tf = FakeTF() + try: + module._on_odometry(_odom()) + assert module._current_pose is None + finally: + module.stop() + + +def test_lookahead_floor_at_low_speed() -> None: assert lookahead_distance(0.1, 1.5, 0.4, 1.5) == 0.4 -def test_lookahead_scales_in_linear_region(): +def test_lookahead_scales_in_linear_region() -> None: assert lookahead_distance(0.5, 1.5, 0.4, 1.5) == 0.75 -def test_lookahead_clamped_at_ceiling(): +def test_lookahead_clamped_at_ceiling() -> None: assert lookahead_distance(2.0, 1.5, 0.4, 1.5) == 1.5 diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py index 905c4f8973..778597cec8 100644 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -14,6 +14,8 @@ from __future__ import annotations +from typing import Any + from reactivex.disposable import Disposable from dimos.core.core import rpc @@ -22,10 +24,19 @@ from dimos.msgs.geometry_msgs.PointStamped import PointStamped from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.navigation.tf_pose import OdomBasePose, base_height_above_ground +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() class GoalRelayConfig(ModuleConfig): - pass + base_frame: str = "base_link" + # Lidar height above the ground while standing. None skips the ground + # correction, leaving it to the planner's start_z_offset_m. Set exactly one + # of the two: both set drops the start pose to the ground twice. + lidar_height: float | None = None class GoalRelay(Module): @@ -35,10 +46,17 @@ class GoalRelay(Module): odometry: In[Odometry] goal: In[PointStamped] + tf: In[TFMessage] start_pose: Out[PoseStamped] goal_pose: Out[PoseStamped] + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._base_pose: OdomBasePose | None = None + self._base_height: float | None = None + self._warned_base_frame = False + @rpc def start(self) -> None: super().start() @@ -46,7 +64,37 @@ def start(self) -> None: self.register_disposable(Disposable(self.goal.subscribe(self._on_goal))) def _on_odometry(self, msg: Odometry) -> None: - self.start_pose.publish(msg.to_pose_stamped()) + if self._base_pose is None: + self._base_pose = OdomBasePose(self.tfbuffer, self.config.base_frame) + start = self._base_pose.resolve(msg) + if start is None: + return + if self.config.lidar_height is not None: + base_height = self._resolve_base_height(msg.child_frame_id, self.config.lidar_height) + if base_height is None: + return + start.position.z -= base_height + self.start_pose.publish(start) + + def _resolve_base_height(self, sensor_frame: str, lidar_height: float) -> float | None: + # The base height comes from subtracting the mount leg, which + # base-frame odometry does not have. + if sensor_frame == self.config.base_frame: + if not self._warned_base_frame: + self._warned_base_frame = True + logger.warning( + "Odometry is stamped at %s, so lidar_height cannot ground-project it. " + "Dropping frames until odometry arrives stamped at a sensor.", + sensor_frame, + ) + return None + if self._base_height is None: + assert self._base_pose is not None + leg = self._base_pose.sensor_to_base(sensor_frame) + if leg is None: + return None + self._base_height = base_height_above_ground(lidar_height, -leg) + return self._base_height def _on_goal(self, point: PointStamped) -> None: self.goal_pose.publish(point.to_pose_stamped()) diff --git a/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py b/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py index 739029361a..c89c9accbf 100644 --- a/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py +++ b/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py @@ -33,6 +33,10 @@ class MLSPlannerNativeConfig(NativeModuleConfig): world_frame: str = "map" voxel_size: float = 0.08 robot_height: float = 0.3 + # Subtracted from the start pose z before snapping to a surface. Leave 0 + # when the publisher already ground-projects via GoalRelay lidar_height. + # Set exactly one of the two: both set drops the start pose twice. + start_z_offset_m: float = 0.0 max_overhead_m: float = 2.0 surface_closing_radius: float = 0.3 diff --git a/dimos/navigation/nav_3d/mls_planner/odom_body_frame.py b/dimos/navigation/nav_3d/mls_planner/odom_body_frame.py deleted file mode 100644 index 2afef1e2f1..0000000000 --- a/dimos/navigation/nav_3d/mls_planner/odom_body_frame.py +++ /dev/null @@ -1,67 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from __future__ import annotations - -from typing import Any - -from pydantic import Field -from reactivex.disposable import Disposable - -from dimos.core.core import rpc -from dimos.core.module import Module, ModuleConfig -from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.nav_msgs.Odometry import Odometry - - -class OdomBodyFrameConfig(ModuleConfig): - # base_link from sensor mount rotation, xyzw. - mount_rotation: list[float] = Field(default_factory=lambda: [0.0, 0.0, 0.0, 1.0]) - body_frame_id: str = "base_link" - - -class OdomBodyFrame(Module): - """Re-express tilted-sensor LIO odometry in the level robot body frame. - - Composes out the fixed mount rotation from the orientation. Position and - twist pass through. - """ - - config: OdomBodyFrameConfig - - odometry: In[Odometry] - body_odometry: Out[Odometry] - - def __init__(self, **kwargs: Any) -> None: - super().__init__(**kwargs) - self._mount_inv = Quaternion(*self.config.mount_rotation).inverse() - - @rpc - def start(self) -> None: - super().start() - self.register_disposable(Disposable(self.odometry.subscribe(self._on_odometry))) - - def _on_odometry(self, msg: Odometry) -> None: - leveled = msg.orientation * self._mount_inv - self.body_odometry.publish( - Odometry( - ts=msg.ts, - frame_id=msg.frame_id, - child_frame_id=self.config.body_frame_id, - pose=Pose(msg.position, leveled), - twist=msg.twist, - ) - ) diff --git a/dimos/navigation/nav_3d/mls_planner/rust/Cargo.lock b/dimos/navigation/nav_3d/mls_planner/rust/Cargo.lock index b82cabedd8..86d8d134a8 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/Cargo.lock +++ b/dimos/navigation/nav_3d/mls_planner/rust/Cargo.lock @@ -602,6 +602,8 @@ version = "0.1.0" dependencies = [ "dimos-lcm", "dimos-module-macros", + "lcm-msgs", + "nalgebra 0.35.0", "serde", "serde_json", "tokio", @@ -927,6 +929,30 @@ dependencies = [ "syn", ] +[[package]] +name = "glam" +version = "0.30.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19fc433e8437a212d1b6f1e68c7824af3aed907da60afa994e7f542d18d12aa9" + +[[package]] +name = "glam" +version = "0.31.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556f6b2ea90b8d15a74e0e7bb41671c9bdf38cd9f78c284d750b9ce58a2b5be7" + +[[package]] +name = "glam" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f70749695b063ecbf6b62949ccccde2e733ec3ecbbd71d467dca4e5c6c97cca0" + +[[package]] +name = "glam" +version = "0.33.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f22fb22f065b308be0d8724e3706c7fa3fc2a6c7d6899df4cad7860e7a75436" + [[package]] name = "hashbrown" version = "0.12.3" @@ -1183,7 +1209,7 @@ dependencies = [ "getrandom 0.2.17", "image", "itertools 0.12.1", - "nalgebra", + "nalgebra 0.32.6", "num", "rand 0.8.6", "rand_distr", @@ -1533,10 +1559,41 @@ dependencies = [ "num-complex", "num-rational", "num-traits", - "simba", + "simba 0.8.1", + "typenum", +] + +[[package]] +name = "nalgebra" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adc43a60c217b0c6ff46e47f26911015ad8d2e5a8be1af668c67e370d99a4346" +dependencies = [ + "approx", + "glam 0.30.10", + "glam 0.31.1", + "glam 0.32.1", + "glam 0.33.2", + "matrixmultiply", + "nalgebra-macros", + "num-complex", + "num-rational", + "num-traits", + "simba 0.10.0", "typenum", ] +[[package]] +name = "nalgebra-macros" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "973e7178a678cfd059ccec50887658d482ce16b0aa9da3888ddeab5cd5eb4889" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "nanorand" version = "0.7.0" @@ -2537,6 +2594,15 @@ dependencies = [ "bytemuck", ] +[[package]] +name = "safe_arch" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f7caad094bd561859bcd467734a720c3c1f5d1f338995351fefe2190c45efed" +dependencies = [ + "bytemuck", +] + [[package]] name = "same-file" version = "1.0.6" @@ -2827,7 +2893,19 @@ dependencies = [ "num-complex", "num-traits", "paste", - "wide", + "wide 0.7.33", +] + +[[package]] +name = "simba" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f45c644a9f3a386f9288625d9f0c1e999e1acf07a37df35d0516c7f199d9cb2" +dependencies = [ + "approx", + "num-complex", + "num-traits", + "wide 1.5.0", ] [[package]] @@ -3646,7 +3724,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ce5da8ecb62bcd8ec8b7ea19f69a51275e91299be594ea5cc6ef7819e16cd03" dependencies = [ "bytemuck", - "safe_arch", + "safe_arch 0.7.4", +] + +[[package]] +name = "wide" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfdfe6a32973f2d1b268b8895845a8a96cac2f0191e72c27cc929036060dbf89" +dependencies = [ + "bytemuck", + "safe_arch 1.0.0", ] [[package]] diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs index 51c79826b0..181b932f58 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs @@ -319,8 +319,7 @@ impl Worker { let Some(start) = *self.latest_start.lock().expect("start mutex") else { return; }; - // Ground-project the sensor pose so the start snaps to the supporting surface. - let start = (start.0, start.1, start.2 - self.config.robot_height); + let start = (start.0, start.1, start.2 - self.config.start_z_offset_m); let goal = { let mut guard = self.active_goal.lock().expect("goal mutex"); let Some(goal) = *guard else { diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs index 2f50236635..314c0adc84 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs @@ -37,6 +37,10 @@ pub struct Config { pub voxel_size: f32, #[validate(range(exclusive_min = 0.0))] pub robot_height: f32, + /// Subtracted from the start pose z before snapping to a surface. 0 when + /// the publisher already ground-projects. + #[validate(range(min = 0.0))] + pub start_z_offset_m: f32, /// Ignore surface more than this far above the sensor. #[validate(range(min = 0.0))] pub max_overhead_m: f32, @@ -572,6 +576,7 @@ mod region_tests { world_frame: String::new(), voxel_size: 0.1, robot_height: 0.5, + start_z_offset_m: 0.0, max_overhead_m: 2.0, surface_closing_radius: 0.3, node_spacing_m: 1.0, diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs index 183231d404..8245cc1d33 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs @@ -901,6 +901,7 @@ mod tests { world_frame: "world".into(), voxel_size: VOXEL, robot_height: Z_TOL, + start_z_offset_m: 0.0, max_overhead_m: 2.0, surface_closing_radius: 0.0, node_spacing_m: 1.0, @@ -927,6 +928,7 @@ mod tests { world_frame: "world".into(), voxel_size: VOXEL, robot_height: Z_TOL, + start_z_offset_m: 0.0, max_overhead_m: 2.0, surface_closing_radius: 0.0, node_spacing_m: 1.0, diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs index 16d25c84e5..a59d8f14a9 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs @@ -86,6 +86,8 @@ impl MLSPlanner { world_frame: String::new(), voxel_size, robot_height, + // Unused here. Only the binary's replan loop projects the start. + start_z_offset_m: 0.0, max_overhead_m, surface_closing_radius, node_spacing_m, diff --git a/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py b/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py new file mode 100644 index 0000000000..67f4b6ff73 --- /dev/null +++ b/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py @@ -0,0 +1,145 @@ +# 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 Any + +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.navigation.nav_3d.mls_planner.goal_relay import GoalRelay +from dimos.protocol.tf.tf import MultiTBuffer + +MOUNT_Z = 0.163 +LIDAR_HEIGHT = 0.45 + + +class FakeTF(MultiTBuffer): + """In-memory tf with the dispose() hook and call counter the module tests need.""" + + def __init__(self) -> None: + super().__init__() + self.gets = 0 + + def get( + self, + parent_frame: str, + child_frame: str, + time_point: float | None = None, + time_tolerance: float | None = None, + *, + forward_tolerance: float = 0.0, + ) -> Transform | None: + self.gets += 1 + return super().get( + parent_frame, + child_frame, + time_point, + time_tolerance, + forward_tolerance=forward_tolerance, + ) + + def dispose(self) -> None: + pass + + +def _mount() -> Transform: + return Transform( + translation=Vector3(0.0, 0.0, MOUNT_Z), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="base_link", + child_frame_id="mid360_link", + ts=1.0, + ) + + +def _odom(z: float = 3.0) -> Odometry: + return Odometry( + ts=1.0, + frame_id="odom", + child_frame_id="mid360_link", + pose=Pose(Vector3(1.0, 2.0, z), Quaternion(0.0, 0.0, 0.0, 1.0)), + ) + + +def _relay(tf: FakeTF, **config: Any) -> tuple[GoalRelay, list[PoseStamped]]: + module = GoalRelay(**config) + module._tf = tf + captured: list[PoseStamped] = [] + module.start_pose.subscribe(captured.append) + return module, captured + + +def test_start_pose_is_ground_projected() -> None: + tf = FakeTF() + tf.receive_transform(_mount()) + module, captured = _relay(tf, lidar_height=LIDAR_HEIGHT) + try: + module._on_odometry(_odom()) + # Base sits MOUNT_Z below the sensor, then drops by the base's height + # above ground (0.45 - MOUNT_Z): together exactly the lidar height. + assert len(captured) == 1 + assert abs(captured[0].position.z - (3.0 - LIDAR_HEIGHT)) < 1e-9 + finally: + module.stop() + + +def test_drops_frames_without_the_mount_tf() -> None: + module, captured = _relay(FakeTF(), lidar_height=LIDAR_HEIGHT) + try: + module._on_odometry(_odom()) + assert captured == [] + finally: + module.stop() + + +def test_base_frame_odometry_is_dropped_rather_than_over_projected() -> None: + tf = FakeTF() + tf.receive_transform(_mount()) + module, captured = _relay(tf, lidar_height=LIDAR_HEIGHT) + try: + odom = _odom() + odom.child_frame_id = "base_link" + module._on_odometry(odom) + assert captured == [] + finally: + module.stop() + + +def test_no_lidar_height_skips_the_ground_correction() -> None: + tf = FakeTF() + tf.receive_transform(_mount()) + module, captured = _relay(tf) + try: + module._on_odometry(_odom()) + assert len(captured) == 1 + assert abs(captured[0].position.z - (3.0 - MOUNT_Z)) < 1e-9 + finally: + module.stop() + + +def test_mount_is_looked_up_once() -> None: + tf = FakeTF() + tf.receive_transform(_mount()) + module, captured = _relay(tf, lidar_height=LIDAR_HEIGHT) + try: + module._on_odometry(_odom()) + module._on_odometry(_odom(z=4.0)) + assert len(captured) == 2 + assert abs(captured[1].position.z - (4.0 - LIDAR_HEIGHT)) < 1e-9 + assert tf.gets == 1 + finally: + module.stop() diff --git a/dimos/navigation/nav_3d/mls_planner/test_odom_body_frame.py b/dimos/navigation/nav_3d/mls_planner/test_odom_body_frame.py deleted file mode 100644 index 571b0fbdd0..0000000000 --- a/dimos/navigation/nav_3d/mls_planner/test_odom_body_frame.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -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.msgs.nav_msgs.Odometry import Odometry -from dimos.navigation.nav_3d.mls_planner.odom_body_frame import OdomBodyFrame - - -def _level(mount_rotation, orientation): - """Run one odometry message through the handler and return the output.""" - module = OdomBodyFrame(mount_rotation=list(mount_rotation), body_frame_id="base_link") - try: - captured = [] - module.body_odometry.subscribe(captured.append) - module._on_odometry( - Odometry( - ts=1.0, - frame_id="odom", - child_frame_id="mid360_link", - pose=Pose(Vector3(1.0, 2.0, 3.0), orientation), - ) - ) - return captured[0] - finally: - module.stop() - - -def test_composes_out_the_mount_pitch(): - # A level body reads its own mount tilt as the sensor's world orientation, so - # composing the mount out returns identity. - mount = Quaternion.from_euler(Vector3(0.0, 0.3, 0.0)) - out = _level(mount.to_tuple(), mount) - assert out.orientation.angle_to(Quaternion(0.0, 0.0, 0.0, 1.0)) < 1e-5 - - -def test_preserves_body_yaw_under_mount_tilt(): - # A body yawed by a known angle keeps that yaw after the mount is composed out. - mount = Quaternion.from_euler(Vector3(0.0, 0.3, 0.0)) - body = Quaternion.from_euler(Vector3(0.0, 0.0, 0.7)) - out = _level(mount.to_tuple(), body * mount) - assert out.orientation.angle_to(body) < 1e-5 - - -def test_relabels_child_frame_and_passes_position_through(): - out = _level([0.0, 0.0, 0.0, 1.0], Quaternion(0.0, 0.0, 0.0, 1.0)) - assert out.child_frame_id == "base_link" - assert out.position.to_tuple() == (1.0, 2.0, 3.0) diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index 5ad3088aca..ffea4626b9 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -29,11 +29,17 @@ from dimos.mapping.ray_tracing.transformer import RayTraceMap from dimos.memory2.store.sqlite import SqliteStore +from dimos.memory2.tf import StreamTF from dimos.memory2.transform import FnTransformer from dimos.memory2.type.observation import Observation +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation from dimos.navigation.nav_3d.mls_planner.mls_planner import MLSPlanner +from dimos.navigation.tf_pose import base_height_above_ground +from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT, ROBOT_LENGTH, ROBOT_WIDTH from dimos.utils.data import resolve_named_path if TYPE_CHECKING: @@ -41,16 +47,16 @@ TIMELINE = "ts" -# Body-frame axis-triad length for the odometry transform (m). -ODOM_AXIS_LEN = 0.5 -# Arrow radius as a fraction of the triad length. +AXIS_LEN = 0.5 AXIS_RADIUS_RATIO = 25 -# The travelled trail. Blue, but light enough to read against the map's turbo -# lows, and clear of PATH_PALETTE so it never reads as a planned path. -ODOM_PATH_COLOR = [80, 160, 255] +# Mount frames as recorded on the tf stream. +BASE_FRAME = "base_link" +SENSOR_FRAME = "mid360_link" -# Distinct path colors for overlaid configurations, config 0 first. +SENSOR_PATH_COLOR = [80, 160, 255] + +# Different colors for each path when running with multiple configs PATH_PALETTE = [ [0, 255, 0], [255, 0, 255], @@ -131,21 +137,82 @@ def _log_path_wp(waypoints: NDArray[np.float32] | None, entity: str, color: list rr.log(entity, rr.LineStrips3D([points], colors=[color], radii=0.05)) +def _base_from_sensor(store: SqliteStore) -> Transform | None: + """Sensor to robot base link transform from the recorded tf stream.""" + tf = StreamTF.from_store(store) + if tf is None: + print("no tf stream in the recording; skipping the base_link triad") + return None + return tf.get(SENSOR_FRAME, BASE_FRAME) + + +def _base_pose(pose: tuple[float, ...], ts: float, base_from_sensor: Transform) -> Transform: + """World to robot base link.""" + px, py, pz, qx, qy, qz, qw = pose + sensor = Transform( + translation=Vector3(px, py, pz), + rotation=Quaternion(qx, qy, qz, qw), + frame_id="world", + child_frame_id=base_from_sensor.frame_id, + ts=ts, + ) + return sensor + base_from_sensor + + +def _plan_start( + pose: tuple[float, ...], + ts: float, + base_from_sensor: Transform | None, + base_height: float, + robot_height: float, +) -> tuple[tuple[float, float, float], Transform | None]: + """Ground-projected planner start, plus the base pose when tf has the mount. + + Without a tf stream the start is the sensor pose dropped by the robot height. + """ + px, py, pz, *_ = pose + if base_from_sensor is None: + return (float(px), float(py), float(pz) - robot_height), None + base = _base_pose(pose, ts, base_from_sensor) + start = ( + float(base.translation.x), + float(base.translation.y), + float(base.translation.z) - base_height, + ) + return start, base + + def _log_odometry( - pose: tuple[float, ...], ts: float, trail: list[tuple[float, float, float]] + pose: tuple[float, ...], + ts: float, + trail: list[tuple[float, float, float]], + base: Transform | None, ) -> None: - """Log the odometry pose as a moving body-frame transform and the growing trail.""" + """Trace the sensor moving throughout the scene.""" import rerun as rr px, py, pz, qx, qy, qz, qw = pose rr.set_time(TIMELINE, timestamp=ts) rr.log( - "world/odom", + "world/mid360_link", rr.Transform3D(translation=[px, py, pz], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw])), ) trail.append((px, py, pz)) if len(trail) > 1: - rr.log("world/odom_path", rr.LineStrips3D([trail], colors=[ODOM_PATH_COLOR], radii=0.015)) + rr.log( + "world/mid360_path", rr.LineStrips3D([trail], colors=[SENSOR_PATH_COLOR], radii=0.015) + ) + if base is None: + return + rr.log( + "world/base_link", + rr.Transform3D( + translation=[base.translation.x, base.translation.y, base.translation.z], + quaternion=rr.Quaternion( + xyzw=[base.rotation.x, base.rotation.y, base.rotation.z, base.rotation.w] + ), + ), + ) def _clearance_colors(clearance: NDArray[np.float32], clamp_m: float) -> NDArray[np.uint8]: @@ -324,7 +391,7 @@ def _process_frame( ray_obs: Observation[PointCloud2], planners: list[tuple[str, list[int], MLSPlanner]], goal: tuple[float, float, float], - robot_height: float, + start: tuple[float, float, float], render_voxel: float, clearance_clamp: float, hard_clearance: float, @@ -335,8 +402,7 @@ def _process_frame( assert ray_obs.pose_tuple is not None bounds = ray_obs.tags["region_bounds"] - px, py, pz, *_ = ray_obs.pose_tuple - start = (float(px), float(py), float(pz) - robot_height) + _, _, pz, *_ = ray_obs.pose_tuple ox, oy, radius, z_min, z_max = bounds pts = ray_obs.data.points_f32() rr.set_time(TIMELINE, timestamp=ray_obs.ts) @@ -407,7 +473,9 @@ def main( help="Min occupied neighbors a surface voxel needs to be emitted; " "0 emits all, higher drops isolated returns", ), - robot_height: float = typer.Option(0.3, "--robot-height", help="Robot height (m)"), + robot_height: float = typer.Option( + ROBOT_HEIGHT, "--robot-height", help="Robot height, ground to tallest point / lidar (m)" + ), max_overhead: float = typer.Option( 2.0, "--max-overhead", help="Ignore surface more than this far above the sensor (m)" ), @@ -518,39 +586,71 @@ def main( rr.log("world/goal", rr.Points3D([goal], colors=[[255, 0, 0]], radii=0.1), static=True) - # Static XYZ axis triad in the odometry body frame (world/odom transform). - rr.log( - "world/odom/axes", - rr.Arrows3D( - origins=[[0.0, 0.0, 0.0]] * 3, - vectors=[ - [ODOM_AXIS_LEN, 0.0, 0.0], - [0.0, ODOM_AXIS_LEN, 0.0], - [0.0, 0.0, ODOM_AXIS_LEN], - ], - colors=[[255, 0, 0], [0, 255, 0], [0, 0, 255]], - radii=ODOM_AXIS_LEN / AXIS_RADIUS_RATIO, - ), - static=True, + base_from_sensor = _base_from_sensor(store) + base_height = ( + base_height_above_ground(robot_height, base_from_sensor.inverse()) + if base_from_sensor is not None + else 0.0 ) - odom_trail: list[tuple[float, float, float]] = [] + entities = ["world/mid360_link/axes"] + ( + ["world/base_link/axes"] if base_from_sensor else [] + ) + for entity in entities: + rr.log( + entity, + rr.Arrows3D( + origins=[[0.0, 0.0, 0.0]] * 3, + vectors=[ + [AXIS_LEN, 0.0, 0.0], + [0.0, AXIS_LEN, 0.0], + [0.0, 0.0, AXIS_LEN], + ], + colors=[[255, 0, 0], [0, 255, 0], [0, 0, 255]], + radii=AXIS_LEN / AXIS_RADIUS_RATIO, + ), + static=True, + ) + if base_from_sensor is not None: + rr.log( + "world/base_link/outline", + rr.Boxes3D( + half_sizes=[ROBOT_LENGTH / 2, ROBOT_WIDTH / 2, robot_height / 2], + colors=[(0, 255, 127)], + ), + static=True, + ) + # wall_clearance is the planner's proxy for the robot radius. + rr.log( + "world/base_link/clearance", + rr.Cylinders3D( + lengths=[robot_height], + radii=[wall_clearance], + colors=[(255, 120, 120, 80)], + fill_mode="solid", + ), + static=True, + ) + sensor_trail: list[tuple[float, float, float]] = [] try: frame = 0 for ray_obs in ray_pipeline: if ray_obs.pose_tuple is None: continue + start, base = _plan_start( + ray_obs.pose_tuple, ray_obs.ts, base_from_sensor, base_height, robot_height + ) ref_timing = _process_frame( ray_obs, planners, goal, - robot_height, + start, render_voxel, clearance_clamp, ref_clearance, crop, ) - _log_odometry(ray_obs.pose_tuple, ray_obs.ts, odom_trail) + _log_odometry(ray_obs.pose_tuple, ray_obs.ts, sensor_trail, base) frame += 1 print( f"frame={frame} configs={len(planners)} " diff --git a/dimos/navigation/test_tf_pose.py b/dimos/navigation/test_tf_pose.py new file mode 100644 index 0000000000..1ded4b6980 --- /dev/null +++ b/dimos/navigation/test_tf_pose.py @@ -0,0 +1,139 @@ +# 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 dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.navigation.tf_pose import OdomBasePose, base_height_above_ground +from dimos.protocol.tf.tf import MultiTBuffer + +IDENTITY = Quaternion(0.0, 0.0, 0.0, 1.0) +MOUNT_Z = 0.163 +LIDAR_HEIGHT = 0.45 + + +class CountingTF(MultiTBuffer): + def __init__(self) -> None: + super().__init__() + self.gets = 0 + + def get( + self, + parent_frame: str, + child_frame: str, + time_point: float | None = None, + time_tolerance: float | None = None, + *, + forward_tolerance: float = 0.0, + ) -> Transform | None: + self.gets += 1 + return super().get( + parent_frame, + child_frame, + time_point, + time_tolerance, + forward_tolerance=forward_tolerance, + ) + + +def _mount(z: float = MOUNT_Z, pitch: float = 0.0) -> Transform: + return Transform( + translation=Vector3(0.0, 0.0, z), + rotation=Quaternion.from_euler(Vector3(0.0, pitch, 0.0)), + frame_id="base_link", + child_frame_id="mid360_link", + ts=1.0, + ) + + +def _odom(orientation: Quaternion = IDENTITY) -> Odometry: + return Odometry( + ts=1.0, + frame_id="odom", + child_frame_id="mid360_link", + pose=Pose(Vector3(1.0, 2.0, 3.0), orientation), + ) + + +def test_translates_to_base_frame() -> None: + tf = MultiTBuffer() + tf.receive_transform(_mount()) + pose = OdomBasePose(tf, "base_link").resolve(_odom()) + assert pose is not None + assert pose.frame_id == "odom" + assert pose.ts == 1.0 + assert abs(pose.position.x - 1.0) < 1e-9 + assert abs(pose.position.y - 2.0) < 1e-9 + assert abs(pose.position.z - (3.0 - MOUNT_Z)) < 1e-9 + + +def test_composes_out_the_mount_pitch() -> None: + # A level body reads its own mount tilt as the sensor's world orientation, so + # composing the mount out returns identity. + mount = _mount(pitch=0.3) + tf = MultiTBuffer() + tf.receive_transform(mount) + pose = OdomBasePose(tf, "base_link").resolve(_odom(orientation=mount.rotation)) + assert pose is not None + assert pose.orientation.angle_to(IDENTITY) < 1e-5 + + +def test_preserves_body_yaw_under_mount_tilt() -> None: + mount = _mount(pitch=0.3) + body = Quaternion.from_euler(Vector3(0.0, 0.0, 0.7)) + tf = MultiTBuffer() + tf.receive_transform(mount) + pose = OdomBasePose(tf, "base_link").resolve(_odom(orientation=body * mount.rotation)) + assert pose is not None + assert pose.orientation.angle_to(body) < 1e-5 + + +def test_drops_frames_until_the_mount_leg_arrives() -> None: + tf = MultiTBuffer() + resolver = OdomBasePose(tf, "base_link") + assert resolver.resolve(_odom()) is None + tf.receive_transform(_mount()) + resolver._next_lookup = 0.0 + assert resolver.resolve(_odom()) is not None + + +def test_missing_leg_lookups_are_throttled() -> None: + tf = CountingTF() + resolver = OdomBasePose(tf, "base_link") + assert resolver.resolve(_odom()) is None + assert resolver.resolve(_odom()) is None + assert tf.gets == 1 + + +def test_mount_leg_is_looked_up_once() -> None: + tf = CountingTF() + tf.receive_transform(_mount()) + resolver = OdomBasePose(tf, "base_link") + assert resolver.resolve(_odom()) is not None + assert resolver.resolve(_odom()) is not None + assert tf.gets == 1 + + +def test_base_frame_odometry_passes_through() -> None: + resolver = OdomBasePose(MultiTBuffer(), "base_link") + msg = Odometry(ts=1.0, frame_id="odom", child_frame_id="base_link") + pose = resolver.resolve(msg) + assert pose is not None + assert pose.frame_id == "odom" + + +def test_base_height_above_ground() -> None: + assert abs(base_height_above_ground(LIDAR_HEIGHT, _mount()) - (LIDAR_HEIGHT - MOUNT_Z)) < 1e-9 diff --git a/dimos/navigation/tf_pose.py b/dimos/navigation/tf_pose.py new file mode 100644 index 0000000000..a74c5cbaaf --- /dev/null +++ b/dimos/navigation/tf_pose.py @@ -0,0 +1,85 @@ +# 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. + +"""Resolve odometry into base-frame poses using the static mount tf.""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING + +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + from dimos.msgs.nav_msgs.Odometry import Odometry + from dimos.protocol.tf.tf import TFLookup + +logger = setup_logger() + + +def base_height_above_ground(lidar_height: float, base_to_sensor: Transform) -> float: + """Height of the base frame origin above the ground while standing.""" + return lidar_height - base_to_sensor.translation.z + + +class OdomBasePose: + """Turn odometry messages into the base-frame pose they imply.""" + + # While the leg is missing, retry the lookup at most this often. The buffer + # warns on every miss, so per-message retries would flood the log. + RETRY_PERIOD_S = 1.0 + + def __init__(self, tf: TFLookup, base_frame: str) -> None: + self._tf = tf + self.base_frame = base_frame + self._legs: dict[str, Transform] = {} + self._waiting = False + self._next_lookup = 0.0 + + def resolve(self, msg: Odometry) -> PoseStamped | None: + """The base pose for one message. None until tf has the mount leg.""" + if msg.child_frame_id == self.base_frame: + return msg.to_pose_stamped() + leg = self.sensor_to_base(msg.child_frame_id) + if leg is None: + return None + odom = Transform.from_pose(msg.child_frame_id, msg.to_pose_stamped()) + return (odom + leg).to_pose(ts=msg.ts) + + def sensor_to_base(self, sensor_frame: str) -> Transform | None: + """The cached static sensor -> base leg. Logs once per outage, not per message.""" + if sensor_frame == self.base_frame: + return Transform.identity() + leg = self._legs.get(sensor_frame) + if leg is None: + if self._waiting and time.monotonic() < self._next_lookup: + return None + leg = self._tf.get(sensor_frame, self.base_frame) + if leg is None: + self._next_lookup = time.monotonic() + self.RETRY_PERIOD_S + if not self._waiting: + self._waiting = True + logger.warning( + "No %s -> %s transform on tf yet, dropping odometry until it arrives.", + sensor_frame, + self.base_frame, + ) + return None + if self._waiting: + self._waiting = False + logger.info("Got the %s -> %s transform, resuming.", sensor_frame, self.base_frame) + self._legs[sensor_frame] = leg + return leg diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 81d9138547..ecf75ac543 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -244,7 +244,6 @@ "object-tracker2-d": "dimos.perception.experimental.object_tracker_2d.ObjectTracker2D", "object-tracker3-d": "dimos.perception.experimental.object_tracker_3d.ObjectTracker3D", "object-tracking": "dimos.perception.experimental.object_tracker.ObjectTracking", - "odom-body-frame": "dimos.navigation.nav_3d.mls_planner.odom_body_frame.OdomBodyFrame", "osm-skill": "dimos.agents.skills.osm.OsmSkill", "path-follower": "dimos.navigation.cmu_nav.modules.path_follower.path_follower.PathFollower", "path-following-coordinator": "dimos.control.path_following_coordinator.PathFollowingCoordinator", diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py index e77cdb2d1d..92b84714e6 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py @@ -69,7 +69,7 @@ def _default_recording_dir() -> Path: unitree_go2_mid360_record = autoconnect( MovementManager.blueprint(), - GO2Connection.blueprint().remappings( + GO2Connection.blueprint(publish_tf=False).remappings( [ (GO2Connection, "lidar", "go2_lidar"), (GO2Connection, "odom", "go2_odom"), diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_mls_htc.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_mls_htc.py index 6c4de75e18..88c8f8321c 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_mls_htc.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_mls_htc.py @@ -82,6 +82,8 @@ def _render_path(msg: Any) -> Any: world_frame="world", voxel_size=voxel_size, robot_height=go2_lidar_height, + # The start pose is raw go2 odometry, so the planner ground-projects it. + start_z_offset_m=go2_lidar_height, wall_clearance_m=0.2, wall_buffer_m=0.75, wall_buffer_weight=100.0, diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 0280be7f57..2c03b51142 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -16,7 +16,6 @@ """3d navigation on Go2 with ray tracing and MLS planning""" from datetime import datetime -import math import os from pathlib import Path from typing import Any @@ -36,21 +35,16 @@ from dimos.navigation.movement_manager.movement_manager import MovementManager from dimos.navigation.nav_3d.mls_planner.goal_relay import GoalRelay from dimos.navigation.nav_3d.mls_planner.mls_planner_native import MLSPlannerNative -from dimos.navigation.nav_3d.mls_planner.odom_body_frame import OdomBodyFrame from dimos.navigation.nav_3d.mls_planner.viz import planner_visual_override from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import rerun_config from dimos.robot.unitree.go2.connection import GO2Connection -from dimos.robot.unitree.go2.go2_mid360_static_transforms import ( - MID360_PITCH_DOWN, - base_link_from_mid360, -) +from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT, ROBOT_LENGTH, ROBOT_WIDTH +from dimos.robot.unitree.go2.go2_mid360_static_transforms import Go2Mid360StaticTf from dimos.visualization.vis_module import vis_module voxel_size = 0.08 # Raise above 0 to draw what the planner searched over (surface, nodes, weighted edges). planner_viz_hz = 0.0 -# base_link <- lidar mount rotation, so nav reads odometry in the level body frame. -_sensor_mount_rotation = list(base_link_from_mid360().rotation.to_tuple()) # Body-frame axis-triad length (m). _axis_len = 0.5 @@ -100,13 +94,13 @@ def _render_path(msg: Any) -> Any: def _static_robot_body(rr: Any) -> list[Any]: - """Go2-shaped box on pointlio's sensor frame, counter-rotated for the lidar pitch.""" + """Go2-shaped box on the body frame.""" return [ - rr.Boxes3D(half_sizes=[0.35, 0.155, 0.2], colors=[(0, 255, 127)]), - rr.Transform3D( - parent_frame="tf#/mid360_link", - rotation=rr.RotationAxisAngle(axis=(0, 1, 0), degrees=-math.degrees(MID360_PITCH_DOWN)), + rr.Boxes3D( + half_sizes=[ROBOT_LENGTH / 2, ROBOT_WIDTH / 2, ROBOT_HEIGHT / 2], + colors=[(0, 255, 127)], ), + rr.Transform3D(parent_frame="tf#/base_link"), ] @@ -125,7 +119,7 @@ def _axis_triad(rr: Any) -> Any: def _static_body_axes(rr: Any) -> Any: - """XYZ triad on the leveled robot body (child of the counter-rotated box).""" + """XYZ triad on the robot body (child of the box).""" return _axis_triad(rr) @@ -144,10 +138,8 @@ def _static_sensor_axes(rr: Any) -> list[Any]: }, # Ring buffer replayed to a connecting viewer. Small so connect catches up fast. "memory_limit": "64MB", - # base_link tf comes from the go2 internal odometry, which is not the map - # frame. Anchor the robot box to pointlio's mid360_link frame instead and hide - # the camera frustum that rides base_link. The box lives on its own entity: - # a static transform on world/tf/mid360_link itself would override the live tf. + # The robot box hangs off base_link. It lives on its own entity: a static + # transform on world/tf/base_link would override the live tf. "static": { "world/robot_body": _static_robot_body, "world/robot_body/axes": _static_body_axes, @@ -168,7 +160,11 @@ def _static_sensor_axes(rr: Any) -> list[Any]: vis_module(viewer_backend=global_config.viewer, rerun_config=_nav_rerun_config), # "mcf" for stair traversal GO2Connection.blueprint( - lidar=False, camera=False, motion_mode="mcf", odom_frame_id="go2_odom" + lidar=False, + camera=False, + motion_mode="mcf", + odom_frame_id="go2_odom", + publish_tf=False, ).remappings( [ (GO2Connection, "lidar", "lidar_l1"), @@ -176,9 +172,7 @@ def _static_sensor_axes(rr: Any) -> list[Any]: ] ), PointLio.blueprint(), - # Level pointlio's tilted-sensor odometry into the body frame so the follower - # steers on a true heading. The ray tracer keeps the raw sensor odometry. - OdomBodyFrame.blueprint(mount_rotation=_sensor_mount_rotation), + Go2Mid360StaticTf.blueprint(), RayTracingVoxelMap.blueprint( voxel_size=voxel_size, emit_every=1, @@ -192,7 +186,7 @@ def _static_sensor_axes(rr: Any) -> list[Any]: MLSPlannerNative.blueprint( world_frame="odom", voxel_size=voxel_size, - robot_height=0.3, + robot_height=ROBOT_HEIGHT, surface_closing_radius=0.3, wall_clearance_m=0.1, wall_buffer_m=0.75, @@ -201,16 +195,13 @@ def _static_sensor_axes(rr: Any) -> list[Any]: step_penalty_weight=4.0, viz_publish_hz=planner_viz_hz, ).remappings([(MLSPlannerNative, "global_map", "global_map_unused")]), - GoalRelay.blueprint(), - BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6).remappings( - [(BasicPathFollower, "odometry", "body_odometry")] - ), + GoalRelay.blueprint(lidar_height=ROBOT_HEIGHT), + BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6), MovementManager.blueprint(), ).global_config(n_workers=10, robot_model="unitree_go2", obstacle_avoidance=False) -# The nav blueprint leaves PointLio on its default lidar / odometry topics, so -# remap the recorder's ports onto them. Streams are recorded under the port -# names pointlio_lidar / pointlio_odometry regardless of the topic. +# PointLio keeps its default topics here, so point the recorder's ports at them. +# Streams are recorded under the port names regardless of the topic. if _RECORD: unitree_go2_nav_3d = autoconnect( unitree_go2_nav_3d, diff --git a/dimos/robot/unitree/go2/blueprints/test_tf_topology.py b/dimos/robot/unitree/go2/blueprints/test_tf_topology.py new file mode 100644 index 0000000000..eef20aa954 --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/test_tf_topology.py @@ -0,0 +1,73 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Every Go2 blueprint publishes a tf tree, not a tf graph. + +tf gives each frame one parent. Two publishers writing the same child frame make +the buffer a graph instead, and a lookup then resolves by hop count rather than by +which source is authoritative, so the wrong odometry can win silently. Each +blueprint keeps the invariant a different way: nav_3d turns GO2Connection's tf off +so the static mount tree owns base_link, and the static tree is rooted at +mid360_link so it never writes the frame PointLio owns. +""" + +import pytest + +from dimos.core.coordination.blueprints import Blueprint +from dimos.hardware.sensors.lidar.pointlio.module import PointLio +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_mid360_record import ( + unitree_go2_mid360_record, +) +from dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_nav_3d import ( + unitree_go2_nav_3d, +) +from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.robot.unitree.go2.go2_mid360_static_transforms import ( + Go2Mid360StaticTf, + mount_transforms, +) + +BLUEPRINTS = [unitree_go2_nav_3d, unitree_go2_mid360_record] + + +def _tf_children_by_publisher(blueprint: Blueprint) -> dict[str, set[str]]: + """Child frames each tf publisher the blueprint actually enables will write.""" + odom = PoseStamped(ts=1.0, frame_id="go2_odom") + children: dict[str, set[str]] = {} + for atom in blueprint.blueprints: + if atom.module is GO2Connection and atom.kwargs.get("publish_tf", True): + children["GO2Connection"] = {t.child_frame_id for t in GO2Connection._odom_to_tf(odom)} + if atom.module is Go2Mid360StaticTf: + children["Go2Mid360StaticTf"] = {t.child_frame_id for t in mount_transforms()} + if atom.module is PointLio: + sensor_frame = atom.kwargs.get("sensor_frame_id", "mid360_link") + children["PointLio"] = {sensor_frame} + return children + + +@pytest.mark.parametrize("blueprint", BLUEPRINTS) +def test_no_frame_has_two_tf_parents(blueprint: Blueprint) -> None: + by_publisher = _tf_children_by_publisher(blueprint) + assert by_publisher, "blueprint publishes no tf, so this asserts nothing" + claimed: set[str] = set() + for publisher, frames in by_publisher.items(): + clash = claimed & frames + assert not clash, f"{publisher} also writes {sorted(clash)}" + claimed |= frames + + +def test_static_tree_does_not_write_the_pointlio_frame() -> None: + """Rooting the mount tree at mid360_link is what keeps it off PointLio's edge.""" + assert "mid360_link" not in {t.child_frame_id for t in mount_transforms()} diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index ee9224bbdb..b285273f23 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -75,6 +75,9 @@ class ConnectionConfig(ModuleConfig): # TF parent frame of the internal odometry (odom_frame_id -> base_link). # Rename (e.g. "go2_odom") when another odom source owns the tree root odom_frame_id: str = "world" + # Turn off where another module owns the base_link edge. The odom port + # keeps publishing either way. + publish_tf: bool = True class Go2ConnectionProtocol(Protocol): @@ -397,8 +400,9 @@ def _odom_to_tf(cls, odom: PoseStamped, prefix: str = "") -> list[Transform]: def _publish_tf(self, msg: PoseStamped) -> None: msg.frame_id = self.config.odom_frame_id - transforms = self._odom_to_tf(msg, prefix=self.config.frame_id_prefix or "") - self.tf.publish(TFMessage(*transforms)) + if self.config.publish_tf: + transforms = self._odom_to_tf(msg, prefix=self.config.frame_id_prefix or "") + self.tf.publish(TFMessage(*transforms)) if self.odom.transport: self.odom.publish(msg) diff --git a/dimos/robot/unitree/go2/constants.py b/dimos/robot/unitree/go2/constants.py new file mode 100644 index 0000000000..51d06da0f8 --- /dev/null +++ b/dimos/robot/unitree/go2/constants.py @@ -0,0 +1,21 @@ +# 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. + +"""Physical constants for the Unitree Go2.""" + +# Robot footprint in meters. Length is forward x, width is left y. +ROBOT_LENGTH = 0.6858 +ROBOT_WIDTH = 0.3175 +# Ground to the tallest point. +ROBOT_HEIGHT = 0.45 diff --git a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py index 358bcafc9d..cf42173e45 100644 --- a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py +++ b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py @@ -14,16 +14,19 @@ """Static mount frames for the Go2 + Mid-360 + front-camera rig. -Published continuously onto tf while recording (see :class:`Go2Mid360StaticTf`) so the -mount geometry lands in the recording's tf stream and companion streams (camera, go2 -lidar) can be anchored to ``base_link``. +Published continuously onto tf (see :class:`Go2Mid360StaticTf`) so the mount geometry +lands in the tf stream and companion streams (camera, go2 lidar) can be anchored to +``base_link``. Mount geometry (measured on the physical rig) --------------------------------------------- - base_link -> front_camera: 32.7cm forward, ~4.3cm up (URDF front_camera mount). -- front_camera -> mid360_link: lidar is 3.2cm back, 12cm up, pitched 44 deg down. +- front_camera -> mid360_link: lidar is 3.2cm back, 12cm up, pitched 60 deg down. - front_camera -> camera_optical: the standard ROS optical rotation (x-right, y-down, z-forward). + +The published tree is rooted at mid360_link so the static edges stay off the entities +the live odom -> mid360_link edge writes. The tf buffer composes either direction. """ from __future__ import annotations @@ -37,7 +40,7 @@ frames_to_edge_transforms, ) -MID360_PITCH_DOWN = math.radians(44.0) +MID360_PITCH_DOWN = math.radians(60.0) # rpy that maps a sensor frame to its optical frame (z-forward, x-right, y-down) OPTICAL_RPY = (-math.pi / 2, 0.0, -math.pi / 2) @@ -50,14 +53,14 @@ ] -def base_link_from_mid360() -> Transform: - """Composed base_link -> mid360_link transform from the static mount tree.""" +def mount_transforms() -> list[Transform]: + """The mount tree as published: rooted at mid360_link.""" edges = {t.child_frame_id: t for t in frames_to_edge_transforms(FRAMES)} - return edges["front_camera"] + edges["mid360_link"] + return [-edges["mid360_link"], -edges["front_camera"], edges["camera_optical"]] class Go2Mid360StaticTf(StaticTfPublisher): """Publishes the Go2/Mid-360 mount tree onto tf on a fixed interval.""" def transforms(self) -> list[Transform]: - return frames_to_edge_transforms(FRAMES) + return mount_transforms() diff --git a/dimos/robot/unitree/go2/test_connection.py b/dimos/robot/unitree/go2/test_connection.py index 533254fe2e..79e99fd402 100644 --- a/dimos/robot/unitree/go2/test_connection.py +++ b/dimos/robot/unitree/go2/test_connection.py @@ -18,6 +18,7 @@ dimos/robot/unitree/test_connection.py; this pins the go2-local routing. """ +from collections.abc import Callable, Iterator from types import SimpleNamespace from unittest.mock import MagicMock @@ -66,6 +67,44 @@ def test_odom_to_tf_unprefixed_by_default() -> None: ) +@pytest.fixture +def connection(stub_webrtc: MagicMock) -> Iterator[Callable[[bool], GO2Connection]]: + """Build GO2Connections with the tf and odom ports stubbed, and stop them after.""" + built: list[GO2Connection] = [] + + def build(publish_tf: bool) -> GO2Connection: + conn = GO2Connection( + g=GlobalConfig(robot_ip="127.0.0.1"), + publish_tf=publish_tf, + odom_frame_id="go2_odom", + ) + conn.tf = MagicMock() + conn.odom = MagicMock() + built.append(conn) + return conn + + yield build + for conn in built: + conn.stop() + + +def test_publish_tf_off_keeps_odometry_on_its_port( + connection: Callable[[bool], GO2Connection], +) -> None: + """Turning tf off hands the base_link edge to another publisher, not the odom port.""" + conn = connection(publish_tf=False) + conn._publish_tf(PoseStamped(ts=1.0, frame_id="ignored")) + assert conn.tf.publish.call_count == 0 + assert conn.odom.publish.call_count == 1 + + +def test_publish_tf_on_by_default(connection: Callable[[bool], GO2Connection]) -> None: + conn = connection(publish_tf=True) + conn._publish_tf(PoseStamped(ts=1.0, frame_id="ignored")) + assert conn.tf.publish.call_count == 1 + assert conn.odom.publish.call_count == 1 + + def test_odom_to_tf_prefixed() -> None: """.namespace() sets frame_id_prefix: robot-local frames get prefixed, the odom parent frame stays global so all robots hang off one tree root.""" diff --git a/dimos/robot/unitree/go2/test_go2_mid360_static_transforms.py b/dimos/robot/unitree/go2/test_go2_mid360_static_transforms.py new file mode 100644 index 0000000000..7cf999cd0f --- /dev/null +++ b/dimos/robot/unitree/go2/test_go2_mid360_static_transforms.py @@ -0,0 +1,64 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The published mount tree composes back to the measured rig geometry. + +mount_transforms() inverts two of the four FRAMES edges to root the tree at +mid360_link, so the geometry a consumer reads back is not the geometry written +in FRAMES. These pin the composed result, which is what nav actually uses. +""" + +import math + +from dimos.protocol.tf.tf import MultiTBuffer +from dimos.robot.unitree.go2.go2_mid360_static_transforms import ( + MID360_PITCH_DOWN, + mount_transforms, +) + +# base_link -> mid360_link, summed down the FRAMES chain. +MOUNT_X = 0.32715 - 0.032 +MOUNT_Z = 0.04297 + 0.12 + + +def _buffer() -> MultiTBuffer: + buffer = MultiTBuffer() + buffer.receive_transform(*mount_transforms()) + return buffer + + +def test_mount_height_survives_the_inversion() -> None: + """The lidar sits MOUNT_Z above base_link, the offset every ground projection uses.""" + leg = _buffer().get("mid360_link", "base_link") + assert leg is not None + base_to_sensor = -leg + assert abs(base_to_sensor.translation.z - MOUNT_Z) < 1e-6 + assert abs(base_to_sensor.translation.x - MOUNT_X) < 1e-6 + + +def test_mount_pitch_survives_the_inversion() -> None: + """A sign flip here steers the follower off-heading rather than failing loudly.""" + leg = _buffer().get("mid360_link", "base_link") + assert leg is not None + pitch = (-leg).rotation.euler.y + assert abs(pitch - MID360_PITCH_DOWN) < 1e-6 + assert abs(math.degrees(pitch) - 60.0) < 1e-6 + + +def test_camera_optical_hangs_off_base_link() -> None: + """The tree is rooted at mid360_link, so the camera edge is reachable by composition.""" + optical = _buffer().get("base_link", "camera_optical") + assert optical is not None + assert abs(optical.translation.x - 0.32715) < 1e-6 + assert abs(optical.translation.z - 0.04297) < 1e-6 diff --git a/dimos/robot/unitree/go2/zenoh/blueprints.py b/dimos/robot/unitree/go2/zenoh/blueprints.py index 05fb406254..708c400023 100644 --- a/dimos/robot/unitree/go2/zenoh/blueprints.py +++ b/dimos/robot/unitree/go2/zenoh/blueprints.py @@ -26,44 +26,31 @@ ``DanLocalPlanner`` + ``DanHolonomicTC`` pair from ``unitree-go2-mls-htc``. """ -import math from typing import Any from dimos.core.coordination.blueprints import autoconnect from dimos.core.global_config import global_config from dimos.mapping.ray_tracing.module import RayTracingVoxelMap -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.navigation.basic_path_follower.module import BasicPathFollower from dimos.navigation.dannav.holonomic_tc.module import DanHolonomicTC from dimos.navigation.dannav.local_planner.module import DanLocalPlanner from dimos.navigation.movement_manager.movement_manager import MovementManager from dimos.navigation.nav_3d.mls_planner.goal_relay import GoalRelay from dimos.navigation.nav_3d.mls_planner.mls_planner_native import MLSPlannerNative -from dimos.navigation.nav_3d.mls_planner.odom_body_frame import OdomBodyFrame from dimos.navigation.nav_3d.mls_planner.viz import planner_visual_override +from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT from dimos.robot.unitree.go2.zenoh.zenohconnection import GO2Zenoh from dimos.visualization.vis_module import vis_module voxel_size = 0.08 # Raise above 0 (2.0 works) to draw what the planner searched over: surface, nodes and -# cost-coloured edges. Drives both its publishing and the rerun overrides. +# cost-colored edges. Drives both its publishing and the rerun overrides. planner_viz_hz = 2.0 -# Feeds both the static tf GO2Zenoh publishes and the rotation that levels its odometry — -# they must agree or nav steers off-heading. Verified against Point-LIO's own attitude. +# GO2Zenoh publishes this mount onto tf, where nav reads its odometry corrections. MID360_MOUNT_RPY_DEG = (-60.0, 0.0, -90.0) -def _mount_rotation() -> list[float]: - """base_link <- lidar rotation, so nav reads odometry in the level body frame. - - base_link -> front_camera carries no rotation, so this is just the mount rpy above. - """ - rpy = Vector3(*(math.radians(d) for d in MID360_MOUNT_RPY_DEG)) - return list(Quaternion.from_euler(rpy).to_tuple()) - - def _camera_info_to_pinhole(camera_info: Any) -> Any: """Log the pinhole onto the video's entity instead of camera_info's own. @@ -149,7 +136,7 @@ def _rerun_config(visual_override: dict[str, Any] | None = None) -> dict[str, An _mls_planner = MLSPlannerNative.blueprint( world_frame="odom", voxel_size=voxel_size, - robot_height=0.3, + robot_height=ROBOT_HEIGHT, surface_closing_radius=0.3, wall_clearance_m=0.1, wall_buffer_m=0.75, @@ -183,26 +170,15 @@ def _rerun_config(visual_override: dict[str, Any] | None = None) -> dict[str, An go2_zenoh_nav = autoconnect( go2_zenoh_raycaster, _mls_planner, - OdomBodyFrame.blueprint(mount_rotation=_mount_rotation()), - GoalRelay.blueprint(), - BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6).remappings( - [(BasicPathFollower, "odometry", "body_odometry")] - ), + GoalRelay.blueprint(lidar_height=ROBOT_HEIGHT), + BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6), MovementManager.blueprint(), ).global_config(transport="zenoh", n_workers=8, robot_model="unitree_go2") -# The nav stack with BasicPathFollower swapped for the DanLocalPlanner + DanHolonomicTC -# pair from unitree-go2-mls-htc. The raw planner stream moves to planner_path; the gate -# forwards committed paths on path, so world/planner_path is muted in rerun. go2_zenoh_htc = autoconnect( go2_zenoh_raycaster, - OdomBodyFrame.blueprint(mount_rotation=_mount_rotation()), _mls_planner.remappings([(MLSPlannerNative, "path", "planner_path")]), - # Fed the leveled odometry, so its start_pose doubles as the body-frame PoseStamped - # the Dan modules consume — mirroring mls_htc, where planner start and follower odom - # are the same topic. - GoalRelay.blueprint().remappings([(GoalRelay, "odometry", "body_odometry")]), - # Setting resample_spacing_m to > 0.0 will smooth out jagged paths returned by MLSP + GoalRelay.blueprint(lidar_height=ROBOT_HEIGHT), DanLocalPlanner.blueprint(resample_spacing_m=0.1).remappings( [(DanLocalPlanner, "odom", "start_pose")] ), diff --git a/examples/native-modules/rust/Cargo.lock b/examples/native-modules/rust/Cargo.lock index cc40f75244..071a544889 100644 --- a/examples/native-modules/rust/Cargo.lock +++ b/examples/native-modules/rust/Cargo.lock @@ -62,6 +62,15 @@ version = "1.0.103" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + [[package]] name = "arc-swap" version = "1.9.1" @@ -187,6 +196,12 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" + [[package]] name = "byteorder" version = "1.5.0" @@ -532,6 +547,8 @@ version = "0.1.0" dependencies = [ "dimos-lcm", "dimos-module-macros", + "lcm-msgs", + "nalgebra", "serde", "serde_json", "tokio", @@ -868,6 +885,30 @@ dependencies = [ "syn", ] +[[package]] +name = "glam" +version = "0.30.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19fc433e8437a212d1b6f1e68c7824af3aed907da60afa994e7f542d18d12aa9" + +[[package]] +name = "glam" +version = "0.31.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556f6b2ea90b8d15a74e0e7bb41671c9bdf38cd9f78c284d750b9ce58a2b5be7" + +[[package]] +name = "glam" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f70749695b063ecbf6b62949ccccde2e733ec3ecbbd71d467dca4e5c6c97cca0" + +[[package]] +name = "glam" +version = "0.33.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f22fb22f065b308be0d8724e3706c7fa3fc2a6c7d6899df4cad7860e7a75436" + [[package]] name = "hashbrown" version = "0.12.3" @@ -1343,6 +1384,16 @@ dependencies = [ "regex-automata", ] +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.0" @@ -1376,6 +1427,37 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "nalgebra" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adc43a60c217b0c6ff46e47f26911015ad8d2e5a8be1af668c67e370d99a4346" +dependencies = [ + "approx", + "glam 0.30.10", + "glam 0.31.1", + "glam 0.32.1", + "glam 0.33.2", + "matrixmultiply", + "nalgebra-macros", + "num-complex", + "num-rational", + "num-traits", + "simba", + "typenum", +] + +[[package]] +name = "nalgebra-macros" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "973e7178a678cfd059ccec50887658d482ce16b0aa9da3888ddeab5cd5eb4889" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "nanorand" version = "0.7.0" @@ -1457,6 +1539,15 @@ dependencies = [ "zeroize", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + [[package]] name = "num-conv" version = "0.2.2" @@ -1483,6 +1574,17 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -1807,7 +1909,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls", - "socket2 0.5.10", + "socket2 0.6.3", "thiserror 2.0.18", "tokio", "tracing", @@ -1846,7 +1948,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.5.10", + "socket2 0.6.3", "tracing", "windows-sys 0.52.0", ] @@ -1931,6 +2033,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "rcgen" version = "0.14.8" @@ -2151,7 +2259,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2183,6 +2291,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "safe_arch" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f7caad094bd561859bcd467734a720c3c1f5d1f338995351fefe2190c45efed" +dependencies = [ + "bytemuck", +] + [[package]] name = "same-file" version = "1.0.6" @@ -2474,6 +2591,18 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "simba" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f45c644a9f3a386f9288625d9f0c1e999e1acf07a37df35d0516c7f199d9cb2" +dependencies = [ + "approx", + "num-complex", + "num-traits", + "wide", +] + [[package]] name = "simd-adler32" version = "0.3.9" @@ -3249,6 +3378,16 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "wide" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfdfe6a32973f2d1b268b8895845a8a96cac2f0191e72c27cc929036060dbf89" +dependencies = [ + "bytemuck", + "safe_arch", +] + [[package]] name = "winapi" version = "0.3.9" @@ -3271,7 +3410,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/examples/native-modules/rust/Cargo.toml b/examples/native-modules/rust/Cargo.toml index 6bfee2b74f..712e4fef45 100644 --- a/examples/native-modules/rust/Cargo.toml +++ b/examples/native-modules/rust/Cargo.toml @@ -11,6 +11,14 @@ path = "src/ping.rs" name = "pong" path = "src/pong.rs" +[[bin]] +name = "tf_listener" +path = "src/tf_listener.rs" + +[[bin]] +name = "tf_broadcaster" +path = "src/tf_broadcaster.rs" + [dependencies] dimos-module = { path = "../../../native/rust/dimos-module" } lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } diff --git a/examples/native-modules/rust/src/tf_broadcaster.rs b/examples/native-modules/rust/src/tf_broadcaster.rs new file mode 100644 index 0000000000..da9c18f75f --- /dev/null +++ b/examples/native-modules/rust/src/tf_broadcaster.rs @@ -0,0 +1,51 @@ +// 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. + +use std::time::{SystemTime, UNIX_EPOCH}; + +use dimos_module::nalgebra::Isometry3; +use dimos_module::{run_with_transport, Module, Tf, Transform}; +use tokio::time::{interval, Duration}; + +#[derive(Module)] +#[module(setup = start_broadcast)] +struct TfBroadcaster { + #[tf] + tf: Tf, +} + +impl TfBroadcaster { + async fn start_broadcast(&mut self) { + let tf = self.tf.clone(); + tokio::spawn(async move { + let mut ticker = interval(Duration::from_millis(100)); + loop { + ticker.tick().await; + let ts = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock before epoch") + .as_secs_f64(); + let t = Transform::new("c", "d", ts, Isometry3::translation(0.5, 0.0, 0.0)); + if tf.publish(&[t]).await.is_err() { + break; + } + } + }); + } +} + +#[tokio::main] +async fn main() { + run_with_transport::().await; +} diff --git a/examples/native-modules/rust/src/tf_listener.rs b/examples/native-modules/rust/src/tf_listener.rs new file mode 100644 index 0000000000..8d7efd1790 --- /dev/null +++ b/examples/native-modules/rust/src/tf_listener.rs @@ -0,0 +1,54 @@ +// 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. + +use dimos_module::{run_with_transport, Module, Tf}; +use tokio::time::{interval, Duration}; + +#[derive(Module)] +#[module(setup = start_lookup)] +struct TfListener { + #[tf] + tf: Tf, +} + +impl TfListener { + async fn start_lookup(&mut self) { + let tf = self.tf.clone(); + tokio::spawn(async move { + let mut ticker = interval(Duration::from_millis(500)); + loop { + ticker.tick().await; + match tf.get_latest("a", "d") { + Some(t) => { + let p = t.translation(); + tracing::info!( + parent = %t.parent, + child = %t.child, + x = p.x, + y = p.y, + z = p.z, + "tf lookup", + ); + } + None => tracing::info!("a -> d not available yet"), + } + } + }); + } +} + +#[tokio::main] +async fn main() { + run_with_transport::().await; +} diff --git a/examples/native-modules/rust_tf.py b/examples/native-modules/rust_tf.py new file mode 100644 index 0000000000..362f59d9c5 --- /dev/null +++ b/examples/native-modules/rust_tf.py @@ -0,0 +1,121 @@ +# 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. + +"""Python and Rust each publish part of a tf chain, a Rust #[tf] module reads it back. + +Python publishes a -> b -> c, a Rust broadcaster publishes c -> d, and a Rust +listener composes the full a -> d transform. + +Run with: + python examples/native-modules/rust_tf.py +""" + +from __future__ import annotations + +import asyncio +import math +from pathlib import Path +import time + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.coordination.module_coordinator import ModuleCoordinator +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.native_module import NativeModule, NativeModuleConfig +from dimos.core.stream import IO +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.tf2_msgs.TFMessage import TFMessage + +_RUST_DIR = Path(__file__).parent / "rust" +_EXAMPLES = _RUST_DIR / "target" / "release" +_BUILD = "cargo build --release" + + +class TfProducer(Module): + """Publishes a time-varying a -> b -> c transform chain onto /tf.""" + + tf: IO[TFMessage] + + _running: bool = False + + @rpc + def start(self) -> None: + super().start() + self._running = True + self.spawn(self._publish_loop()) + + async def _publish_loop(self) -> None: + start = time.time() + while self._running: + t = time.time() - start + now = time.time() + self.tfbuffer.publish( + Transform( + translation=Vector3(0.0, math.cos(t), math.sin(t)), + frame_id="a", + child_frame_id="b", + ts=now, + ), + Transform( + translation=Vector3(1.0, 0.0, 0.0), + frame_id="b", + child_frame_id="c", + ts=now, + ), + ) + await asyncio.sleep(0.1) + + @rpc + def stop(self) -> None: + self._running = False + super().stop() + + +class TfListenerConfig(NativeModuleConfig): + executable: str = str(_EXAMPLES / "tf_listener") + build_command: str = _BUILD + cwd: str = str(_RUST_DIR) + stdin_config: bool = True + + +class TfListenerModule(NativeModule): + """Rust module that looks up a -> d and logs it. + + Expect to see (1.5, cos(t), sin(t)) + """ + + config: TfListenerConfig + tf: IO[TFMessage] + + +class TfBroadcasterConfig(NativeModuleConfig): + executable: str = str(_EXAMPLES / "tf_broadcaster") + build_command: str = _BUILD + cwd: str = str(_RUST_DIR) + stdin_config: bool = True + + +class TfBroadcasterModule(NativeModule): + """Rust module that publishes the c -> d transform.""" + + config: TfBroadcasterConfig + tf: IO[TFMessage] + + +if __name__ == "__main__": + bp = autoconnect( + TfProducer.blueprint(), TfBroadcasterModule.blueprint(), TfListenerModule.blueprint() + ).global_config(viewer="none") + ModuleCoordinator.build(bp).loop() diff --git a/native/rust/Cargo.lock b/native/rust/Cargo.lock index 1e95f02742..8364b02486 100644 --- a/native/rust/Cargo.lock +++ b/native/rust/Cargo.lock @@ -62,6 +62,15 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + [[package]] name = "arc-swap" version = "1.9.1" @@ -187,6 +196,12 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" + [[package]] name = "byteorder" version = "1.5.0" @@ -533,6 +548,7 @@ dependencies = [ "dimos-lcm", "dimos-module-macros", "lcm-msgs", + "nalgebra", "serde", "serde_json", "tokio", @@ -858,6 +874,30 @@ dependencies = [ "syn", ] +[[package]] +name = "glam" +version = "0.30.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19fc433e8437a212d1b6f1e68c7824af3aed907da60afa994e7f542d18d12aa9" + +[[package]] +name = "glam" +version = "0.31.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556f6b2ea90b8d15a74e0e7bb41671c9bdf38cd9f78c284d750b9ce58a2b5be7" + +[[package]] +name = "glam" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f70749695b063ecbf6b62949ccccde2e733ec3ecbbd71d467dca4e5c6c97cca0" + +[[package]] +name = "glam" +version = "0.33.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f22fb22f065b308be0d8724e3706c7fa3fc2a6c7d6899df4cad7860e7a75436" + [[package]] name = "hashbrown" version = "0.12.3" @@ -1333,6 +1373,16 @@ dependencies = [ "regex-automata", ] +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.0" @@ -1366,6 +1416,37 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "nalgebra" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adc43a60c217b0c6ff46e47f26911015ad8d2e5a8be1af668c67e370d99a4346" +dependencies = [ + "approx", + "glam 0.30.10", + "glam 0.31.1", + "glam 0.32.1", + "glam 0.33.2", + "matrixmultiply", + "nalgebra-macros", + "num-complex", + "num-rational", + "num-traits", + "simba", + "typenum", +] + +[[package]] +name = "nalgebra-macros" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "973e7178a678cfd059ccec50887658d482ce16b0aa9da3888ddeab5cd5eb4889" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "nanorand" version = "0.7.0" @@ -1447,6 +1528,15 @@ dependencies = [ "zeroize", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + [[package]] name = "num-conv" version = "0.2.2" @@ -1473,6 +1563,17 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -1921,6 +2022,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "rcgen" version = "0.14.8" @@ -2173,6 +2280,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "safe_arch" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f7caad094bd561859bcd467734a720c3c1f5d1f338995351fefe2190c45efed" +dependencies = [ + "bytemuck", +] + [[package]] name = "same-file" version = "1.0.6" @@ -2464,6 +2580,18 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "simba" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f45c644a9f3a386f9288625d9f0c1e999e1acf07a37df35d0516c7f199d9cb2" +dependencies = [ + "approx", + "num-complex", + "num-traits", + "wide", +] + [[package]] name = "simd-adler32" version = "0.3.9" @@ -3260,6 +3388,16 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "wide" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfdfe6a32973f2d1b268b8895845a8a96cac2f0191e72c27cc929036060dbf89" +dependencies = [ + "bytemuck", + "safe_arch", +] + [[package]] name = "winapi" version = "0.3.9" diff --git a/native/rust/README.md b/native/rust/README.md index 8b89359944..4c31fc1d81 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -26,6 +26,9 @@ struct MyModule { #[output(encode = Twist::encode)] out: Output, + #[io(decode = Twist::decode, encode = Twist::encode)] + shared: Io, + #[config] config: MyConfig, } @@ -37,6 +40,9 @@ impl MyModule { // processing function expected by cmd: Input async fn handle_cmd(&mut self, msg: Twist) { /* ... */ } + // processing function expected by shared: Io + async fn handle_shared(&mut self, msg: Twist) { /* ... */ } + // teardown / clean up logic async fn on_stop(&mut self) { /* ... */ } } @@ -57,7 +63,9 @@ Every transport is compiled into the binary. `run_with_transport` opens the one - `#[module(setup = fn, teardown = fn)]`: on the struct. Both optional. Names methods on `Self`. `setup` runs once before the input dispatch loop starts (use it to spawn background tasks or initialize resources); `teardown` runs once after the loop exits (use it for cleanup). - `#[input(decode = fn, handler = fn)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. - `#[output(encode = fn)]`: on a field of type `Output`. `encode` is required. +- `#[io(decode = fn, encode = fn, handler = fn)]`: on a field of type `Io`, a port that publishes to and subscribes on one topic. `decode` and `encode` are required; `handler` defaults to `handle_`. The transports deliver a message back to its own sender, so the handler also sees what the module publishes. Use `#[output]` instead when the module only publishes. - `#[config]`: on one field. The type must be defined with `#[native_config]` (see [Config](#config)). At most one per struct. If absent, `Config` defaults to `dimos_module::NoConfig`. +- `#[tf]`: on a field of type `Tf`. Subscribes to the `tf` topic, answers transform queries, and publishes transforms (see [Transforms](#transforms)). No arguments. - Unattributed fields are initialized via `Default::default()` and treated as module state. ## Config @@ -100,6 +108,63 @@ At runtime `run()` enforces the mapping on the Python payload: deserialization r Field name = port name. Ports map to topics via the stdin JSON; unmapped ports fall back to `/{port}`. +## Transforms + +A `#[tf]` field gives a module a view of the transform graph, the Rust counterpart to Python's `tf.get()` and `tf.publish()`. It subscribes to the `tf` topic (mapped like any other port, default `/tf`), buffers each `parent -> child` edge it sees, and answers queries by composing transforms along the shortest path through the graph. + +```rust +#[derive(Module)] +struct VoxelMap { + #[input(decode = PointCloud2::decode)] + lidar: Input, + #[tf] + tf: Tf, +} + +impl VoxelMap { + async fn handle_lidar(&mut self, cloud: PointCloud2) { + // De-rotate a scan from the lidar's mount frame into the robot base frame. + if let Some(t) = self.tf.get_latest("base_link", "mid360_link") { + let point_in_base = t.rotation() * point_in_lidar + t.translation(); + } + } +} +``` + +`Tf` is a cheap-to-clone handle; the graph fills in the background as `tf` messages arrive. `get_latest(parent, child)` is the common case. For a query against a particular stamp, `lookup(parent, child)` starts one that `.at(time)` points at the sample nearest that stamp and `.tolerance(secs)` bounds how far that sample may sit from it, finished with `.get()`: + +```rust +let at_scan = self.tf.lookup("map", "base_link").at(scan_ts).tolerance(0.1).get(); +``` + +A message and the transform it needs arrive on separate topics, so the transform for a given stamp is often merely late. `.within(duration)` replaces `.get()` to wait for one, returning as soon as the lookup succeeds or `None` at the deadline: + +```rust +let at_scan = self.tf.lookup("odom", &cloud.header.frame_id) + .at(scan_ts) + .tolerance(0.02) + .within(Duration::from_millis(200)) + .await; +``` + +`.tolerance()` and `.within()` are different clocks. Tolerance bounds how far the chosen sample may sit from `.at()` in message stamps — accuracy. `.within()` bounds how long to wait in wall time — patience. Always set a tolerance when waiting, or the lookup is satisfied by anything inside the buffer window and returns a stale transform immediately. + +`.within()` suspends the caller, and awaiting it inside a `handle_*` method parks that module's whole dispatch loop, so every other topic it subscribes to stops being served until it returns. Prefer `.get()` there; move waiting onto its own task when the wait may be long. + +Either way the result is `None` when no path connects the frames or no sample falls within the tolerance. It exposes its `nalgebra` parts via `translation()` (a `Vector3`) and `rotation()` (a `UnitQuaternion`). Lookups are nearest-in-time, not interpolated. + +A result composed over several edges carries the stamp of the stalest edge on the path, so `ts` reads as the age of the whole answer rather than of one hop. A chain mixing a live edge with a static one is only as fresh as the live edge, in either direction. + +`publish` sends transforms onto the same `tf` topic, the counterpart to Python's `tf.publish()`. Published transforms also feed the module's own graph, so a lookup right after the publish sees them. Build the isometry from `dimos_module::nalgebra`, re-exported so the version matches the SDK's types: + +```rust +use dimos_module::nalgebra::Isometry3; +use dimos_module::Transform; + +let iso = Isometry3::translation(0.5, 0.0, 0.0); +self.tf.publish(&[Transform::new("base_link", "gripper", ts, iso)]).await?; +``` + ## What `#[derive(Module)]` generates Just for reference, in the example above the macro expands to: diff --git a/native/rust/dimos-module-macros/src/lib.rs b/native/rust/dimos-module-macros/src/lib.rs index 5644693af9..8679b2289c 100644 --- a/native/rust/dimos-module-macros/src/lib.rs +++ b/native/rust/dimos-module-macros/src/lib.rs @@ -17,7 +17,7 @@ use proc_macro2::TokenStream as TokenStream2; use quote::{format_ident, quote}; use syn::{parse_macro_input, Data, DeriveInput, Field, Fields, Ident, Path, Type}; -#[proc_macro_derive(Module, attributes(input, output, config, module))] +#[proc_macro_derive(Module, attributes(input, output, io, config, tf, module))] pub fn derive_module(input: TokenStream) -> TokenStream { let input = parse_macro_input!(input as DeriveInput); match expand(input) { @@ -174,10 +174,24 @@ fn is_option(ty: &Type) -> bool { matches!(ty, Type::Path(p) if p.path.segments.last().is_some_and(|s| s.ident == "Option")) } +const ONE_ATTR_ONLY: &str = "field has multiple module attributes; only one of #[input], \ + #[output], #[io], #[config], #[tf] is allowed"; + enum FieldKind { - Input { decode: Path, handler: Ident }, - Output { encode: Path }, + Input { + decode: Path, + handler: Ident, + }, + Output { + encode: Path, + }, + Io { + decode: Path, + encode: Path, + handler: Ident, + }, Config, + Tf, State, } @@ -272,24 +286,30 @@ fn expand(input: DeriveInput) -> syn::Result { FieldKind::Output { encode } => { quote!(#name: builder.output(#name_str, #encode)) } + FieldKind::Io { decode, encode, .. } => { + quote!(#name: builder.io(#name_str, #decode, #encode)) + } FieldKind::Config => quote!(#name: config), + FieldKind::Tf => quote!(#name: builder.tf()), FieldKind::State => quote!(#name: ::core::default::Default::default()), } }); - let input_fields: Vec<&ClassifiedField> = classified + // Every port that receives messages gets an arm in the select! loop. + let handled_fields: Vec<(&Ident, &Ident)> = classified .iter() - .filter(|f| matches!(f.kind, FieldKind::Input { .. })) + .filter_map(|f| match &f.kind { + FieldKind::Input { handler, .. } | FieldKind::Io { handler, .. } => { + Some((f.name, handler)) + } + _ => None, + }) .collect(); - let handle_body = if input_fields.is_empty() { + let handle_body = if handled_fields.is_empty() { quote!(::std::future::pending::<()>().await) } else { - let handle_arms = input_fields.iter().map(|f| { - let FieldKind::Input { handler, .. } = &f.kind else { - unreachable!() - }; - let name = f.name; + let handle_arms = handled_fields.iter().map(|(name, handler)| { quote!( ::core::option::Option::Some(msg) = self.#name.recv() => { self.#handler(msg).await @@ -353,10 +373,7 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { let path = attr.path(); if path.is_ident("input") { if found.is_some() { - return Err(syn::Error::new_spanned( - attr, - "field has multiple module attributes; only one of #[input], #[output], #[config] is allowed", - )); + return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); } let mut decode: Option = None; let mut handler: Option = None; @@ -378,10 +395,7 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { found = Some(FieldKind::Input { decode, handler }); } else if path.is_ident("output") { if found.is_some() { - return Err(syn::Error::new_spanned( - attr, - "field has multiple module attributes; only one of #[input], #[output], #[config] is allowed", - )); + return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); } let mut encode: Option = None; attr.parse_nested_meta(|meta| { @@ -398,14 +412,47 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { syn::Error::new_spanned(attr, "#[output] requires `encode = ...`") })?; found = Some(FieldKind::Output { encode }); + } else if path.is_ident("io") { + if found.is_some() { + return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); + } + let mut decode: Option = None; + let mut encode: Option = None; + let mut handler: Option = None; + attr.parse_nested_meta(|meta| { + if meta.path.is_ident("decode") { + decode = Some(meta.value()?.parse()?); + } else if meta.path.is_ident("encode") { + encode = Some(meta.value()?.parse()?); + } else if meta.path.is_ident("handler") { + handler = Some(meta.value()?.parse()?); + } else { + return Err(meta.error( + "unrecognized #[io] argument; expected `decode = ...`, `encode = ...` or `handler = ...`", + )); + } + Ok(()) + })?; + let decode = decode + .ok_or_else(|| syn::Error::new_spanned(attr, "#[io] requires `decode = ...`"))?; + let encode = encode + .ok_or_else(|| syn::Error::new_spanned(attr, "#[io] requires `encode = ...`"))?; + let handler = handler.unwrap_or_else(|| format_ident!("handle_{}", name)); + found = Some(FieldKind::Io { + decode, + encode, + handler, + }); } else if path.is_ident("config") { if found.is_some() { - return Err(syn::Error::new_spanned( - attr, - "field has multiple module attributes; only one of #[input], #[output], #[config] is allowed", - )); + return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); } found = Some(FieldKind::Config); + } else if path.is_ident("tf") { + if found.is_some() { + return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); + } + found = Some(FieldKind::Tf); } } diff --git a/native/rust/dimos-module/Cargo.toml b/native/rust/dimos-module/Cargo.toml index e11df3f02b..fd45eb17e2 100644 --- a/native/rust/dimos-module/Cargo.toml +++ b/native/rust/dimos-module/Cargo.toml @@ -8,6 +8,8 @@ license = "Apache-2.0" [dependencies] dimos-lcm = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } dimos-module-macros = { version = "=0.1.0", path = "../dimos-module-macros" } +lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +nalgebra = "0.35.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time", "signal", "io-std", "io-util"] } serde = { version = "1", features = ["derive"] } serde_json = "1" @@ -17,7 +19,6 @@ validator = { version = "0.20", features = ["derive"] } zenoh = { version = "1", features = ["unstable"] } [dev-dependencies] -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } tracing-test = "0.2" [[example]] diff --git a/native/rust/dimos-module/src/lib.rs b/native/rust/dimos-module/src/lib.rs index 56f17ec8b8..8a97eddfe2 100644 --- a/native/rust/dimos-module/src/lib.rs +++ b/native/rust/dimos-module/src/lib.rs @@ -12,18 +12,27 @@ // See the License for the specific language governing permissions and // limitations under the License. +// #[derive(Module)] emits ::dimos_module paths, so tests in this crate that +// derive Module need the crate to be nameable from inside itself. +#[cfg(test)] +extern crate self as dimos_module; + pub mod lcm; pub mod log; pub mod module; +pub mod tf; pub mod transport; pub mod zenoh; pub use dimos_module_macros::{native_config, Module}; pub use lcm::LcmTransport; -pub use module::{run, Builder, Input, Module, ModuleConfig, NativeConfig, NoConfig, Output}; +pub use module::{run, Builder, Input, Io, Module, ModuleConfig, NativeConfig, NoConfig, Output}; +pub use tf::{Lookup, Tf, Transform}; pub use transport::Transport; pub use zenoh::ZenohTransport; +pub use nalgebra; + // Re-export LcmOptions so callers don't need to depend on dimos-lcm directly. pub use dimos_lcm::LcmOptions; diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index c02077a130..dc953bd24b 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -125,14 +125,41 @@ pub struct Output { impl Output { pub async fn publish(&self, msg: &T) -> io::Result<()> { - let data = (self.encode)(msg); - self.sender - .send(data) - .await - .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "background task gone")) + publish_encoded(&self.sender, (self.encode)(msg)).await + } +} + +/// A port that publishes to and subscribes on the same topic. +/// +/// The transports deliver a message back to its own sender, so an `Io` port +/// sees the whole topic including its own publishes. +pub struct Io { + pub topic: String, + receiver: mpsc::Receiver, + encode: fn(&T) -> Vec, + sender: mpsc::Sender>, +} + +impl Io { + pub async fn recv(&mut self) -> Option { + self.receiver.recv().await + } + + pub async fn publish(&self, msg: &T) -> io::Result<()> { + publish_encoded(&self.sender, (self.encode)(msg)).await } } +pub(crate) async fn publish_encoded( + sender: &mpsc::Sender>, + data: Vec, +) -> io::Result<()> { + sender + .send(data) + .await + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "background task gone")) +} + /// Extract `(topics, config)` from an already-parsed config object. `run` /// parses the line once and also reads `qos` from it, so this takes the value. fn parse_config_value( @@ -275,60 +302,130 @@ pub trait Module: Sized + Send + 'static { pub struct Builder { topics: HashMap, + // Every port the module asked for a topic, matched against topics after build. + requested: BTreeSet, routes: HashMap>>, // One publish queue per output channel, drained by its own worker. outputs: Vec<(String, mpsc::Receiver>)>, + tf: Option, } impl Builder { pub(crate) fn new(topics: HashMap) -> Self { Self { topics, + requested: BTreeSet::new(), routes: HashMap::new(), outputs: Vec::new(), + tf: None, } } - fn topic_for(&self, port: &str) -> String { + fn topic_for(&mut self, port: &str) -> String { + self.requested.insert(port.to_string()); self.topics .get(port) .cloned() .unwrap_or_else(|| format!("/{port}")) } - pub fn input( + // A mismatch is dead wiring: an unclaimed topic reaches no port, and an + // unsent one leaves the port on a fallback name nothing else publishes to. + pub(crate) fn enforce_topics_match_ports(&self) -> io::Result<()> { + let provided: BTreeSet<&String> = self.topics.keys().collect(); + let requested: BTreeSet<&String> = self.requested.iter().collect(); + if provided == requested { + return Ok(()); + } + let missing: Vec<&&String> = requested.difference(&provided).collect(); + let unexpected: Vec<&&String> = provided.difference(&requested).collect(); + Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "topics do not match module ports: missing {missing:?}, unexpected {unexpected:?}" + ), + )) + } + + fn add_route( &mut self, - port: &str, + topic: &str, decode: fn(&[u8]) -> io::Result, - ) -> Input { - let topic = self.topic_for(port); + ) -> mpsc::Receiver { let (tx, rx) = mpsc::channel(INPUT_CHANNEL_CAPACITY); self.routes - .entry(topic.clone()) + .entry(topic.to_string()) .or_default() .push(Box::new(TypedRoute { - topic: topic.clone(), + topic: topic.to_string(), decode, sender: tx, drop_count: AtomicU64::new(0), last_log_ns: AtomicU64::new(0), })); - Input { - topic, - receiver: rx, - } + rx + } + + fn add_publisher(&mut self, topic: &str) -> mpsc::Sender> { + let (tx, rx) = mpsc::channel(PUBLISH_CHANNEL_CAPACITY); + self.outputs.push((topic.to_string(), rx)); + tx + } + + pub fn input( + &mut self, + port: &str, + decode: fn(&[u8]) -> io::Result, + ) -> Input { + let topic = self.topic_for(port); + let receiver = self.add_route(&topic, decode); + Input { topic, receiver } } pub fn output(&mut self, port: &str, encode: fn(&T) -> Vec) -> Output { let topic = self.topic_for(port); - let (tx, rx) = mpsc::channel(PUBLISH_CHANNEL_CAPACITY); - self.outputs.push((topic.clone(), rx)); + let sender = self.add_publisher(&topic); Output { topic, encode, - sender: tx, + sender, } } + + /// A port that both subscribes and publishes on one topic. + pub fn io( + &mut self, + port: &str, + decode: fn(&[u8]) -> io::Result, + encode: fn(&T) -> Vec, + ) -> Io { + let topic = self.topic_for(port); + let receiver = self.add_route(&topic, decode); + let sender = self.add_publisher(&topic); + Io { + topic, + receiver, + encode, + sender, + } + } + + /// A handle that answers transform queries and publishes on the `tf` topic. + /// + /// The graph fills in the background as `tf` messages arrive. Repeated calls + /// share one graph. + pub fn tf(&mut self) -> crate::tf::Tf { + if let Some(tf) = &self.tf { + return tf.clone(); + } + let topic = self.topic_for("tf"); + let sender = self.add_publisher(&topic); + let (tf, route) = + crate::tf::tf_subscription(topic.clone(), crate::tf::DEFAULT_TF_WINDOW_SECS, sender); + self.routes.entry(topic).or_default().push(route); + self.tf = Some(tf.clone()); + tf + } } /// Subscribe each channel on the transport, dispatching its messages to that @@ -423,6 +520,7 @@ where let mut builder = Builder::new(topics); let mut module = M::build(&mut builder, config); + builder.enforce_topics_match_ports()?; subscribe_routes(&transport, builder.routes).await?; // Kept alive until teardown so the subscriptions stay live. @@ -739,13 +837,13 @@ mod tests { #[test] fn unmapped_port_falls_back_to_slash_port() { - let builder = builder_with_topics(&[]); + let mut builder = builder_with_topics(&[]); assert_eq!(builder.topic_for("cmd_vel"), "/cmd_vel"); } #[test] fn mapped_port_uses_given_topic() { - let builder = builder_with_topics(&[("cmd_vel", "/robot/cmd_vel")]); + let mut builder = builder_with_topics(&[("cmd_vel", "/robot/cmd_vel")]); assert_eq!(builder.topic_for("cmd_vel"), "/robot/cmd_vel"); } @@ -770,6 +868,108 @@ mod tests { assert_eq!(output.topic, "/robot/cmd_vel"); } + #[test] + fn topics_matching_ports_exactly_pass() { + let mut builder = builder_with_topics(&[("cmd", "/robot/cmd"), ("odom", "/robot/odom")]); + builder.input("cmd", |b| Ok(b.to_vec())); + builder.output("odom", |b: &Vec| b.clone()); + builder.enforce_topics_match_ports().expect("exact match"); + } + + #[test] + fn a_port_the_coordinator_never_sent_is_rejected() { + let mut builder = builder_with_topics(&[("cmd", "/robot/cmd")]); + builder.input("cmd", |b| Ok(b.to_vec())); + builder.output("odom", |b: &Vec| b.clone()); + let err = builder + .enforce_topics_match_ports() + .expect_err("odom has no topic"); + assert!(err.to_string().contains("missing [\"odom\"]"), "{err}"); + } + + #[test] + fn a_topic_no_port_claimed_is_rejected() { + let mut builder = builder_with_topics(&[("cmd", "/robot/cmd"), ("stale", "/robot/stale")]); + builder.input("cmd", |b| Ok(b.to_vec())); + let err = builder + .enforce_topics_match_ports() + .expect_err("stale is unclaimed"); + assert!(err.to_string().contains("unexpected [\"stale\"]"), "{err}"); + } + + #[test] + fn a_tf_field_claims_the_tf_topic() { + let mut builder = builder_with_topics(&[("tf", "/tf#tf2_msgs.TFMessage")]); + builder.tf(); + builder.enforce_topics_match_ports().expect("tf claimed"); + } + + #[test] + fn a_module_with_no_ports_and_no_topics_passes() { + let builder = builder_with_topics(&[]); + builder + .enforce_topics_match_ports() + .expect("nothing to match"); + } + + #[test] + fn io_uses_mapped_topic() { + let mut builder = builder_with_topics(&[("cmd", "/robot/cmd")]); + let io = builder.io("cmd", |b| Ok(b.to_vec()), |b: &Vec| b.clone()); + assert_eq!(io.topic, "/robot/cmd"); + } + + #[test] + fn io_registers_one_route_and_one_publisher_on_the_same_topic() { + let mut builder = builder_with_topics(&[("cmd", "/robot/cmd")]); + let _io = builder.io("cmd", |b| Ok(b.to_vec()), |b: &Vec| b.clone()); + assert_eq!(builder.routes.get("/robot/cmd").map(Vec::len), Some(1)); + assert_eq!(builder.outputs.len(), 1); + assert_eq!(builder.outputs[0].0, "/robot/cmd"); + } + + #[tokio::test] + async fn io_receives_on_its_route_and_publishes_to_its_queue() { + let mut builder = builder_with_topics(&[("cmd", "/robot/cmd")]); + let mut io = builder.io("cmd", |b| Ok(b.to_vec()), |b: &Vec| b.clone()); + + builder.routes["/robot/cmd"][0].try_dispatch(b"inbound"); + assert_eq!(io.recv().await.expect("inbound message"), b"inbound"); + + io.publish(&b"outbound".to_vec()).await.expect("publish"); + let (_, rx) = &mut builder.outputs[0]; + assert_eq!(rx.recv().await.expect("published bytes"), b"outbound"); + } + + #[tokio::test] + async fn io_publish_errors_when_the_publish_worker_is_gone() { + let mut builder = builder_with_topics(&[]); + let io = builder.io("cmd", |b| Ok(b.to_vec()), |b: &Vec| b.clone()); + builder.outputs.clear(); + let err = io + .publish(&b"x".to_vec()) + .await + .expect_err("publish should fail with no worker"); + assert_eq!(err.kind(), io::ErrorKind::BrokenPipe); + } + + #[test] + fn tf_uses_mapped_topic() { + let mut builder = builder_with_topics(&[("tf", "/robot/tf")]); + builder.tf(); + assert!(builder.routes.contains_key("/robot/tf")); + assert_eq!(builder.outputs[0].0, "/robot/tf"); + } + + #[test] + fn repeated_tf_calls_share_one_graph() { + let mut builder = builder_with_topics(&[("tf", "/tf")]); + builder.tf(); + builder.tf(); + assert_eq!(builder.outputs.len(), 1); + assert_eq!(builder.routes.get("/tf").map(Vec::len), Some(1)); + } + // recv/publish concurrency #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -988,4 +1188,52 @@ mod tests { assert_eq!(route.drop_count.load(Ordering::Relaxed), 1); assert!(logs_contain("handler was full")); } + + // Exercises the code #[derive(Module)] generates for an #[io] field. + mod derive_io { + use super::*; + use crate::Io; + + struct Msg(Vec); + + fn decode(bytes: &[u8]) -> io::Result { + Ok(Msg(bytes.to_vec())) + } + + fn encode(msg: &Msg) -> Vec { + msg.0.clone() + } + + #[derive(crate::Module)] + struct Echo { + #[io(decode = decode, encode = encode)] + cmd: Io, + } + + impl Echo { + async fn handle_cmd(&mut self, msg: Msg) { + if msg.0 == b"ping" { + self.cmd + .publish(&Msg(b"pong".to_vec())) + .await + .expect("publish"); + } + } + } + + #[tokio::test] + async fn io_field_is_wired_to_its_handler_and_can_publish() { + let mut builder = Builder::new(topics(&[("cmd", "/robot/cmd")])); + let mut echo = Echo::build(&mut builder, NoConfig); + + builder.routes["/robot/cmd"][0].try_dispatch(b"ping"); + // Dropping the routes closes the sender, so handle() drains and returns. + builder.routes.clear(); + echo.handle().await; + + let (topic, rx) = &mut builder.outputs[0]; + assert_eq!(topic, "/robot/cmd"); + assert_eq!(rx.recv().await.expect("handler reply"), b"pong"); + } + } } diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs new file mode 100644 index 0000000000..15ffcd5a26 --- /dev/null +++ b/native/rust/dimos-module/src/tf.rs @@ -0,0 +1,1211 @@ +// Copyright 2026 Dimensional Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Transform client for native modules. +//! +//! Each `/tf` edge is buffered per `(parent, child)`, and [`Tf::lookup`] composes +//! the shortest path through the frame graph. Lookups are nearest-in-time within +//! a tolerance, not interpolated. [`Tf::publish`] sends transforms onto the same +//! topic and feeds the local graph. + +use std::collections::{HashMap, HashSet, VecDeque}; +use std::io; +use std::sync::atomic::AtomicU64; +use std::sync::{Arc, Mutex, RwLock}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use nalgebra::{Isometry3, Quaternion, Translation3, UnitQuaternion, Vector3}; +use tokio::sync::{mpsc, Notify}; +use tracing::warn; + +use crate::module::Route; + +/// How many seconds of history each edge keeps. +pub(crate) const DEFAULT_TF_WINDOW_SECS: f64 = 10.0; + +const WARN_INTERVAL: Duration = Duration::from_secs(1); + +fn now_secs() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs_f64()) + .unwrap_or(0.0) +} + +/// A rigid transform from `parent` to `child` at a point in time. +/// +/// It maps a point expressed in `child` coordinates into `parent` coordinates. +#[derive(Clone, Debug)] +pub struct Transform { + pub parent: String, + pub child: String, + pub ts: f64, + iso: Isometry3, +} + +impl Transform { + pub fn new( + parent: impl Into, + child: impl Into, + ts: f64, + iso: Isometry3, + ) -> Self { + Self { + parent: parent.into(), + child: child.into(), + ts, + iso, + } + } + + pub fn translation(&self) -> Vector3 { + self.iso.translation.vector + } + + pub fn rotation(&self) -> UnitQuaternion { + self.iso.rotation + } + + fn inverse(&self) -> Transform { + Transform { + parent: self.child.clone(), + child: self.parent.clone(), + ts: self.ts, + iso: self.iso.inverse(), + } + } + + fn compose(&self, other: &Transform) -> Transform { + Transform { + parent: self.parent.clone(), + child: other.child.clone(), + ts: self.ts, + iso: self.iso * other.iso, + } + } +} + +struct Sample { + ts: f64, + iso: Isometry3, +} + +// One edge's time-sorted history, capped to a fixed-duration window. +struct TBuffer { + window_secs: f64, + samples: VecDeque, +} + +impl TBuffer { + fn new(window_secs: f64) -> Self { + Self { + window_secs, + samples: VecDeque::new(), + } + } + + fn add(&mut self, ts: f64, iso: Isometry3) { + // A stamp a whole window behind the newest is a clock reset, not jitter. + if self + .samples + .back() + .is_some_and(|s| ts < s.ts - self.window_secs) + { + self.samples.clear(); + } + let pos = self.samples.partition_point(|s| s.ts <= ts); + self.samples.insert(pos, Sample { ts, iso }); + // Anchored to the newest sample so a late message cannot widen the window. + let newest = self.samples.back().map_or(ts, |s| s.ts); + self.prune(newest - self.window_secs); + } + + fn prune(&mut self, min_ts: f64) { + let drop_to = self.samples.partition_point(|s| s.ts < min_ts); + for _ in 0..drop_to { + self.samples.pop_front(); + } + } + + fn last(&self) -> Option<&Sample> { + self.samples.back() + } + + // Nearest sample in time. On a tie, prefer the later sample. Returns None + // when the closest sample is further than tolerance from ts. + fn find_closest(&self, ts: f64, tolerance: Option) -> Option<&Sample> { + let pos = self.samples.partition_point(|s| s.ts < ts); + let prev = pos.checked_sub(1).and_then(|i| self.samples.get(i)); + let next = self.samples.get(pos); + let best = match (prev, next) { + (Some(p), Some(n)) => { + if (n.ts - ts).abs() <= (ts - p.ts).abs() { + n + } else { + p + } + } + (Some(p), None) => p, + (None, Some(n)) => n, + (None, None) => return None, + }; + match tolerance { + Some(tol) if (best.ts - ts).abs() > tol => None, + _ => Some(best), + } + } + + fn sample( + &self, + parent: &str, + child: &str, + time: Option, + tolerance: Option, + ) -> Option { + let s = match time { + None => self.last()?, + Some(t) => self.find_closest(t, Some(tolerance.unwrap_or(self.window_secs)))?, + }; + Some(Transform { + parent: parent.to_string(), + child: child.to_string(), + ts: s.ts, + iso: s.iso, + }) + } +} + +/// The transform graph: one [`TBuffer`] per `(parent, child)` edge. +struct MultiTBuffer { + window_secs: f64, + buffers: HashMap<(String, String), TBuffer>, +} + +impl MultiTBuffer { + fn new(window_secs: f64) -> Self { + Self { + window_secs, + buffers: HashMap::new(), + } + } + + fn receive(&mut self, parent: &str, child: &str, ts: f64, iso: Isometry3) { + let window_secs = self.window_secs; + self.buffers + .entry((parent.to_string(), child.to_string())) + .or_insert_with(|| TBuffer::new(window_secs)) + .add(ts, iso); + } + + fn connections(&self, frame: &str) -> Vec { + let mut out = Vec::new(); + for (parent, child) in self.buffers.keys() { + if parent == frame { + out.push(child.clone()); + } + if child == frame { + out.push(parent.clone()); + } + } + out + } + + fn edge( + &self, + parent: &str, + child: &str, + time: Option, + tolerance: Option, + ) -> Option { + if parent == child { + return Some(Transform { + parent: parent.to_string(), + child: child.to_string(), + ts: time.unwrap_or_else(now_secs), + iso: Isometry3::identity(), + }); + } + if let Some(buf) = self.buffers.get(&(parent.to_string(), child.to_string())) { + return buf.sample(parent, child, time, tolerance); + } + if let Some(buf) = self.buffers.get(&(child.to_string(), parent.to_string())) { + return buf + .sample(child, parent, time, tolerance) + .map(|t| t.inverse()); + } + None + } + + fn get( + &self, + parent: &str, + child: &str, + time: Option, + tolerance: Option, + ) -> Option { + if let Some(direct) = self.edge(parent, child, time, tolerance) { + return Some(direct); + } + let path = self.bfs(parent, child, time, tolerance)?; + // A composition is only as fresh as its stalest edge. + let oldest = path.iter().map(|t| t.ts).fold(f64::INFINITY, f64::min); + let mut steps = path.into_iter(); + let first = steps.next()?; + let mut composed = steps.fold(first, |acc, step| acc.compose(&step)); + composed.ts = oldest; + Some(composed) + } + + fn bfs( + &self, + parent: &str, + child: &str, + time: Option, + tolerance: Option, + ) -> Option> { + let mut queue: VecDeque<(String, Vec)> = VecDeque::new(); + queue.push_back((parent.to_string(), Vec::new())); + let mut visited: HashSet = HashSet::new(); + visited.insert(parent.to_string()); + + while let Some((frame, path)) = queue.pop_front() { + if frame == child { + return Some(path); + } + for next in self.connections(&frame) { + if !visited.contains(&next) { + if let Some(edge) = self.edge(&frame, &next, time, tolerance) { + visited.insert(next.clone()); + let mut extended = path.clone(); + extended.push(edge); + queue.push_back((next, extended)); + } + } + } + } + None + } +} + +// The graph plus the signal that it changed. Every write notifies, so a writer +// cannot leave a waiter asleep on a transform that has already landed. +struct Graph { + buffer: RwLock, + changed: Notify, + // Keyed per pair: warn_throttled! throttles per call site, and one site + // serves every lookup, so a missing pair would mute all the others. + warned: Mutex>, +} + +impl Graph { + fn new(window_secs: f64) -> Self { + Self { + buffer: RwLock::new(MultiTBuffer::new(window_secs)), + changed: Notify::new(), + warned: Mutex::new(HashMap::new()), + } + } + + fn should_warn(&self, parent: &str, child: &str) -> bool { + let mut warned = self.warned.lock().expect("tf warn map lock poisoned"); + let last = warned + .entry((parent.to_string(), child.to_string())) + .or_default(); + crate::log::check_and_record(last, WARN_INTERVAL.as_nanos() as u64) + } + + fn update(&self, edits: impl FnOnce(&mut MultiTBuffer)) { + edits(&mut self.buffer.write().expect("tf buffer lock poisoned")); + self.changed.notify_waiters(); + } + + fn get( + &self, + parent: &str, + child: &str, + time: Option, + tolerance: Option, + ) -> Option { + self.buffer + .read() + .expect("tf buffer lock poisoned") + .get(parent, child, time, tolerance) + } +} + +/// A cheap-to-clone handle for querying and publishing transforms. +/// +/// Obtain one from `Builder::tf` (or a `#[tf]` field on a `#[derive(Module)]` +/// struct). The graph is filled in the background as `/tf` messages arrive. +#[derive(Clone)] +pub struct Tf { + graph: Arc, + sender: mpsc::Sender>, +} + +impl Tf { + /// Start a lookup of the transform from `parent` to `child`. + /// + /// Refine it with [`Lookup::at`] and [`Lookup::tolerance`], then finish with + /// [`Lookup::get`]. Use [`Tf::get_latest`] when no refinement is needed. + /// + /// ```ignore + /// let at_scan = tf.lookup("map", "base_link").at(scan_ts).tolerance(0.1).get(); + /// ``` + pub fn lookup<'a>(&'a self, parent: &'a str, child: &'a str) -> Lookup<'a> { + Lookup { + tf: self, + parent, + child, + time: None, + tolerance: None, + } + } + + /// The latest transform from `parent` to `child`. + /// + /// Shorthand for `lookup(parent, child).get()`. + pub fn get_latest(&self, parent: &str, child: &str) -> Option { + self.lookup(parent, child).get() + } + + /// Publish transforms on the `tf` topic. + /// + /// The transforms also feed the local graph, so a lookup right after sees + /// them without waiting for the transport round trip. + pub async fn publish(&self, transforms: &[Transform]) -> io::Result<()> { + self.graph.update(|buffer| { + for t in transforms { + buffer.receive(&t.parent, &t.child, t.ts, t.iso); + } + }); + let msg = lcm_msgs::tf2_msgs::TFMessage { + transforms: transforms.iter().map(to_stamped).collect(), + }; + crate::module::publish_encoded(&self.sender, msg.encode()).await + } +} + +/// A transform lookup being built. Created by [`Tf::lookup`]. +pub struct Lookup<'a> { + tf: &'a Tf, + parent: &'a str, + child: &'a str, + time: Option, + tolerance: Option, +} + +impl Lookup<'_> { + /// Take the sample nearest `time` rather than the latest one. + pub fn at(mut self, time: f64) -> Self { + self.time = Some(time); + self + } + + /// Bound how far, in seconds, the chosen sample may sit from [`Lookup::at`]. + pub fn tolerance(mut self, tolerance: f64) -> Self { + self.tolerance = Some(tolerance); + self + } + + fn resolve(&self) -> Option { + self.tf + .graph + .get(self.parent, self.child, self.time, self.tolerance) + } + + // A lookup that resolves to nothing is otherwise invisible: the caller sees + // None and the buffer says nothing about which frames or stamp missed. + fn warn_unresolved(&self) { + if !self.tf.graph.should_warn(self.parent, self.child) { + return; + } + warn!( + parent = %self.parent, + child = %self.child, + at = self.time.unwrap_or_else(now_secs), + tolerance = self.tolerance.unwrap_or(f64::NAN), + "No transform found between frames", + ); + } + + /// Resolve the lookup against the transforms buffered so far. + /// + /// `None` when no path connects the frames, or when the nearest sample is + /// outside the tolerance. + pub fn get(self) -> Option { + let found = self.resolve(); + if found.is_none() { + self.warn_unresolved(); + } + found + } + + /// Resolve the lookup, waiting up to `timeout` for a late transform. + /// + /// Returns as soon as the lookup succeeds, or `None` at the deadline. + /// Awaiting this inside a `handle_*` method parks the module's whole + /// dispatch loop, so prefer [`Lookup::get`] there and move a long wait onto + /// a task of its own. + pub async fn within(self, timeout: Duration) -> Option { + let deadline = tokio::time::Instant::now() + timeout; + loop { + // Registered before the resolve below, so a transform landing between + // the two still wakes this waiter instead of it sleeping out the + // whole timeout. + let changed = self.tf.graph.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + + if let Some(transform) = self.resolve() { + return Some(transform); + } + // Only the deadline warns. An intermediate miss is the normal state + // of a wait, not a failure. + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + self.warn_unresolved(); + return None; + } + if tokio::time::timeout(remaining, changed).await.is_err() { + self.warn_unresolved(); + return None; + } + } + } +} + +fn to_stamped(t: &Transform) -> lcm_msgs::geometry_msgs::TransformStamped { + let mut sec = t.ts.floor(); + let mut nsec = ((t.ts - sec) * 1e9).round(); + if nsec >= 1e9 { + sec += 1.0; + nsec -= 1e9; + } + let p = t.iso.translation.vector; + let q = t.iso.rotation; + lcm_msgs::geometry_msgs::TransformStamped { + header: lcm_msgs::std_msgs::Header { + seq: 0, + stamp: lcm_msgs::std_msgs::Time { + sec: sec as i32, + nsec: nsec as i32, + }, + frame_id: t.parent.clone(), + }, + child_frame_id: t.child.clone(), + transform: lcm_msgs::geometry_msgs::Transform { + translation: lcm_msgs::geometry_msgs::Vector3 { + x: p.x, + y: p.y, + z: p.z, + }, + rotation: lcm_msgs::geometry_msgs::Quaternion { + x: q.i, + y: q.j, + z: q.k, + w: q.w, + }, + }, + } +} + +// Decodes /tf messages into the shared graph. Registered as a Route so the +// module's existing recv loop dispatches tf traffic to it. +struct TfRoute { + topic: String, + graph: Arc, +} + +impl Route for TfRoute { + fn try_dispatch(&self, data: &[u8]) { + let msg = match lcm_msgs::tf2_msgs::TFMessage::decode(data) { + Ok(msg) => msg, + Err(e) => { + crate::error_throttled!( + Duration::from_secs(1), + topic = %self.topic, + error = %e, + "tf decode error" + ); + return; + } + }; + self.graph.update(|buffer| { + for st in &msg.transforms { + let t = &st.transform.translation; + let q = &st.transform.rotation; + // Normalizing a zero-norm quaternion yields a NaN rotation. + let Some(rotation) = + UnitQuaternion::try_new(Quaternion::new(q.w, q.x, q.y, q.z), 1e-9) + else { + crate::error_throttled!( + Duration::from_secs(1), + topic = %self.topic, + parent = %st.header.frame_id, + child = %st.child_frame_id, + "tf rotation is not a valid quaternion" + ); + continue; + }; + let iso = Isometry3::from_parts(Translation3::new(t.x, t.y, t.z), rotation); + let ts = st.header.stamp.sec as f64 + st.header.stamp.nsec as f64 * 1e-9; + buffer.receive(&st.header.frame_id, &st.child_frame_id, ts, iso); + } + }); + } +} + +// Builds the shared graph plus the handle and the route that feeds it. The +// sender carries published messages to the tf topic's publish worker. +pub(crate) fn tf_subscription( + topic: String, + window_secs: f64, + sender: mpsc::Sender>, +) -> (Tf, Box) { + let graph = Arc::new(Graph::new(window_secs)); + let tf = Tf { + graph: Arc::clone(&graph), + sender, + }; + let route = Box::new(TfRoute { topic, graph }); + (tf, route) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::f64::consts::PI; + + fn tf_with(window_secs: f64) -> (Tf, MultiHandle) { + let (tf, _rx, handle) = tf_with_publish(window_secs); + (tf, handle) + } + + fn tf_with_publish(window_secs: f64) -> (Tf, mpsc::Receiver>, MultiHandle) { + let graph = Arc::new(Graph::new(window_secs)); + let (tx, rx) = mpsc::channel(8); + ( + Tf { + graph: Arc::clone(&graph), + sender: tx, + }, + rx, + MultiHandle { graph }, + ) + } + + // Test-only writer that bypasses LCM and pushes edges straight into the graph. + struct MultiHandle { + graph: Arc, + } + + impl MultiHandle { + fn add(&self, parent: &str, child: &str, ts: f64, xyz: (f64, f64, f64), yaw: f64) { + let iso = Isometry3::from_parts( + Translation3::new(xyz.0, xyz.1, xyz.2), + UnitQuaternion::from_euler_angles(0.0, 0.0, yaw), + ); + self.graph + .update(|buffer| buffer.receive(parent, child, ts, iso)); + } + } + + #[test] + fn direct_edge() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("base_link", "arm", 1.0, (1.0, -1.0, 0.0), 0.0); + let t = tf.get_latest("base_link", "arm").unwrap(); + assert!((t.translation().x - 1.0).abs() < 1e-9); + assert!((t.translation().y + 1.0).abs() < 1e-9); + assert_eq!(t.parent, "base_link"); + assert_eq!(t.child, "arm"); + } + + #[test] + fn reverse_edge_returns_inverse() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("base_link", "arm", 1.0, (1.0, 2.0, 3.0), 0.0); + let inv = tf.get_latest("arm", "base_link").unwrap(); + assert!((inv.translation().x + 1.0).abs() < 1e-9); + assert!((inv.translation().y + 2.0).abs() < 1e-9); + assert!((inv.translation().z + 3.0).abs() < 1e-9); + assert_eq!(inv.parent, "arm"); + assert_eq!(inv.child, "base_link"); + } + + #[test] + fn composes_ros_example_chain() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("base_link", "arm", 1.0, (1.0, -1.0, 0.0), PI / 6.0); + h.add("arm", "end_effector", 1.0, (1.0, 1.0, 0.0), 0.0); + let t = tf.get_latest("base_link", "end_effector").unwrap(); + assert!( + (t.translation().x - 1.366).abs() < 1e-3, + "{}", + t.translation().x + ); + assert!( + (t.translation().y - 0.366).abs() < 1e-3, + "{}", + t.translation().y + ); + assert_eq!(t.parent, "base_link"); + assert_eq!(t.child, "end_effector"); + } + + #[test] + fn composes_multi_hop_chain() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("world", "robot", 1.0, (1.0, 2.0, 3.0), 0.0); + h.add("robot", "sensor", 1.0, (0.5, 0.0, 0.2), PI / 2.0); + let t = tf.get_latest("world", "sensor").unwrap(); + assert!((t.translation().x - 1.5).abs() < 1e-3); + assert!((t.translation().y - 2.0).abs() < 1e-3); + assert!((t.translation().z - 3.2).abs() < 1e-3); + } + + #[test] + fn composed_stamp_is_the_stalest_edge() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("world", "robot", 700.0, (1.0, 0.0, 0.0), 0.0); + h.add("robot", "sensor", 1000.0, (0.5, 0.0, 0.0), 0.0); + assert_eq!(tf.get_latest("world", "sensor").unwrap().ts, 700.0); + // Both directions, so the answer does not depend on which end is queried. + assert_eq!(tf.get_latest("sensor", "world").unwrap().ts, 700.0); + } + + #[test] + fn missing_path_returns_none() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("world", "robot", 1.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.get_latest("world", "unconnected").is_none()); + } + + #[test] + fn identity_for_same_frame() { + let (tf, _h) = tf_with(DEFAULT_TF_WINDOW_SECS); + // No query time: identity is stamped now, not the epoch. + let t = tf.get_latest("base_link", "base_link").unwrap(); + assert!((t.translation().norm()).abs() < 1e-12); + assert!( + t.ts > 0.0, + "identity ts should be a fresh stamp, got {}", + t.ts + ); + // Explicit query time is echoed back. + let at = tf.lookup("base_link", "base_link").at(42.0).get().unwrap(); + assert!((at.ts - 42.0).abs() < 1e-9); + } + + #[test] + fn time_query_picks_nearest_sample() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("a", "b", 10.0, (1.0, 0.0, 0.0), 0.0); + h.add("a", "b", 20.0, (2.0, 0.0, 0.0), 0.0); + let near_10 = tf.lookup("a", "b").at(11.0).get().unwrap(); + assert!((near_10.translation().x - 1.0).abs() < 1e-9); + let near_20 = tf.lookup("a", "b").at(18.0).get().unwrap(); + assert!((near_20.translation().x - 2.0).abs() < 1e-9); + } + + #[test] + fn time_query_outside_tolerance_returns_none() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("a", "b", 10.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.lookup("a", "b").at(50.0).tolerance(1.0).get().is_none()); + assert!(tf.lookup("a", "b").at(10.5).tolerance(1.0).get().is_some()); + } + + #[test] + fn time_query_beyond_the_window_returns_none_without_a_tolerance() { + let (tf, h) = tf_with(10.0); + h.add("a", "b", 100.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.lookup("a", "b").at(50.0).get().is_none()); + } + + #[test] + fn time_query_inside_the_window_resolves_without_a_tolerance() { + let (tf, h) = tf_with(10.0); + h.add("a", "b", 100.0, (1.0, 0.0, 0.0), 0.0); + let t = tf + .lookup("a", "b") + .at(95.0) + .get() + .expect("within the window"); + assert!((t.translation().x - 1.0).abs() < 1e-9); + } + + // An explicit tolerance is the caller opting into staleness, so it widens + // past the window rather than being clamped by it. + #[test] + fn an_explicit_tolerance_reaches_past_the_window() { + let (tf, h) = tf_with(10.0); + h.add("a", "b", 100.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.lookup("a", "b").at(50.0).tolerance(60.0).get().is_some()); + } + + // The window bounds queries against a stamp, not the latest sample. With no + // requested time, the newest edge is returned however old it is. + #[test] + fn latest_is_not_bounded_by_the_window() { + let (tf, h) = tf_with(10.0); + h.add("a", "b", 100.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.get_latest("a", "b").is_some()); + } + + #[tokio::test] + async fn within_returns_without_waiting_when_already_buffered() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("a", "b", 5.0, (1.0, 0.0, 0.0), 0.0); + let t = tf + .lookup("a", "b") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_secs(30)) + .await + .expect("already available"); + assert!((t.translation().x - 1.0).abs() < 1e-9); + } + + // Returns as soon as the transform lands, not at the deadline. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn within_resolves_when_the_transform_arrives_late() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(30)).await; + h.add("a", "b", 5.0, (1.0, 0.0, 0.0), 0.0); + }); + let started = tokio::time::Instant::now(); + let t = tf + .lookup("a", "b") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_secs(30)) + .await + .expect("arrived inside the budget"); + assert!( + started.elapsed() < Duration::from_secs(5), + "waited {:?}, should have returned on arrival", + started.elapsed() + ); + assert!((t.translation().x - 1.0).abs() < 1e-9); + } + + #[tokio::test] + async fn within_times_out_when_nothing_arrives() { + let (tf, _h) = tf_with(DEFAULT_TF_WINDOW_SECS); + let started = tokio::time::Instant::now(); + let t = tf + .lookup("a", "b") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_millis(50)) + .await; + assert!(t.is_none()); + assert!( + started.elapsed() >= Duration::from_millis(50), + "returned early" + ); + } + + // The waiter is on a -> c, but what lands is b -> c. Waking only on the + // queried edge would sleep through the composition becoming possible. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn within_wakes_when_a_later_edge_completes_the_chain() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("a", "b", 5.0, (1.0, 0.0, 0.0), 0.0); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(30)).await; + h.add("b", "c", 5.0, (2.0, 0.0, 0.0), 0.0); + }); + let t = tf + .lookup("a", "c") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_secs(30)) + .await + .expect("chain completed inside the budget"); + assert!( + (t.translation().x - 3.0).abs() < 1e-9, + "{}", + t.translation().x + ); + } + + // A waiting handler is woken by the transport's dispatch task, which its own + // stall cannot block. Guards against the wait deadlocking against tf intake. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn within_is_woken_by_a_transform_dispatched_through_the_route() { + let (tx, _rx) = mpsc::channel(8); + let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_WINDOW_SECS, tx); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(30)).await; + route.try_dispatch(&stamped_message("a", "b", 5.0, 1.0)); + }); + let t = tf + .lookup("a", "b") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_secs(30)) + .await + .expect("route dispatch woke the waiter"); + assert!((t.translation().x - 1.0).abs() < 1e-9); + } + + fn stamped_message(parent: &str, child: &str, ts: f64, x: f64) -> Vec { + rotated_message(parent, child, ts, x, (0.0, 0.0, 0.0, 1.0)) + } + + fn rotated_message( + parent: &str, + child: &str, + ts: f64, + x: f64, + quat: (f64, f64, f64, f64), + ) -> Vec { + use lcm_msgs::geometry_msgs::{ + Quaternion as LQuat, Transform as LTransform, TransformStamped, Vector3 as LVec3, + }; + use lcm_msgs::std_msgs::{Header, Time}; + let (x_q, y_q, z_q, w_q) = quat; + lcm_msgs::tf2_msgs::TFMessage { + transforms: vec![TransformStamped { + header: Header { + seq: 0, + stamp: Time { + sec: ts as i32, + nsec: 0, + }, + frame_id: parent.to_string(), + }, + child_frame_id: child.to_string(), + transform: LTransform { + translation: LVec3 { x, y: 0.0, z: 0.0 }, + rotation: LQuat { + x: x_q, + y: y_q, + z: z_q, + w: w_q, + }, + }, + }], + } + .encode() + } + + #[test] + fn a_zero_rotation_on_the_wire_is_dropped_rather_than_stored_as_nan() { + let (tx, _rx) = mpsc::channel(4); + let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_WINDOW_SECS, tx); + route.try_dispatch(&rotated_message("a", "b", 5.0, 1.0, (0.0, 0.0, 0.0, 0.0))); + assert!(tf.get_latest("a", "b").is_none()); + + route.try_dispatch(&stamped_message("a", "b", 6.0, 1.0)); + let t = tf.get_latest("a", "b").expect("valid rotation is accepted"); + assert!(t.rotation().coords.iter().all(|c| c.is_finite())); + } + + #[test] + #[tracing_test::traced_test] + fn a_lookup_that_finds_nothing_warns_with_the_frames() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("world", "robot", 1.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.get_latest("world", "gripper").is_none()); + assert!(logs_contain("No transform found between frames")); + assert!(logs_contain("gripper")); + } + + #[test] + #[tracing_test::traced_test] + fn repeated_misses_warn_once_per_frame_pair() { + let (tf, _h) = tf_with(DEFAULT_TF_WINDOW_SECS); + for _ in 0..5 { + assert!(tf.get_latest("world", "gripper").is_none()); + } + logs_assert(|lines: &[&str]| { + match lines.iter().filter(|l| l.contains("gripper")).count() { + 1 => Ok(()), + n => Err(format!("expected 1 warning for the repeated pair, got {n}")), + } + }); + + // A pair that is throttled must not mute an unrelated one. + assert!(tf.get_latest("world", "camera").is_none()); + assert!(logs_contain("camera")); + } + + #[test] + #[tracing_test::traced_test] + fn a_resolved_lookup_stays_quiet() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("world", "robot", 1.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.get_latest("world", "robot").is_some()); + assert!(!logs_contain("No transform found between frames")); + } + + // A wait in progress is not a failure, so only the deadline warns. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[tracing_test::traced_test] + async fn within_warns_only_once_it_gives_up() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(20)).await; + h.add("a", "b", 5.0, (1.0, 0.0, 0.0), 0.0); + }); + assert!(tf + .lookup("a", "b") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_secs(30)) + .await + .is_some()); + assert!(!logs_contain("No transform found between frames")); + + assert!(tf + .lookup("a", "missing") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_millis(20)) + .await + .is_none()); + assert!(logs_contain("No transform found between frames")); + } + + #[test] + fn prunes_samples_outside_window() { + let mut buf = TBuffer::new(5.0); + buf.add(1.0, Isometry3::identity()); + buf.add(2.0, Isometry3::identity()); + buf.add(10.0, Isometry3::identity()); + assert_eq!(buf.samples.len(), 1); + assert!((buf.last().unwrap().ts - 10.0).abs() < 1e-9); + } + + #[test] + fn a_late_sample_does_not_spare_ones_the_window_has_aged_out() { + let mut buf = TBuffer::new(5.0); + buf.add(10.0, Isometry3::identity()); + buf.add(11.0, Isometry3::identity()); + // Late, but still inside the window. + buf.add(7.0, Isometry3::identity()); + assert_eq!(buf.samples.len(), 3); + + buf.add(20.0, Isometry3::identity()); + assert_eq!(buf.samples.len(), 1); + assert!((buf.last().unwrap().ts - 20.0).abs() < 1e-9); + } + + #[test] + fn a_clock_reset_drops_the_pre_jump_samples() { + let mut buf = TBuffer::new(5.0); + for i in 0..20 { + buf.add(1000.0 + i as f64, Isometry3::identity()); + } + for i in 0..20 { + buf.add(100.0 + i as f64, Isometry3::identity()); + } + assert!( + buf.samples.len() <= 6, + "buffer grew to {}", + buf.samples.len() + ); + assert!((buf.last().unwrap().ts - 119.0).abs() < 1e-9); + } + + #[test] + fn tf_route_decodes_into_graph() { + use lcm_msgs::geometry_msgs::{ + Quaternion as LQuat, Transform as LTransform, Vector3 as LVec3, + }; + use lcm_msgs::std_msgs::{Header, Time}; + use lcm_msgs::tf2_msgs::TFMessage; + + let (tx, _rx) = mpsc::channel(8); + let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_WINDOW_SECS, tx); + let msg = TFMessage { + transforms: vec![lcm_msgs::geometry_msgs::TransformStamped { + header: Header { + seq: 0, + stamp: Time { + sec: 5, + nsec: 500_000_000, + }, + frame_id: "base_link".to_string(), + }, + child_frame_id: "mid360_link".to_string(), + transform: LTransform { + translation: LVec3 { + x: 0.1, + y: 0.2, + z: 0.3, + }, + rotation: LQuat { + x: 0.0, + y: 0.0, + z: 0.0, + w: 1.0, + }, + }, + }], + }; + route.try_dispatch(&msg.encode()); + + let t = tf.get_latest("base_link", "mid360_link").unwrap(); + assert!((t.translation().x - 0.1).abs() < 1e-9); + assert!((t.translation().y - 0.2).abs() < 1e-9); + assert!((t.translation().z - 0.3).abs() < 1e-9); + assert!((t.ts - 5.5).abs() < 1e-9); + } + + #[tokio::test] + async fn publish_feeds_local_graph() { + let (tf, _rx, _h) = tf_with_publish(DEFAULT_TF_WINDOW_SECS); + let iso = Isometry3::from_parts( + Translation3::new(1.0, 2.0, 3.0), + UnitQuaternion::from_euler_angles(0.0, 0.0, PI / 2.0), + ); + tf.publish(&[Transform::new("map", "base_link", 7.0, iso)]) + .await + .unwrap(); + + let t = tf.get_latest("map", "base_link").unwrap(); + assert!((t.translation().x - 1.0).abs() < 1e-9); + assert!((t.translation().y - 2.0).abs() < 1e-9); + assert!((t.translation().z - 3.0).abs() < 1e-9); + assert!((t.ts - 7.0).abs() < 1e-9); + } + + #[tokio::test] + async fn publish_round_trips_through_route() { + let (tf_out, mut rx, _h) = tf_with_publish(DEFAULT_TF_WINDOW_SECS); + let iso = Isometry3::from_parts( + Translation3::new(0.5, -0.5, 0.25), + UnitQuaternion::from_euler_angles(0.0, 0.0, PI / 6.0), + ); + tf_out + .publish(&[Transform::new("a", "b", 3.25, iso)]) + .await + .unwrap(); + let bytes = rx.recv().await.unwrap(); + + let (tx, _rx2) = mpsc::channel(8); + let (tf_in, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_WINDOW_SECS, tx); + route.try_dispatch(&bytes); + + let t = tf_in.get_latest("a", "b").unwrap(); + assert!((t.translation().x - 0.5).abs() < 1e-9); + assert!((t.translation().y + 0.5).abs() < 1e-9); + assert!((t.translation().z - 0.25).abs() < 1e-9); + assert!((t.ts - 3.25).abs() < 1e-9); + let (_, _, yaw) = t.rotation().euler_angles(); + assert!((yaw - PI / 6.0).abs() < 1e-9); + } + + #[test] + fn stamp_rounding_does_not_overflow_nsec() { + let st = to_stamped(&Transform::new( + "a", + "b", + 1.9999999999, + Isometry3::identity(), + )); + assert_eq!(st.header.stamp.sec, 2); + assert_eq!(st.header.stamp.nsec, 0); + } + + #[test] + fn add_out_of_order_keeps_samples_sorted() { + let mut buf = TBuffer::new(DEFAULT_TF_WINDOW_SECS); + buf.add(3.0, Isometry3::identity()); + buf.add(1.0, Isometry3::identity()); + buf.add(2.0, Isometry3::identity()); + assert!((buf.last().unwrap().ts - 3.0).abs() < 1e-9); + let s = buf.find_closest(1.9, None).unwrap(); + assert!((s.ts - 2.0).abs() < 1e-9); + } + + #[test] + fn tie_prefers_the_later_sample() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("a", "b", 10.0, (1.0, 0.0, 0.0), 0.0); + h.add("a", "b", 12.0, (2.0, 0.0, 0.0), 0.0); + let t = tf.lookup("a", "b").at(11.0).get().unwrap(); + assert!((t.translation().x - 2.0).abs() < 1e-9); + } + + // Two routes to d: three hops through b, c and two through x. BFS must + // compose the two-hop route. + #[test] + fn bfs_takes_the_fewest_hops_on_a_branching_graph() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("a", "b", 1.0, (1.0, 0.0, 0.0), 0.0); + h.add("b", "c", 1.0, (1.0, 0.0, 0.0), 0.0); + h.add("c", "d", 1.0, (1.0, 0.0, 0.0), 0.0); + h.add("a", "x", 1.0, (10.0, 0.0, 0.0), 0.0); + h.add("x", "d", 1.0, (1.0, 0.0, 0.0), 0.0); + let t = tf.get_latest("a", "d").unwrap(); + assert!( + (t.translation().x - 11.0).abs() < 1e-9, + "{}", + t.translation().x + ); + } + + // Inverse of a rotated edge is t' = -R^T t, the classic sign/order trap. + #[test] + fn reverse_edge_inverts_rotation_and_translation() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("base_link", "arm", 1.0, (1.0, 2.0, 3.0), PI / 2.0); + let inv = tf.get_latest("arm", "base_link").unwrap(); + assert!( + (inv.translation().x + 2.0).abs() < 1e-9, + "{:?}", + inv.translation() + ); + assert!((inv.translation().y - 1.0).abs() < 1e-9); + assert!((inv.translation().z + 3.0).abs() < 1e-9); + let (_, _, yaw) = inv.rotation().euler_angles(); + assert!((yaw + PI / 2.0).abs() < 1e-9); + } + + #[test] + fn composed_chain_accumulates_rotation() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("a", "b", 1.0, (0.0, 0.0, 0.0), PI / 6.0); + h.add("b", "c", 1.0, (0.0, 0.0, 0.0), PI / 6.0); + let t = tf.get_latest("a", "c").unwrap(); + let (_, _, yaw) = t.rotation().euler_angles(); + assert!((yaw - PI / 3.0).abs() < 1e-9); + } + + #[test] + fn dispatch_of_undecodable_bytes_leaves_the_graph_empty() { + let (tx, _rx) = mpsc::channel(8); + let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_WINDOW_SECS, tx); + route.try_dispatch(b"garbage"); + assert!(tf.get_latest("a", "b").is_none()); + } + + #[tokio::test] + async fn publish_errors_when_the_background_task_is_gone() { + let (tf, rx, _h) = tf_with_publish(DEFAULT_TF_WINDOW_SECS); + drop(rx); + let err = tf + .publish(&[Transform::new("a", "b", 1.0, Isometry3::identity())]) + .await + .unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::BrokenPipe); + } +}