diff --git a/dimos/mapping/cli/map.py b/dimos/mapping/cli/map.py index ec0197b1e6b..6082ebd6d6f 100644 --- a/dimos/mapping/cli/map.py +++ b/dimos/mapping/cli/map.py @@ -15,6 +15,9 @@ from __future__ import annotations from collections.abc import Callable, Iterable +from copy import deepcopy +from importlib import resources +from itertools import chain import math from pathlib import Path import subprocess @@ -25,9 +28,8 @@ if TYPE_CHECKING: from dimos.memory.stream import Stream from dimos.memory.type.observation import Observation - from dimos.msgs.geometry_msgs.Transform import Transform - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos.messages.geometry_msgs.msg import TransformStamped + from dimos.messages.sensor_msgs.msg import CameraInfo, Image, PointCloud2 from dimos.navigation.go2.loop_closure.pgo import PoseGraph PATH_THICKNESS = 0.01 @@ -40,6 +42,19 @@ _WORLD_FRAMES = ("world", "map", "odom") +def _camera_info_static() -> CameraInfo: + """Load the packaged calibration without importing a robot connection.""" + from dimos.message_utils.camera_info import camera_info_from_yaml + from dimos.message_utils.time import time_from_seconds + from dimos.messages.std_msgs.msg import Header + + source = resources.files("dimos.robot.unitree.go2").joinpath("front_camera_720.yaml") + with resources.as_file(source) as path: + return camera_info_from_yaml( + path, header=Header(stamp=time_from_seconds(0), frame_id="camera_optical") + ) + + def _detect_world(tf_buf: Any, cloud_frame: str, ts: float) -> str | None: """Pick the first conventional world frame that resolves the cloud frame via tf.""" if cloud_frame in _WORLD_FRAMES: @@ -109,7 +124,7 @@ def _accumulate( block_count: int, device: str, graph: PoseGraph | None = None, - register: Callable[[Observation[Any]], Transform | None] | None = None, + register: Callable[[Observation[Any]], TransformStamped | None] | None = None, carve_columns: bool = False, progress_cb: Callable[[Observation[Any]], None] | None = None, ) -> PointCloud2 | None: @@ -123,17 +138,21 @@ def _accumulate( Disposal of the underlying ``VoxelGrid`` is handled by ``VoxelMapTransformer``. """ from dimos.mapping.voxels.module import VoxelMapTransformer + from dimos.message_utils.geometry import compose_transform_values + from dimos.message_utils.pointcloud import transform_cloud + from dimos.messages.geometry_msgs.msg import TransformStamped + from dimos.messages.std_msgs.msg import Header def prepared() -> Iterable[Observation[PointCloud2]]: for obs in obs_iter: if progress_cb is not None: progress_cb(obs) - if len(obs.data) == 0: + if obs.data.width * obs.data.height == 0: continue # sensor->world via `register`, unless the clouds are already # world-registered. graph adds the PGO correction on top # (correction ∘ tf), applied after the registration. - tf: Transform | None = None + tf: TransformStamped | None = None if register is not None: tf = register(obs) if tf is None: @@ -142,28 +161,44 @@ def prepared() -> Iterable[Observation[PointCloud2]]: if obs.pose_tuple is None: continue correction = graph.correction_at(obs.ts) - tf = correction if tf is None else correction + tf - yield obs if tf is None else obs.derive(data=obs.data.transform(tf)) + tf = TransformStamped( + header=Header( + stamp=deepcopy(obs.data.header.stamp), frame_id=correction.header.frame_id + ), + child_frame_id=obs.data.header.frame_id, + transform=correction.transform + if tf is None + else compose_transform_values(correction.transform, tf.transform), + ) + yield obs if tf is None else obs.derive(data=transform_cloud(obs.data, tf)) + frames = iter(prepared()) + first = next(frames, None) + if first is None: + return None vmt = VoxelMapTransformer( + frame_id=first.data.header.frame_id, emit_every=0, # batch mode: emit once on exhaustion voxel_size=voxel, block_count=block_count, device=device, carve_columns=carve_columns, ) - result = next(iter(vmt(iter(prepared()))), None) + result = next(iter(vmt(chain([first], frames))), None) return result.data if result is not None else None def _denoise(cloud: PointCloud2 | None) -> PointCloud2 | None: """Statistical outlier removal via o3d; drops sparse floaters, keeps colors.""" - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + import open3d as o3d # type: ignore[import-untyped] - if cloud is None or len(cloud.pointcloud.points) < 20: + from dimos.message_utils.pointcloud import pointcloud_to_open3d, select_points + + if cloud is None or cloud.width * cloud.height < 20: return cloud - clean, _ = cloud.pointcloud_tensor.remove_statistical_outliers(nb_neighbors=20, std_ratio=2.0) - return PointCloud2(pointcloud=clean, frame_id=cloud.frame_id, ts=cloud.ts) + tensor = o3d.t.geometry.PointCloud.from_legacy(pointcloud_to_open3d(cloud)) + _, mask = tensor.remove_statistical_outliers(nb_neighbors=20, std_ratio=2.0) + return select_points(cloud, mask.numpy()) def _log_reconstruction( @@ -184,13 +219,22 @@ def _log_reconstruction( import rerun.blueprint as rrb from dimos.memory.vis.color import Color - from dimos.msgs.geometry_msgs.Transform import Transform + from dimos.message_utils.time import time_from_seconds + from dimos.messages.geometry_msgs.msg import Transform, TransformStamped, Vector3 + from dimos.messages.std_msgs.msg import Header + from dimos.visualization.rerun.cdr import cloud_archetype rr.send_blueprint(rrb.Blueprint(rrb.Spatial3DView(origin="world"))) if global_map is not None: rr.log( "world/raw_map/pointcloud", - global_map.to_rerun(voxel_size=voxel / 2, bottom_cutoff=bottom_cutoff), + cloud_archetype( + global_map, + voxel_size=voxel / 2, + bottom_cutoff=bottom_cutoff, + mode="spheres", + height_colors=True, + ), static=True, ) if path: @@ -202,13 +246,25 @@ def _log_reconstruction( if pgo_map is not None: rr.log( "world/pgo_map/pointcloud", - pgo_map.to_rerun(voxel_size=voxel / 2, bottom_cutoff=bottom_cutoff), + cloud_archetype( + pgo_map, + voxel_size=voxel / 2, + bottom_cutoff=bottom_cutoff, + mode="spheres", + height_colors=True, + ), static=True, ) if full_pgo_map is not None: rr.log( "world/full_pgo_map/pointcloud", - full_pgo_map.to_rerun(voxel_size=voxel / 2, bottom_cutoff=bottom_cutoff), + cloud_archetype( + full_pgo_map, + voxel_size=voxel / 2, + bottom_cutoff=bottom_cutoff, + mode="spheres", + height_colors=True, + ), static=True, ) if pgo_path: @@ -225,8 +281,16 @@ def _log_reconstruction( if graph is not None and graph.loops: loop_strips = [ [ - (lc.source.translation.x, lc.source.translation.y, lc.source.translation.z), - (lc.target.translation.x, lc.target.translation.y, lc.target.translation.z), + ( + lc.source.transform.translation.x, + lc.source.transform.translation.y, + lc.source.transform.translation.z, + ), + ( + lc.target.transform.translation.x, + lc.target.transform.translation.y, + lc.target.transform.translation.z, + ), ] for lc in graph.loops ] @@ -269,23 +333,28 @@ def _log_reconstruction( pgo_centers: list[tuple[float, float, float]] = [] pgo_quats: list[tuple[float, float, float, float]] = [] for d in marker_dets: - raw_tf = Transform( - translation=d.data.center, - rotation=d.data.orientation, - frame_id="world", + raw_tf = TransformStamped( + header=Header(stamp=time_from_seconds(d.ts), frame_id="world"), child_frame_id=f"marker_{d.data.marker_id}", - ts=d.ts, + transform=Transform( + translation=Vector3(d.data.center.x, d.data.center.y, d.data.center.z), + rotation=deepcopy(d.data.orientation), + ), ) corrected = graph.correct(raw_tf) pgo_centers.append( - (corrected.translation.x, corrected.translation.y, corrected.translation.z) + ( + corrected.transform.translation.x, + corrected.transform.translation.y, + corrected.transform.translation.z, + ) ) pgo_quats.append( ( - corrected.rotation.x, - corrected.rotation.y, - corrected.rotation.z, - corrected.rotation.w, + corrected.transform.rotation.x, + corrected.transform.rotation.y, + corrected.transform.rotation.z, + corrected.transform.rotation.w, ) ) _log_markers( @@ -328,7 +397,7 @@ def main( export: bool = typer.Option( False, "--export", - help="Export PGO map to ./.pc2.lcm in cwd (implies --pgo)", + help="Export PGO map to ./.pc2.cdr in cwd (implies --pgo)", ), full_pgo: bool = typer.Option( False, @@ -418,17 +487,20 @@ def main( ), ) -> None: """Rebuild a voxel map from a recorded SQLite dataset, write a .rrd, and open it in rerun.""" + from dimos_cdrgen.registry import encode import rerun as rr from dimos.memory.cli.dataset import open_store, resolve_dataset from dimos.memory.transform import QualityWindow, SpeedLimit from dimos.memory.utils.progress import progress - from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos.message_utils.camera_info import camera_info_from_yaml + from dimos.message_utils.image import image_sharpness + from dimos.message_utils.time import time_from_seconds + from dimos.messages.geometry_msgs.msg import Quaternion, Transform, Vector3 + from dimos.messages.sensor_msgs.msg import Image, PointCloud2 + from dimos.messages.std_msgs.msg import Header from dimos.navigation.go2.loop_closure.pgo import PGO from dimos.perception.fiducial.marker_transformer import DetectMarkers - from dimos.robot.unitree.go2.connection import BASE_TO_OPTICAL, _camera_info_static from dimos.visualization.rerun.init import rerun_init db_path = resolve_dataset(dataset) @@ -454,7 +526,7 @@ def main( tf_buf = StreamTF.from_store(store) # Streams are homogeneous: read the cloud frame from the first observation. first_obs = next(iter(lidar), None) - cloud_frame: str | None = first_obs.data.frame_id if first_obs is not None else None + cloud_frame: str | None = first_obs.data.header.frame_id if first_obs is not None else None world = frame if world is None and first_obs is not None and cloud_frame is not None: @@ -473,7 +545,7 @@ def main( # Registration: sensor-frame clouds get a per-frame tf lookup lifting them # into the world frame (frames with no tf answer are dropped); clouds # already stamped with the world frame accumulate verbatim (register=None). - register: Callable[[Observation[Any]], Transform | None] | None = None + register: Callable[[Observation[Any]], TransformStamped | None] | None = None if first_obs is not None and cloud_frame is not None and cloud_frame != world: # Fail fast when registration is impossible: probe the first cloud's # timestamp (unbounded tolerance — "possible at all", not "in range"). @@ -490,8 +562,10 @@ def main( print(f"registering clouds {world!r} ← {cloud_frame!r} via tf") buf = tf_buf - def _register(obs: Observation[Any]) -> Transform | None: - return buf.get(world, obs.data.frame_id, time_point=obs.ts, time_tolerance=tf_tolerance) + def _register(obs: Observation[Any]) -> TransformStamped | None: + return buf.get( + world, obs.data.header.frame_id, time_point=obs.ts, time_tolerance=tf_tolerance + ) register = _register elif cloud_frame is not None: @@ -504,12 +578,15 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: tf = register(obs) if tf is None: return None - return (tf.translation.x, tf.translation.y, tf.translation.z) - pose = obs.pose + return ( + tf.transform.translation.x, + tf.transform.translation.y, + tf.transform.translation.z, + ) + pose = obs.pose_tuple # Reject placeholder poses: zero translation OR uninitialized rotation. - # Same condition as pgo_keyframes so dedup and PGO see the same frames. - if pose is not None and not (pose.position.is_zero() or pose.orientation.is_zero()): - return (pose.position.x, pose.position.y, pose.position.z) + if pose is not None and any(pose[:3]) and any(pose[3:]): + return (pose[0], pose[1], pose[2]) return None # Spatial dedup: bucket frames by 3D cell using the trajectory position, @@ -553,7 +630,11 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: graph = lidar.tap(bar).transform(PGO()).last().data pgo_path = [ - (kf.optimized.translation.x, kf.optimized.translation.y, kf.optimized.translation.z) + ( + kf.optimized.transform.translation.x, + kf.optimized.transform.translation.y, + kf.optimized.transform.translation.z, + ) for kf in graph.keyframes ] @@ -618,8 +699,21 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: store.stream(image_pose).from_time(seek or None).to_time(duration) ) print(f"re-posing color_image from {image_pose!r} + camera optical mount") - color_image = pose_fill(color_image, src_pose, tolerance=0.1, mount=BASE_TO_OPTICAL) - cam_info = CameraInfo.from_yaml(str(camera_info)) if camera_info else _camera_info_static() + color_image = pose_fill( + color_image, + src_pose, + tolerance=0.1, + mount=Transform( + translation=Vector3(0.3, 0, 0), rotation=Quaternion(-0.5, 0.5, -0.5, 0.5) + ), + ) + cam_info = ( + camera_info_from_yaml( + camera_info, header=Header(stamp=time_from_seconds(0), frame_id="camera_optical") + ) + if camera_info + else _camera_info_static() + ) xf = DetectMarkers( camera_info=cam_info, marker_length_m=marker_size, @@ -631,7 +725,7 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: # than the limits. Defaults match replay_marker.py so positions agree. with progress(n_images, "detecting markers") as bar: pipeline: Stream[Image] = color_image.tap(bar).transform( - QualityWindow(lambda img: img.sharpness, window=marker_quality_window) + QualityWindow(image_sharpness, window=marker_quality_window) ) if marker_max_speed > 0: pipeline = pipeline.transform( @@ -677,14 +771,15 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: subprocess.Popen(["rerun", str(out)]) if export and pgo_map is not None: - out_path = Path.cwd() / f"{db_path.stem}.pc2.lcm" + out_path = Path.cwd() / f"{db_path.stem}.pc2.cdr" print(f"exporting PGO twopass map to {out_path}...") - out_path.write_bytes(pgo_map.lcm_encode()) + out_path.write_bytes(encode(pgo_map)) print(f"wrote {out_path}") print() print("load back with:") - print(" from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2") - print(f' pcd = PointCloud2.lcm_decode(open("{out_path.name}", "rb").read())') + print(" from dimos.messages.sensor_msgs.msg import PointCloud2") + print(" from dimos_cdrgen.registry import decode") + print(f' pcd = decode(open("{out_path.name}", "rb").read(), PointCloud2)') if __name__ == "__main__": diff --git a/dimos/mapping/cli/pose_fill.py b/dimos/mapping/cli/pose_fill.py index 38465ff59fb..d220a2d10d3 100644 --- a/dimos/mapping/cli/pose_fill.py +++ b/dimos/mapping/cli/pose_fill.py @@ -40,7 +40,7 @@ from dimos.memory.backend import Backend from dimos.memory.stream import Stream from dimos.memory.type.observation import Observation - from dimos.msgs.geometry_msgs.Transform import Transform + from dimos.messages.geometry_msgs.msg import Transform def pose_fill( @@ -54,10 +54,10 @@ def pose_fill( Pairs are formed by :meth:`Stream.align` (nearest ``|Δts| <= tolerance``); primaries with no match in tolerance are dropped. The pose is read from the - matched pose observation's *payload* (``PoseStamped`` / ``Odometry`` / any - object exposing ``.position`` + ``.orientation``) because pose-message - streams carry the pose in the value, not the indexed pose columns. The - target stream's payload stays lazy. + matched pose observation's payload: generated ``Pose``, ``Transform`` or + their standard stamped/covariance/odometry forms. Pose-message streams + carry the pose in the value, not the indexed pose columns. The target + stream's payload stays lazy. *mount* composes a static child transform onto each source pose (``world_base ∘ mount``) before attaching — e.g. ``base_link → @@ -65,9 +65,8 @@ def pose_fill( image pose. """ from dimos.memory.type.observation import _to_tuple - 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.message_utils.geometry import compose_transform_values + from dimos.messages.geometry_msgs.msg import Quaternion, Transform, Vector3 def _fill(pair_obs: Observation[Any]) -> Observation[Any]: primary, secondary = cast( @@ -77,7 +76,7 @@ def _fill(pair_obs: Observation[Any]) -> Observation[Any]: return primary.with_pose(secondary.data) x, y, z, qx, qy, qz, qw = cast("tuple[float, ...]", _to_tuple(secondary.data)) world_base = Transform(translation=Vector3(x, y, z), rotation=Quaternion(qx, qy, qz, qw)) - return primary.with_pose(world_base + mount) + return primary.with_pose(compose_transform_values(world_base, mount)) return stream.align(pose_stream, tolerance=tolerance).map(_fill) @@ -112,48 +111,57 @@ def pose_fill_db( "collide on ts). Delete it (and any -wal/-shm sidecars) first." ) - src = SqliteStore(path=str(src_path), must_exist=True) - dest = SqliteStore(path=str(dest_path)) - names = streams if streams is not None else src.list_streams() - if target not in names: - raise ValueError(f"target stream {target!r} not in {names}") - if pose_source not in src.list_streams(): - raise ValueError(f"pose_source stream {pose_source!r} not found in source dataset") - - written: dict[str, int] = {} - for name in names: - src_b = cast("Backend[Any]", src.stream(name)._source) - dest_b = cast("Backend[Any]", dest.stream(name, src_b.data_type, codec=src_b.codec)._source) - - # Re-posed observations for the target (ids preserved from source, so - # the source blob still resolves); plain metadata for everything else. - if name == target: - rows: Iterable[Observation[Any]] = pose_fill( - src.stream(name).from_time(seek or None).to_time(duration).order_by("ts"), - src.stream(pose_source).from_time(seek or None).to_time(duration).order_by("ts"), - tolerance=tolerance, + with ( + SqliteStore(path=str(src_path), must_exist=True) as src, + SqliteStore(path=str(dest_path)) as dest, + ): + names = streams if streams is not None else src.list_streams() + if target not in names: + raise ValueError(f"target stream {target!r} not in {names}") + if pose_source not in src.list_streams(): + raise ValueError(f"pose_source stream {pose_source!r} not found in source dataset") + + written: dict[str, int] = {} + for name in names: + src_b = cast("Backend[Any]", src.stream(name)._source) + dest_b = cast( + "Backend[Any]", dest.stream(name, src_b.data_type, codec=src_b.codec)._source ) - else: - rows = src.stream(name).from_time(seek or None).to_time(duration).order_by("ts") - - scalar = src_b.data_type in (int, float) - n = 0 - for obs in rows: - blob = ( - None if scalar or src_b.blob_store is None else src_b.blob_store.get(name, obs.id) - ) - row_id = dest_b.metadata_store.insert(obs) - if blob is not None: - assert dest_b.blob_store is not None - dest_b.blob_store.put(name, row_id, blob) - n += 1 - commit = getattr(dest_b.metadata_store, "commit", None) - if commit is not None: - commit() - written[name] = n - - src.stop() - dest.stop() + + # Re-posed observations for the target (ids preserved from source, so + # the source blob still resolves); plain metadata for everything else. + if name == target: + rows: Iterable[Observation[Any]] = pose_fill( + src.stream(name).from_time(seek or None).to_time(duration).order_by("ts"), + src.stream(pose_source) + .from_time(seek or None) + .to_time(duration) + .order_by("ts"), + tolerance=tolerance, + ) + else: + rows = src.stream(name).from_time(seek or None).to_time(duration).order_by("ts") + + scalar = src_b.data_type in (int, float) + n = 0 + for obs in rows: + blob = ( + None + if scalar or src_b.blob_store is None + else src_b.blob_store.get(name, obs.id) + ) + row_id = dest_b.metadata_store.insert(obs) + if blob is not None: + assert dest_b.blob_store is not None + dest_b.blob_store.put(name, row_id, blob) + n += 1 + commit = getattr(dest_b.metadata_store, "commit", None) + if commit is not None: + commit() + written[name] = n + + src.stop() + dest.stop() return written diff --git a/dimos/mapping/cli/replay.py b/dimos/mapping/cli/replay.py index 9f30e89cfd2..9affca8e1e5 100644 --- a/dimos/mapping/cli/replay.py +++ b/dimos/mapping/cli/replay.py @@ -41,7 +41,7 @@ if TYPE_CHECKING: from dimos.memory.stream import Stream - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos.messages.sensor_msgs.msg import PointCloud2 TIMELINE = "ts" @@ -64,6 +64,8 @@ def _log_clouds( """ import rerun as rr + from dimos.visualization.rerun.cdr import cloud_archetype + n = total if total is not None else stream.count() with progress(n, label) as bar: for obs in stream: @@ -71,7 +73,13 @@ def _log_clouds( rr.set_time(TIMELINE, timestamp=obs.ts) rr.log( entity, - obs.data.to_rerun(voxel_size=voxel, mode=point_mode, bottom_cutoff=bottom_cutoff), + cloud_archetype( + obs.data, + voxel_size=voxel, + mode=point_mode, + bottom_cutoff=bottom_cutoff, + height_colors=True, + ), ) @@ -176,14 +184,14 @@ def main( """Dump a recording to .rrd (lidar clouds + camera frames) and open it in rerun.""" import rerun as rr + from dimos.mapping.cli.map import _camera_info_static from dimos.mapping.voxels.module import VoxelMapTransformer from dimos.memory.cli.dataset import open_store, resolve_dataset, stream_payload_types from dimos.memory.transform import throttle - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - from dimos.msgs.nav_msgs.Odometry import Odometry - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation - from dimos.robot.unitree.go2.connection import _camera_info_static + from dimos.messages.geometry_msgs.msg import PoseStamped + from dimos.messages.nav_msgs.msg import Odometry + from dimos.messages.sensor_msgs.msg import Image, PointCloud2 + from dimos.visualization.rerun.cdr import camera_pinhole, cloud_archetype, image_archetype src_path = resolve_dataset(dataset) store = open_store(src_path) @@ -201,11 +209,10 @@ def main( rr.init("dimos map_rrd", recording_id=src_path.stem) rr.save(str(out)) - register_colormap_annotation("turbo") # Static pinhole on the camera entity; per-frame Transform3D goes on the # same entity. Image is the child so it projects through the pinhole. - pinhole = cam_info.to_rerun() + pinhole = camera_pinhole(cam_info) assert not isinstance(pinhole, list) rr.log("world/camera", pinhole, static=True) @@ -268,8 +275,12 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: ).last() rr.log( f"world/{name}_map", - final.data.to_rerun( - voxel_size=voxel / 4, mode=point_mode, bottom_cutoff=bottom_cutoff + cloud_archetype( + final.data, + voxel_size=voxel / 4, + mode=point_mode, + bottom_cutoff=bottom_cutoff, + height_colors=True, ), static=True, ) @@ -339,7 +350,7 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: translation=[x, y, z], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]) ), ) - rr.log("world/camera/image", img_obs.data.to_rerun()) + rr.log("world/camera/image", image_archetype(img_obs.data)) print(f"wrote {out}") if no_gui: diff --git a/dimos/mapping/cli/replay_marker.py b/dimos/mapping/cli/replay_marker.py index d2e72a2bc47..94a1b91c655 100644 --- a/dimos/mapping/cli/replay_marker.py +++ b/dimos/mapping/cli/replay_marker.py @@ -64,13 +64,14 @@ def main( """Dump an AprilTag detection replay to .rrd and open it in rerun.""" import rerun as rr + from dimos.mapping.cli.map import _camera_info_static from dimos.memory.cli.dataset import open_store, resolve_dataset from dimos.memory.transform import QualityWindow, SpeedLimit from dimos.memory.vis.color import Color - from dimos.msgs.sensor_msgs.Image import Image - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos.message_utils.image import image_sharpness + from dimos.messages.sensor_msgs.msg import Image, PointCloud2 from dimos.perception.fiducial.marker_transformer import DetectMarkers - from dimos.robot.unitree.go2.connection import _camera_info_static + from dimos.visualization.rerun.cdr import camera_pinhole, image_archetype db_path = resolve_dataset(dataset) store = open_store(db_path) @@ -83,7 +84,7 @@ def main( # Static pinhole on the camera entity; per-frame Transform3D goes on the # same entity. Image is the child so it projects through the pinhole. - pinhole = cam_info.to_rerun() + pinhole = camera_pinhole(cam_info) assert not isinstance(pinhole, list) rr.log("world/camera", pinhole, static=True) @@ -116,7 +117,7 @@ def main( translation=[x, y, z], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]) ), ) - rr.log("world/camera/image", img_obs.data.to_rerun()) + rr.log("world/camera/image", image_archetype(img_obs.data)) # Clear any prior detection box at this ts; pass 3 will overwrite # this with the actual detection iff one fires for this frame. rr.log( @@ -129,7 +130,7 @@ def main( # Pass 3: marker detections (filtered same way as `dimos map`). xf = DetectMarkers(camera_info=cam_info, marker_length_m=marker_size) pipeline: Stream[Any] = color_image.transform( - QualityWindow(lambda img: img.sharpness, window=quality_window) + QualityWindow(image_sharpness, window=quality_window) ) if marker_max_speed > 0: pipeline = pipeline.transform( @@ -195,7 +196,7 @@ def main( smoothing_window=smoothing_window, ) pipeline_tracked: Stream[Any] = color_image.transform( - QualityWindow(lambda img: img.sharpness, window=quality_window) + QualityWindow(image_sharpness, window=quality_window) ) if marker_max_speed > 0: pipeline_tracked = pipeline_tracked.transform( diff --git a/dimos/mapping/cli/test_cli.py b/dimos/mapping/cli/test_cli.py index 53674915c6d..74b06b148d4 100644 --- a/dimos/mapping/cli/test_cli.py +++ b/dimos/mapping/cli/test_cli.py @@ -29,9 +29,15 @@ import traceback from types import SimpleNamespace +from dimos_cdrgen.registry import encode +import numpy as np import pytest from typer.testing import CliRunner +from dimos.message_utils.pointcloud import pointcloud_from_xyz +from dimos.message_utils.time import time_from_seconds +from dimos.messages.std_msgs.msg import Header + # These drive the real CLI against an LFS recording (CPU voxel accumulation, # multi-second runs) — self-hosted runner only. pytestmark = pytest.mark.self_hosted @@ -152,7 +158,7 @@ def test_pose_fill_snippet(dataset: str, tmp_path: Path) -> None: assert out.exists() from dimos.memory.store.sqlite import SqliteStore - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos.messages.sensor_msgs.msg import PointCloud2 store = SqliteStore(path=str(out), must_exist=True) try: @@ -226,15 +232,12 @@ def test_global_snippet(dataset: str, tmp_path: Path) -> None: def test_view_pc2(tmp_path: Path) -> None: - import numpy as np - - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 - - pc2 = tmp_path / "map.pc2.lcm" - cloud = PointCloud2.from_numpy( - np.random.default_rng(0).random((256, 3)), frame_id="world", timestamp=0.0 + pc2 = tmp_path / "map.pc2.cdr" + cloud = pointcloud_from_xyz( + np.random.default_rng(0).random((256, 3)), + header=Header(stamp=time_from_seconds(0), frame_id="world"), ) - pc2.write_bytes(cloud.lcm_encode()) + pc2.write_bytes(encode(cloud)) out = tmp_path / "view.rrd" res = _run("view", str(pc2), "--out", str(out)) diff --git a/dimos/mapping/cli/test_generated_cli.py b/dimos/mapping/cli/test_generated_cli.py new file mode 100644 index 00000000000..6161acdcfff --- /dev/null +++ b/dimos/mapping/cli/test_generated_cli.py @@ -0,0 +1,331 @@ +# 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. + + +import cv2 +from dimos_cdrgen.registry import decode, encode +import numpy as np +import pytest +import rerun as rr +from typer.testing import CliRunner + +from dimos.cli.commands.map import map_app +from dimos.mapping.cli.map import _accumulate, _camera_info_static, _denoise +from dimos.memory.cli import dataset +from dimos.memory.store.sqlite import SqliteStore +from dimos.memory.type.observation import Observation +from dimos.message_utils.image import image_from_array +from dimos.message_utils.pointcloud import ( + pointcloud_from_xyz, + pointcloud_from_xyz_rgb, + pointcloud_stamps, + pointcloud_xyz, +) +from dimos.message_utils.time import time_from_nanoseconds, time_from_seconds +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos.messages.sensor_msgs.msg import Image, PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage +from dimos.navigation.go2.loop_closure.pgo import Keyframe, PoseGraph +from dimos.visualization.rerun.cdr import cloud_archetype + + +@pytest.fixture +def recorded_dataset(tmp_path): + path = tmp_path / "canonical.db" + pixels = np.full((720, 1280, 3), 255, dtype=np.uint8) + tile = cv2.aruco.generateImageMarker( + cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_APRILTAG_36h11), 7, 200 + ) + pixels[260:460, 540:740] = tile[:, :, None] + with SqliteStore(path=str(path)) as store: + lidar = store.stream("lidar", PointCloud2) + images = store.stream("color_image", Image) + odom = store.stream("odom", PoseStamped) + for index in range(3): + stamp = time_from_nanoseconds(1700000000123456789 + index * 1000000000) + header = Header(stamp=stamp, frame_id="world") + cloud = pointcloud_from_xyz( + np.array([[index + 1.025, 0.025, 0.025]], dtype=np.float32), header=header + ) + pose = (float(index + 1), 0, 0, 0, 0, 0.05, 0.998749217771909) + lidar.append(cloud, ts=10.0 + index, pose=pose) + odom.append( + PoseStamped( + header=header, + pose=Pose(position=Point(*pose[:3]), orientation=Quaternion(*pose[3:])), + ), + ts=10.0 + index, + pose=pose, + ) + header.frame_id = "camera_optical" + images.append( + image_from_array(pixels, encoding="rgb8", header=header), ts=10.0 + index, pose=pose + ) + return path + + +@pytest.fixture +def recording_calls(mocker): + log = mocker.spy(rr, "log") + opened = mocker.spy(dataset, "open_store") + try: + yield log + finally: + recording = rr.get_global_data_recording() + if recording is not None: + recording.flush() + rr.disconnect() + for store in opened.spy_return_list: + store.stop() + + +def invoke(*arguments): + result = CliRunner().invoke(map_app, list(arguments)) + assert result.exit_code == 0, (result.output, result.exception) + return result + + +def test_global_cli_exports_cdr_and_view_reads_it( + recorded_dataset, recording_calls, monkeypatch, tmp_path +): + monkeypatch.chdir(tmp_path) + invoke( + "global", + str(recorded_dataset), + "--device", + "CPU:0", + "--export", + "--full-pgo", + "--no-gui", + "--out", + str(tmp_path / "map.rrd"), + ) + exported = tmp_path / "canonical.pc2.cdr" + cloud = decode(exported.read_bytes(), PointCloud2) + assert cloud.width == 3 and cloud.height == 1 + assert cloud.header.frame_id == "world_corrected" + assert cloud.header.stamp == time_from_nanoseconds(1700000002123456789) + np.testing.assert_allclose( + np.sort(pointcloud_xyz(cloud)[:, 0]), [1.025, 2.025, 3.025], atol=1e-6 + ) + invoke("view", str(exported), "--out", str(tmp_path / "view.rrd")) + visible = [ + call.args[1] + for call in recording_calls.call_args_list + if call.args[0] == "world/canonical/pointcloud" + ] + assert len(visible) == 1 + np.testing.assert_allclose( + np.sort(np.asarray(visible[0].positions.as_arrow_array().to_pylist())[:, 0]), + [1.025, 2.025, 3.025], + atol=1e-6, + ) + + +@pytest.mark.parametrize("mode", ["points", "boxes", "spheres"]) +def test_replay_cli_logs_cdr_clouds_images_and_maps( + recorded_dataset, recording_calls, tmp_path, mode +): + output = tmp_path / f"replay-{mode}.rrd" + invoke( + "replay", + str(recorded_dataset), + "--no-gui", + "--out", + str(output), + "--map", + "--map-final", + "--map-device", + "CPU:0", + "--map-emit-every", + "1", + "--point-mode", + mode, + ) + paths = [call.args[0] for call in recording_calls.call_args_list] + assert paths.count("world/lidar") == 3 + assert paths.count("world/lidar_voxels") == 3 + assert paths.count("world/camera/image") == 3 + assert "world/lidar_map" in paths + recording = rr.get_global_data_recording() + assert recording is not None + recording.flush() + assert output.stat().st_size > 1000 + + +def test_replay_marker_cli_detects_actual_apriltag(recorded_dataset, recording_calls, tmp_path): + invoke( + "replay-marker", + str(recorded_dataset), + "--no-gui", + "--out", + str(tmp_path / "markers.rrd"), + "--marker-max-speed", + "0", + "--quality-window", + "0", + "--smoothing-window", + "0", + ) + detections = [ + call.args[1] + for call in recording_calls.call_args_list + if call.args[0].startswith("world/markers/det_") + ] + assert len(detections) == 3 + assert all(item.labels.as_arrow_array().to_pylist() == ["id=7"] for item in detections) + + +def test_registration_then_drift_correction_preserves_header_and_input(): + header = Header(frame_id="lidar", stamp=time_from_nanoseconds(1700000000123456789)) + cloud = pointcloud_from_xyz(np.array([[0, 0, 0]], dtype=np.float32), header=header) + before = encode(cloud) + + def transform(x, frame, child): + return TransformStamped( + header=Header(frame_id=frame, stamp=time_from_seconds(10)), + child_frame_id=child, + transform=Transform(translation=Vector3(x, 0, 0), rotation=Quaternion(0, 0, 0, 1)), + ) + + graph = PoseGraph( + keyframes=( + Keyframe( + ts=10, + local=transform(0, "world_raw", "body"), + optimized=transform(2, "world_corrected", "body"), + ), + ) + ) + obs = Observation(ts=10, _data=cloud, pose=(1, 0, 0, 0, 0, 0, 1)) + output = _accumulate( + [obs], + voxel=0.05, + block_count=128, + device="CPU:0", + graph=graph, + register=lambda _: transform(1, "world", "lidar"), + ) + assert output is not None + np.testing.assert_allclose(pointcloud_xyz(output), [[3.025, 0.025, 0.025]], atol=1e-6) + assert output.header.frame_id == "world_corrected" and output.header.stamp == header.stamp + assert encode(cloud) == before + + +def test_denoise_retains_point_capture_times_and_exact_header(): + rng = np.random.default_rng(7) + points = np.concatenate([rng.normal(0, 0.03, (100, 3)), [[100, 100, 100]]]) + header = Header(frame_id="world", stamp=time_from_nanoseconds(1700000000123456789)) + cloud = pointcloud_from_xyz(points, header=header, stamps=np.arange(101, dtype=np.float64)) + before = encode(cloud) + output = _denoise(cloud) + assert output is not None and output.width < cloud.width + assert np.linalg.norm(pointcloud_xyz(output), axis=1).max() < 1 + assert output.header == header and output.fields == cloud.fields + assert np.max(pointcloud_stamps(output)) < 100 + assert encode(cloud) == before + + +@pytest.mark.parametrize("mode", ["points", "boxes", "spheres"]) +def test_cloud_rendering_filters_geometry_and_own_colors_together(mode): + cloud = pointcloud_from_xyz_rgb( + np.array([[1, 2, -1], [3, 4, 2]], dtype=np.float32), + np.array([[255, 0, 0], [0, 255, 0]], dtype=np.uint8), + header=Header(frame_id="world", stamp=time_from_seconds(0)), + ) + before = encode(cloud) + value = cloud_archetype(cloud, voxel_size=0.2, mode=mode, bottom_cutoff=0) + assert value.colors.as_arrow_array().to_pylist() == [0x00FF00FF] + positions = value.centers if mode == "boxes" else value.positions + np.testing.assert_array_equal(positions.as_arrow_array().to_pylist(), [[3, 4, 2]]) + assert encode(cloud) == before + + +def test_packaged_camera_calibration_is_a_generated_value(): + camera = _camera_info_static() + assert camera.__msgtype__ == "sensor_msgs/msg/CameraInfo" + assert camera.width == 1280 and camera.height == 720 + assert camera.header.frame_id == "camera_optical" + + +def test_global_markers_compose_external_pose_stream_and_pgo( + recorded_dataset, recording_calls, tmp_path +): + invoke( + "global", + str(recorded_dataset), + "--device", + "CPU:0", + "--pgo", + "--no-gui", + "--markers", + "--marker-max-speed", + "0", + "--marker-quality-window", + "0", + "--marker-smoothing", + "0", + "--image-pose", + "odom", + "--out", + str(tmp_path / "posed-markers.rrd"), + ) + paths = [call.args[0] for call in recording_calls.call_args_list] + assert "world/raw_map/markers/fill" in paths and "world/pgo_map/markers/fill" in paths + fill = next( + call.args[1] + for call in recording_calls.call_args_list + if call.args[0] == "world/pgo_map/markers/fill" + ) + assert len(fill.centers.as_arrow_array()) == 3 + + +def test_global_cli_registers_cdr_sensor_cloud_using_recorded_tf(recording_calls, tmp_path): + path = tmp_path / "sensor.db" + header = Header(frame_id="lidar", stamp=time_from_nanoseconds(1700000000123456789)) + cloud = pointcloud_from_xyz(np.array([[1.025, 0.025, 0.025]], dtype=np.float32), header=header) + transform = TransformStamped( + header=Header(frame_id="map", stamp=time_from_seconds(10)), + child_frame_id="lidar", + transform=Transform(translation=Vector3(2, 0, 0), rotation=Quaternion(0, 0, 0, 1)), + ) + with SqliteStore(path=str(path)) as store: + store.stream("lidar", PointCloud2).append(cloud, ts=10) + store.stream("tf", TFMessage).append(TFMessage(transforms=[transform]), ts=10) + invoke( + "global", + str(path), + "--device", + "CPU:0", + "--no-gui", + "--out", + str(tmp_path / "registered.rrd"), + ) + points = next( + call.args[1] + for call in recording_calls.call_args_list + if call.args[0] == "world/raw_map/pointcloud" + ) + np.testing.assert_allclose( + points.positions.as_arrow_array().to_pylist(), [[3.025, 0.025, 0.025]], atol=1e-6 + ) diff --git a/dimos/mapping/cli/test_pose_fill.py b/dimos/mapping/cli/test_pose_fill.py index 2953d9e13e6..85149c8ebab 100644 --- a/dimos/mapping/cli/test_pose_fill.py +++ b/dimos/mapping/cli/test_pose_fill.py @@ -20,11 +20,28 @@ from __future__ import annotations -from dimos.mapping.cli.pose_fill import pose_fill +from copy import deepcopy +import math + +from dimos_cdrgen.registry import encode +import numpy as np + +from dimos.mapping.cli.pose_fill import pose_fill, pose_fill_db from dimos.memory.store.memory import MemoryStore -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.memory.store.sqlite import SqliteStore +from dimos.message_utils.geometry import quaternion_from_euler +from dimos.message_utils.pointcloud import pointcloud_from_xyz +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + Vector3, +) +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header def test_pose_fill_attaches_nearest_pose() -> None: @@ -33,7 +50,12 @@ def test_pose_fill_attaches_nearest_pose() -> None: target = store.stream("image", str) poses = store.stream("odom", Transform) target.append("img0", ts=0.0) - poses.append(Transform(translation=Vector3(1.0, 0.0, 0.0)), ts=0.001) + poses.append( + Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), rotation=Quaternion(x=0, y=0, z=0, w=1) + ), + ts=0.001, + ) out = pose_fill(target, poses, tolerance=0.05).to_list() @@ -50,8 +72,15 @@ def test_pose_fill_mount_composes_static_child_transform() -> None: poses = store.stream("odom", Transform) target.append("img0", ts=0.0) # Base pose 1m forward in x; mount offsets 1m in y (identity rotations). - poses.append(Transform(translation=Vector3(1.0, 0.0, 0.0)), ts=0.001) - mount = Transform(translation=Vector3(0.0, 1.0, 0.0), rotation=Quaternion()) + poses.append( + Transform( + translation=Vector3(x=1.0, y=0.0, z=0.0), rotation=Quaternion(x=0, y=0, z=0, w=1) + ), + ts=0.001, + ) + mount = Transform( + translation=Vector3(x=0, y=1, z=0), rotation=Quaternion(x=0, y=0, z=0, w=1) + ) out = pose_fill(target, poses, tolerance=0.05, mount=mount).to_list() @@ -59,3 +88,57 @@ def test_pose_fill_mount_composes_static_child_transform() -> None: assert out[0].pose_tuple is not None # Identity rotations → composed translation is the component-wise sum. assert out[0].pose_tuple[:3] == (1.0, 1.0, 0.0) + + +def test_stamped_pose_mount_rotates_the_offset_and_keeps_payload(): + header = Header(frame_id="world", stamp=time_from_nanoseconds(1700000000123456789)) + pose = PoseStamped( + header=header, + pose=Pose( + position=Point(x=1, y=2, z=3), orientation=quaternion_from_euler(0, 0, math.pi / 2) + ), + ) + mount = Transform(translation=Vector3(x=1, y=0, z=0), rotation=Quaternion(x=0, y=0, z=0, w=1)) + with MemoryStore() as store: + target = store.stream("target", str) + source = store.stream("pose", PoseStamped) + target.append("frame", ts=10.0) + source.append(pose, ts=10.01) + before = encode(pose) + [result] = pose_fill(target, source, mount=mount).to_list() + assert result.data == "frame" + np.testing.assert_allclose(result.pose_tuple[:3], [1, 3, 3], atol=1e-12) + assert encode(pose) == before + + +def test_sqlite_pose_fill_copies_original_cdr_bytes_and_exact_header(tmp_path): + source_path = tmp_path / "source.db" + destination = tmp_path / "posed.db" + header = Header(frame_id="lidar", stamp=time_from_nanoseconds(1700000000123456789)) + cloud = pointcloud_from_xyz(np.array([[1, 2, 3]], dtype=np.float32), header=header) + pose_header = deepcopy(header) + pose_header.frame_id = "odom" + pose = PoseStamped( + header=pose_header, + pose=Pose(position=Point(x=4, y=5, z=6), orientation=Quaternion(x=0, y=0, z=0, w=1)), + ) + with SqliteStore(path=str(source_path)) as store: + store.stream("lidar", PointCloud2).append(cloud, ts=10.0) + store.stream("odom", PoseStamped).append(pose, ts=10.001) + written = pose_fill_db(source_path, destination, target="lidar", pose_source="odom") + assert written == {"lidar": 1, "odom": 1} + with ( + SqliteStore(path=str(source_path), must_exist=True) as original, + SqliteStore(path=str(destination), must_exist=True) as store, + ): + result = store.stream("lidar", PointCloud2).first() + assert result.pose_tuple == (4.0, 5.0, 6.0, 0.0, 0.0, 0.0, 1.0) + assert result.ts == 10.0 + assert encode(result.data) == encode(cloud) + assert result.data.header.stamp.nanosec == 123456789 + assert encode(store.stream("odom", PoseStamped).first().data) == encode(pose) + before = original.stream("lidar", PointCloud2) + after = store.stream("lidar", PointCloud2) + assert before._source.blob_store.get( + "lidar", before.first().id + ) == after._source.blob_store.get("lidar", result.id) diff --git a/dimos/mapping/cli/view.py b/dimos/mapping/cli/view.py index a76eccc173d..87c353ba627 100644 --- a/dimos/mapping/cli/view.py +++ b/dimos/mapping/cli/view.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""`dimos map view`: open a `.pc2.lcm` cloud - `dimos map global --export` - in rerun.""" +"""`dimos map view`: open a `.pc2.cdr` cloud - `dimos map global --export` - in rerun.""" from __future__ import annotations @@ -22,7 +22,7 @@ def main( - path: Path = typer.Argument(..., help="Point cloud to view (.pc2.lcm)"), + path: Path = typer.Argument(..., help="Point cloud to view (.pc2.cdr)"), voxel: float = typer.Option(0.05, "--voxel", help="Rendered point size (m)"), bottom_cutoff: float | None = typer.Option( None, @@ -33,14 +33,16 @@ def main( None, "--out", help="Write a .rrd instead of opening the viewer" ), ) -> None: - """Load an LCM-encoded PointCloud2 and display it in rerun.""" + """Load an CDR-encoded PointCloud2 and display it in rerun.""" + from dimos_cdrgen.registry import decode import rerun as rr - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos.messages.sensor_msgs.msg import PointCloud2 + from dimos.visualization.rerun.cdr import cloud_archetype from dimos.visualization.rerun.init import rerun_init - cloud = PointCloud2.lcm_decode(path.read_bytes()) - print(f"{path}: {len(cloud.pointcloud.points)} points in {cloud.frame_id!r}") + cloud = decode(path.read_bytes(), PointCloud2) + print(f"{path}: {cloud.width * cloud.height} points in {cloud.header.frame_id!r}") rerun_init("dimos map view") if out is not None: @@ -49,7 +51,13 @@ def main( rr.spawn() rr.log( f"world/{path.name.split('.')[0]}/pointcloud", - cloud.to_rerun(voxel_size=voxel / 2, bottom_cutoff=bottom_cutoff), + cloud_archetype( + cloud, + voxel_size=voxel / 2, + bottom_cutoff=bottom_cutoff, + mode="spheres", + height_colors=True, + ), static=True, ) if out is not None: diff --git a/dimos/mapping/costmapper.py b/dimos/mapping/costmapper.py index 31b94b164f5..556e916fe00 100644 --- a/dimos/mapping/costmapper.py +++ b/dimos/mapping/costmapper.py @@ -27,8 +27,9 @@ HeightCostConfig, OccupancyConfig, ) -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.occupancy import occupancy_view +from dimos.messages.nav_msgs.msg import OccupancyGrid +from dimos.messages.sensor_msgs.msg import PointCloud2 from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -90,19 +91,22 @@ def stop(self) -> None: def _calculate_costmap(self, msg: PointCloud2) -> OccupancyGrid: occupancy_function = OCCUPANCY_ALGOS[self.config.algo] grid = occupancy_function(msg, **asdict(self.config.config)) + if not isinstance(grid, OccupancyGrid): + raise TypeError("CostMapper algorithms must return generated OccupancyGrid values") self._apply_initial_safe_radius(grid) return grid def _apply_initial_safe_radius(self, grid: OccupancyGrid) -> None: radius_meters = self.config.initial_safe_radius_meters - if radius_meters <= 0 or grid.grid.size == 0: + if radius_meters <= 0 or grid.data.size == 0: return - resolution = grid.resolution - origin_x = grid.origin.position.x - origin_y = grid.origin.position.y + resolution = grid.info.resolution + origin_x = grid.info.origin.position.x + origin_y = grid.info.origin.position.y - rows, columns = np.ogrid[: grid.grid.shape[0], : grid.grid.shape[1]] + cells = occupancy_view(grid).copy() + rows, columns = np.ogrid[: cells.shape[0], : cells.shape[1]] cell_world_x = columns * resolution + origin_x cell_world_y = rows * resolution + origin_y distance_squared_meters = cell_world_x**2 + cell_world_y**2 @@ -111,4 +115,5 @@ def _apply_initial_safe_radius(self, grid: OccupancyGrid) -> None: # the disc. Avoids floating-point boundary flakiness from radius/resolution. effective_radius_meters = radius_meters + resolution * 0.5 safe_mask = distance_squared_meters <= effective_radius_meters**2 - grid.grid[safe_mask] = 0 + cells[safe_mask] = 0 + grid.data = cells.reshape(-1) diff --git a/dimos/mapping/occupancy/extrude_occupancy.py b/dimos/mapping/occupancy/extrude_occupancy.py index 799319cbf6d..e9b89504991 100644 --- a/dimos/mapping/occupancy/extrude_occupancy.py +++ b/dimos/mapping/occupancy/extrude_occupancy.py @@ -16,13 +16,14 @@ import numpy as np from numpy.typing import NDArray -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.mapping.occupancy.operations import _grid_values, _GridT +from dimos.message_utils.occupancy import CostValues # Rectangle type: (x, y, width, height) Rect = tuple[int, int, int, int] -def identify_convex_shapes(occupancy_grid: OccupancyGrid) -> list[Rect]: +def identify_convex_shapes(occupancy_grid: _GridT) -> list[Rect]: """Identify occupied zones and decompose them into convex rectangles. This function finds all occupied cells in the occupancy grid and @@ -36,7 +37,7 @@ def identify_convex_shapes(occupancy_grid: OccupancyGrid) -> list[Rect]: Returns: List of rectangles as (x, y, width, height) tuples in grid coords. """ - grid = occupancy_grid.grid + grid = _grid_values(occupancy_grid) # Create binary mask of occupied cells (treat UNKNOWN as OCCUPIED) occupied_mask = ((grid == CostValues.OCCUPIED) | (grid == CostValues.UNKNOWN)).astype( @@ -166,7 +167,7 @@ def _largest_rect_in_histogram( def generate_mujoco_scene( - occupancy_grid: OccupancyGrid, + occupancy_grid: _GridT, ) -> str: """Generate a MuJoCo scene XML from an occupancy grid. @@ -184,9 +185,9 @@ def generate_mujoco_scene( # Get rectangles from the occupancy grid rectangles = identify_convex_shapes(occupancy_grid) - resolution = occupancy_grid.resolution - origin_x = occupancy_grid.origin.position.x - origin_y = occupancy_grid.origin.position.y + resolution = occupancy_grid.info.resolution + origin_x = occupancy_grid.info.origin.position.x + origin_y = occupancy_grid.info.origin.position.y # Build XML xml_lines = [ @@ -208,7 +209,7 @@ def generate_mujoco_scene( # Add each rectangle as a box geom for i, (gx, gy, gw, gh) in enumerate(rectangles): # Convert grid coordinates to world coordinates - # Grid origin is top-left, world origin is at occupancy_grid.origin + # Grid origin is top-left, world origin is at occupancy_grid.info.origin # gx, gy are in grid cells, need to convert to meters world_x = origin_x + (gx + gw / 2) * resolution world_y = origin_y + (gy + gh / 2) * resolution diff --git a/dimos/mapping/occupancy/gradient.py b/dimos/mapping/occupancy/gradient.py index 802b52a9faf..8265b0d6d1f 100644 --- a/dimos/mapping/occupancy/gradient.py +++ b/dimos/mapping/occupancy/gradient.py @@ -17,7 +17,8 @@ import numpy as np from scipy import ndimage -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.mapping.occupancy.operations import _grid_values, _GridT, _with_grid +from dimos.message_utils.occupancy import CostValues if TYPE_CHECKING: from numpy.typing import NDArray @@ -27,8 +28,8 @@ def gradient( - occupancy_grid: OccupancyGrid, obstacle_threshold: int = 50, max_distance: float = 2.0 -) -> OccupancyGrid: + occupancy_grid: _GridT, obstacle_threshold: int = 50, max_distance: float = 2.0 +) -> _GridT: """Create a gradient OccupancyGrid for path planning. Creates a gradient where free space has value 0 and values increase near obstacles. @@ -49,19 +50,19 @@ def gradient( """ # Remember which cells are unknown - unknown_mask = occupancy_grid.grid == CostValues.UNKNOWN + unknown_mask = _grid_values(occupancy_grid) == CostValues.UNKNOWN # Create binary obstacle map # Consider cells >= threshold as obstacles (1), everything else as free (0) # Unknown cells are not considered obstacles for distance calculation - obstacle_map = (occupancy_grid.grid >= obstacle_threshold).astype(np.float32) + obstacle_map = (_grid_values(occupancy_grid) >= obstacle_threshold).astype(np.float32) # Compute distance transform (distance to nearest obstacle in cells) # Unknown cells are treated as if they don't exist for distance calculation distance_cells = cast("NDArray[np.float64]", ndimage.distance_transform_edt(1 - obstacle_map)) # Convert to meters and clip to max distance - distance_meters = np.clip(distance_cells * occupancy_grid.resolution, 0, max_distance) + distance_meters = np.clip(distance_cells * occupancy_grid.info.resolution, 0, max_distance) # Invert and scale to 0-100 range # Far from obstacles (max_distance) -> 0 @@ -78,20 +79,14 @@ def gradient( gradient_data[unknown_mask] = CostValues.UNKNOWN # Create new OccupancyGrid with gradient - gradient_grid = OccupancyGrid( - grid=gradient_data, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + gradient_grid = _with_grid(occupancy_grid, gradient_data) return gradient_grid def voronoi_gradient( - occupancy_grid: OccupancyGrid, obstacle_threshold: int = 50, max_distance: float = 2.0 -) -> OccupancyGrid: + occupancy_grid: _GridT, obstacle_threshold: int = 50, max_distance: float = 2.0 +) -> _GridT: """Create a Voronoi-based gradient OccupancyGrid for path planning. Unlike the regular gradient which can result in suboptimal paths in narrow @@ -122,23 +117,17 @@ def voronoi_gradient( - 100: At obstacles """ # Remember which cells are unknown - unknown_mask = occupancy_grid.grid == CostValues.UNKNOWN + unknown_mask = _grid_values(occupancy_grid) == CostValues.UNKNOWN # Create binary obstacle map - obstacle_map = (occupancy_grid.grid >= obstacle_threshold).astype(np.float32) + obstacle_map = (_grid_values(occupancy_grid) >= obstacle_threshold).astype(np.float32) # Check if there are any obstacles if not np.any(obstacle_map): # No obstacles - everything is free - gradient_data = np.zeros_like(occupancy_grid.grid, dtype=np.int8) + gradient_data = np.zeros_like(_grid_values(occupancy_grid), dtype=np.int8) gradient_data[unknown_mask] = CostValues.UNKNOWN - return OccupancyGrid( - grid=gradient_data, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + return _with_grid(occupancy_grid, gradient_data) # Label connected obstacle regions (clusters) # This groups all cells of the same wall/obstacle together @@ -192,7 +181,7 @@ def voronoi_gradient( gradient_values[obstacle_map > 0] = CostValues.OCCUPIED # Apply max_distance clipping - cells beyond max_distance from obstacles get cost 0 - max_distance_cells = max_distance / occupancy_grid.resolution + max_distance_cells = max_distance / occupancy_grid.info.resolution gradient_values[distance_cells > max_distance_cells] = 0 # Convert to int8 @@ -201,10 +190,4 @@ def voronoi_gradient( # Preserve unknown cells gradient_data[unknown_mask] = CostValues.UNKNOWN - return OccupancyGrid( - grid=gradient_data, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + return _with_grid(occupancy_grid, gradient_data) diff --git a/dimos/mapping/occupancy/inflation.py b/dimos/mapping/occupancy/inflation.py index 610157ddf62..3c0ed695e6a 100644 --- a/dimos/mapping/occupancy/inflation.py +++ b/dimos/mapping/occupancy/inflation.py @@ -15,10 +15,11 @@ import numpy as np from scipy import ndimage -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.mapping.occupancy.operations import _grid_values, _GridT, _with_grid +from dimos.message_utils.occupancy import CostValues -def simple_inflate(occupancy_grid: OccupancyGrid, radius: float) -> OccupancyGrid: +def simple_inflate(occupancy_grid: _GridT, radius: float) -> _GridT: """Inflate obstacles by a given radius (binary inflation). Args: radius: Inflation radius in meters @@ -26,10 +27,10 @@ def simple_inflate(occupancy_grid: OccupancyGrid, radius: float) -> OccupancyGri New OccupancyGrid with inflated obstacles """ # Convert radius to grid cells - cell_radius = int(np.ceil(radius / occupancy_grid.resolution)) + cell_radius = int(np.ceil(radius / occupancy_grid.info.resolution)) # Get grid as numpy array - grid_array = occupancy_grid.grid + grid_array = _grid_values(occupancy_grid) # Create circular kernel for binary inflation y, x = np.ogrid[-cell_radius : cell_radius + 1, -cell_radius : cell_radius + 1] @@ -44,10 +45,4 @@ def simple_inflate(occupancy_grid: OccupancyGrid, radius: float) -> OccupancyGri result_grid[inflated] = CostValues.OCCUPIED # Create new OccupancyGrid with inflated data using numpy constructor - return OccupancyGrid( - grid=result_grid, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + return _with_grid(occupancy_grid, result_grid) diff --git a/dimos/mapping/occupancy/operations.py b/dimos/mapping/occupancy/operations.py index 8609ac634b3..a82397c06f5 100644 --- a/dimos/mapping/occupancy/operations.py +++ b/dimos/mapping/occupancy/operations.py @@ -12,15 +12,54 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy +from typing import TypeVar + import numpy as np +from numpy.typing import NDArray from scipy import ndimage -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.message_utils.occupancy import CostValues, occupancy_view, world_to_grid +from dimos.messages.geometry_msgs.msg import Point +from dimos.messages.nav_msgs.msg import OccupancyGrid as CdrOccupancyGrid +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid + +_GridT = TypeVar("_GridT", OccupancyGrid, CdrOccupancyGrid) + + +def _grid_values(message: OccupancyGrid | CdrOccupancyGrid) -> NDArray[np.int8]: + if isinstance(message, CdrOccupancyGrid): + return occupancy_view(message) + return message.grid + + +def _grid_coordinates( + message: OccupancyGrid | CdrOccupancyGrid, point: Point | Vector3 +) -> tuple[float, float]: + if isinstance(message, CdrOccupancyGrid): + return world_to_grid(message, Point(x=point.x, y=point.y, z=point.z)) + result = message.world_to_grid((point.x, point.y, point.z)) + return result.x, result.y + +def _with_grid(message: _GridT, cells: NDArray[np.int8]) -> _GridT: + if isinstance(message, CdrOccupancyGrid): + if cells.shape != (message.info.height, message.info.width): + raise ValueError("Replacement occupancy cells must retain map dimensions") + result = deepcopy(message) + result.data = np.array(cells, dtype=np.int8, copy=True).reshape(-1) + return result + return OccupancyGrid( + grid=cells, + resolution=message.resolution, + origin=message.origin, + frame_id=message.frame_id, + ts=message.ts, + ) -def smooth_occupied( - occupancy_grid: OccupancyGrid, min_neighbor_fraction: float = 0.4 -) -> OccupancyGrid: + +def smooth_occupied(occupancy_grid: _GridT, min_neighbor_fraction: float = 0.4) -> _GridT: """Smooth occupied zones by removing unsupported protrusions. Removes occupied cells that don't have sufficient neighboring occupied @@ -33,7 +72,7 @@ def smooth_occupied( Returns: New OccupancyGrid with smoothed occupied zones """ - grid_array = occupancy_grid.grid + grid_array = _grid_values(occupancy_grid) occupied_mask = grid_array >= CostValues.OCCUPIED # Count occupied neighbors for each cell (8-connectivity). @@ -49,16 +88,10 @@ def smooth_occupied( result_grid = grid_array.copy() result_grid[unsupported] = CostValues.FREE - return OccupancyGrid( - grid=result_grid, - resolution=occupancy_grid.resolution, - origin=occupancy_grid.origin, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + return _with_grid(occupancy_grid, result_grid) -def overlay_occupied(base: OccupancyGrid, overlay: OccupancyGrid) -> OccupancyGrid: +def overlay_occupied(base: _GridT, overlay: _GridT) -> _GridT: """Overlay occupied zones from one grid onto another. Marks cells as occupied in the base grid wherever they are occupied @@ -70,19 +103,13 @@ def overlay_occupied(base: OccupancyGrid, overlay: OccupancyGrid) -> OccupancyGr Returns: New OccupancyGrid with combined occupied zones """ - if base.grid.shape != overlay.grid.shape: + if _grid_values(base).shape != _grid_values(overlay).shape: raise ValueError( - f"Grid shapes must match: base {base.grid.shape} vs overlay {overlay.grid.shape}" + f"Grid shapes must match: base {_grid_values(base).shape} vs overlay {_grid_values(overlay).shape}" ) - result_grid = base.grid.copy() - overlay_occupied_mask = overlay.grid >= CostValues.OCCUPIED + result_grid = _grid_values(base).copy() + overlay_occupied_mask = _grid_values(overlay) >= CostValues.OCCUPIED result_grid[overlay_occupied_mask] = CostValues.OCCUPIED - return OccupancyGrid( - grid=result_grid, - resolution=base.resolution, - origin=base.origin, - frame_id=base.frame_id, - ts=base.ts, - ) + return _with_grid(base, result_grid) diff --git a/dimos/mapping/occupancy/path_map.py b/dimos/mapping/occupancy/path_map.py index 69a1f937385..a1f5e29ddba 100644 --- a/dimos/mapping/occupancy/path_map.py +++ b/dimos/mapping/occupancy/path_map.py @@ -14,17 +14,16 @@ from dimos.mapping.occupancy.gradient import GradientStrategy, gradient, voronoi_gradient from dimos.mapping.occupancy.inflation import simple_inflate -from dimos.mapping.occupancy.operations import overlay_occupied, smooth_occupied +from dimos.mapping.occupancy.operations import _GridT, overlay_occupied, smooth_occupied from dimos.mapping.occupancy.types import NavigationStrategy -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid def make_navigation_map( - occupancy_grid: OccupancyGrid, + occupancy_grid: _GridT, robot_width: float, strategy: NavigationStrategy, gradient_strategy: GradientStrategy, -) -> OccupancyGrid: +) -> _GridT: half_width = robot_width / 2 gradient_distance = 1.5 diff --git a/dimos/mapping/occupancy/path_mask.py b/dimos/mapping/occupancy/path_mask.py index b71372bbf62..a21d86a106f 100644 --- a/dimos/mapping/occupancy/path_mask.py +++ b/dimos/mapping/occupancy/path_mask.py @@ -15,13 +15,17 @@ import numpy as np from numpy.typing import NDArray -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +from dimos.mapping.occupancy.operations import _grid_coordinates, _grid_values +from dimos.message_utils.occupancy import CostValues +from dimos.messages.geometry_msgs.msg import PoseStamped as CdrPoseStamped +from dimos.messages.nav_msgs.msg import OccupancyGrid as CdrOccupancyGrid, Path as CdrPath +from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid from dimos.msgs.nav_msgs.Path import Path def make_path_mask( - occupancy_grid: OccupancyGrid, - path: Path, + occupancy_grid: OccupancyGrid | CdrOccupancyGrid, + path: Path | CdrPath, robot_width: float, pose_index: int = 0, max_length: float = float("inf"), @@ -44,9 +48,9 @@ def make_path_mask( """ import cv2 - mask = np.zeros((occupancy_grid.height, occupancy_grid.width), dtype=np.uint8) + mask = np.zeros((occupancy_grid.info.height, occupancy_grid.info.width), dtype=np.uint8) - line_width_pixels = max(1, int(robot_width / occupancy_grid.resolution)) + line_width_pixels = max(1, int(robot_width / occupancy_grid.info.resolution)) poses = path.poses if len(poses) < pose_index + 2: @@ -55,8 +59,9 @@ def make_path_mask( # Draw lines between consecutive points cumulative_length = 0.0 for i in range(pose_index, len(poses) - 1): - pos1 = poses[i].position - pos2 = poses[i + 1].position + first, second = poses[i], poses[i + 1] + pos1 = first.pose.position if isinstance(first, CdrPoseStamped) else first.position + pos2 = second.pose.position if isinstance(second, CdrPoseStamped) else second.position segment_length = np.sqrt( (pos2.x - pos1.x) ** 2 + (pos2.y - pos1.y) ** 2 + (pos2.z - pos1.z) ** 2 @@ -67,11 +72,11 @@ def make_path_mask( cumulative_length += segment_length - grid_pt1 = occupancy_grid.world_to_grid(pos1) - grid_pt2 = occupancy_grid.world_to_grid(pos2) + grid_pt1 = _grid_coordinates(occupancy_grid, pos1) + grid_pt2 = _grid_coordinates(occupancy_grid, pos2) - pt1 = (round(grid_pt1.x), round(grid_pt1.y)) - pt2 = (round(grid_pt2.x), round(grid_pt2.y)) + pt1 = (round(grid_pt1[0]), round(grid_pt1[1])) + pt2 = (round(grid_pt2[0]), round(grid_pt2[1])) cv2.line(mask, pt1, pt2, (255.0,), thickness=line_width_pixels) @@ -82,7 +87,7 @@ def make_path_mask( if total_points == 0: return bool_mask - occupied_mask = occupancy_grid.grid >= CostValues.OCCUPIED + occupied_mask = _grid_values(occupancy_grid) >= CostValues.OCCUPIED occupied_in_path = bool_mask & occupied_mask occupied_count = np.sum(occupied_in_path) diff --git a/dimos/mapping/occupancy/path_resampling.py b/dimos/mapping/occupancy/path_resampling.py index 1f50fad7cdf..292273b438d 100644 --- a/dimos/mapping/occupancy/path_resampling.py +++ b/dimos/mapping/occupancy/path_resampling.py @@ -13,23 +13,75 @@ # limitations under the License. +from copy import deepcopy import math +from typing import TypeVar, cast import numpy as np from scipy.ndimage import uniform_filter1d +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose as CdrPose, + PoseStamped as CdrPoseStamped, + Quaternion as CdrQuaternion, +) +from dimos.messages.nav_msgs.msg import Path as CdrPath 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.Vector3 import Vector3 from dimos.msgs.nav_msgs.Path import Path from dimos.utils.logging_config import setup_logger -from dimos.utils.transform_utils import euler_to_quaternion logger = setup_logger() -def _add_orientations_to_path(path: Path, goal_orientation: Quaternion) -> None: +_PathT = TypeVar("_PathT", Path, CdrPath) +_Stamped = PoseStamped | CdrPoseStamped +_Quaternion = Quaternion | CdrQuaternion + + +def _pose(stamped: _Stamped) -> PoseStamped | CdrPose: + return stamped.pose if isinstance(stamped, CdrPoseStamped) else stamped + + +def _new_pose( + path: Path | CdrPath, position: tuple[float, float, float], orientation: _Quaternion +) -> _Stamped: + if isinstance(path, CdrPath): + x, y, z = position + return CdrPoseStamped( + header=deepcopy(path.header), + pose=CdrPose( + position=Point(x=float(x), y=float(y), z=float(z)), + orientation=CdrQuaternion( + x=orientation.x, y=orientation.y, z=orientation.z, w=orientation.w + ), + ), + ) + return PoseStamped( + frame_id=path.frame_id, + position=position, + orientation=Quaternion(orientation.x, orientation.y, orientation.z, orientation.w), + ) + + +def _new_path(path: _PathT, poses: list[_Stamped]) -> _PathT: + if isinstance(path, CdrPath): + return CdrPath(header=deepcopy(path.header), poses=cast("list[CdrPoseStamped]", poses)) + return Path(frame_id=path.frame_id, poses=cast("list[PoseStamped]", poses)) + + +def _set_orientation(stamped: _Stamped, orientation: _Quaternion) -> None: + if isinstance(stamped, CdrPoseStamped): + stamped.pose.orientation = CdrQuaternion( + x=orientation.x, y=orientation.y, z=orientation.z, w=orientation.w + ) + else: + stamped.orientation = Quaternion(orientation.x, orientation.y, orientation.z, orientation.w) + + +def _add_orientations_to_path(path: Path | CdrPath, goal_orientation: _Quaternion) -> None: """Add orientations to path poses based on direction of movement. Args: @@ -44,8 +96,8 @@ def _add_orientations_to_path(path: Path, goal_orientation: Quaternion) -> None: # Calculate orientations for all poses except the last one for i in range(len(path.poses) - 1): - current_pose = path.poses[i] - next_pose = path.poses[i + 1] + current_pose = _pose(path.poses[i]) + next_pose = _pose(path.poses[i + 1]) # Calculate direction to next point dx = next_pose.position.x - current_pose.position.x @@ -55,24 +107,29 @@ def _add_orientations_to_path(path: Path, goal_orientation: Quaternion) -> None: yaw = math.atan2(dy, dx) # Convert to quaternion (roll=0, pitch=0, yaw) - orientation = euler_to_quaternion(Vector3(0, 0, yaw)) - current_pose.orientation = orientation + orientation = Quaternion(0, 0, math.sin(yaw / 2), math.cos(yaw / 2)) + _set_orientation(path.poses[i], orientation) # Set last pose orientation identity_quat = Quaternion(0, 0, 0, 1) - if goal_orientation != identity_quat: + if (goal_orientation.x, goal_orientation.y, goal_orientation.z, goal_orientation.w) != ( + 0, + 0, + 0, + 1, + ): # Use the provided goal orientation if it's not the identity - path.poses[-1].orientation = goal_orientation + _set_orientation(path.poses[-1], goal_orientation) elif len(path.poses) > 1: # Use the previous pose's orientation - path.poses[-1].orientation = path.poses[-2].orientation + _set_orientation(path.poses[-1], _pose(path.poses[-2]).orientation) else: # Single pose with identity goal orientation - path.poses[-1].orientation = identity_quat + _set_orientation(path.poses[-1], identity_quat) # TODO: replace goal_pose with just goal_orientation -def simple_resample_path(path: Path, goal_pose: Pose, spacing: float) -> Path: +def simple_resample_path(path: _PathT, goal_pose: Pose | CdrPose, spacing: float) -> _PathT: """Resample a path to have approximately uniform spacing between poses. Args: @@ -82,21 +139,21 @@ def simple_resample_path(path: Path, goal_pose: Pose, spacing: float) -> Path: Returns: A new Path with resampled poses """ - if len(path) < 2 or spacing <= 0: + if len(path.poses) < 2 or spacing <= 0: return path - resampled = [] - resampled.append(path.poses[0]) + resampled: list[_Stamped] = [] + resampled.append(deepcopy(path.poses[0])) accumulated_distance = 0.0 for i in range(1, len(path.poses)): current = path.poses[i] - prev = path.poses[i - 1] + prev: _Stamped = path.poses[i - 1] # Calculate segment distance - dx = current.x - prev.x - dy = current.y - prev.y + dx = _pose(current).position.x - _pose(prev).position.x + dy = _pose(current).position.y - _pose(prev).position.y segment_length = (dx**2 + dy**2) ** 0.5 if segment_length < 1e-10: @@ -114,13 +171,9 @@ def simple_resample_path(path: Path, goal_pose: Pose, spacing: float) -> Path: break # Create new pose - new_x = prev.x + dir_x * dist_along - new_y = prev.y + dir_y * dist_along - new_pose = PoseStamped( - frame_id=path.frame_id, - position=[new_x, new_y, 0.0], - orientation=prev.orientation, # Keep same orientation - ) + new_x = _pose(prev).position.x + dir_x * dist_along + new_y = _pose(prev).position.y + dir_y * dist_along + new_pose = _new_pose(path, (new_x, new_y, 0.0), _pose(prev).orientation) resampled.append(new_pose) # Update for next iteration @@ -133,10 +186,13 @@ def simple_resample_path(path: Path, goal_pose: Pose, spacing: float) -> Path: # Add last pose if not already there if len(path.poses) > 1: last = path.poses[-1] - if not resampled or (resampled[-1].x != last.x or resampled[-1].y != last.y): - resampled.append(last) + if not resampled or ( + _pose(resampled[-1]).position.x != _pose(last).position.x + or _pose(resampled[-1]).position.y != _pose(last).position.y + ): + resampled.append(deepcopy(last)) - ret = Path(frame_id=path.frame_id, poses=resampled) + ret = _new_path(path, resampled) _add_orientations_to_path(ret, goal_pose.orientation) @@ -144,8 +200,8 @@ def simple_resample_path(path: Path, goal_pose: Pose, spacing: float) -> Path: def smooth_resample_path( - path: Path, goal_pose: Pose, spacing: float, smoothing_window: int = 100 -) -> Path: + path: _PathT, goal_pose: Pose | CdrPose, spacing: float, smoothing_window: int = 100 +) -> _PathT: """Resample a path with smoothing to reduce jagged corners and abrupt turns. This produces smoother paths than simple_resample_path by: @@ -165,21 +221,17 @@ def smooth_resample_path( """ if len(path.poses) == 1: - p = path.poses[0].position + p = _pose(path.poses[0]).position o = goal_pose.orientation - new_pose = PoseStamped( - frame_id=path.frame_id, - position=[p.x, p.y, p.z], - orientation=[o.x, o.y, o.z, o.w], - ) - return Path(frame_id=path.frame_id, poses=[new_pose]) + new_pose = _new_pose(path, (p.x, p.y, p.z), o) + return _new_path(path, [new_pose]) - if len(path) < 2 or spacing <= 0: + if len(path.poses) < 2 or spacing <= 0: return path # Extract x, y coordinates from path - xs = np.array([p.x for p in path.poses]) - ys = np.array([p.y for p in path.poses]) + xs = np.array([_pose(p).position.x for p in path.poses]) + ys = np.array([_pose(p).position.y for p in path.poses]) # Remove duplicate consecutive points diffs = np.sqrt(np.diff(xs) ** 2 + np.diff(ys) ** 2) @@ -243,16 +295,14 @@ def smooth_resample_path( sampled_y = np.interp(sample_distances, arc_length_smooth, ys_smooth) # Create resampled poses - resampled = [] + resampled: list[_Stamped] = [] for i in range(len(sampled_x)): - new_pose = PoseStamped( - frame_id=path.frame_id, - position=[float(sampled_x[i]), float(sampled_y[i]), 0.0], - orientation=Quaternion(0, 0, 0, 1), + new_pose = _new_pose( + path, (float(sampled_x[i]), float(sampled_y[i]), 0.0), Quaternion(0, 0, 0, 1) ) resampled.append(new_pose) - ret = Path(frame_id=path.frame_id, poses=resampled) + ret = _new_path(path, resampled) _add_orientations_to_path(ret, goal_pose.orientation) diff --git a/dimos/mapping/occupancy/test_generated_operations.py b/dimos/mapping/occupancy/test_generated_operations.py new file mode 100644 index 00000000000..a3617caa21b --- /dev/null +++ b/dimos/mapping/occupancy/test_generated_operations.py @@ -0,0 +1,124 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from copy import deepcopy + +from dimos_cdrgen.registry import decode, encode +import numpy as np +import pytest + +from dimos.mapping.occupancy.extrude_occupancy import generate_mujoco_scene, identify_convex_shapes +from dimos.mapping.occupancy.gradient import gradient, voronoi_gradient +from dimos.mapping.occupancy.inflation import simple_inflate +from dimos.mapping.occupancy.operations import overlay_occupied, smooth_occupied +from dimos.message_utils.occupancy import occupancy_from_array, occupancy_view +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.nav_msgs.msg import OccupancyGrid +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid as LegacyGrid + + +@pytest.fixture(params=[(0, 0), (-1, 999999999), (1700000000, 123456789)]) +def grids(request): + cells = np.zeros((7, 7), dtype=np.int8) + cells[0, 0] = -1 + cells[3, 3] = cells[4, 4] = 100 + generated = occupancy_from_array( + cells, + header=Header(frame_id="map", stamp=Time(sec=request.param[0], nanosec=request.param[1])), + resolution=0.25, + ) + generated.info.map_load_time = Time(sec=7, nanosec=9) + return generated, LegacyGrid(grid=cells, resolution=0.25, frame_id="map", ts=0) + + +@pytest.mark.parametrize( + "operation,kwargs", + [ + (gradient, {"max_distance": 0.5}), + (voronoi_gradient, {"max_distance": 0.5}), + (smooth_occupied, {}), + (simple_inflate, {"radius": 0.25}), + ], +) +def test_existing_algorithms_preserve_input_identity_and_exact_generated_metadata( + grids, operation, kwargs +): + generated, legacy = grids + original = deepcopy(generated) + + actual = operation(generated, **kwargs) + old = operation(legacy, **kwargs) + received = decode(encode(actual), OccupancyGrid) + + assert type(actual) is OccupancyGrid + assert type(old) is LegacyGrid + np.testing.assert_array_equal(occupancy_view(received), old.grid) + assert received.header == original.header + assert received.info.map_load_time == Time(sec=7, nanosec=9) + assert received.info.origin == original.info.origin + assert occupancy_view(received)[0, 0] == -1 + np.testing.assert_array_equal(generated.data, original.data) + assert generated.header == original.header + assert generated.info == original.info + actual.header.stamp.nanosec = 0 + actual.info.origin.position.x = 99 + actual.data[0] = 42 + np.testing.assert_array_equal(generated.data, original.data) + assert generated.header == original.header + assert generated.info == original.info + assert legacy.ts == 0 + + +def test_gradient_retains_unknown_obstacles_and_exact_distance_cost(grids): + generated, _ = grids + + actual = gradient(generated, max_distance=0.5) + + assert occupancy_view(actual)[0, 0] == -1 + assert occupancy_view(actual)[3, 3] == 100 + assert occupancy_view(actual)[3, 2] == 50 + assert occupancy_view(actual)[1, 1] == 0 + + +def test_overlay_preserves_base_metadata_and_rejects_different_dimensions(grids): + base, _ = grids + cells = np.zeros((7, 7), dtype=np.int8) + cells[2, 5] = 100 + overlay = occupancy_from_array( + cells, header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), resolution=0.25 + ) + + actual = overlay_occupied(base, overlay) + + assert occupancy_view(actual)[2, 5] == 100 + assert occupancy_view(actual)[3, 3] == 100 + assert actual.header == base.header + assert actual.info.map_load_time == base.info.map_load_time + other = occupancy_from_array( + np.zeros((2, 2), dtype=np.int8), header=overlay.header, resolution=0.25 + ) + with pytest.raises(ValueError, match="Grid shapes must match"): + overlay_occupied(base, other) + + +def test_generated_and_legacy_extrusion_produce_identical_offline_scene(grids): + generated, legacy = grids + + actual = generate_mujoco_scene(generated) + + assert actual == generate_mujoco_scene(legacy) + assert identify_convex_shapes(generated) == identify_convex_shapes(legacy) + assert '' in actual + assert 'name="wall_0"' in actual diff --git a/dimos/mapping/occupancy/test_generated_path_resampling.py b/dimos/mapping/occupancy/test_generated_path_resampling.py new file mode 100644 index 00000000000..c8c80792148 --- /dev/null +++ b/dimos/mapping/occupancy/test_generated_path_resampling.py @@ -0,0 +1,178 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from copy import deepcopy +import math + +from dimos_cdrgen.registry import decode, encode +import numpy as np +import pytest + +from dimos.mapping.occupancy.path_resampling import simple_resample_path, smooth_resample_path +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.geometry_msgs.Pose import Pose as LegacyPose +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped as LegacyStamped +from dimos.msgs.nav_msgs.Path import Path as LegacyPath + + +def make_path(points, stamp): + header = Header(frame_id="map", stamp=time_from_nanoseconds(stamp)) + return Path( + header=deepcopy(header), + poses=[ + PoseStamped( + header=deepcopy(header), + pose=Pose( + position=Point(x=x, y=y, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1) + ), + ) + for x, y in points + ], + ) + + +def values(path): + return np.array( + [ + [ + p.pose.position.x, + p.pose.position.y, + p.pose.position.z, + p.pose.orientation.x, + p.pose.orientation.y, + p.pose.orientation.z, + p.pose.orientation.w, + ] + for p in path.poses + ] + ) + + +@pytest.mark.parametrize("method", [simple_resample_path, smooth_resample_path]) +@pytest.mark.parametrize("stamp", [0, -1, 1700000000123456789]) +@pytest.mark.parametrize("goal_yaw", [0, math.pi / 2]) +def test_generated_resampling_keeps_exact_metadata_and_matches_existing_numerics( + method, stamp, goal_yaw +): + points = [(0.0, 0.0), (0.4, 0.0), (0.4, 0.4), (0.8, 0.4)] + source = decode(encode(make_path(points, stamp)), Path) + before = deepcopy(source) + quaternion = Quaternion(x=0, y=0, z=math.sin(goal_yaw / 2), w=math.cos(goal_yaw / 2)) + goal = Pose(position=Point(x=0.8, y=0.4, z=0), orientation=quaternion) + legacy = LegacyPath( + frame_id="map", + poses=[ + LegacyStamped(frame_id="map", position=[x, y, 0], orientation=[0, 0, 0, 1]) + for x, y in points + ], + ) + legacy_goal = LegacyPose(0.8, 0.4, 0, 0, 0, quaternion.z, quaternion.w) + + result = method(source, goal, 0.1) + expected = method(legacy, legacy_goal, 0.1) + + assert type(result) is Path + assert result.header == source.header + assert all(p.header == source.header for p in result.poses) + np.testing.assert_allclose( + values(result), + [ + [ + p.position.x, + p.position.y, + p.position.z, + p.orientation.x, + p.orientation.y, + p.orientation.z, + p.orientation.w, + ] + for p in expected.poses + ], + rtol=0, + atol=1e-14, + ) + np.testing.assert_allclose(values(result)[[0, -1], :3], [[0, 0, 0], [0.8, 0.4, 0]]) + assert decode(encode(result), Path) == result + assert source == before + result.poses[0].pose.position.x = 100 + result.poses[0].header.frame_id = "changed" + result.header.frame_id = "changed" + assert source == before + assert result.poses[-1].header == source.header + + +@pytest.mark.parametrize("method", [simple_resample_path, smooth_resample_path]) +@pytest.mark.parametrize( + "points,spacing", + [ + ([], 0.1), + ([(1.0, 2.0)], 0.1), + ([(0.0, 0.0), (1.0, 1.0)], 0), + ([(0.0, 0.0), (1.0, 1.0)], -1), + ], +) +def test_generated_resampling_handles_empty_single_repeated_and_disabled_paths( + method, points, spacing +): + path = make_path(points, 0) + goal = Pose(position=Point(x=1, y=2, z=0), orientation=Quaternion(x=0, y=0, z=1, w=0)) + + result = method(path, goal, spacing) + + assert result.header == path.header + assert len(result.poses) == len(points) + assert [(p.pose.position.x, p.pose.position.y, p.pose.position.z) for p in result.poses] == [ + (x, y, 0) for x, y in points + ] + + +@pytest.mark.parametrize("method,count", [(simple_resample_path, 1), (smooth_resample_path, 2)]) +def test_generated_repeated_coordinates_retain_existing_method_behavior(method, count): + path = make_path([(1.0, 2.0), (1.0, 2.0)], 0) + goal = Pose(position=Point(x=1, y=2, z=0), orientation=Quaternion(x=0, y=0, z=1, w=0)) + result = method(path, goal, 0.1) + assert len(result.poses) == count + assert [(p.pose.position.x, p.pose.position.y) for p in result.poses] == [(1.0, 2.0)] * count + + +def test_generated_single_smooth_pose_uses_goal_orientation_and_copies_metadata(): + path = make_path([(1.0, 2.0)], -1) + before = deepcopy(path) + goal = Pose(position=Point(x=1, y=2, z=0), orientation=Quaternion(x=0, y=0, z=1, w=0)) + result = smooth_resample_path(path, goal, 0.1) + assert result.poses[0].pose.orientation == goal.orientation + assert result.header == path.header and result.poses[0].header == path.header + assert path == before + + +def test_simple_resampling_preserves_existing_endpoint_source_headers_without_aliasing(): + path = make_path([(0.0, 0.0), (0.25, 0.0)], -1) + path.poses[0].header.stamp = time_from_nanoseconds(0) + path.poses[-1].header.stamp = time_from_nanoseconds(17) + before = deepcopy(path) + goal = Pose(position=Point(x=0.25, y=0, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1)) + + result = simple_resample_path(path, goal, 0.1) + + assert [p.header.stamp for p in result.poses] == [ + path.poses[0].header.stamp, + path.header.stamp, + path.header.stamp, + path.poses[-1].header.stamp, + ] + assert path == before diff --git a/dimos/mapping/occupancy/test_generated_rendering.py b/dimos/mapping/occupancy/test_generated_rendering.py new file mode 100644 index 00000000000..2c1db5f466e --- /dev/null +++ b/dimos/mapping/occupancy/test_generated_rendering.py @@ -0,0 +1,189 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from copy import deepcopy +import math + +from dimos_cdrgen.registry import decode, encode +import numpy as np +import pytest + +from dimos.mapping.occupancy.path_mask import make_path_mask +from dimos.mapping.occupancy.visualizations import generate_rgba_texture, visualize_occupancy_grid +from dimos.mapping.occupancy.visualize_path import visualize_path +from dimos.message_utils.image import image_view +from dimos.message_utils.occupancy import grid_to_world, occupancy_from_array +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.sensor_msgs.msg import Image +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped as LegacyStamped +from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid as LegacyGrid +from dimos.msgs.nav_msgs.Path import Path as LegacyPath + + +def make_values(stamp, rotated=False): + cells = np.zeros((9, 9), dtype=np.int8) + cells[0, 0] = -1 + cells[7:9, 7:9] = 100 + cells[5, 5] = 50 + grid = occupancy_from_array( + cells, + header=Header(frame_id="map", stamp=time_from_nanoseconds(stamp)), + resolution=0.25, + origin=Pose( + position=Point(x=0, y=0, z=0), + orientation=Quaternion( + x=0, + y=0, + z=math.sin(math.pi / 4) if rotated else 0, + w=math.cos(math.pi / 4) if rotated else 1, + ), + ), + ) + path = Path( + header=deepcopy(grid.header), + poses=[ + PoseStamped( + header=deepcopy(grid.header), + pose=Pose( + position=grid_to_world(grid, coordinate), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) + for coordinate in [(1, 1), (4, 1), (4, 4)] + ], + ) + old_grid = LegacyGrid(cells.copy(), resolution=0.25, frame_id="map", ts=0) + old_path = LegacyPath( + frame_id="map", + poses=[ + LegacyStamped( + frame_id="map", position=[x * 0.25, y * 0.25, 0], orientation=[0, 0, 0, 1] + ) + for x, y in [(1, 1), (4, 1), (4, 4)] + ], + ) + return grid, path, old_grid, old_path + + +@pytest.mark.parametrize("palette", ["rainbow", "turbo"]) +@pytest.mark.parametrize("stamp", [0, -1, 1700000000123456789]) +def test_generated_map_pixels_match_existing_palette_and_path_overlay(palette, stamp): + grid, path, old_grid, old_path = make_values(stamp) + before = encode(grid) + + result = visualize_occupancy_grid(grid, palette, path) + expected = visualize_occupancy_grid(old_grid, palette, old_path) + + assert type(result) is Image and result.encoding == "bgr8" + assert result.header == grid.header + np.testing.assert_array_equal(image_view(result), expected.data) + decoded = decode(encode(result), Image) + np.testing.assert_array_equal(image_view(decoded), expected.data) + result.header.frame_id = "changed" + result.data[0] = 255 + assert encode(grid) == before + + +@pytest.mark.parametrize("stamp", [0, -1, 1700000000123456789]) +def test_generated_robot_footprints_match_existing_pixels_and_preserve_header(stamp): + grid, path, old_grid, old_path = make_values(stamp) + result = visualize_path(grid, path, 0.2, 0.4, scale=4) + expected = visualize_path(old_grid, old_path, 0.2, 0.4, scale=4) + assert result.header == grid.header and result.encoding == "bgr8" + assert image_view(result).shape == (36, 36, 3) + np.testing.assert_array_equal(image_view(result), expected.data) + + +@pytest.mark.parametrize( + "pose_index,max_length", [(0, float("inf")), (0, 0.8), (1, float("inf")), (2, 0)] +) +@pytest.mark.parametrize("rotated", [False, True]) +def test_generated_path_mask_keeps_existing_clipping_and_rotated_origin_coordinates( + pose_index, max_length, rotated +): + grid, path, old_grid, old_path = make_values(0, rotated) + result = make_path_mask(grid, path, 0.2, pose_index, max_length) + expected = make_path_mask(old_grid, old_path, 0.2, pose_index, max_length) + np.testing.assert_array_equal(result, expected) + + +@pytest.mark.parametrize("colormap", [None, "viridis"]) +@pytest.mark.parametrize("opacity", [0.0, 0.5, 1.0]) +@pytest.mark.parametrize("cost_range", [None, (0, 50)]) +def test_generated_rgba_texture_matches_existing_colormap_opacity_and_range( + colormap, opacity, cost_range +): + grid, _, old_grid, _ = make_values(-1) + result = generate_rgba_texture(grid, colormap, opacity, cost_range, background="#102030") + expected = generate_rgba_texture(old_grid, colormap, opacity, cost_range, background="#102030") + np.testing.assert_array_equal(result, expected) + np.testing.assert_array_equal(result[0, 0], [0, 0, 0, 0]) + + +def test_generated_default_rgba_values_follow_existing_foxglove_palette(): + grid = occupancy_from_array( + np.array([[-1, 0, 100]], dtype=np.int8), + resolution=1, + header=Header(frame_id="map", stamp=time_from_nanoseconds(0)), + ) + np.testing.assert_array_equal( + generate_rgba_texture(grid), [[[0, 0, 0, 0], [72, 73, 129, 255], [0, 0, 0, 255]]] + ) + + +@pytest.mark.parametrize("occupied_x", [1, 20]) +def test_generated_mask_removes_small_occupied_edge_fraction(occupied_x): + cells = np.zeros((5, 41), dtype=np.int8) + cells[2, occupied_x] = 100 + grid = occupancy_from_array( + cells, resolution=1, header=Header(frame_id="map", stamp=time_from_nanoseconds(0)) + ) + path = Path( + header=deepcopy(grid.header), + poses=[ + PoseStamped( + header=deepcopy(grid.header), + pose=Pose( + position=Point(x=x, y=2, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1) + ), + ) + for x in [0, 40] + ], + ) + expected = np.zeros((5, 41), dtype=bool) + expected[2, :] = True + expected[2, occupied_x] = False + np.testing.assert_array_equal(make_path_mask(grid, path, 1), expected) + + +def test_generated_mask_rejects_unsafe_occupied_fraction(): + grid, path, _, _ = make_values(0) + cells = np.array(grid.data, copy=True).reshape(9, 9) + cells[1, 2] = 100 + grid.data = cells.reshape(-1) + with pytest.raises(ValueError, match="More than 5%"): + make_path_mask(grid, path, 0.2) + + +def test_generated_turbo_handles_unknown_int8_cells_with_exact_existing_colors(): + grid = occupancy_from_array( + np.array([[-1, 0, 100]], dtype=np.int8), + resolution=1, + header=Header(frame_id="map", stamp=time_from_nanoseconds(0)), + ) + image = visualize_occupancy_grid(grid, "turbo") + np.testing.assert_array_equal(image_view(image), [[[28, 24, 34], [27, 23, 35], [0, 0, 0]]]) diff --git a/dimos/mapping/occupancy/visualizations.py b/dimos/mapping/occupancy/visualizations.py index 560707ccae4..4c7914eafd7 100644 --- a/dimos/mapping/occupancy/visualizations.py +++ b/dimos/mapping/occupancy/visualizations.py @@ -12,13 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy from functools import lru_cache -from typing import Literal, TypeAlias +from typing import Literal, TypeAlias, overload import matplotlib.pyplot as plt import numpy as np from numpy.typing import NDArray +from dimos.mapping.occupancy.operations import _grid_coordinates, _grid_values, _GridT +from dimos.message_utils.image import image_from_array +from dimos.messages.geometry_msgs.msg import PoseStamped as CdrPoseStamped +from dimos.messages.nav_msgs.msg import OccupancyGrid as CdrOccupancyGrid, Path as CdrPath +from dimos.messages.sensor_msgs.msg import Image as CdrImage from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid from dimos.msgs.nav_msgs.Path import Path from dimos.msgs.sensor_msgs.Image import Image, ImageFormat @@ -26,38 +32,60 @@ Palette: TypeAlias = Literal["rainbow", "turbo"] +def _map_image( + grid: OccupancyGrid | CdrOccupancyGrid, pixels: NDArray[np.uint8] +) -> Image | CdrImage: + if isinstance(grid, CdrOccupancyGrid): + return image_from_array(pixels, encoding="bgr8", header=deepcopy(grid.header)) + return Image(data=pixels, format=ImageFormat.BGR, frame_id=grid.frame_id, ts=grid.ts) + + +@overload +def visualize_occupancy_grid( + occupancy_grid: CdrOccupancyGrid, palette: Palette, path: Path | CdrPath | None = None +) -> CdrImage: ... + + +@overload def visualize_occupancy_grid( - occupancy_grid: OccupancyGrid, palette: Palette, path: Path | None = None -) -> Image: + occupancy_grid: OccupancyGrid, palette: Palette, path: Path | CdrPath | None = None +) -> Image: ... + + +def visualize_occupancy_grid( + occupancy_grid: OccupancyGrid | CdrOccupancyGrid, + palette: Palette, + path: Path | CdrPath | None = None, +) -> Image | CdrImage: match palette: case "rainbow": - bgr_image = rainbow_image(occupancy_grid.grid) + bgr_image = rainbow_image(_grid_values(occupancy_grid)) case "turbo": - bgr_image = turbo_image(occupancy_grid.grid) + bgr_image = turbo_image(_grid_values(occupancy_grid)) case _: raise NotImplementedError() if path is not None and len(path.poses) > 0: _draw_path(occupancy_grid, bgr_image, path) - return Image( - data=bgr_image, - format=ImageFormat.BGR, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + return _map_image(occupancy_grid, bgr_image) -def _draw_path(occupancy_grid: OccupancyGrid, bgr_image: NDArray[np.uint8], path: Path) -> None: +def _draw_path( + occupancy_grid: OccupancyGrid | CdrOccupancyGrid, + bgr_image: NDArray[np.uint8], + path: Path | CdrPath, +) -> None: import cv2 points = [] for pose in path.poses: - grid_coord = occupancy_grid.world_to_grid([pose.x, pose.y, pose.z]) - pixel_x = int(grid_coord.x) - pixel_y = int(grid_coord.y) + position = pose.pose.position if isinstance(pose, CdrPoseStamped) else pose.position + x, y = _grid_coordinates(occupancy_grid, position) + pixel_x = int(x) + pixel_y = int(y) - if 0 <= pixel_x < occupancy_grid.width and 0 <= pixel_y < occupancy_grid.height: + if 0 <= pixel_x < occupancy_grid.info.width and 0 <= pixel_y < occupancy_grid.info.height: points.append((pixel_x, pixel_y)) if len(points) > 1: @@ -109,7 +137,8 @@ def turbo_image(grid: NDArray[np.int8]) -> NDArray[np.uint8]: # Map grid values to lookup indices # Values: -1 -> 255, 0-100 -> 0-100, clipped to valid range - lookup_indices = np.where(grid == -1, 255, np.clip(grid, 0, 100)).astype(np.uint8) + lookup_indices = np.clip(grid, 0, 100).astype(np.uint8) + lookup_indices[grid == -1] = 255 # Create BGR image using lookup table (vectorized operation) return color_lut[lookup_indices] @@ -140,7 +169,7 @@ def _interpolate_turbo(t: float) -> tuple[int, int, int]: def generate_rgba_texture( - grid: OccupancyGrid, + grid: _GridT, colormap: str | None = None, opacity: float = 1.0, cost_range: tuple[int, int] | None = None, @@ -159,6 +188,7 @@ def generate_rgba_texture( RGBA numpy array of shape (height, width, 4). Note: NOT flipped - caller handles orientation. """ + cells = _grid_values(grid) if background is not None: bg = background.lstrip("#") bg_rgb = np.array([int(bg[i : i + 2], 16) for i in (0, 2, 4)], dtype=np.float32) @@ -166,18 +196,18 @@ def generate_rgba_texture( bg_rgb = np.array([0, 0, 0], dtype=np.float32) if cost_range is not None: - in_range_mask = (grid.grid >= cost_range[0]) & (grid.grid <= cost_range[1]) + in_range_mask = (cells >= cost_range[0]) & (cells <= cost_range[1]) else: in_range_mask = None if colormap is not None: cmap = plt.get_cmap(colormap) - grid_float = grid.grid.astype(np.float32) + grid_float = cells.astype(np.float32) - vis = np.zeros((grid.height, grid.width, 4), dtype=np.uint8) + vis = np.zeros((grid.info.height, grid.info.width, 4), dtype=np.uint8) - free_mask = grid.grid == 0 - occupied_mask = grid.grid > 0 + free_mask = cells == 0 + occupied_mask = cells > 0 if np.any(free_mask): fg = np.array(cmap(0.0)[:3]) * 255 @@ -193,21 +223,21 @@ def generate_rgba_texture( vis[occupied_mask, :3] = blended.astype(np.uint8) vis[occupied_mask, 3] = 255 - unknown_mask = grid.grid == -1 + unknown_mask = cells == -1 vis[unknown_mask] = 0 if in_range_mask is not None: - out_of_range = ~in_range_mask & (grid.grid != -1) + out_of_range = ~in_range_mask & (cells != -1) vis[out_of_range, :3] = bg_rgb.astype(np.uint8) vis[out_of_range, 3] = 255 return vis # Default: Foxglove-style coloring - vis = np.zeros((grid.height, grid.width, 4), dtype=np.uint8) + vis = np.zeros((grid.info.height, grid.info.width, 4), dtype=np.uint8) - free_mask = grid.grid == 0 - occupied_mask = grid.grid > 0 + free_mask = cells == 0 + occupied_mask = cells > 0 fg_free = np.array([72, 73, 129], dtype=np.float32) blended_free = fg_free * opacity + bg_rgb * (1 - opacity) @@ -215,18 +245,18 @@ def generate_rgba_texture( vis[free_mask, 3] = 255 if np.any(occupied_mask): - costs = grid.grid[occupied_mask].astype(np.float32) + costs = cells[occupied_mask].astype(np.float32) factor = (1 - costs / 100).clip(0, 1) fg_occ = np.column_stack([72 * factor, 73 * factor, 129 * factor]) blended_occ = fg_occ * opacity + bg_rgb * (1 - opacity) vis[occupied_mask, :3] = blended_occ.astype(np.uint8) vis[occupied_mask, 3] = 255 - unknown_mask = grid.grid == -1 + unknown_mask = cells == -1 vis[unknown_mask] = 0 if in_range_mask is not None: - out_of_range = ~in_range_mask & (grid.grid != -1) + out_of_range = ~in_range_mask & (cells != -1) vis[out_of_range, :3] = bg_rgb.astype(np.uint8) vis[out_of_range, 3] = 255 diff --git a/dimos/mapping/occupancy/visualize_path.py b/dimos/mapping/occupancy/visualize_path.py index 4d3de0f0a0e..e50a7c8c617 100644 --- a/dimos/mapping/occupancy/visualize_path.py +++ b/dimos/mapping/occupancy/visualize_path.py @@ -12,47 +12,78 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import overload + import numpy as np -from dimos.mapping.occupancy.visualizations import visualize_occupancy_grid +from dimos.mapping.occupancy.operations import _grid_coordinates +from dimos.mapping.occupancy.visualizations import _map_image, visualize_occupancy_grid +from dimos.message_utils.geometry import yaw +from dimos.message_utils.image import image_view +from dimos.messages.geometry_msgs.msg import PoseStamped as CdrPoseStamped +from dimos.messages.nav_msgs.msg import OccupancyGrid as CdrOccupancyGrid, Path as CdrPath +from dimos.messages.sensor_msgs.msg import Image as CdrImage from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.Image import Image, ImageFormat +from dimos.msgs.sensor_msgs.Image import Image + + +@overload +def visualize_path( + occupancy_grid: CdrOccupancyGrid, + path: Path | CdrPath, + robot_width: float, + robot_length: float, + thickness: int = 1, + scale: int = 8, +) -> CdrImage: ... +@overload def visualize_path( occupancy_grid: OccupancyGrid, - path: Path, + path: Path | CdrPath, + robot_width: float, + robot_length: float, + thickness: int = 1, + scale: int = 8, +) -> Image: ... + + +def visualize_path( + occupancy_grid: OccupancyGrid | CdrOccupancyGrid, + path: Path | CdrPath, robot_width: float, robot_length: float, thickness: int = 1, scale: int = 8, -) -> Image: +) -> Image | CdrImage: import cv2 image = visualize_occupancy_grid(occupancy_grid, "rainbow") - bgr = image.data + bgr = image_view(image) if isinstance(image, CdrImage) else image.data bgr = cv2.resize( bgr, (bgr.shape[1] * scale, bgr.shape[0] * scale), interpolation=cv2.INTER_NEAREST, - ) + ).astype(np.uint8, copy=False) # Convert robot dimensions from meters to grid cells, then to scaled pixels - resolution = occupancy_grid.resolution + resolution = occupancy_grid.info.resolution robot_width_px = int((robot_width / resolution) * scale) robot_length_px = int((robot_length / resolution) * scale) # Draw robot rectangle at each path point for pose in path.poses: # Convert world coordinates to grid coordinates - grid_coord = occupancy_grid.world_to_grid([pose.x, pose.y, pose.z]) - cx = int(grid_coord.x * scale) - cy = int(grid_coord.y * scale) + position = pose.pose.position if isinstance(pose, CdrPoseStamped) else pose.position + x, y = _grid_coordinates(occupancy_grid, position) + cx = int(x * scale) + cy = int(y * scale) # Get yaw angle from pose orientation - yaw = pose.yaw + angle = yaw(pose.pose.orientation) if isinstance(pose, CdrPoseStamped) else pose.yaw # Define rectangle corners centered at origin (length along x, width along y) half_length = robot_length_px / 2 @@ -68,8 +99,8 @@ def visualize_path( ) # Rotate corners by yaw angle - cos_yaw = np.cos(yaw) - sin_yaw = np.sin(yaw) + cos_yaw = np.cos(angle) + sin_yaw = np.sin(angle) rotation_matrix = np.array([[cos_yaw, -sin_yaw], [sin_yaw, cos_yaw]]) rotated_corners = corners @ rotation_matrix.T @@ -81,9 +112,4 @@ def visualize_path( pts = rotated_corners.astype(np.int32).reshape((-1, 1, 2)) cv2.polylines(bgr, [pts], isClosed=True, color=(0, 0, 0), thickness=thickness) - return Image( - data=bgr, - format=ImageFormat.BGR, - frame_id=occupancy_grid.frame_id, - ts=occupancy_grid.ts, - ) + return _map_image(occupancy_grid, bgr) diff --git a/dimos/mapping/odometry_hist.py b/dimos/mapping/odometry_hist.py index f368a7fed1b..4f3dcd13f8a 100644 --- a/dimos/mapping/odometry_hist.py +++ b/dimos/mapping/odometry_hist.py @@ -15,6 +15,7 @@ from __future__ import annotations from collections import deque +from copy import deepcopy import math from typing import Any @@ -22,13 +23,17 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.nav_msgs.Path import Path +from dimos.message_utils.time import to_nanoseconds +from dimos.messages.geometry_msgs.msg import PoseStamped +from dimos.messages.nav_msgs.msg import Odometry, Path +from dimos.msgs.nav_msgs.Path import Path as LegacyPath +from dimos.visualization.rerun.cdr import path_archetype -def path_at_true_height(path: Path) -> Any: +def path_at_true_height(path: Path | LegacyPath) -> Any: """The default z lift clears a costmap a bare odometry demo has none of.""" + if isinstance(path, Path): + return path_archetype(path, z_offset=0.0, radii=0.02) return path.to_rerun(z_offset=0.0, radii=0.02) @@ -50,37 +55,38 @@ class OdometryHist(Module): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._poses: deque[PoseStamped] = deque(maxlen=self.config.max_poses) - # None, not 0.0: a recording stamped from zero would suppress the first path. - self._last_publish_ts: float | None = None + # None distinguishes an unpublished history from a recording starting at zero. + self._last_publish_ns: int | None = None self._unpublished = False async def handle_odometry(self, msg: Odometry) -> None: - position = msg.pose.position + position = msg.pose.pose.position + stamp_ns = to_nanoseconds(msg.header.stamp) point = (position.x, position.y, position.z) - frame_id = self.config.frame_id or msg.frame_id + frame_id = self.config.frame_id or msg.header.frame_id previous = self._poses[-1] if self._poses else None if ( previous is None - or math.dist((previous.x, previous.y, previous.z), point) >= self.config.min_step_meters - ): - orientation = msg.pose.orientation - self._poses.append( - PoseStamped( - ts=msg.ts, - frame_id=frame_id, - position=list(point), - orientation=[orientation.x, orientation.y, orientation.z, orientation.w], - ) + or math.dist( + (previous.pose.position.x, previous.pose.position.y, previous.pose.position.z), + point, ) + >= self.config.min_step_meters + ): + header = deepcopy(msg.header) + header.frame_id = frame_id + self._poses.append(PoseStamped(header=header, pose=deepcopy(msg.pose.pose))) self._unpublished = True if not self._unpublished: return - if self._last_publish_ts is not None: - elapsed = msg.ts - self._last_publish_ts + if self._last_publish_ns is not None: + elapsed = stamp_ns - self._last_publish_ns # A replay restart moves the stamp backwards; publish rather than wait it out. - if 0.0 <= elapsed < self.config.min_publish_interval_seconds: + if 0 <= elapsed < self.config.min_publish_interval_seconds * 1_000_000_000: return - self._last_publish_ts = msg.ts + self._last_publish_ns = stamp_ns self._unpublished = False - self.odom_hist.publish(Path(ts=msg.ts, frame_id=frame_id, poses=list(self._poses))) + header = deepcopy(msg.header) + header.frame_id = frame_id + self.odom_hist.publish(Path(header=header, poses=deepcopy(list(self._poses)))) diff --git a/dimos/mapping/pointclouds/occupancy.py b/dimos/mapping/pointclouds/occupancy.py index d723b901543..20b7246680f 100644 --- a/dimos/mapping/pointclouds/occupancy.py +++ b/dimos/mapping/pointclouds/occupancy.py @@ -14,13 +14,20 @@ from __future__ import annotations +from copy import deepcopy from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Protocol, TypeVar +from typing import TYPE_CHECKING, Any, Protocol, TypeVar, overload import numpy as np +from numpy.typing import NDArray from scipy import ndimage -from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.message_utils.occupancy import occupancy_from_array +from dimos.message_utils.pointcloud import pointcloud_xyz +from dimos.messages.geometry_msgs.msg import Point, Pose, Quaternion +from dimos.messages.nav_msgs.msg import OccupancyGrid as CdrGrid +from dimos.messages.sensor_msgs.msg import PointCloud2 as CdrCloud +from dimos.msgs.geometry_msgs.Pose import Pose as LegacyPose from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid if TYPE_CHECKING: @@ -37,6 +44,49 @@ class OccupancyConfig: frame_id: str | None = None +def _cloud_points(cloud: PointCloud2 | CdrCloud) -> NDArray[np.float64]: + if isinstance(cloud, CdrCloud): + points = pointcloud_xyz(cloud) + return points[np.isfinite(points).all(axis=1)] + points, _ = cloud.as_numpy() + return points.astype(np.float64) + + +def _legacy_stamp(cloud: PointCloud2 | CdrCloud) -> float: + if isinstance(cloud, CdrCloud): + return 0.0 # Generated output uses its original Header, never this float. + return cloud.ts if cloud.ts is not None else 0.0 + + +def _grid_message( + cloud: PointCloud2 | CdrCloud, + config: OccupancyConfig, + cells: NDArray[np.int8], + origin: tuple[float, float, float] = (0.0, 0.0, 0.0), + *, + legacy_stamp: float | None = None, +) -> OccupancyGrid | CdrGrid: + if isinstance(cloud, CdrCloud): + header = deepcopy(cloud.header) + header.frame_id = config.frame_id or header.frame_id + return occupancy_from_array( + cells, + header=header, + resolution=config.resolution, + origin=Pose( + position=Point(x=origin[0], y=origin[1], z=origin[2]), + orientation=Quaternion(w=1, x=0, y=0, z=0), + ), + ) + return OccupancyGrid( + grid=cells, + resolution=config.resolution, + origin=LegacyPose(*origin), + frame_id=config.frame_id or cloud.frame_id, + ts=legacy_stamp, + ) + + ConfigT = TypeVar("ConfigT", bound=OccupancyConfig, covariant=True) @@ -54,7 +104,7 @@ def __call__(self, cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: ... # Populated after function definitions below -OCCUPANCY_ALGOS: dict[str, Callable[..., OccupancyGrid]] = {} +OCCUPANCY_ALGOS: dict[str, Callable[..., OccupancyGrid | CdrGrid]] = {} @dataclass(frozen=True) @@ -67,7 +117,15 @@ class HeightCostConfig(OccupancyConfig): smoothing: float = 1.0 -def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: +@overload +def height_cost_occupancy(cloud: CdrCloud, **kwargs: Any) -> CdrGrid: ... + + +@overload +def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: ... + + +def height_cost_occupancy(cloud: PointCloud2 | CdrCloud, **kwargs: Any) -> OccupancyGrid | CdrGrid: """Create a costmap based on terrain slope (rate of change of height). Costs are assigned based on the gradient magnitude of the terrain height. @@ -83,17 +141,11 @@ def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: OccupancyGrid with costs 0-100 based on terrain slope, -1 for unknown """ cfg = HeightCostConfig(**kwargs) - points, _ = cloud.as_numpy() - points = points.astype(np.float64) # Upcast to avoid float32 rounding - ts = cloud.ts if hasattr(cloud, "ts") and cloud.ts is not None else 0.0 + points = _cloud_points(cloud) + ts = _legacy_stamp(cloud) if len(points) == 0: - return OccupancyGrid( - width=1, - height=1, - resolution=cfg.resolution, - frame_id=cfg.frame_id or cloud.frame_id, - ) + return _grid_message(cloud, cfg, np.full((1, 1), -1, dtype=np.int8)) # Find bounds of the point cloud in X-Y plane (use all points) min_x = np.min(points[:, 0]) @@ -113,11 +165,7 @@ def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: height = int(np.ceil((max_y - min_y) / cfg.resolution)) # Create origin pose - origin = Pose() - origin.position.x = min_x - origin.position.y = min_y - origin.position.z = 0.0 - origin.orientation.w = 1.0 + origin = (float(min_x), float(min_y), 0.0) # Step 1: Build min and max height maps for each cell # Initialize with NaN to track which cells have observations @@ -204,13 +252,7 @@ def height_cost_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: else: cost = np.full((height, width), -1, dtype=np.int8) - return OccupancyGrid( - grid=cost, - resolution=cfg.resolution, - origin=origin, - frame_id=cfg.frame_id or cloud.frame_id, - ts=ts, - ) + return _grid_message(cloud, cfg, cost, origin, legacy_stamp=ts) @dataclass(frozen=True) @@ -223,7 +265,15 @@ class GeneralOccupancyConfig(OccupancyConfig): # can remove, just needs pulling out of unitree type/map.py -def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: +@overload +def general_occupancy(cloud: CdrCloud, **kwargs: Any) -> CdrGrid: ... + + +@overload +def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: ... + + +def general_occupancy(cloud: PointCloud2 | CdrCloud, **kwargs: Any) -> OccupancyGrid | CdrGrid: """Create an OccupancyGrid from a PointCloud2 message. Args: @@ -235,16 +285,10 @@ def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: OccupancyGrid with occupied cells where points were projected """ cfg = GeneralOccupancyConfig(**kwargs) - points, _ = cloud.as_numpy() - points = points.astype(np.float64) # Upcast to avoid float32 rounding + points = _cloud_points(cloud) if len(points) == 0: - return OccupancyGrid( - width=1, - height=1, - resolution=cfg.resolution, - frame_id=cfg.frame_id or cloud.frame_id, - ) + return _grid_message(cloud, cfg, np.full((1, 1), -1, dtype=np.int8)) # Filter points by height for obstacles obstacle_mask = (points[:, 2] >= cfg.min_height) & (points[:, 2] <= cfg.max_height) @@ -272,11 +316,7 @@ def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: height = int(np.ceil((max_y - min_y) / cfg.resolution)) # Create origin pose (bottom-left corner of the grid) - origin = Pose() - origin.position.x = min_x - origin.position.y = min_y - origin.position.z = 0.0 - origin.orientation.w = 1.0 # No rotation + origin = (float(min_x), float(min_y), 0.0) # Initialize grid (all unknown) grid = np.full((height, width), -1, dtype=np.int8) @@ -328,15 +368,9 @@ def general_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: # Create and return OccupancyGrid # Get timestamp from cloud if available - ts = cloud.ts if hasattr(cloud, "ts") and cloud.ts is not None else 0.0 + ts = _legacy_stamp(cloud) - return OccupancyGrid( - grid=grid, - resolution=cfg.resolution, - origin=origin, - frame_id=cfg.frame_id or cloud.frame_id, - ts=ts, - ) + return _grid_message(cloud, cfg, grid, origin, legacy_stamp=ts) @dataclass(frozen=True) @@ -353,7 +387,15 @@ class SimpleOccupancyConfig(OccupancyConfig): smoothing: float = 1.0 -def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: +@overload +def simple_occupancy(cloud: CdrCloud, **kwargs: Any) -> CdrGrid: ... + + +@overload +def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: ... + + +def simple_occupancy(cloud: PointCloud2 | CdrCloud, **kwargs: Any) -> OccupancyGrid | CdrGrid: """Create a simple occupancy grid with morphological closing. Args: @@ -365,16 +407,10 @@ def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: OccupancyGrid with occupied/free cells """ cfg = SimpleOccupancyConfig(**kwargs) - points, _ = cloud.as_numpy() - points = points.astype(np.float64) # Upcast to avoid float32 rounding + points = _cloud_points(cloud) if len(points) == 0: - return OccupancyGrid( - width=1, - height=1, - resolution=cfg.resolution, - frame_id=cfg.frame_id or cloud.frame_id, - ) + return _grid_message(cloud, cfg, np.full((1, 1), -1, dtype=np.int8)) # Find bounds of the point cloud in X-Y plane min_x = float(np.min(points[:, 0])) - 1.0 @@ -387,11 +423,7 @@ def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: height = int(np.ceil((max_y - min_y) / cfg.resolution)) # Create origin pose (bottom-left corner of the grid) - origin = Pose() - origin.position.x = min_x - origin.position.y = min_y - origin.position.z = 0.0 - origin.orientation.w = 1.0 + origin = (float(min_x), float(min_y), 0.0) # Initialize grid (all unknown) grid = np.full((height, width), -1, dtype=np.int8) @@ -411,15 +443,9 @@ def simple_occupancy(cloud: PointCloud2, **kwargs: Any) -> OccupancyGrid: cfg.max_height, ) - ts = cloud.ts if hasattr(cloud, "ts") and cloud.ts is not None else 0.0 + ts = _legacy_stamp(cloud) - return OccupancyGrid( - grid=grid, - resolution=cfg.resolution, - origin=origin, - frame_id=cfg.frame_id or cloud.frame_id, - ts=ts, - ) + return _grid_message(cloud, cfg, grid, origin, legacy_stamp=ts) # Populate algorithm registry diff --git a/dimos/mapping/ray_tracing/module.py b/dimos/mapping/ray_tracing/module.py index 16bd93fb535..190bcae2147 100644 --- a/dimos/mapping/ray_tracing/module.py +++ b/dimos/mapping/ray_tracing/module.py @@ -18,10 +18,10 @@ from typing import TYPE_CHECKING from dimos.core.native_module import NativeModule, NativeModuleConfig -from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.core.stream import IO, In, Out +from dimos.messages.dimos_msgs.msg import PointCloudRegion, RegionBounds +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.tf2_msgs.msg import TFMessage from dimos.spec import mapping # Max stamp gap between a cloud and the transform used to register it (s). @@ -97,18 +97,19 @@ class RayTracingVoxelMap(NativeModule, mapping.GlobalPointcloud): # An externally loaded map cloud, placed by the latest tf world_frame -> # cloud frame_id. Only the first cloud seeds the map. loaded_map: In[PointCloud2] - tf: In[TFMessage] + tf: IO[TFMessage] + tf_static: In[TFMessage] global_map: Out[PointCloud2] local_map: Out[PointCloud2] local_map_fine: Out[PointCloud2] - region_bounds: Out[PoseStamped] + region_bounds: Out[RegionBounds] # One region of a seeded map as it lands, gated like local_map. seed_bounds # carries its cylinder like region_bounds. - seed_map: Out[PointCloud2] - seed_bounds: Out[PoseStamped] + seed_map: Out[PointCloudRegion] + seed_bounds: Out[RegionBounds] # The map for viewers, one region cell per message with the cell packed in - # the header seq. - map_regions: Out[PointCloud2] + # the explicit region_id. + map_regions: Out[PointCloudRegion] if TYPE_CHECKING: diff --git a/dimos/mapping/ray_tracing/rust/Cargo.lock b/dimos/mapping/ray_tracing/rust/Cargo.lock index a2abd4749a3..03b0f3222e2 100644 --- a/dimos/mapping/ray_tracing/rust/Cargo.lock +++ b/dimos/mapping/ray_tracing/rust/Cargo.lock @@ -255,6 +255,15 @@ dependencies = [ "shlex", ] +[[package]] +name = "cdr" +version = "0.2.4" +source = "git+https://github.com/TomCC7/cdr-rs?rev=38a4368255d408faf374726653b39d4a4e9735cb#38a4368255d408faf374726653b39d4a4e9735cb" +dependencies = [ + "byteorder", + "serde", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -585,6 +594,17 @@ dependencies = [ "subtle", ] +[[package]] +name = "dimos-generated-messages" +version = "0.0.14" +dependencies = [ + "cdr", + "heapless", + "ros2msg", + "serde", + "serde-big-array", +] + [[package]] name = "dimos-lcm" version = "0.1.0" @@ -597,12 +617,13 @@ dependencies = [ [[package]] name = "dimos-module" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#fc6a4e58da89bb073a38dce5ac63728c6559d3a6" +version = "0.2.0" dependencies = [ + "cdr", + "dimos-generated-messages", "dimos-lcm", "dimos-module-macros", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549)", + "lcm-msgs", "nalgebra", "rayon", "serde", @@ -610,6 +631,7 @@ dependencies = [ "tokio", "tracing", "tracing-subscriber", + "transforms", "url", "validator", "zenoh", @@ -618,7 +640,6 @@ dependencies = [ [[package]] name = "dimos-module-macros" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#fc6a4e58da89bb073a38dce5ac63728c6559d3a6" dependencies = [ "proc-macro2", "quote", @@ -632,8 +653,8 @@ version = "0.1.0" dependencies = [ "ahash", "arrayvec", + "dimos-generated-messages", "dimos-module", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen)", "nalgebra", "rayon", "serde", @@ -983,6 +1004,15 @@ version = "0.33.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e762b9b188634fc508d4ec54b62ea3d352c43fd431d18d05d46f559f217f4725" +[[package]] +name = "hash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" +dependencies = [ + "byteorder", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -1024,6 +1054,22 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] + +[[package]] +name = "heapless" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" +dependencies = [ + "hash32", + "serde", + "stable_deref_trait", +] [[package]] name = "heck" @@ -1430,14 +1476,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "lcm-msgs" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" -dependencies = [ - "byteorder", -] - [[package]] name = "lcm-msgs" version = "0.1.0" @@ -2022,6 +2060,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -2438,6 +2486,23 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "ros2msg" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fe291be5cb6d4b8e6bf9c9de80c40bd1cb09ecac50deafa2fdcae0027e04988" +dependencies = [ + "heck", + "pest", + "pest_derive", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "syn 2.0.119", + "thiserror", +] + [[package]] name = "rustc-hash" version = "2.1.3" @@ -2678,6 +2743,15 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-big-array" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11fc7cc2c76d73e0f27ee52abbd64eec84d46f370c88371120433196934e4b7f" +dependencies = [ + "serde", +] + [[package]] name = "serde_core" version = "1.0.229" @@ -3359,6 +3433,18 @@ dependencies = [ "tracing-serde", ] +[[package]] +name = "transforms" +version = "2.2.0" +source = "git+https://github.com/TomCC7/transforms?rev=5ae9563de0c02638c04e393aadfad28444a4398e#5ae9563de0c02638c04e393aadfad28444a4398e" +dependencies = [ + "approx", + "hashbrown 0.17.1", + "libm", + "serde", + "thiserror", +] + [[package]] name = "twox-hash" version = "1.6.3" diff --git a/dimos/mapping/ray_tracing/rust/Cargo.toml b/dimos/mapping/ray_tracing/rust/Cargo.toml index 8d248fd1dbe..1c0e389348b 100644 --- a/dimos/mapping/ray_tracing/rust/Cargo.toml +++ b/dimos/mapping/ray_tracing/rust/Cargo.toml @@ -21,22 +21,23 @@ python = "dimos.mapping.ray_tracing.module:RayTracingVoxelMap" threads = 2 [package.metadata.dimos.module.ray_tracing.inputs] -lidar = "sensor_msgs.PointCloud2" -loaded_map = "sensor_msgs.PointCloud2" -tf = "tf2_msgs.TFMessage" +lidar = "sensor_msgs/msg/PointCloud2" +loaded_map = "sensor_msgs/msg/PointCloud2" +tf = "tf2_msgs/msg/TFMessage" +tf_static = "tf2_msgs/msg/TFMessage" [package.metadata.dimos.module.ray_tracing.outputs] -global_map = "sensor_msgs.PointCloud2" -local_map = "sensor_msgs.PointCloud2" -local_map_fine = "sensor_msgs.PointCloud2" -seed_map = "sensor_msgs.PointCloud2" -seed_bounds = "geometry_msgs.PoseStamped" -region_bounds = "geometry_msgs.PoseStamped" -map_regions = "sensor_msgs.PointCloud2" +global_map = "sensor_msgs/msg/PointCloud2" +local_map = "sensor_msgs/msg/PointCloud2" +local_map_fine = "sensor_msgs/msg/PointCloud2" +seed_map = "dimos_msgs/msg/PointCloudRegion" +seed_bounds = "dimos_msgs/msg/RegionBounds" +region_bounds = "dimos_msgs/msg/RegionBounds" +map_regions = "dimos_msgs/msg/PointCloudRegion" [dependencies] -dimos-module = { git = "https://github.com/dimensionalOS/dimos", rev = "fc6a4e58da89bb073a38dce5ac63728c6559d3a6", version = "=0.1.0" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-module = { path = "../../../../native/rust/dimos-module" } +dimos-generated-messages = { path = "../../../../build/message-source/rust" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } serde = { version = "1", features = ["derive"] } ahash = "0.8" diff --git a/dimos/mapping/ray_tracing/rust/crate-hashes.json b/dimos/mapping/ray_tracing/rust/crate-hashes.json index 187d28513c6..3ee4f926e23 100644 --- a/dimos/mapping/ray_tracing/rust/crate-hashes.json +++ b/dimos/mapping/ray_tracing/rust/crate-hashes.json @@ -1,7 +1,6 @@ { - "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#lcm-msgs@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=", + "git+https://github.com/TomCC7/cdr-rs?rev=38a4368255d408faf374726653b39d4a4e9735cb#cdr@0.2.4": "sha256-akg/bLBY3wQ1YRr2bijE5ZqB211Cst5lKjELLrgiG6o=", + "git+https://github.com/TomCC7/transforms?rev=5ae9563de0c02638c04e393aadfad28444a4398e#transforms@2.2.0": "sha256-Kh8aVjj7zQ+Xd1QHOcAhSJbubg7azaWFSMpaI/ZAgM4=", "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549#dimos-lcm@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=", - "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549#lcm-msgs@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=", - "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#dimos-module-macros@0.1.0": "sha256-ECU2xYh/j+a3XwyhO3UR/GPruo5jNzM3n/9S5BIvVrA=", - "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#dimos-module@0.1.0": "sha256-ECU2xYh/j+a3XwyhO3UR/GPruo5jNzM3n/9S5BIvVrA=" + "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549#lcm-msgs@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=" } \ No newline at end of file diff --git a/dimos/mapping/ray_tracing/rust/flake.nix b/dimos/mapping/ray_tracing/rust/flake.nix index 13dbb3a9ffb..50a8978f535 100644 --- a/dimos/mapping/ray_tracing/rust/flake.nix +++ b/dimos/mapping/ray_tracing/rust/flake.nix @@ -15,7 +15,32 @@ pkgs = nixpkgs.legacyPackages.${system}; name = "dimos-voxel-ray-tracing"; - src = nix-filter.lib { root = ./.; exclude = [ "target" "build" "result" "__pycache__" ]; }; + moduleSource = nix-filter.lib { root = ./.; exclude = [ "target" "build" "result" "__pycache__" ]; }; + sharedArchive = pkgs.fetchurl { + url = "https://codeload.github.com/dimensionalOS/dimos/tar.gz/a9365487b4323151ea85041c6a944b509c0c3268"; + sha256 = "e339487dd2281980fcde7ab03c46dfd2e274491e0f73e66d9c6cbcaeea8902b2"; + }; + sharedSource = pkgs.runCommand "dimos-ray-tracing-cdr-source-input" { + nativeBuildInputs = [ pkgs.gnutar pkgs.gzip ]; + } '' + mkdir -p $out + tar -xzf ${sharedArchive} --strip-components=1 -C $out + ''; + messageSource = (import "${sharedSource}/native/message-source.nix" { inherit pkgs; }).source; + src = pkgs.runCommand "ray-tracing-cdr-source" { } '' + cp -r ${moduleSource} $out + chmod -R u+w $out + mkdir -p $out/sdk + cp -r ${sharedSource}/native/rust/dimos-module ${sharedSource}/native/rust/dimos-module-macros $out/sdk/ + cp -r ${messageSource}/rust $out/messages + substituteInPlace $out/sdk/dimos-module/Cargo.toml \ + --replace-fail '../../../build/message-source/rust' '../../messages' + substituteInPlace $out/Cargo.toml \ + --replace-fail '../../../../native/rust/dimos-module' 'sdk/dimos-module' \ + --replace-fail '../../../../build/message-source/rust' 'messages' \ + --replace-fail '[workspace]' '[workspace] + exclude = ["sdk/dimos-module", "sdk/dimos-module-macros", "messages"]' + ''; generated = crate2nix.tools.${system}.generatedCargoNix { inherit name src; }; diff --git a/dimos/mapping/ray_tracing/rust/src/lib.rs b/dimos/mapping/ray_tracing/rust/src/lib.rs index 34edbda1ca8..d06ec2b9248 100644 --- a/dimos/mapping/ray_tracing/rust/src/lib.rs +++ b/dimos/mapping/ray_tracing/rust/src/lib.rs @@ -16,3 +16,5 @@ pub mod mapper; pub mod module; pub mod region_viz; pub mod voxel_ray_tracer; + +mod wire; diff --git a/dimos/mapping/ray_tracing/rust/src/module.rs b/dimos/mapping/ray_tracing/rust/src/module.rs index 0bb2fc6d4e9..ce80f6ff662 100644 --- a/dimos/mapping/ray_tracing/rust/src/module.rs +++ b/dimos/mapping/ray_tracing/rust/src/module.rs @@ -21,13 +21,21 @@ use crate::voxel_ray_tracer::{ partition_seed, region_at, region_of, Cell, ChunkKey, Config, Cylinder, SeedPartition, SeedRegion, }; -use dimos_module::pointcloud::extract_xyz; +use crate::wire::extract_xyz; +use dimos_generated_messages::builtin_interfaces::msg::time::Time; +use dimos_generated_messages::dimos_msgs::msg::{ + point_cloud_region::PointCloudRegion, region_bounds::RegionBounds, +}; +use dimos_generated_messages::geometry_msgs::msg::point::Point; +use dimos_generated_messages::sensor_msgs::msg::{ + point_cloud2::PointCloud2, point_field::PointField, +}; +use dimos_generated_messages::std_msgs::msg::header::Header; + +use dimos_module::cdr; use dimos_module::{ - debug_throttled, error_throttled, warn_throttled, Input, Module, Output, Tf, Transform, + debug_throttled, error_throttled, warn_throttled, Input, Module, Output, Tf, TransformStamped, }; -use lcm_msgs::geometry_msgs::{Point, Pose as PoseMsg, PoseStamped, Quaternion}; -use lcm_msgs::sensor_msgs::{PointCloud2, PointField}; -use lcm_msgs::std_msgs::{Header, Time}; use tokio::sync::mpsc; use tokio::sync::mpsc::error::TryRecvError; use tokio::task::JoinHandle; @@ -53,42 +61,41 @@ const SEED_STARVED_WARN_AFTER: Duration = Duration::from_secs(5); #[derive(Module)] #[module(name = "ray_tracing", setup = spawn_worker, teardown = stop_worker)] pub struct RayTracingVoxelMap { - #[input(decode = PointCloud2::decode, handler = on_lidar)] + #[input(decode = dimos_module::cdr::decode, handler = on_lidar)] lidar: Input, // An externally loaded map cloud. Only the first one seeds the map. - #[input(decode = PointCloud2::decode, handler = on_loaded_map)] + #[input(decode = dimos_module::cdr::decode, handler = on_loaded_map)] loaded_map: Input, #[tf] tf: Tf, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] global_map: Output, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] local_map: Output, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] local_map_fine: Output, - // Cylinder bounds of the local map. Position is the center, orientation holds - // radius, z_min, z_max. Stamped like local_map so consumers pair them. - #[output(encode = PoseStamped::encode)] - region_bounds: Output, + // Explicit cylinder dimensions, stamped like local_map so consumers pair them. + #[output(encode = cdr::encode)] + region_bounds: Output, // The map for viewers, one region cell per message with the cell packed in - // the header seq. - #[output(encode = PointCloud2::encode)] - map_regions: Output, + // the explicit region_id. + #[output(encode = cdr::encode)] + map_regions: Output, // One region of a seeded map as it lands, gated like local_map. seed_bounds // carries its cylinder like region_bounds. - #[output(encode = PointCloud2::encode)] - seed_map: Output, + #[output(encode = cdr::encode)] + seed_map: Output, - #[output(encode = PoseStamped::encode)] - seed_bounds: Output, + #[output(encode = cdr::encode)] + seed_bounds: Output, #[config] config: Config, @@ -310,17 +317,17 @@ struct Worker { global_map: Output, local_map: Output, local_map_fine: Output, - region_bounds: Output, - map_regions: Output, - seed_map: Output, - seed_bounds: Output, + region_bounds: Output, + map_regions: Output, + seed_map: Output, + seed_bounds: Output, } impl Worker { async fn run(mut self) { let mut state = State { mapper: Mapper::new(self.config.clone()), - last_frame_stamp: Time::default(), + last_frame_stamp: Time { sec: 0, nanosec: 0 }, seed: SeedState::Idle, viz: RegionSweep::default(), }; @@ -362,19 +369,20 @@ impl Worker { /// Fold one lidar frame into the map and publish whatever is due. async fn ingest_frame(&self, state: &mut State, msg: PointCloud2, received: Instant) { - // Register with the transform nearest the cloud stamp, waiting briefly + // Register with the transform at the cloud stamp, waiting briefly // for one still in flight rather than dropping the cloud. let stamp = time_secs(&msg.header.stamp); let tolerance = self.config.tf_match_tolerance_s; let lookup = self .tf .lookup(&self.config.world_frame, &msg.header.frame_id) - .at(stamp) - .tolerance(tolerance); + .at(msg.header.stamp.clone()); let latest = self .tf .get_latest(&self.config.world_frame, &msg.header.frame_id) - .map(|latest| latest.ts); + .map(|latest| { + latest.header.stamp.sec as f64 + latest.header.stamp.nanosec as f64 * 1e-9 + }); if is_stale(stamp, latest, self.config.max_cloud_age_s) { warn_throttled!(Duration::from_secs(5), cloud_frame = %msg.header.frame_id, "Skipped a cloud older than max_cloud_age_s: the map is behind and catching up."); return; @@ -391,7 +399,7 @@ impl Worker { stamp, world_frame = %self.config.world_frame, cloud_frame = %msg.header.frame_id, - "No transform within tolerance of the cloud stamp, dropped a cloud.", + "TF history cannot resolve the cloud stamp, dropped a cloud.", ); return; }; @@ -437,7 +445,7 @@ impl Worker { // Bounds pair with local_map by stamp, so publish them on its cadence. if let Some(c) = region { - let bounds_msg = bounds_to_pose(&c, out_frame_id, stamp.clone()); + let bounds_msg = bounds_message(&c, out_frame_id, stamp.clone()); publish_bounds(&self.region_bounds, &bounds_msg).await; } @@ -466,7 +474,7 @@ impl Worker { } /// Publish the map region cells due this viz tick, the cell in each - /// header seq. + /// explicit region_id. async fn publish_map_regions(&self, state: &mut State, stamp: &Time) { let due = tokio::task::block_in_place(|| { map_regions_due( @@ -478,9 +486,13 @@ impl Worker { ) }); for (cell, points) in due { - let mut cloud = points_to_cloud(&points, &self.config.world_frame, stamp.clone()); - cloud.header.seq = pack_cell(cell); - publish_cloud(&self.map_regions, &cloud).await; + let cloud = points_to_cloud(&points, &self.config.world_frame, stamp.clone()); + let region = PointCloudRegion { + header: cloud.header.clone(), + region_id: pack_cell(cell), + cloud, + }; + publish_cloud(&self.map_regions, ®ion).await; } } @@ -563,7 +575,7 @@ impl Worker { } /// Publish one seeded region as the map now holds it. Cloud and bounds - /// carry the region number in their header seq, which consumers pair on. + /// carry the region number in their explicit region_id, which consumers pair on. async fn publish_seed_region( &self, mapper: &Mapper, @@ -580,19 +592,23 @@ impl Worker { } /// One seeded region as its bounds and cloud messages, both stamped alike -/// with the region number in their header seq. +/// with the region number in their explicit region_id. fn seed_region_messages( cylinder: &Cylinder, points: &[f32], seq: i32, frame_id: &str, stamp: &Time, -) -> (PoseStamped, PointCloud2) { - let mut bounds = bounds_to_pose(cylinder, frame_id, stamp.clone()); - bounds.header.seq = seq; - let mut cloud = points_to_cloud(points, frame_id, stamp.clone()); - cloud.header.seq = seq; - (bounds, cloud) +) -> (RegionBounds, PointCloudRegion) { + let mut bounds = bounds_message(cylinder, frame_id, stamp.clone()); + bounds.region_id = seq; + let cloud = points_to_cloud(points, frame_id, stamp.clone()); + let region = PointCloudRegion { + header: cloud.header.clone(), + region_id: seq, + cloud, + }; + (bounds, region) } /// The regions due this viz tick with their points: those whose chunks @@ -657,35 +673,29 @@ fn prepare_seed( const LOADED_MAP_TF_WAIT_TIMEOUT: Duration = Duration::from_secs(2); fn time_secs(t: &Time) -> f64 { - t.sec as f64 + t.nsec as f64 * 1e-9 + t.sec as f64 + t.nanosec as f64 * 1e-9 } -/// A cylinder as the bounds message: position is the center, orientation -/// holds radius, z_min, z_max. -fn bounds_to_pose(c: &Cylinder, frame_id: &str, stamp: Time) -> PoseStamped { - PoseStamped { +/// Explicit cylinder bounds, stamped like the paired cloud. +fn bounds_message(c: &Cylinder, frame_id: &str, stamp: Time) -> RegionBounds { + RegionBounds { header: Header { - seq: 0, stamp, frame_id: frame_id.to_string(), }, - pose: PoseMsg { - position: Point { - x: c.cx as f64, - y: c.cy as f64, - z: 0.0, - }, - orientation: Quaternion { - x: c.radius as f64, - y: c.z_min as f64, - z: c.z_max as f64, - w: 0.0, - }, + region_id: 0, + center: Point { + x: c.cx as f64, + y: c.cy as f64, + z: 0.0, }, + radius: c.radius as f64, + z_min: c.z_min as f64, + z_max: c.z_max as f64, } } -async fn publish_bounds(out: &Output, msg: &PoseStamped) { +async fn publish_bounds(out: &Output, msg: &RegionBounds) { if let Err(e) = out.publish(msg).await { error_throttled!( Duration::from_secs(1), @@ -697,21 +707,25 @@ async fn publish_bounds(out: &Output, msg: &PoseStamped) { } /// The f32 pose of a transform, for registering clouds. -fn tf_to_pose(t: &Transform) -> Pose { - let translation = t.translation().cast::(); - let rotation = t.rotation().cast::(); +fn tf_to_pose(t: &TransformStamped) -> Pose { + let translation = &t.transform.translation; + let rotation = &t.transform.rotation; Pose { - position: (translation.x, translation.y, translation.z), + position: ( + translation.x as f32, + translation.y as f32, + translation.z as f32, + ), orientation: ( - rotation.coords.x, - rotation.coords.y, - rotation.coords.z, - rotation.coords.w, + rotation.x as f32, + rotation.y as f32, + rotation.z as f32, + rotation.w as f32, ), } } -fn write_point(data: &mut Vec, n: &mut i32, x: f32, y: f32, z: f32) { +fn write_point(data: &mut Vec, n: &mut u32, x: f32, y: f32, z: f32) { data.extend_from_slice(&x.to_le_bytes()); data.extend_from_slice(&y.to_le_bytes()); data.extend_from_slice(&z.to_le_bytes()); @@ -719,16 +733,15 @@ fn write_point(data: &mut Vec, n: &mut i32, x: f32, y: f32, z: f32) { *n += 1; } -fn make_cloud(data: Vec, n: i32, frame_id: &str, stamp: Time) -> PointCloud2 { - let make_field = |name: &str, off: i32| PointField { +fn make_cloud(data: Vec, n: u32, frame_id: &str, stamp: Time) -> PointCloud2 { + let make_field = |name: &str, off: u32| PointField { name: name.into(), offset: off, - datatype: PointField::FLOAT32 as u8, + datatype: PointField::FLOAT32, count: 1, }; PointCloud2 { header: Header { - seq: 0, stamp, frame_id: frame_id.into(), }, @@ -751,14 +764,14 @@ fn make_cloud(data: Vec, n: i32, frame_id: &str, stamp: Time) -> PointCloud2 /// Pack flat (x, y, z) triples into an LCM cloud message. fn points_to_cloud(points: &[f32], frame_id: &str, stamp: Time) -> PointCloud2 { let mut data = Vec::with_capacity((points.len() / 3) * 16); - let mut n: i32 = 0; + let mut n: u32 = 0; for p in points.as_chunks::<3>().0 { write_point(&mut data, &mut n, p[0], p[1], p[2]); } make_cloud(data, n, frame_id, stamp) } -async fn publish_cloud(out: &Output, cloud: &PointCloud2) { +async fn publish_cloud(out: &Output, cloud: &T) { if let Err(e) = out.publish(cloud).await { error_throttled!( Duration::from_secs(1), @@ -847,8 +860,13 @@ mod tests { Translation3::new(10.0, 0.0, 0.0), UnitQuaternion::from_axis_angle(&Vector3::z_axis(), std::f64::consts::FRAC_PI_2), ); - let t = Transform::new("odom", "map", 0.0, iso); - let cloud = points_to_cloud(&[3.5, 0.0, 0.5], "map", Time::default()); + let t = dimos_module::tf::from_isometry( + "odom", + "map", + Time { sec: 0, nanosec: 0 }, + iso, + ); + let cloud = points_to_cloud(&[3.5, 0.0, 0.5], "map", Time { sec: 0, nanosec: 0 }); let part = prepare_seed(&cloud, tf_to_pose(&t), 1.0, (0.0, 0.0, 0.0), 4.0) .expect("loaded map must decode"); @@ -859,7 +877,7 @@ mod tests { assert_eq!(created, 1); let region = mapper.local_points(&part.regions[0].cylinder.bounds()); - let seeded = points_to_cloud(®ion, "odom", Time::default()); + let seeded = points_to_cloud(®ion, "odom", Time { sec: 0, nanosec: 0 }); assert!(cloud_points(&seeded).contains(&voxel_center(10, 3, 0))); } @@ -895,14 +913,14 @@ mod tests { } fn lidar_job(seq: i32) -> Job { - let mut cloud = PointCloud2::default(); - cloud.header.seq = seq; + let mut cloud = points_to_cloud(&[], "", Time { sec: 0, nanosec: 0 }); + cloud.header.stamp.nanosec = seq as u32; Job::Lidar(cloud, Instant::now()) } fn lidar_seq(turn: Turn) -> Option { match turn { - Turn::Job(Job::Lidar(cloud, _)) => Some(cloud.header.seq), + Turn::Job(Job::Lidar(cloud, _)) => Some(cloud.header.stamp.nanosec as i32), _ => None, } } @@ -973,18 +991,63 @@ mod tests { } #[test] - fn seed_region_messages_pair_on_seq_and_stamp() { + fn cdr_seed_regions_keep_explicit_ids_and_epoch_nanoseconds() { + let stamp = Time { + sec: 1_700_000_000, + nanosec: 123_456_789, + }; + let cylinder = Cylinder { + cx: 1., + cy: -2., + radius: 3., + z_min: -1., + z_max: 4., + }; + for id in [i32::MIN, 0, i32::MAX] { + let (bounds, region) = + seed_region_messages(&cylinder, &[1., 2., 3.], id, "odom", &stamp); + for little in [false, true] { + let bounds_bytes = dimos_module::cdr::encode_endian(&bounds, little).unwrap(); + let cloud_bytes = dimos_module::cdr::encode_endian(®ion, little).unwrap(); + let received_bounds: RegionBounds = + dimos_module::cdr::decode(&bounds_bytes).unwrap(); + let received_cloud: PointCloudRegion = + dimos_module::cdr::decode(&cloud_bytes).unwrap(); + assert_eq!( + (received_bounds.region_id, received_cloud.region_id), + (id, id) + ); + assert_eq!(received_bounds.header, received_cloud.header); + assert_eq!(received_cloud.cloud.header.stamp, stamp); + assert_eq!( + ( + received_bounds.center.y, + received_bounds.radius, + received_bounds.z_min + ), + (-2., 3., -1.) + ); + assert_eq!( + extract_xyz(&received_cloud.cloud).unwrap(), + vec![[1., 2., 3.]] + ); + } + } + } + + #[test] + fn seed_region_messages_pair_on_region_id_and_stamp() { let part = two_region_seed(); - let stamp = Time { sec: 7, nsec: 9 }; + let stamp = Time { sec: 7, nanosec: 9 }; for (index, region) in part.regions.iter().enumerate() { let seq = index as i32 + 1; let (bounds, cloud) = seed_region_messages(®ion.cylinder, &[1.5, 1.5, 0.5], seq, "odom", &stamp); - assert_eq!((bounds.header.seq, cloud.header.seq), (seq, seq)); + assert_eq!((bounds.region_id, cloud.region_id), (seq, seq)); assert_eq!(bounds.header.stamp, stamp); assert_eq!(cloud.header.stamp, stamp); - assert_eq!(bounds.pose.position.x, region.cylinder.cx as f64); - assert_eq!(cloud.width, 1); + assert_eq!(bounds.center.x, region.cylinder.cx as f64); + assert_eq!(cloud.cloud.width, 1); } } @@ -1045,7 +1108,7 @@ mod tests { assert_eq!(occupied, vec![(3, -2, 1)]); // Encode and decode the occupied voxel center exactly as the port would. - let cloud = points_to_cloud(&[3.5, -1.5, 1.5], "world", Time::default()); + let cloud = points_to_cloud(&[3.5, -1.5, 1.5], "world", Time { sec: 0, nanosec: 0 }); let Ok(points) = extract_xyz(&cloud) else { panic!("occupied cloud must decode"); }; @@ -1072,12 +1135,12 @@ mod tests { let global = points_to_cloud( &emit_points(&map, 1.0, None, &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); let local = points_to_cloud( &emit_points(&map, 1.0, Some(&cylinder), &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); assert!(cloud_points(&global).contains(&voxel_center(0, 0, 0))); assert!(cloud_points(&local).contains(&voxel_center(0, 0, 0))); @@ -1097,12 +1160,12 @@ mod tests { let global = points_to_cloud( &emit_points(&map, 1.0, None, &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); let local = points_to_cloud( &emit_points(&map, 1.0, Some(&cylinder), &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); assert!(cloud_points(&global).contains(&voxel_center(5, 0, 0))); assert!(!cloud_points(&local).contains(&voxel_center(5, 0, 0))); @@ -1123,12 +1186,12 @@ mod tests { let global = points_to_cloud( &emit_points(&map, 1.0, None, &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); let local = points_to_cloud( &emit_points(&map, 1.0, Some(&cylinder), &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); assert!(cloud_points(&global).contains(&voxel_center(0, 0, 5))); assert!(!cloud_points(&local).contains(&voxel_center(0, 0, 5))); @@ -1151,12 +1214,12 @@ mod tests { let global = points_to_cloud( &emit_points(&map, 1.0, None, &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); let local = points_to_cloud( &emit_points(&map, 1.0, Some(&cylinder), &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); assert!(cloud_points(&global).contains(&voxel_center(1, 0, 0))); assert!(cloud_points(&global).contains(&voxel_center(10, 10, 10))); @@ -1189,7 +1252,7 @@ mod tests { let local = points_to_cloud( &emit_points(&map, 1.0, Some(&cylinder), &live), "world", - Time::default(), + Time { sec: 0, nanosec: 0 }, ); let pts = cloud_points(&local); assert!(pts.contains(&voxel_center(1, 1, 0)), "dense patch kept"); diff --git a/dimos/mapping/ray_tracing/rust/src/wire.rs b/dimos/mapping/ray_tracing/rust/src/wire.rs new file mode 100644 index 00000000000..b00c503fde7 --- /dev/null +++ b/dimos/mapping/ray_tracing/rust/src/wire.rs @@ -0,0 +1,90 @@ +// 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. + +//! Generated CDR and TF boundary adapters for the numeric voxel mapper. + +use dimos_generated_messages::sensor_msgs::msg::{ + point_cloud2::PointCloud2, point_field::PointField, +}; + +use dimos_module::pointcloud::{extract_xyz_rows, ExtractError}; + +pub fn extract_xyz(cloud: &PointCloud2) -> Result, ExtractError> { + let offset = |name: &str| { + cloud + .fields + .iter() + .find(|f| f.name == name && f.datatype == PointField::FLOAT32 && f.count == 1) + .map(|f| f.offset as usize) + .ok_or(ExtractError("missing scalar float32 XYZ field")) + }; + extract_xyz_rows( + &cloud.data, + [cloud.width as usize, cloud.height as usize], + [cloud.point_step as usize, cloud.row_step as usize], + [offset("x")?, offset("y")?, offset("z")?], + cloud.is_bigendian, + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use dimos_generated_messages::builtin_interfaces::msg::time::Time; + use dimos_generated_messages::std_msgs::msg::header::Header; + + #[test] + fn organized_cloud_reader_keeps_endian_and_row_padding_and_filters_missing_points() { + let mut data = vec![0; 80]; + for (offset, point) in [ + (0, [1_f32, 2., 3.]), + (16, [4., 5., 6.]), + (40, [7., 8., 9.]), + (56, [f32::NAN, 0., 0.]), + ] { + for (i, value) in point.iter().enumerate() { + data[offset + i * 4..offset + i * 4 + 4].copy_from_slice(&value.to_be_bytes()); + } + } + let mut cloud = PointCloud2 { + header: Header { + stamp: Time { sec: 0, nanosec: 0 }, + frame_id: "lidar".into(), + }, + height: 2, + width: 2, + fields: ["x", "y", "z"] + .into_iter() + .enumerate() + .map(|(i, n)| PointField { + name: n.into(), + offset: (i * 4) as u32, + datatype: PointField::FLOAT32, + count: 1, + }) + .collect(), + is_bigendian: true, + point_step: 16, + row_step: 40, + data, + is_dense: false, + }; + assert_eq!( + extract_xyz(&cloud).unwrap(), + vec![[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]] + ); + cloud.row_step = 12; + assert!(extract_xyz(&cloud).is_err()); + } +} diff --git a/dimos/mapping/ray_tracing/test_generated_native.py b/dimos/mapping/ray_tracing/test_generated_native.py new file mode 100644 index 00000000000..19ca2d53073 --- /dev/null +++ b/dimos/mapping/ray_tracing/test_generated_native.py @@ -0,0 +1,226 @@ +# 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 contextlib import ExitStack +from copy import deepcopy +import json +import os +from pathlib import Path +import socket +from threading import Event, Lock +import time + +import numpy as np +import pytest +import zenoh + +from dimos.core.global_config import global_config +from dimos.core.transport import ZenohTransport +from dimos.mapping.ray_tracing.module import RayTracingVoxelMap +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.message_utils.time import time_from_nanoseconds, to_nanoseconds +from dimos.messages.dimos_msgs.msg import PointCloudRegion, RegionBounds +from dimos.messages.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage +from dimos.protocol.pubsub.impl.zenohpubsub import Topic as ZenohTopic +from dimos.protocol.rpc.pubsubrpc import LCMRPC +from dimos.protocol.service.zenohservice import ZenohConfig + + +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +def test_native_mapper_declares_canonical_cloud_and_region_ports(): + module = RayTracingVoxelMap(rpc_transport=MemoryRPC) + try: + assert module.lidar.type is PointCloud2 and module.global_map.type is PointCloud2 + assert ( + module.map_regions.type is PointCloudRegion and module.seed_map.type is PointCloudRegion + ) + assert module.region_bounds.type is RegionBounds and module.seed_bounds.type is RegionBounds + assert module.tf.type is TFMessage + finally: + module.stop() + + +@pytest.mark.native_e2e +def test_actual_native_mapper_cdr_registered_clouds_keyed_regions_and_seed_pair(monkeypatch): + executable = os.environ.get("DIMOS_NATIVE_RAY_TRACING_BINARY") + assert executable, "Build and set the native binary before running native_e2e" + binary = Path(executable).resolve() + assert binary.is_file() + monkeypatch.setattr(global_config, "transport", "zenoh") + with socket.socket() as reservation: + reservation.bind(("127.0.0.1", 0)) + endpoint = f"tcp/127.0.0.1:{reservation.getsockname()[1]}" + config = zenoh.Config() + for key, value in { + "mode": "router", + "listen/endpoints": [endpoint], + "scouting/multicast/enabled": False, + "scouting/gossip/enabled": False, + }.items(): + config.insert_json5(key, json.dumps(value)) + with ExitStack() as cleanup: + cleanup.enter_context(zenoh.open(config)) + session = ZenohConfig( + mode="client", + connect=[endpoint], + listen=[], + scouting=False, + multicast=False, + gossip=False, + connect_timeout=5, + ) + module = RayTracingVoxelMap( + rpc_transport=MemoryRPC, + executable=str(binary), + source_dir=None, + build_command=None, + session=session, + voxel_size=0.1, + support_min=0, + min_health=0, + max_health=1, + emit_every=1, + global_emit_every=1, + viz_emit_every=1, + viz_cell_m=4, + world_frame="odom", + ) + names = [ + "lidar", + "loaded_map", + "tf", + "tf_static", + "global_map", + "local_map", + "local_map_fine", + "region_bounds", + "map_regions", + "seed_map", + "seed_bounds", + ] + wires = {} + for name in names: + port = getattr(module, name) + wire = ZenohTransport( + ZenohTopic("native/ray/" + name, port.type), **session.model_dump() + ) + port.transport = wire + wires[name] = wire + cleanup.callback(wire.stop) + cleanup.callback(module.stop) + live_names = ["global_map", "region_bounds", "map_regions"] + seed_names = ["seed_map", "seed_bounds"] + received = {name: {} for name in live_names + seed_names} + live_ready, seed_ready = Event(), Event() + lock = Lock() + + def receive(name, value): + stamp = to_nanoseconds(value.header.stamp) + group = seed_names if name in seed_names else live_names + key = (stamp, value.region_id) if name in seed_names else stamp + with lock: + received[name].setdefault(key, value) + if set.intersection(*(set(received[n]) for n in group)): + (seed_ready if name in seed_names else live_ready).set() + + for name in received: + cleanup.callback(wires[name].subscribe(lambda value, name=name: receive(name, value))) + module.start() + epoch = 1700000000123456789 + sequence = 0 + source_stamps = set() + + def tf_and_header(): + nonlocal sequence + header = Header(frame_id="odom", stamp=time_from_nanoseconds(epoch + sequence)) + source_stamps.add(epoch + sequence) + sequence += 1 + wires["tf_static"].publish( + TFMessage( + transforms=[ + TransformStamped( + header=Header(frame_id="odom", stamp=time_from_nanoseconds(0)), + child_frame_id="lidar", + transform=Transform( + translation=Vector3(x=0, y=0, z=0), + rotation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) + ] + ) + ) + wires["tf"].publish( + TFMessage( + transforms=[ + TransformStamped( + header=deepcopy(header), + child_frame_id=child, + transform=Transform( + translation=Vector3(x=x, y=0, z=0), + rotation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) + for child, x in [("map", 1)] + ] + ) + ) + return header + + deadline = time.monotonic() + 15 + while not live_ready.is_set() and time.monotonic() < deadline: + header = tf_and_header() + header.frame_id = "lidar" + wires["lidar"].publish( + pointcloud_from_xyz( + np.array([[4.15, -0.85, 0.25], [4.15, -0.85, 0.35]], dtype=np.float32), + header=header, + ) + ) + live_ready.wait(0.15) + assert live_ready.is_set(), "No matching generated cloud/region/bounds batch" + with lock: + stamp = max(set.intersection(*(set(received[n]) for n in live_names))) + cloud, bounds, region = [received[n][stamp] for n in live_names] + assert stamp in source_stamps + assert type(cloud) is PointCloud2 and cloud.header.frame_id == "odom" + assert bounds.center.x == bounds.center.y == 0 and bounds.radius > 4 + assert region.region_id == (1 << 16) | 0xFFFF + assert region.header == region.cloud.header + assert len(pointcloud_xyz(region.cloud)) == 2 + np.testing.assert_allclose(pointcloud_xyz(cloud)[:, 0], 4.15, atol=0.05) + + deadline = time.monotonic() + 15 + while not seed_ready.is_set() and time.monotonic() < deadline: + header = tf_and_header() + header.frame_id = "map" + wires["loaded_map"].publish( + pointcloud_from_xyz(np.array([[7.15, 2.15, 0.25]], dtype=np.float32), header=header) + ) + seed_ready.wait(0.15) + assert seed_ready.is_set(), "No paired canonical seed region and bounds" + with lock: + key = max(set.intersection(*(set(received[n]) for n in seed_names))) + seeded, seed_bounds = [received[n][key] for n in seed_names] + assert key[0] in source_stamps and seeded.region_id == seed_bounds.region_id + assert seeded.header == seed_bounds.header == seeded.cloud.header + assert seeded.cloud.header.frame_id == "odom" + np.testing.assert_allclose(pointcloud_xyz(seeded.cloud)[:, 0], 8.15, atol=0.05) + assert module._process is not None and module._process.poll() is None diff --git a/dimos/mapping/ray_tracing/test_transformer.py b/dimos/mapping/ray_tracing/test_transformer.py index 378c1058fc5..9bf496f8f1b 100644 --- a/dimos/mapping/ray_tracing/test_transformer.py +++ b/dimos/mapping/ray_tracing/test_transformer.py @@ -22,16 +22,27 @@ pytest.importorskip("dimos_voxel_ray_tracing") +from dimos_cdrgen.registry import encode + from dimos.mapping.ray_tracing.module import TF_MATCH_TOLERANCE_S from dimos.mapping.ray_tracing.transformer import RayTraceMap, pose_from_tf from dimos.mapping.ray_tracing.voxel_map import VoxelRayMapper from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tf import StreamTF from dimos.memory.type.observation import Observation -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.message_utils.geometry import quaternion_from_euler +from dimos.message_utils.pointcloud import pointcloud_from_xyz +from dimos.message_utils.time import time_from_nanoseconds, time_from_seconds, to_nanoseconds +from dimos.messages.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage + + +def _cloud(points, frame_id="lidar", timestamp=0.0): + return pointcloud_from_xyz( + points, header=Header(frame_id=frame_id, stamp=time_from_seconds(timestamp)) + ) def _default_margin() -> float: @@ -46,7 +57,7 @@ def _obs( id=0, ts=ts, pose=pose, - _data=PointCloud2.from_numpy(points), + _data=_cloud(points, timestamp=ts), ) @@ -62,11 +73,27 @@ def test_emit_every_n_yields_on_cadence_and_flushes_remainder() -> None: results = list(RayTraceMap(emit_every=3)(iter(obs))) assert [r.tags["frame_count"] for r in results] == [3, 6, 7] + assert [to_nanoseconds(r.data.header.stamp) for r in results] == [ + 2000000000, + 5000000000, + 6000000000, + ] + + +@pytest.mark.parametrize("stamp_ns", [-1, 0, 1700000000123456789]) +def test_accumulation_keeps_exact_source_stamp_and_does_not_mutate_input(stamp_ns) -> None: + source = _obs(_cube(), ts=stamp_ns / 1e9, pose=(0, 0, 0)) + source.data.header.stamp = time_from_nanoseconds(stamp_ns) + before = encode(source.data) + [result] = list(RayTraceMap()(iter([source]))) + assert result.data.header.frame_id == "world" + assert to_nanoseconds(result.data.header.stamp) == stamp_ns + assert encode(source.data) == before def test_poseless_obs_are_skipped() -> None: points = _cube() - poseless = Observation(id=1, ts=0.0, pose=None, _data=PointCloud2.from_numpy(points)) + poseless = Observation(id=1, ts=0.0, pose=None, _data=_cloud(points, timestamp=0.0)) posed = _obs(points, ts=1.0, pose=(0.0, 0.0, 0.0)) results = list(RayTraceMap()(iter([poseless, posed]))) @@ -119,7 +146,7 @@ def test_registers_sensor_frame_cloud_by_pose() -> None: id=0, ts=1.0, pose=(5.0, 0.0, 2.0, 0.0, s, 0.0, s), - _data=PointCloud2.from_numpy(point), + _data=_cloud(point), ) [emitted] = list(rtm(iter([obs]))) @@ -147,15 +174,16 @@ def test_reused_instance_continues_the_map() -> None: def test_pose_from_tf_attaches_the_pose_at_the_stamp(tmp_path: Path) -> None: ts = 10.0 with SqliteStore(path=str(tmp_path / "tf.db")) as store: - edge = Transform( - frame_id="world", + edge = TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(ts)), child_frame_id="lidar", - translation=Vector3(1.0, 2.0, 3.0), - ts=ts, + transform=Transform( + translation=Vector3(x=1, y=2, z=3), rotation=quaternion_from_euler(0, 0, 0) + ), ) - store.stream("tf", TFMessage).append(TFMessage(edge), ts=ts, pose=None) + store.stream("tf", TFMessage).append(TFMessage(transforms=[edge]), ts=ts, pose=None) tf = StreamTF(store.stream("tf", TFMessage)) - cloud = PointCloud2.from_numpy(_cube(), frame_id="lidar") + cloud = _cloud(_cube(), frame_id="lidar") matched = Observation(id=0, ts=ts, _data=cloud) far = Observation(id=1, ts=ts + 10 * TF_MATCH_TOLERANCE_S, _data=cloud) diff --git a/dimos/mapping/ray_tracing/test_viz.py b/dimos/mapping/ray_tracing/test_viz.py index 6c8d0e0eaa3..acc04ae3b57 100644 --- a/dimos/mapping/ray_tracing/test_viz.py +++ b/dimos/mapping/ray_tracing/test_viz.py @@ -15,7 +15,18 @@ import numpy as np from dimos.mapping.ray_tracing.viz import render_map_region, voxel_map_points -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import pointcloud_from_xyz +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.dimos_msgs.msg import PointCloudRegion +from dimos.messages.std_msgs.msg import Header + + +def _region(points): + header = Header(frame_id="world", stamp=time_from_nanoseconds(1700000000123456789)) + return PointCloudRegion( + header=header, region_id=0, cloud=pointcloud_from_xyz(points, header=header) + ) + HEIGHT_RANGE = (-1.0, 3.0) @@ -30,15 +41,15 @@ def test_the_height_ramp_spans_the_range_and_clamps_outside_it() -> None: def test_map_regions_land_static_on_their_own_cell_entities() -> None: - region = PointCloud2.from_numpy(_at_heights([-1.0, 3.0])) - region.seq = (1 << 16) | (-2 & 0xFFFF) + region = _region(_at_heights([-1.0, 3.0])) + region.region_id = (1 << 16) | (-2 & 0xFFFF) (path, arch, static) = render_map_region(region, 0.1, HEIGHT_RANGE)[0] assert path == "world/map_regions/1_-2" and static assert len(arch.positions.as_arrow_array()) == 2 assert arch.class_ids.as_arrow_array().to_pylist() == [0, 255] - emptied = PointCloud2.from_numpy(np.zeros((0, 3), dtype=np.float32)) - emptied.seq = region.seq + emptied = _region(np.zeros((0, 3), dtype=np.float32)) + emptied.region_id = region.region_id (path, arch, static) = render_map_region(emptied, 0.1, HEIGHT_RANGE)[0] assert path == "world/map_regions/1_-2" and static assert len(arch.positions.as_arrow_array()) == 0 diff --git a/dimos/mapping/ray_tracing/transformer.py b/dimos/mapping/ray_tracing/transformer.py index 96bd197fd64..a569fef333e 100644 --- a/dimos/mapping/ray_tracing/transformer.py +++ b/dimos/mapping/ray_tracing/transformer.py @@ -14,12 +14,16 @@ from __future__ import annotations +from copy import deepcopy from typing import TYPE_CHECKING, Any +import numpy as np + from dimos.mapping.ray_tracing.module import TF_MATCH_TOLERANCE_S from dimos.mapping.ray_tracing.voxel_map import VoxelRayMapper from dimos.memory.transform import FnTransformer, Transformer -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.messages.sensor_msgs.msg import PointCloud2 from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -36,9 +40,12 @@ def pose_from_tf(tf: StreamTF, world_frame: str) -> FnTransformer[PointCloud2, P def attach(obs: Observation[PointCloud2]) -> Observation[PointCloud2]: t = tf.get( - world_frame, obs.data.frame_id, time_point=obs.ts, time_tolerance=TF_MATCH_TOLERANCE_S + world_frame, + obs.data.header.frame_id, + time_point=obs.ts, + time_tolerance=TF_MATCH_TOLERANCE_S, ) - return obs.with_pose(t) + return obs.with_pose(None if t is None else t.transform) return FnTransformer(attach) @@ -76,16 +83,13 @@ def _make_obs( last_obs: Observation[PointCloud2], count: int, ) -> Observation[PointCloud2]: - import open3d as o3d # type: ignore[import-untyped] - import open3d.core as o3c # type: ignore[import-untyped] - tags = {**last_obs.tags, "frame_count": count} cx, cy, radius, z_min, z_max = mapper.take_local_bounds() positions = mapper.local_map((cx, cy, 0.0), radius, z_min, z_max) tags["region_bounds"] = (cx, cy, radius, z_min, z_max) - pcd = o3d.t.geometry.PointCloud() - pcd.point["positions"] = o3c.Tensor.from_numpy(positions) - cloud = PointCloud2(pointcloud=pcd, frame_id="world", ts=last_obs.ts) + header = deepcopy(last_obs.data.header) + header.frame_id = "world" + cloud = pointcloud_from_xyz(positions, header=header) return last_obs.derive(data=cloud, tags=tags) def __call__( @@ -101,7 +105,9 @@ def __call__( logger.warning("RayTraceMap: obs %s has no pose. dropped a cloud", obs.id) continue x, y, z, qx, qy, qz, qw = obs.pose_tuple - mapper.add_frame(obs.data.points_f32(), (x, y, z), (qx, qy, qz, qw)) + mapper.add_frame( + np.asarray(pointcloud_xyz(obs.data), dtype=np.float32), (x, y, z), (qx, qy, qz, qw) + ) last_obs = obs count += 1 diff --git a/dimos/mapping/ray_tracing/utils/loaded_map.py b/dimos/mapping/ray_tracing/utils/loaded_map.py index 2ed096246c9..bca670dfb0c 100644 --- a/dimos/mapping/ray_tracing/utils/loaded_map.py +++ b/dimos/mapping/ray_tracing/utils/loaded_map.py @@ -22,7 +22,8 @@ from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tf import StreamTF from dimos.memory.type.observation import Observation -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import pointcloud_xyz, transform_cloud +from dimos.messages.sensor_msgs.msg import PointCloud2 LOADED_MAP_STREAM = "loaded_map" @@ -41,10 +42,10 @@ def place_loaded_map( loaded_map: Observation[PointCloud2], tf: StreamTF, world_frame: str, ts: float ) -> NDArray[np.float32]: """World-frame points of the loaded map, placed by the transform at ts.""" - placement = tf.get(world_frame, loaded_map.data.frame_id, time_point=ts) + placement = tf.get(world_frame, loaded_map.data.header.frame_id, time_point=ts) if placement is None: raise RuntimeError( - f"no {world_frame}->{loaded_map.data.frame_id} transform at ts={ts:.3f} " + f"no {world_frame}->{loaded_map.data.header.frame_id} transform at ts={ts:.3f} " "to place the loaded map" ) - return loaded_map.data.transform(placement).points_f32() + return np.asarray(pointcloud_xyz(transform_cloud(loaded_map.data, placement)), dtype=np.float32) diff --git a/dimos/mapping/ray_tracing/utils/test_loaded_map.py b/dimos/mapping/ray_tracing/utils/test_loaded_map.py index 8e83b46b104..08151d29be4 100644 --- a/dimos/mapping/ray_tracing/utils/test_loaded_map.py +++ b/dimos/mapping/ray_tracing/utils/test_loaded_map.py @@ -24,9 +24,13 @@ from dimos.mapping.ray_tracing.utils.loaded_map import first_loaded_map, place_loaded_map from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tf import StreamTF -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.message_utils.geometry import transform_from_matrix +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.message_utils.time import time_from_seconds +from dimos.messages.geometry_msgs.msg import TransformStamped +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage pytestmark = [pytest.mark.skipif_aarch64, pytest.mark.skipif_macos] @@ -40,18 +44,21 @@ def store(tmp_path: Path) -> Iterator[SqliteStore]: def _cloud(x: float, ts: float) -> PointCloud2: - return PointCloud2.from_numpy( - np.array([[x, 0.0, 0.0]], dtype=np.float32), frame_id="map", timestamp=ts + return pointcloud_from_xyz( + np.array([[x, 0, 0]], dtype=np.float32), + header=Header(frame_id="map", stamp=time_from_seconds(ts)), ) -def _world_to_map(yaw_deg: float, x: float, y: float, ts: float) -> Transform: +def _world_to_map(yaw_deg: float, x: float, y: float, ts: float) -> TransformStamped: m = np.eye(4) m[:3, :3] = Rotation.from_euler("z", yaw_deg, degrees=True).as_matrix() m[0, 3], m[1, 3] = x, y - tf = Transform.from_matrix(m, frame_id="world", child_frame_id="map") - tf.ts = ts - return tf + return TransformStamped( + header=Header(frame_id="world", stamp=time_from_seconds(ts)), + child_frame_id="map", + transform=transform_from_matrix(m), + ) def test_first_loaded_map_is_none_for_a_missing_or_empty_stream(store: SqliteStore) -> None: @@ -69,17 +76,18 @@ def test_first_loaded_map_is_the_earliest_by_ts(store: SqliteStore) -> None: assert first is not None assert first.ts == T0 - np.testing.assert_allclose(first.data.points_f32(), [[1.0, 0.0, 0.0]]) + np.testing.assert_allclose(pointcloud_xyz(first.data), [[1.0, 0.0, 0.0]]) def test_place_loaded_map_moves_points_into_the_world(store: SqliteStore) -> None: store.stream("tf", TFMessage).append( - TFMessage(_world_to_map(90.0, 1.0, 2.0, T0)), ts=T0, pose=None + TFMessage(transforms=[_world_to_map(90.0, 1.0, 2.0, T0)]), ts=T0, pose=None ) tf = StreamTF(store.stream("tf", TFMessage)) points = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32) loaded = store.stream("loaded_map", PointCloud2).append( - PointCloud2.from_numpy(points, frame_id="map", timestamp=T0), ts=T0 + pointcloud_from_xyz(points, header=Header(frame_id="map", stamp=time_from_seconds(T0))), + ts=T0, ) placed = place_loaded_map(loaded, tf, "world", T0) diff --git a/dimos/mapping/ray_tracing/viz.py b/dimos/mapping/ray_tracing/viz.py index d74facb6e40..a13cd0ea351 100644 --- a/dimos/mapping/ray_tracing/viz.py +++ b/dimos/mapping/ray_tracing/viz.py @@ -20,13 +20,15 @@ import numpy as np +from dimos.message_utils.pointcloud import pointcloud_xyz from dimos.visualization.rerun.bridge import RerunEntry, keyed_by_seq, region_entity if TYPE_CHECKING: from numpy.typing import NDArray from rerun._baseclasses import Archetype - from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 + from dimos.messages.dimos_msgs.msg import PointCloudRegion + from dimos.messages.sensor_msgs.msg import PointCloud2 from dimos.visualization.rerun.bridge import RerunMulti # rerun is imported inside the renderers: it is heavy and only viewers load it. @@ -56,12 +58,14 @@ def voxel_map_points( def render_voxel_map( msg: PointCloud2, voxel_size: float, height_range: tuple[float, float] ) -> Archetype: - return voxel_map_points(msg.points_f32(), voxel_size, height_range) + return voxel_map_points( + np.asarray(pointcloud_xyz(msg), dtype=np.float32), voxel_size, height_range + ) @keyed_by_seq def render_map_region( - msg: PointCloud2, + msg: PointCloudRegion, voxel_size: float, height_range: tuple[float, float], z_band: tuple[float, float] | None = None, @@ -70,11 +74,11 @@ def render_map_region( Points outside z_band (world frame) are dropped, to show one storey. """ - pts = msg.points_f32() + pts = np.asarray(pointcloud_xyz(msg.cloud), dtype=np.float32) if z_band is not None: pts = pts[(pts[:, 2] >= z_band[0]) & (pts[:, 2] < z_band[1])] cell = voxel_map_points(pts, voxel_size, height_range) - return [RerunEntry(region_entity(MAP_REGIONS_ENTITY, msg.seq), cell, static=True)] + return [RerunEntry(region_entity(MAP_REGIONS_ENTITY, msg.region_id), cell, static=True)] def log_loaded_map(points: NDArray[np.float32]) -> None: diff --git a/dimos/mapping/relocalization/blueprints.py b/dimos/mapping/relocalization/blueprints.py index e7acc0a469f..6a5fbe5d965 100644 --- a/dimos/mapping/relocalization/blueprints.py +++ b/dimos/mapping/relocalization/blueprints.py @@ -37,10 +37,11 @@ from dimos.mapping.relocalization.lidar.module import LocalMapRelocalization from dimos.memory.replay import ReplayStream from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.tf2_msgs.msg import TFMessage from dimos.utils.data import resolve_named_path from dimos.utils.logging_config import setup_logger +from dimos.visualization.rerun.cdr import cloud_archetype from dimos.visualization.vis_module import vis_module logger = setup_logger() @@ -87,9 +88,9 @@ def start(self) -> None: self.register_disposable(tf.observable().subscribe(self.tf.publish)) -def _fine_points(cloud: Any) -> Any: +def _fine_points(cloud: PointCloud2) -> Any: """The premap is millimetre-scale; draw it at that size, not the 5 cm default.""" - return cloud.to_rerun(voxel_size=0.0015) + return cloud_archetype(cloud, voxel_size=0.0015, mode="spheres", height_colors=True) # Off until asked for. diff --git a/dimos/mapping/relocalization/go2/module.py b/dimos/mapping/relocalization/go2/module.py index 0bec9eae10d..4744178aa03 100644 --- a/dimos/mapping/relocalization/go2/module.py +++ b/dimos/mapping/relocalization/go2/module.py @@ -14,14 +14,24 @@ """Go2 relocalization: lidar relocalization plus the premap merged into the live map.""" +from copy import deepcopy + +import numpy as np from reactivex import combine_latest from dimos.core.core import rpc from dimos.core.stream import In, Out from dimos.mapping.relocalization.lidar.module import LidarConfig, LidarWindowRelocalization from dimos.mapping.voxels.grid import VoxelGrid -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import ( + concatenate_clouds, + pointcloud_from_xyz, + pointcloud_xyz, + transform_cloud, +) +from dimos.message_utils.time import to_seconds +from dimos.messages.geometry_msgs.msg import TransformStamped +from dimos.messages.sensor_msgs.msg import PointCloud2 from dimos.utils.reactive import backpressure @@ -51,17 +61,30 @@ def start(self) -> None: ).subscribe(self._on_merge_input) ) - def _on_merge_input(self, pair: tuple[PointCloud2, Transform]) -> None: + def _on_merge_input(self, pair: tuple[PointCloud2, TransformStamped]) -> None: local, tf = pair assert self.premap is not None - premap_in_world = self.premap.transform(tf) + premap_in_world = transform_cloud(self.premap, tf) if self.config.use_carving: - grid = VoxelGrid(carve_columns=True, frame_id=local.frame_id, show_startup_log=False) + grid = VoxelGrid( + carve_columns=True, frame_id=local.header.frame_id, show_startup_log=False + ) try: - grid.add_frame(premap_in_world) - grid.add_frame(local) - self.merged_map.publish(grid.get_global_pointcloud2()) + grid.add_points( + np.asarray(pointcloud_xyz(premap_in_world), dtype=np.float32), + timestamp=to_seconds(premap_in_world.header.stamp), + ) + grid.add_points( + np.asarray(pointcloud_xyz(local), dtype=np.float32), + timestamp=to_seconds(local.header.stamp), + ) + self.merged_map.publish( + pointcloud_from_xyz( + grid.get_global_pointcloud().point["positions"].numpy(), + header=deepcopy(local.header), + ) + ) finally: grid.dispose() else: - self.merged_map.publish(local + premap_in_world) + self.merged_map.publish(concatenate_clouds(local, premap_in_world)) diff --git a/dimos/mapping/relocalization/lidar/module.py b/dimos/mapping/relocalization/lidar/module.py index 35ffde47898..c87a21a1eed 100644 --- a/dimos/mapping/relocalization/lidar/module.py +++ b/dimos/mapping/relocalization/lidar/module.py @@ -16,7 +16,6 @@ from collections import deque from functools import reduce -import operator import time from typing import Any @@ -30,7 +29,8 @@ RelocalizeConfig, ) from dimos.mapping.relocalization.module import Config, RelocalizationModule -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import concatenate_clouds, pointcloud_to_open3d +from dimos.messages.sensor_msgs.msg import PointCloud2 from dimos.utils.logging_config import setup_logger from dimos.utils.reactive import backpressure @@ -51,7 +51,7 @@ def window( ops.do_action(held.append), ops.throttle_first(interval), ops.filter(lambda _: len(held) >= cfg.min_frames), - ops.map(lambda _: reduce(operator.add, held)), + ops.map(lambda _: reduce(concatenate_clouds, held)), ) @@ -91,7 +91,9 @@ def start(self) -> None: return # Downsampling the premap and computing its normals and FPFH is the # pipeline's dominant cost, so it is a startup cost, not a per-fix one. - self._relocalizer = LidarRelocalizer(self.premap.pointcloud, self.config.relocalize) + self._relocalizer = LidarRelocalizer( + pointcloud_to_open3d(self.premap), self.config.relocalize + ) self.register_disposable( backpressure( self.clouds().pipe( @@ -108,20 +110,20 @@ def _maybe_log_skip(self, msg: PointCloud2) -> None: now = time.monotonic() if now - self._last_skip_log > 5.0: logger.warning( - f"relocalize skipped: n_pts={len(msg)} " + f"relocalize skipped: n_pts={msg.width * msg.height} " f"< min_local_points={self.config.min_local_points}" ) self._last_skip_log = now def _has_enough_points(self, msg: PointCloud2) -> bool: - return len(msg) >= self.config.min_local_points + return msg.width * msg.height >= self.config.min_local_points def _relocalize(self, msg: PointCloud2) -> None: assert self._relocalizer is not None t0 = time.monotonic() try: tf = self._relocalizer.relocalize( - msg.pointcloud, self.config.world_frame, self.config.map_frame + pointcloud_to_open3d(msg), self.config.world_frame, self.config.map_frame ) except Exception: logger.exception("relocalize() failed") @@ -129,11 +131,11 @@ def _relocalize(self, msg: PointCloud2) -> None: dt = time.monotonic() - t0 if tf is None: logger.info( - f"relocalize lidar: refused after {dt:.1f}s n_pts={len(msg)} " + f"relocalize lidar: refused after {dt:.1f}s n_pts={msg.width * msg.height} " f"(below fitness_threshold={self.config.relocalize.fitness_threshold})" ) return - logger.info(f"relocalize lidar: time_cost={dt:.1f}s n_pts={len(msg)}") + logger.info(f"relocalize lidar: time_cost={dt:.1f}s n_pts={msg.width * msg.height}") self.submit(tf, "lidar") diff --git a/dimos/mapping/relocalization/lidar/relocalize.py b/dimos/mapping/relocalization/lidar/relocalize.py index 587346e1240..a6c43f60643 100644 --- a/dimos/mapping/relocalization/lidar/relocalize.py +++ b/dimos/mapping/relocalization/lidar/relocalize.py @@ -35,11 +35,16 @@ from __future__ import annotations +from collections.abc import Callable from typing import TYPE_CHECKING, NamedTuple import numpy as np +from packaging.version import Version -from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.message_utils.geometry import inverse_transform, transform_from_matrix +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import TransformStamped +from dimos.messages.std_msgs.msg import Header from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger @@ -65,7 +70,7 @@ class _Prepared(NamedTuple): class RelocAttempt(NamedTuple): """One decided attempt: the ``world -> map`` fix if accepted, and what it was scored on.""" - fix: Transform | None + fix: TransformStamped | None result: RegistrationResult @@ -187,11 +192,16 @@ def normals(pcd: PointCloud, voxel: float) -> PointCloud: return pcd coarse = normals(cloud.voxel_down_sample(cfg.voxel_coarse), cfg.voxel_coarse) - fpfh = reg.compute_fpfh_feature( - coarse, - o3d.geometry.KDTreeSearchParamHybrid( - radius=cfg.voxel_coarse * cfg.fpfh_radius_factor, max_nn=cfg.fpfh_max_nn - ), + search = o3d.geometry.KDTreeSearchParamHybrid( + radius=cfg.voxel_coarse * cfg.fpfh_radius_factor, max_nn=cfg.fpfh_max_nn + ) + # Open3D 0.20 requires indices; older supported versions compute all + # points with the two-argument API. Both paths describe every point. + features: Callable[..., Feature] = reg.compute_fpfh_feature + fpfh = ( + features(coarse, search, indices=list(range(len(coarse.points)))) + if Version(o3d.__version__) >= Version("0.20") + else features(coarse, search) ) fine = normals(cloud.voxel_down_sample(cfg.voxel_fine), cfg.voxel_fine) return _Prepared(coarse=coarse, fpfh=fpfh, fine=fine) @@ -259,7 +269,7 @@ def refine(coarse_T: np.ndarray) -> RegistrationResult: def relocalize( self, local_map: PointCloud, world_frame: str, map_frame: str - ) -> Transform | None: + ) -> TransformStamped | None: """The ``world_frame -> map_frame`` transform, or ``None`` when nothing was good enough.""" return self.attempt(local_map, world_frame, map_frame).fix @@ -269,7 +279,9 @@ def attempt(self, local_map: PointCloud, world_frame: str, map_frame: str) -> Re logger.info(f"align: fitness={result.fitness:.3f} rmse={result.inlier_rmse:.3f}") if result.fitness < self.config.fitness_threshold: return RelocAttempt(None, result) - placement = Transform.from_matrix( - np.asarray(result.transformation), frame_id=map_frame, child_frame_id=world_frame + placement = TransformStamped( + header=Header(frame_id=map_frame, stamp=time_from_nanoseconds(0)), + child_frame_id=world_frame, + transform=transform_from_matrix(np.asarray(result.transformation)), ) - return RelocAttempt(placement.inverse(), result) + return RelocAttempt(inverse_transform(placement), result) diff --git a/dimos/mapping/relocalization/lidar/replay.py b/dimos/mapping/relocalization/lidar/replay.py index 7845e0a3b23..6367ffa31ca 100644 --- a/dimos/mapping/relocalization/lidar/replay.py +++ b/dimos/mapping/relocalization/lidar/replay.py @@ -20,10 +20,13 @@ from __future__ import annotations +from copy import deepcopy import math from pathlib import Path from typing import TYPE_CHECKING, NamedTuple +from dimos_cdrgen.registry import decode +import numpy as np import typer from dimos.mapping.ray_tracing.module import RayTracingVoxelMapConfig @@ -34,12 +37,16 @@ from dimos.mapping.relocalization.module import yaw_deg from dimos.memory.store.sqlite import SqliteStore from dimos.memory.tf import StreamTF -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.message_utils.geometry import vector_array +from dimos.message_utils.pointcloud import pointcloud_to_open3d, pointcloud_xyz, transform_cloud +from dimos.message_utils.time import time_from_seconds +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.tf2_msgs.msg import TFMessage from dimos.utils.data import resolve_named_path +from dimos.visualization.rerun.cdr import register_colormap_annotation if TYPE_CHECKING: - from dimos.msgs.geometry_msgs.Transform import Transform + from dimos.messages.geometry_msgs.msg import TransformStamped # rerun is imported inside the functions that log: it is heavy and only this command loads it. @@ -55,29 +62,33 @@ class ReplayAttempt(NamedTuple): t_s: float fitness: float - fix: Transform | None + fix: TransformStamped | None class Replay(NamedTuple): attempts: list[ReplayAttempt] - fix: Transform | None + fix: TransformStamped | None fix_ts: float - recorded: Transform | None + recorded: TransformStamped | None -def fix_error(fix: Transform, recorded: Transform) -> tuple[float, float]: +def fix_error(fix: TransformStamped, recorded: TransformStamped) -> tuple[float, float]: """Yaw in degrees and translation in meters between two fixes of the same frame pair.""" dyaw = (yaw_deg(fix) - yaw_deg(recorded) + 180.0) % 360.0 - 180.0 - return dyaw, (fix.translation - recorded.translation).length() + return dyaw, float( + np.linalg.norm( + vector_array(fix.transform.translation) - vector_array(recorded.transform.translation) + ) + ) -def recorded_fix(store: SqliteStore, world_frame: str, map_frame: str) -> Transform | None: +def recorded_fix(store: SqliteStore, world_frame: str, map_frame: str) -> TransformStamped | None: """The first world -> map edge the live run published, if the recording has one.""" if "tf" not in store.list_streams(): return None for obs in store.stream("tf", TFMessage).order_by("ts"): for tf in obs.data.transforms: - if (tf.frame_id, tf.child_frame_id) == (world_frame, map_frame): + if (tf.header.frame_id, tf.child_frame_id) == (world_frame, map_frame): return tf return None @@ -131,11 +142,11 @@ def replay( ray = RayTraceMap(voxel_size=voxel_size) frames = lidar.transform(pose_from_tf(tf, world_frame)).transform(ray) - relocalizer = LidarRelocalizer(premap.pointcloud, PRESETS[preset]) + relocalizer = LidarRelocalizer(pointcloud_to_open3d(premap), PRESETS[preset]) recorded = recorded_fix(store, world_frame, MAP_FRAME) attempts: list[ReplayAttempt] = [] - fix: Transform | None = None + fix: TransformStamped | None = None fix_ts = 0.0 next_attempt = 0.0 stop_at = math.inf @@ -147,11 +158,18 @@ def replay( if obs.ts >= stop_at: break rr.set_time(TIMELINE, timestamp=obs.ts) - rr.log("world/local_map", voxel_map_points(obs.data.points_f32(), voxel_size)) - if fix is not None or obs.ts < next_attempt or len(obs.data) < min_local_points: + rr.log( + "world/local_map", + voxel_map_points(np.asarray(pointcloud_xyz(obs.data), dtype=np.float32), voxel_size), + ) + if ( + fix is not None + or obs.ts < next_attempt + or obs.data.width * obs.data.height < min_local_points + ): continue next_attempt = obs.ts + reloc_interval - fix_attempt = relocalizer.attempt(obs.data.pointcloud, world_frame, MAP_FRAME) + fix_attempt = relocalizer.attempt(pointcloud_to_open3d(obs.data), world_frame, MAP_FRAME) attempts.append(ReplayAttempt(obs.ts - t0, fix_attempt.result.fitness, fix_attempt.fix)) rr.log(FITNESS_SERIES, rr.Scalars(fix_attempt.result.fitness)) _print_attempt(attempts[-1], recorded) @@ -159,12 +177,12 @@ def replay( continue fix, fix_ts = fix_attempt.fix, obs.ts stop_at = obs.ts + after_s - log_loaded_map(premap.transform(fix).points_f32()) + log_loaded_map(np.asarray(pointcloud_xyz(transform_cloud(premap, fix)), dtype=np.float32)) if recorded is not None: rr.log( "world/recorded_map", rr.Points3D( - premap.transform(recorded).points_f32(), + pointcloud_xyz(transform_cloud(premap, recorded)), colors=[RECORDED_MAP_COLOR], radii=PREMAP_POINT_RADIUS, ), @@ -173,24 +191,29 @@ def replay( def write_loaded_map( - store: SqliteStore, premap: PointCloud2, fix: Transform, ts: float, world_frame: str + store: SqliteStore, premap: PointCloud2, fix: TransformStamped, ts: float, world_frame: str ) -> bool: """Append the placed premap to the recording as the stream plan_rrd seeds from.""" stream = store.stream(LOADED_MAP_STREAM, PointCloud2) if stream.count() > 0: print(f"{LOADED_MAP_STREAM} already has {stream.count()} messages, leaving it") return False - placed = PointCloud2(pointcloud=premap.transform(fix).pointcloud, frame_id=world_frame, ts=ts) + placed = transform_cloud(premap, fix) + placed.header = deepcopy(placed.header) + placed.header.frame_id = world_frame + placed.header.stamp = time_from_seconds(ts) stream.append(placed, ts=ts) - print(f"wrote {len(placed)} placed premap points to {LOADED_MAP_STREAM} at {ts:.3f}") + print( + f"wrote {placed.width * placed.height} placed premap points to {LOADED_MAP_STREAM} at {ts:.3f}" + ) return True -def _print_attempt(attempt: ReplayAttempt, recorded: Transform | None) -> None: +def _print_attempt(attempt: ReplayAttempt, recorded: TransformStamped | None) -> None: if attempt.fix is None: print(f"{attempt.t_s:.1f}s refused fitness={attempt.fitness:.3f}") return - t = attempt.fix.translation + t = attempt.fix.transform.translation line = ( f"{attempt.t_s:.1f}s fitness={attempt.fitness:.3f} " f"t=({t.x:.2f}, {t.y:.2f}, {t.z:.2f}) yaw={yaw_deg(attempt.fix):.1f}deg" @@ -203,7 +226,7 @@ def _print_attempt(attempt: ReplayAttempt, recorded: Transform | None) -> None: def main( recording: str = typer.Argument(..., help="Recording .db: bare name (cwd or data/) or path"), - premap: str = typer.Option(..., "--premap", help="Premap .pc2.lcm: bare name or path"), + premap: str = typer.Option(..., "--premap", help="Premap .pc2.cdr: bare name or path"), lidar: str = typer.Option("lidar", "--lidar", help="Lidar stream in the recording"), world_frame: str = typer.Option("odom", "--world-frame", help="Frame the live map is built in"), preset: str = typer.Option(DEFAULT_PRESET, "--preset", help=f"One of {sorted(PRESETS)}"), @@ -232,9 +255,12 @@ def main( ) -> None: """Replay a recording through the relocalizer and compare its fix with the robot's.""" db_path = resolve_named_path(recording, ".db") - premap_cloud = PointCloud2.lcm_decode(resolve_named_path(premap, ".pc2.lcm").read_bytes()) + premap_cloud = decode(resolve_named_path(premap, ".pc2.cdr").read_bytes(), PointCloud2) + premap_cloud.header.frame_id = MAP_FRAME _init_recording(db_path.stem, out) - print(f"relocalizing {db_path.name} against {premap} ({len(premap_cloud)} points, {preset})") + print( + f"relocalizing {db_path.name} against {premap} ({premap_cloud.width * premap_cloud.height} points, {preset})" + ) store = SqliteStore(path=str(db_path)) with store: result = replay( diff --git a/dimos/mapping/relocalization/lidar/test_replay.py b/dimos/mapping/relocalization/lidar/test_replay.py index c3b7b9c8863..3ececf05e5b 100644 --- a/dimos/mapping/relocalization/lidar/test_replay.py +++ b/dimos/mapping/relocalization/lidar/test_replay.py @@ -16,6 +16,7 @@ from pathlib import Path from types import SimpleNamespace +from dimos_cdrgen.registry import encode import numpy as np import pytest from pytest_mock import MockerFixture @@ -25,6 +26,7 @@ pytest.importorskip("dimos_voxel_ray_tracing") from dimos.mapping.ray_tracing.utils.loaded_map import LOADED_MAP_STREAM +from dimos.mapping.relocalization.blueprints import _fine_points from dimos.mapping.relocalization.lidar import replay as replay_module from dimos.mapping.relocalization.lidar.relocalize import PRESETS, RelocAttempt from dimos.mapping.relocalization.lidar.replay import ( @@ -34,21 +36,47 @@ write_loaded_map, ) from dimos.memory.store.sqlite import SqliteStore -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.message_utils.geometry import transform_from_matrix +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.message_utils.time import time_from_seconds, to_seconds +from dimos.messages.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage -def _fix(yaw_deg: float, x: float = 0.0, y: float = 0.0) -> Transform: +@pytest.fixture +def recording(): + with rr.RecordingStream("relocalization-test") as value: + try: + yield value + finally: + rr.disconnect(recording=value) + + +def _cloud(points, *, frame_id, timestamp=0.0): + return pointcloud_from_xyz( + points, header=Header(frame_id=frame_id, stamp=time_from_seconds(timestamp)) + ) + + +def _fix(yaw_deg: float, x: float = 0.0, y: float = 0.0) -> TransformStamped: m = np.eye(4) m[:3, :3] = Rotation.from_euler("z", yaw_deg, degrees=True).as_matrix() m[0, 3], m[1, 3] = x, y - return Transform.from_matrix(m, frame_id="odom", child_frame_id="map") + return TransformStamped( + header=Header(frame_id="odom", stamp=time_from_seconds(0)), + child_frame_id="map", + transform=transform_from_matrix(m), + ) -def _edge(parent: str, child: str, ts: float, x: float = 0.0) -> Transform: - return Transform(frame_id=parent, child_frame_id=child, translation=Vector3(x, 0.0, 0.0), ts=ts) +def _edge(parent: str, child: str, ts: float, x: float = 0.0) -> TransformStamped: + return TransformStamped( + header=Header(frame_id=parent, stamp=time_from_seconds(ts)), + child_frame_id=child, + transform=Transform(translation=Vector3(x, 0, 0), rotation=Quaternion(0, 0, 0, 1)), + ) def test_fix_error_wraps_yaw_and_measures_translation() -> None: @@ -61,16 +89,19 @@ def test_fix_error_wraps_yaw_and_measures_translation() -> None: @pytest.mark.skipif_macos def test_write_loaded_map_writes_once(tmp_path: Path) -> None: pts = np.array([[1, 0, 0], [0, 1, 0]], dtype=np.float32) - premap = PointCloud2.from_numpy(pts, frame_id="map") + premap = _cloud(pts, frame_id="map") fix = _fix(90.0, x=1.0) + originals = encode(premap), encode(fix) with SqliteStore(path=str(tmp_path / "r.db")) as store: assert write_loaded_map(store, premap, fix, 5.0, "odom") is True stream = store.stream(LOADED_MAP_STREAM, PointCloud2) assert stream.count() == 1 written = stream.first() - assert written.data.frame_id == "odom" + assert written.data.header.frame_id == "odom" assert written.ts == 5.0 - np.testing.assert_allclose(written.data.points_f32(), [[1, 1, 0], [0, 0, 0]], atol=1e-6) + assert written.data.header.stamp == time_from_seconds(5.0) + assert (encode(premap), encode(fix)) == originals + np.testing.assert_allclose(pointcloud_xyz(written.data), [[1, 1, 0], [0, 0, 0]], atol=1e-6) assert write_loaded_map(store, premap, fix, 6.0, "odom") is False assert stream.count() == 1 @@ -87,18 +118,18 @@ def test_recorded_fix_is_the_first_world_to_map_edge(tmp_path: Path) -> None: _edge("odom", "map", 2.0, x=2.0), _edge("odom", "map", 3.0, x=3.0), ): - tf.append(TFMessage(edge), ts=edge.ts, pose=None) + tf.append(TFMessage(transforms=[edge]), ts=to_seconds(edge.header.stamp), pose=None) recorded = recorded_fix(store, "odom", "map") assert recorded is not None - assert recorded.translation.x == 2.0 + assert recorded.transform.translation.x == 2.0 assert recorded_fix(store, "odom", "nowhere") is None @pytest.mark.skipif_aarch64 @pytest.mark.skipif_macos def test_replay_spaces_attempts_and_stops_after_the_fix( - tmp_path: Path, mocker: MockerFixture + tmp_path: Path, mocker: MockerFixture, recording ) -> None: fix = _fix(90.0, x=1.0) attempts = [ @@ -111,7 +142,7 @@ def test_replay_spaces_attempts_and_stops_after_the_fix( set_time = mocker.spy(rr, "set_time") span = np.arange(0.0, 1.0, 0.05) dense = np.array([(x, y, 0.5) for x in span for y in span], dtype=np.float32) - premap = PointCloud2.from_numpy(np.eye(3, dtype=np.float32), frame_id="map") + premap = _cloud(np.eye(3, dtype=np.float32), frame_id="map") with SqliteStore(path=str(tmp_path / "r.db")) as store: lidar = store.stream("lidar", PointCloud2) @@ -120,8 +151,8 @@ def test_replay_spaces_attempts_and_stops_after_the_fix( ts = 10.0 + i * 0.5 # The first frame's local map is under min_local_points. points = dense[:3] if i == 0 else dense - tf.append(TFMessage(_edge("odom", "lidar", ts)), ts=ts, pose=None) - lidar.append(PointCloud2.from_numpy(points, frame_id="lidar", timestamp=ts), ts=ts) + tf.append(TFMessage(transforms=[_edge("odom", "lidar", ts)]), ts=ts, pose=None) + lidar.append(_cloud(points, frame_id="lidar", timestamp=ts), ts=ts) result = replay( store, lidar_stream="lidar", @@ -136,7 +167,10 @@ def test_replay_spaces_attempts_and_stops_after_the_fix( to_time=None, ) - relocalizer_class.assert_called_once_with(premap.pointcloud, PRESETS["mid360"]) + relocalizer_class.assert_called_once() + geometry, config = relocalizer_class.call_args.args + np.testing.assert_allclose(np.asarray(geometry.points), pointcloud_xyz(premap)) + assert config is PRESETS["mid360"] assert attempt.call_count == 2 # The sparse frame at 0.0 s is skipped, so the first attempt waits for 0.5 s. assert [a.t_s for a in result.attempts] == [0.5, 1.5] @@ -147,3 +181,12 @@ def test_replay_spaces_attempts_and_stops_after_the_fix( # The frame at 12.0 is inside after_s of the fix, the ones at 12.5 and 13.0 are past it. processed = [call.kwargs["timestamp"] for call in set_time.call_args_list] assert processed == [10.0, 10.5, 11.0, 11.5, 12.0] + + +def test_premap_renderer_preserves_fine_point_size_with_generated_values(): + cloud = _cloud(np.array([[1, 2, 3]], dtype=np.float32), frame_id="map") + before = encode(cloud) + value = _fine_points(cloud) + assert value.positions.as_arrow_array().to_pylist() == [[1, 2, 3]] + np.testing.assert_allclose(value.radii.as_arrow_array().to_pylist(), [0.00075], atol=1e-9) + assert encode(cloud) == before diff --git a/dimos/mapping/relocalization/lidar/tune.md b/dimos/mapping/relocalization/lidar/tune.md index 581de8888af..c6cf07908f8 100644 --- a/dimos/mapping/relocalization/lidar/tune.md +++ b/dimos/mapping/relocalization/lidar/tune.md @@ -30,14 +30,14 @@ Given that we know coordinates from which lidar frame was pulled, when relocaliz The shipped pair was built like this. -Both files are available in LFS so this is only for reproducing them: +The historical pair is available in LFS. This proposal requires a canonical CDR recording and a regenerated `.pc2.cdr` premap; legacy typed recordings and `.pc2.lcm` maps are not decoded automatically. Use local converted or regenerated inputs for this study. Rebuild the map with: dimos map global recording_go2_mid360_2026-05-29_4-45pm-PST_corrected \ --lidar fastlio_lidar --voxel 0.005 --seek 400 --export --no-gui `--seek 400` does the split: the premap is the walk after 400 s, and probes come from before it. `--export` implies `--pgo` and writes -`./.pc2.lcm`. The walk is a loop, so the two halves still cover some +`./.pc2.cdr`. The walk is a loop, so the two halves still cover some of the same ground. Without that they would never match. Register the pair in `DATASETS`: @@ -45,7 +45,7 @@ Register the pair in `DATASETS`: ```python skip "go2-sf-area1": Dataset( recording="recording_go2_mid360_2026-05-29_4-45pm-PST_corrected.db", - premap="recording_go2_mid360_2026-05-29_4-45pm-PST_corrected.pc2.lcm", + premap="recording_go2_mid360_2026-05-29_4-45pm-PST_corrected.pc2.cdr", window=(0.0, 400.0), # premap was built from the scans after 400 s so we use lidar frames before then ), ``` diff --git a/dimos/mapping/relocalization/lidar/tune.py b/dimos/mapping/relocalization/lidar/tune.py index 6237f28b723..e21e41d3a69 100644 --- a/dimos/mapping/relocalization/lidar/tune.py +++ b/dimos/mapping/relocalization/lidar/tune.py @@ -34,11 +34,13 @@ from __future__ import annotations +from copy import deepcopy from functools import lru_cache from itertools import islice import time from typing import TYPE_CHECKING, Any, NamedTuple +from dimos_cdrgen.registry import decode import numpy as np import typer @@ -49,8 +51,17 @@ RelocalizeConfig, ) from dimos.mapping.relocalization.lidar.replay import main as replay_main -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import ( + pointcloud_from_xyz, + pointcloud_to_open3d, + pointcloud_with_xyz, + pointcloud_xyz, +) +from dimos.message_utils.time import time_from_seconds +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header from dimos.utils.data import get_data +from dimos.visualization.rerun.cdr import cloud_archetype if TYPE_CHECKING: from collections.abc import Callable @@ -190,7 +201,7 @@ class Dataset(NamedTuple): # so probes come from before that. "go2-sf-area1": Dataset( recording="recording_go2_mid360_2026-05-29_4-45pm-PST_corrected.db", - premap="recording_go2_mid360_2026-05-29_4-45pm-PST_corrected.pc2.lcm", + premap="recording_go2_mid360_2026-05-29_4-45pm-PST_corrected.pc2.cdr", # Runs the full stretch before the premap's own scans: 0-250 s is # ground the loop revisits after 400 s and the premap therefore # holds, while 250-400 s is a one-time excursion it never saw. @@ -283,7 +294,7 @@ def fixtures(name: str) -> tuple[PointCloud2, Any]: from dimos.memory.store.sqlite import SqliteStore ds = DATASETS[name] - premap = PointCloud2.lcm_decode(get_data(ds.premap).read_bytes()) + premap = decode(get_data(ds.premap).read_bytes(), PointCloud2) store = SqliteStore(path=str(get_data(ds.recording)), must_exist=True) store.start() return premap, store @@ -293,13 +304,13 @@ def fixtures(name: str) -> tuple[PointCloud2, Any]: def premap_index(name: str, voxel: float = 0.2) -> Any: """The premap thinned for nearest-neighbour queries, built once.""" premap, _ = fixtures(name) - return premap.pointcloud.voxel_down_sample(voxel) + return pointcloud_to_open3d(premap).voxel_down_sample(voxel) def premap_distance(name: str, points: np.ndarray) -> np.ndarray: """Per-point distance from ``points`` to the nearest premap point.""" - cloud = PointCloud2.from_numpy(points, frame_id="world", timestamp=0.0) - return np.asarray(cloud.pointcloud.compute_point_cloud_distance(premap_index(name))) + cloud = pointcloud_from_xyz(points, header=Header(frame_id="world", stamp=time_from_seconds(0))) + return np.asarray(pointcloud_to_open3d(cloud).compute_point_cloud_distance(premap_index(name))) def accumulate( @@ -328,7 +339,7 @@ def accumulate( mapper = VoxelRayMapper(voxel_size=voxel, max_range=max_range, emit_every=1) used = 0 - last_ts = t0 + start + last_header = Header(frame_id="world", stamp=time_from_seconds(t0 + start)) # A cloud with no pose is skipped, not counted: n_frames is a number of # registered scans, which is what the window question is about. for obs in islice(frames, n_frames * 4): @@ -337,15 +348,16 @@ def accumulate( if obs.pose_tuple is None: continue x, y, z = obs.pose_tuple[:3] - mapper.add_frame_world(obs.data.points_f32(), (x, y, z)) + mapper.add_frame_world(np.asarray(pointcloud_xyz(obs.data), dtype=np.float32), (x, y, z)) used += 1 - last_ts = obs.ts + last_header = deepcopy(obs.data.header) + last_header.frame_id = "world" if used == n_frames: break if used < n_frames: raise ValueError(f"only {used} of {n_frames} scans had a pose from {start:.0f}s") - return PointCloud2.from_numpy(mapper.global_map(), frame_id="world", timestamp=last_ts) + return pointcloud_from_xyz(mapper.global_map(), header=last_header) @lru_cache(maxsize=512) @@ -386,7 +398,7 @@ def coverage(name: str, cloud: PointCloud2) -> float: whether the aligner answered it: a probe standing where the premap has no points cannot be placed by any parameter setting. """ - distances = premap_distance(name, np.asarray(cloud.pointcloud.points)) + distances = premap_distance(name, pointcloud_xyz(cloud)) return float((distances < COVERAGE_TOL_M).mean()) @@ -467,7 +479,7 @@ def identity_error(T: np.ndarray) -> tuple[float, float, float]: def applied(T: np.ndarray, cloud: PointCloud2) -> np.ndarray: """``cloud``'s points where ``T`` puts them.""" - points: np.ndarray = np.asarray(cloud.pointcloud.points) @ T[:3, :3].T + T[:3, 3] + points: np.ndarray = pointcloud_xyz(cloud) @ T[:3, :3].T + T[:3, 3] return points @@ -495,7 +507,7 @@ def run_probes( premap, _ = fixtures(name) # The premap's preprocessing depends on the config but not the probe, so # it is paid once per trial rather than once per attempt. - relocalizer = LidarRelocalizer(premap.pointcloud, config) + relocalizer = LidarRelocalizer(pointcloud_to_open3d(premap), config) rate = scan_rate(name) probes: list[Probe] = [] @@ -512,7 +524,7 @@ def run_probes( frames = min(max(int(elapsed * rate), min_frames), max_frames) cloud = local_map(name, frames, start, voxel) t0, c0 = time.monotonic(), time.process_time() - result = relocalizer.align(cloud.pointcloud) + result = relocalizer.align(pointcloud_to_open3d(cloud)) dt = time.monotonic() - t0 elapsed += dt wall += dt @@ -525,13 +537,13 @@ def run_probes( break T = np.asarray(result.transformation) - truth_points = np.asarray(cloud.pointcloud.points) + truth_points = pointcloud_xyz(cloud) moved = applied(T, cloud) translation, rotation, tilt = identity_error(T) probes.append( Probe( start=start, - n_points=len(cloud), + n_points=cloud.width * cloud.height, coverage=place_coverage(name, start, voxel), translation_m=translation, rotation_deg=rotation, @@ -621,22 +633,22 @@ def view(name: str, probes: list[Probe], out: str | None) -> None: rr.save(out) else: rr.spawn() - rr.log("world/premap/pointcloud", premap.to_rerun(voxel_size=0.05), static=True) + rr.log( + "world/premap/pointcloud", + cloud_archetype(premap, voxel_size=0.05, mode="spheres", height_colors=True), + static=True, + ) for probe in probes: - aligned = PointCloud2.from_numpy( - applied(probe.transform, probe.cloud), - frame_id=probe.cloud.frame_id, - timestamp=probe.cloud.ts, - ) + aligned = pointcloud_with_xyz(probe.cloud, applied(probe.transform, probe.cloud)) entity = f"world/probe_{probe.start:04.0f}s" rr.log( f"{entity}/truth", - probe.cloud.to_rerun(voxel_size=0.05, colors=[76, 220, 41]), + cloud_archetype(probe.cloud, voxel_size=0.05, mode="spheres", color=(76, 220, 41)), static=True, ) rr.log( f"{entity}/aligned", - aligned.to_rerun(voxel_size=0.05, colors=[231, 76, 60]), + cloud_archetype(aligned, voxel_size=0.05, mode="spheres", color=(231, 76, 60)), static=True, ) if out is not None: @@ -681,7 +693,7 @@ def _register( DatasetOpt = typer.Option(DEFAULT_DATASET, "--dataset", "-d", help="Registered dataset name") RecordingOpt = typer.Option(None, "--recording", help="Recording, resolved through get_data") -PremapOpt = typer.Option(None, "--premap", help="Premap .pc2.lcm, resolved through get_data") +PremapOpt = typer.Option(None, "--premap", help="Premap .pc2.cdr, resolved through get_data") LidarOpt = typer.Option(None, "--lidar", help="Lidar stream in the recording") FromOpt = typer.Option(None, "--from", help="Earliest second probes may come from") ToOpt = typer.Option(None, "--to", help="Latest second probes may come from") @@ -727,7 +739,7 @@ def run( console = Console() table = _table( - f"{name} premap {len(pre):,} pts preset={preset} " + f"{name} premap {pre.width * pre.height:,} pts preset={preset} " f"cutoff={config.fitness_threshold} " f"frames {config.min_frames}-{config.max_frames}", *PROBE_COLUMNS, diff --git a/dimos/mapping/relocalization/module.py b/dimos/mapping/relocalization/module.py index 81b7b81fe18..839092dd458 100644 --- a/dimos/mapping/relocalization/module.py +++ b/dimos/mapping/relocalization/module.py @@ -39,34 +39,40 @@ from __future__ import annotations +from copy import deepcopy import math from typing import Any +from dimos_cdrgen.registry import decode import reactivex as rx from reactivex import Observable, Subject, operators as ops from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import Out -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.message_utils.geometry import quaternion_euler +from dimos.message_utils.time import header_now +from dimos.messages.geometry_msgs.msg import TransformStamped +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.tf2_msgs.msg import TFMessage from dimos.utils.data import resolve_named_path from dimos.utils.logging_config import setup_logger logger = setup_logger() -MAP_SUFFIX = ".pc2.lcm" +MAP_SUFFIX = ".pc2.cdr" -def yaw_deg(tf: Transform) -> float: - return math.degrees(tf.rotation.euler.z) +def yaw_deg(tf: TransformStamped) -> float: + return math.degrees(quaternion_euler(tf.transform.rotation)[2]) -def fix_stream(fixes: Observable[Transform], interval: float) -> Observable[Transform]: +def fix_stream( + fixes: Observable[TransformStamped], interval: float +) -> Observable[TransformStamped]: """Every accepted fix as it lands, then again every ``interval`` s (once only if <= 0).""" - def now_and_again(fix: Transform) -> Observable[Transform]: + def now_and_again(fix: TransformStamped) -> Observable[TransformStamped]: again = rx.interval(interval).pipe(ops.map(lambda _: fix)) if interval > 0 else rx.empty() return rx.concat(rx.of(fix), again) @@ -75,7 +81,7 @@ def now_and_again(fix: Transform) -> Observable[Transform]: class Config(ModuleConfig): # Premap stem or path, e.g. `--map-file=go2_hongkong_office_twopass_map`; - # `.pc2.lcm` is appended if absent. Without one the module runs but never + # `.pc2.cdr` is appended if absent. Without one the module runs but never # attempts a fix. map_file: str | None = None # What the live fixed frame is called, whatever the odometry is in @@ -102,16 +108,14 @@ class RelocalizationModule(Module): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) - self.fixes: Subject[Transform] = Subject() + self.fixes: Subject[TransformStamped] = Subject() self.premap: PointCloud2 | None = None @rpc def start(self) -> None: super().start() self.register_disposable( - fix_stream(self.fixes, self.config.tf_interval).subscribe( - lambda tf: self.tf.publish(TFMessage(tf.now())) - ) + fix_stream(self.fixes, self.config.tf_interval).subscribe(self._publish_fix) ) if not self.config.map_file: logger.info("Relocalization module disabled (no map_file configured)") @@ -119,11 +123,16 @@ def start(self) -> None: self._load_premap(self.config.map_file) logger.info(f"Relocalization module started: map_file={self.config.map_file!r}") + def _publish_fix(self, fix: TransformStamped) -> None: + published = deepcopy(fix) + published.header = header_now(fix.header.frame_id) + self.tf.publish(TFMessage(transforms=[published])) + def _load_premap(self, map_file: str) -> None: path = resolve_named_path(map_file, MAP_SUFFIX) - premap = PointCloud2.lcm_decode(path.read_bytes()) + premap = decode(path.read_bytes(), PointCloud2) - premap.frame_id = self.config.map_frame + premap.header.frame_id = self.config.map_frame self.premap = premap self.register_disposable( @@ -141,19 +150,19 @@ def keep_relocalizing(self) -> bool: """Whether to keep attempting. Implementations gate their input on this.""" return not (self._placed and self.config.relocalize_once) - def submit(self, tf: Transform, source: str = "") -> None: + def submit(self, tf: TransformStamped, source: str = "") -> None: """Publish a ``world_frame -> map`` fix the implementation already decided to believe.""" world, map_frame = self.config.world_frame, self.config.map_frame - assert (tf.frame_id, tf.child_frame_id) == (world, map_frame), ( + assert (tf.header.frame_id, tf.child_frame_id) == (world, map_frame), ( f"relocalize {source}: expected {world!r} -> {map_frame!r}, " - f"got {tf.frame_id!r} -> {tf.child_frame_id!r}" + f"got {tf.header.frame_id!r} -> {tf.child_frame_id!r}" ) logger.info( "relocalize fix", source=source, world=world, map_frame=map_frame, - translation=str(tf.translation), + translation=str(tf.transform.translation), yaw_deg=round(yaw_deg(tf), 1), ) self.fixes.on_next(tf) diff --git a/dimos/mapping/relocalization/test_module.py b/dimos/mapping/relocalization/test_module.py index 34f44ee6d0d..a9bfe9718e4 100644 --- a/dimos/mapping/relocalization/test_module.py +++ b/dimos/mapping/relocalization/test_module.py @@ -12,19 +12,49 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import ExitStack +from threading import Event from types import SimpleNamespace from typing import get_type_hints +from dimos_cdrgen.registry import encode import numpy as np +import open3d as o3d +import pydantic import pytest +from reactivex import Subject +from reactivex.testing import TestScheduler from scipy.spatial.transform import Rotation from dimos.core.stream import In -from dimos.mapping.relocalization.lidar.module import LidarWindowRelocalization -from dimos.mapping.relocalization.module import RelocalizationModule -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.vision_msgs.Detection3DArray import Detection3DArray +from dimos.core.transport import LCMTransport +from dimos.mapping.relocalization.go2.module import Go2Relocalization +from dimos.mapping.relocalization.lidar import relocalize as lidar +from dimos.mapping.relocalization.lidar.module import LidarWindowRelocalization, window +from dimos.mapping.relocalization.lidar.relocalize import MID360 +from dimos.mapping.relocalization.module import RelocalizationModule, fix_stream +from dimos.message_utils.geometry import inverse_transform, transform_from_matrix, transform_matrix +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import TransformStamped +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage +from dimos.messages.vision_msgs.msg import Detection3DArray +from dimos.protocol.rpc.pubsubrpc import LCMRPC + + +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +def _transform(matrix, *, frame_id, child_frame_id): + return TransformStamped( + header=Header(stamp=time_from_nanoseconds(0), frame_id=frame_id), + child_frame_id=child_frame_id, + transform=transform_from_matrix(matrix), + ) @pytest.fixture @@ -33,19 +63,19 @@ def module(): built = [] def build(cls=RelocalizationModule, **config): - m = cls(**config) + m = cls(rpc_transport=MemoryRPC, **config) built.append(m) return m yield build for m in built: - m.dispose() + m.stop() def fixes(m): """Collect the transforms a module accepts.""" got = [] - m.fixes.subscribe(got.append) + m.register_disposable(m.fixes.subscribe(got.append)) return got @@ -53,19 +83,19 @@ def test_submit_publishes_and_checks_frames(module): """submit does not second-guess the fix; the implementation already decided.""" m = module() got = fixes(m) - tf = Transform.from_matrix(np.eye(4), frame_id="world", child_frame_id="map") + tf = _transform(np.eye(4), frame_id="world", child_frame_id="map") m.submit(tf, "x") m.submit(tf, "x") assert got == [tf, tf] # A strategy handing over the placement instead of the frame transform is # caught here rather than publishing a backwards TF. with pytest.raises(AssertionError): - m.submit(Transform.from_matrix(np.eye(4), frame_id="map", child_frame_id="world")) + m.submit(_transform(np.eye(4), frame_id="map", child_frame_id="world")) def test_relocalize_once_stops_after_the_first_fix(module): """What the flag does, either way. Which one is the default is a policy call.""" - tf = Transform.from_matrix(np.eye(4), frame_id="world", child_frame_id="map") + tf = _transform(np.eye(4), frame_id="world", child_frame_id="map") once = module(relocalize_once=True) assert once.keep_relocalizing() and not once.placed @@ -80,44 +110,39 @@ def test_relocalize_once_stops_after_the_first_fix(module): @pytest.mark.parametrize("interval", [3600.0, 0.0]) def test_the_fix_goes_out_the_moment_it_is_accepted(interval): """Not on the next interval tick, and with or without republishing (<= 0 = once per fix).""" - from reactivex import Subject - - from dimos.mapping.relocalization.module import fix_stream fixes, sent = Subject(), [] disposable = fix_stream(fixes, interval=interval).subscribe(sent.append) - tf = Transform.from_matrix(np.eye(4), frame_id="world", child_frame_id="map") + tf = _transform(np.eye(4), frame_id="world", child_frame_id="map") fixes.on_next(tf) assert sent == [tf] disposable.dispose() -def test_premap_defines_the_map_frame_and_waits_for_a_fix(module, tmp_path): +def test_premap_defines_the_map_frame_and_waits_for_a_fix(module, tmp_path, mocker): """Loading is the base's: every strategy reads a premap and publishes it, once placed.""" - path = tmp_path / "somewhere.pc2.lcm" + path = tmp_path / "somewhere.pc2.cdr" path.write_bytes( - PointCloud2.from_numpy(np.zeros((5, 3), dtype=np.float32), timestamp=0.0).lcm_encode() + encode( + pointcloud_from_xyz( + np.zeros((5, 3), dtype=np.float32), + header=Header(stamp=time_from_nanoseconds(123456789), frame_id="original"), + ) + ) ) m = module() - published, disposables = [], [] - # The one collaborator worth faking: a real Out port would publish onto a - # bus nothing in this test is listening to. - m.loaded_map = SimpleNamespace(publish=published.append) - m.register_disposable = disposables.append - + published = mocker.patch.object(m.loaded_map, "publish") m._load_premap(str(path)) - assert m.premap is not None and len(m.premap) == 5 - assert m.premap.frame_id == "map" - assert len(disposables) == 1 # the gated publish - assert published == [] # ... which stays silent until a fix lands - m.submit(Transform.from_matrix(np.eye(4), frame_id="world", child_frame_id="map")) - assert published == [m.premap] # republish_loaded_map=0: once, on that fix - disposables[0].dispose() + assert m.premap is not None and m.premap.width * m.premap.height == 5 + assert m.premap.header.frame_id == "map" + assert m.premap.header.stamp == time_from_nanoseconds(123456789) + published.assert_not_called() + m.submit(_transform(np.eye(4), frame_id="world", child_frame_id="map")) + published.assert_called_once_with(m.premap) def test_relocalizer_refuses_below_its_own_threshold(monkeypatch): """One config surface: the relocalizer holds the knobs and the accept decision.""" - from dimos.mapping.relocalization.lidar import relocalize as lidar placement = np.eye(4) placement[:3, 3] = [3.0, -1.0, 0.0] # the map is 3 m +x of where the robot thought @@ -137,15 +162,12 @@ def relocalizer(threshold): # Accepted: open3d places the live cloud in the map, the TF tree wants the # other direction, and relocalize() is what turns one into the other. tf = relocalizer(0.3).relocalize(None, "world", "map") - assert (tf.frame_id, tf.child_frame_id) == ("world", "map") - np.testing.assert_allclose(tf.to_matrix(), np.linalg.inv(placement), atol=1e-9) + assert (tf.header.frame_id, tf.child_frame_id) == ("world", "map") + np.testing.assert_allclose(transform_matrix(tf.transform), np.linalg.inv(placement), atol=1e-9) def test_no_config_without_naming_a_rig(): """The scales are per-rig, so there is no bare RelocalizeConfig() to fall into.""" - import pydantic - - from dimos.mapping.relocalization.lidar import relocalize as lidar assert lidar.DEFAULT_PRESET in lidar.PRESETS assert lidar.PRESETS["mid360"] is lidar.MID360 @@ -155,35 +177,34 @@ def test_no_config_without_naming_a_rig(): def test_the_match_runs_on_a_window_of_the_last_scans(): """Every scan enters the window; a match sees the last `max_frames` of them.""" - import time - - from reactivex import Subject - - from dimos.mapping.relocalization.lidar.module import window - from dimos.mapping.relocalization.lidar.relocalize import MID360 cfg = MID360.model_copy(update={"min_frames": 2, "max_frames": 3}) scans = Subject() matched = [] - window(scans, cfg, interval=0.001).subscribe(matched.append) - - for i in range(5): - scans.on_next(PointCloud2.from_numpy(np.full((4, 3), i, dtype=np.float32), timestamp=0.0)) - time.sleep(0.01) # clear the throttle, so every scan gets its attempt - - # The first scan is below min_frames; from then on the window is full. - assert [len(c) for c in matched] == [8, 12, 12, 12] - # ... and holds the *last* three scans, not the first. - assert set(matched[-1].points_f32()[:, 0]) == {2.0, 3.0, 4.0} + scheduler = TestScheduler() + disposable = window(scans, cfg, interval=0.001).subscribe(matched.append, scheduler=scheduler) + try: + for i in range(5): + scans.on_next( + pointcloud_from_xyz( + np.full((4, 3), i, dtype=np.float32), + header=Header(stamp=time_from_nanoseconds(i), frame_id="world"), + ) + ) + scheduler.advance_by(0.01) + assert [c.width * c.height for c in matched] == [8, 12, 12, 12] + assert set(pointcloud_xyz(matched[-1])[:, 0]) == {2.0, 3.0, 4.0} + finally: + disposable.dispose() def test_from_matrix_inverse_matches_linalg_inv(): T = np.eye(4) T[:3, :3] = Rotation.from_euler("xyz", [0.1, -0.2, 1.3]).as_matrix() T[:3, 3] = [1.5, -2.0, 0.3] - tf = Transform.from_matrix(T, frame_id="map", child_frame_id="world").inverse() - assert (tf.frame_id, tf.child_frame_id) == ("world", "map") - np.testing.assert_allclose(tf.to_matrix(), np.linalg.inv(T), atol=1e-9) + tf = inverse_transform(_transform(T, frame_id="map", child_frame_id="world")) + assert (tf.header.frame_id, tf.child_frame_id) == ("world", "map") + np.testing.assert_allclose(transform_matrix(tf.transform), np.linalg.inv(T), atol=1e-9) def test_dual_strategy_merges_ports(): @@ -196,3 +217,92 @@ class Dual(LidarWindowRelocalization, FakeImpl): hints = get_type_hints(Dual) assert {"tf", "lidar", "loaded_map", "detections"} <= hints.keys() assert Dual.blueprint() is not None + + +def test_real_cdr_premap_and_fix_ports_keep_metadata_and_own_published_stamp( + module, tmp_path, mocker +): + path = tmp_path / "canonical.pc2.cdr" + cloud = pointcloud_from_xyz( + np.array([[1, 2, 3]], dtype=np.float32), + header=Header(frame_id="exported", stamp=time_from_nanoseconds(123456789)), + ) + payload = encode(cloud) + path.write_bytes(payload) + instance = module(map_file=str(tmp_path / "canonical"), tf_interval=0) + stamp = 1700000000123456789 + mocker.patch("dimos.message_utils.time.time.time_ns", return_value=stamp) + with ExitStack() as cleanup: + received, ready = {}, {} + + def collect(name): + def receive(value): + received[name] = value + ready[name].set() + + return receive + + for name, dtype in (("tf", TFMessage), ("loaded_map", PointCloud2)): + assert getattr(instance, name).type is dtype + transport = LCMTransport("relocalization_" + name, dtype, url="memq://") + cleanup.callback(transport.stop) + getattr(instance, name).transport = transport + ready[name] = Event() + cleanup.callback(transport.subscribe(collect(name))) + cleanup.callback(instance.stop) + instance.start() + assert not received + fix = _transform(np.eye(4), frame_id="world", child_frame_id="map") + before = encode(fix) + instance.submit(fix) + assert ready["tf"].wait(5) and ready["loaded_map"].wait(5) + [sent] = received["tf"].transforms + assert sent.header.stamp == time_from_nanoseconds(stamp) + assert sent.header.frame_id == "world" and sent.child_frame_id == "map" + loaded = received["loaded_map"] + assert loaded.header.frame_id == "map" and loaded.header.stamp == cloud.header.stamp + np.testing.assert_array_equal(pointcloud_xyz(loaded), [[1, 2, 3]]) + assert encode(fix) == before and path.read_bytes() == payload + + +@pytest.mark.parametrize("carving,count", [(True, 2), (False, 4)]) +def test_go2_merge_preserves_map_placement_and_carving_policy(module, mocker, carving, count): + # Exercise the real CPU carving backend independently of host GPU availability. + mocker.patch("open3d.core.cuda.is_available", return_value=False) + instance = module(Go2Relocalization, use_carving=carving) + premap = pointcloud_from_xyz( + np.array([[0.025, 0.025, 0.025], [0.025, 0.025, 0.075]], dtype=np.float32), + header=Header(frame_id="map", stamp=time_from_nanoseconds(100)), + ) + local = pointcloud_from_xyz( + np.array([[1.025, 0.025, 0.125], [2.025, 0.025, 0.025]], dtype=np.float32), + header=Header(frame_id="world", stamp=time_from_nanoseconds(101)), + ) + instance.premap = premap + publish = mocker.patch.object(instance.merged_map, "publish") + matrix = np.eye(4) + matrix[0, 3] = 1 + fix = _transform(matrix, frame_id="world", child_frame_id="map") + originals = encode(premap), encode(local), encode(fix) + instance._on_merge_input((local, fix)) + publish.assert_called_once() + [result] = publish.call_args.args + assert result.width * result.height == count + assert result.header == local.header + points = pointcloud_xyz(result) + assert set(np.round(points[:, 0], 3)) == {1.025, 2.025} + if carving: + np.testing.assert_allclose(points[points[:, 0] < 1.5, 2], [0.125], atol=1e-6) + assert (encode(premap), encode(local), encode(fix)) == originals + + +def test_real_preparation_produces_features_for_every_coarse_point(): + rng = np.random.default_rng(17) + points = rng.uniform(-3, 3, size=(500, 3)) + cloud = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(points)) + relocalizer = lidar.LidarRelocalizer(cloud, lidar.MID360) + prepared = relocalizer._prepare(cloud) + assert len(prepared.coarse.points) > 50 and len(prepared.fine.points) > 50 + assert prepared.fpfh.dimension() == 33 + assert prepared.fpfh.num() == len(prepared.coarse.points) + assert np.isfinite(np.asarray(prepared.fpfh.data)).all() diff --git a/dimos/mapping/test_generated_costmapper.py b/dimos/mapping/test_generated_costmapper.py new file mode 100644 index 00000000000..f1cd0b75148 --- /dev/null +++ b/dimos/mapping/test_generated_costmapper.py @@ -0,0 +1,234 @@ +# 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 contextlib import ExitStack +from copy import deepcopy +from threading import Event + +import cv2 +from dimos_cdrgen.registry import decode, encode +import numpy as np +import pytest +from reactivex.scheduler import ThreadPoolScheduler + +from dimos.core.transport import LCMTransport +from dimos.mapping.costmapper import CostMapper +from dimos.mapping.occupancy.visualizations import rainbow_image +from dimos.mapping.pointclouds.occupancy import ( + SimpleOccupancyConfig, + general_occupancy, + height_cost_occupancy, + simple_occupancy, +) +from dimos.mapping.pointclouds.util import read_pointcloud +from dimos.message_utils.occupancy import occupancy_from_array, occupancy_view +from dimos.message_utils.pointcloud import pointcloud_from_xyz +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.nav_msgs.msg import OccupancyGrid +from dimos.messages.sensor_msgs.msg import PointCloud2, PointField +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid as LegacyGrid +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 as LegacyCloud +from dimos.protocol.rpc.pubsubrpc import LCMRPC +from dimos.utils.data import get_data +from dimos.utils.threadpool import get_max_workers + + +@pytest.fixture(params=[(0, 0), (-1, 999999999), (1700000000, 123456789)]) +def source_header(request): + return Header(frame_id="lidar", stamp=Time(sec=request.param[0], nanosec=request.param[1])) + + +@pytest.mark.parametrize("algorithm", [simple_occupancy, general_occupancy, height_cost_occupancy]) +def test_generated_costmaps_match_existing_math_and_keep_exact_header(source_header, algorithm): + points = np.array( + [[-1, -1, 0], [-0.5, 0, 0.4], [0, 0, 0.4], [0.5, 0.5, 0], [1, 1, 0.2]], dtype=np.float32 + ) + cloud = pointcloud_from_xyz(points, header=source_header) + legacy = LegacyCloud.from_numpy(points, frame_id="lidar", timestamp=0) + + actual = algorithm(cloud, resolution=0.25) + original = algorithm(legacy, resolution=0.25) + received = decode(encode(actual), OccupancyGrid) + + assert type(original) is LegacyGrid + assert type(actual) is OccupancyGrid + assert received.header == source_header + assert received.info.map_load_time == source_header.stamp + assert original.ts == legacy.ts + assert received.info.resolution == original.resolution + assert (received.info.origin.position.x, received.info.origin.position.y) == ( + original.origin.position.x, + original.origin.position.y, + ) + np.testing.assert_array_equal(occupancy_view(received), original.grid) + actual.header.stamp.nanosec = 0 + assert cloud.header == source_header + + +@pytest.mark.parametrize("algorithm", [simple_occupancy, general_occupancy, height_cost_occupancy]) +def test_empty_generated_cloud_retains_source_metadata_and_unknown_map(source_header, algorithm): + cloud = pointcloud_from_xyz(np.empty((0, 3), dtype=np.float32), header=source_header) + + actual = algorithm(cloud, resolution=0.25, frame_id="map") + + assert actual.header.frame_id == "map" + assert actual.header.stamp == source_header.stamp + assert actual.info.map_load_time == source_header.stamp + assert occupancy_view(actual).tolist() == [[-1]] + + +def test_generated_missing_points_are_ignored_before_rasterization(): + header = Header(frame_id="lidar", stamp=Time(sec=0, nanosec=0)) + points = np.array([[0, 0, 0.4], [1, 1, 0], [np.nan, 0, 0], [0, np.inf, 0]], dtype=np.float32) + + actual = simple_occupancy(pointcloud_from_xyz(points, header=header), resolution=0.25) + expected = simple_occupancy(pointcloud_from_xyz(points[:2], header=header), resolution=0.25) + + np.testing.assert_array_equal(actual.data, expected.data) + assert actual.header == header + + +class MemoryRPC(LCMRPC): + """Use the real RPC implementation with an in-memory transport.""" + + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +@pytest.fixture +def mapper_factory(): + with ExitStack() as cleanup: + + def create(**kwargs): + module = CostMapper(rpc_transport=MemoryRPC, **kwargs) + cleanup.callback(module.stop) + return module + + yield create, cleanup + + +def test_initial_safe_radius_uses_owned_cells_without_changing_metadata(mapper_factory): + create, _ = mapper_factory + module = create(initial_safe_radius_meters=0.25) + header = Header(frame_id="map", stamp=Time(sec=1700000000, nanosec=123456789)) + grid = occupancy_from_array(np.full((3, 3), 100, dtype=np.int8), header=header, resolution=0.25) + original = deepcopy(grid.header) + + module._apply_initial_safe_radius(grid) + + assert occupancy_view(grid).tolist() == [[0, 0, 100], [0, 0, 100], [100, 100, 100]] + assert grid.header == original + assert grid.info.map_load_time == original.stamp + + +def test_costmapper_actual_generated_ports_deliver_cdr_without_network( + mapper_factory, owned_scheduler +): + create, cleanup = mapper_factory + global_wire = LCMTransport("costmap_global", PointCloud2, url="memq://") + merged_wire = LCMTransport("costmap_merged", PointCloud2, url="memq://") + output_wire = LCMTransport("costmap_output", OccupancyGrid, url="memq://") + for wire in (global_wire, merged_wire, output_wire): + cleanup.callback(wire.stop) + module = create(algo="simple", config=SimpleOccupancyConfig(resolution=0.25)) + assert module.global_map.type is PointCloud2 + assert module.global_costmap.type is OccupancyGrid + module.global_map.transport = global_wire + module.merged_map.transport = merged_wire + module.global_costmap.transport = output_wire + received = [] + ready = Event() + + def receive(value): + received.append(value) + ready.set() + + cleanup.callback(output_wire.subscribe(receive)) + header = Header(frame_id="lidar", stamp=Time(sec=1700000000, nanosec=123456789)) + cloud = pointcloud_from_xyz(np.array([[0, 0, 0.4], [1, 1, 0]], dtype=np.float32), header=header) + module.start() + + global_wire.broadcast(None, cloud) + + assert ready.wait(5) + assert len(received) == 1 + assert received[0].header == header + assert 100 in received[0].data + + ready.clear() + merged_header = Header(frame_id="merged_map", stamp=Time(sec=1700000000, nanosec=123456790)) + merged = pointcloud_from_xyz( + np.array([[0, 0, 0.4], [1, 1, 0]], dtype=np.float32), header=merged_header + ) + merged_wire.broadcast(None, merged) + assert ready.wait(5) + assert len(received) == 2 + assert received[-1].header == merged_header + + +@pytest.mark.self_hosted +@pytest.mark.parametrize( + "source_name,algorithm,expected_name,grayscale", + [ + ("apartment", simple_occupancy, "occupancy_simple.png", True), + ("big_office.ply", height_cost_occupancy, "big_office_height_cost_occupancy.png", False), + ("big_office.ply", simple_occupancy, "big_office_simple_occupancy.png", False), + ], +) +def test_generated_real_pointcloud_costmap_matches_existing_golden( + source_name, algorithm, expected_name, grayscale +): + source = get_data(source_name) + if source_name == "apartment": + source = source / "sum.ply" + points = np.asarray(read_pointcloud(source).points).astype(" None: self.get_global_pointcloud.invalidate_cache(self) self.get_global_pointcloud2.invalidate_cache(self) - def reproject(self, place: Callable[[PointCloud2], PointCloud2]) -> None: + def add_points(self, points: NDArray[np.float32], *, timestamp: float | None = None) -> None: + """Accumulate finite Nx3 coordinates without requiring a message wrapper. + + The optional per-frame timestamp uses the existing store's seconds unit; + stream adapters keep their exact source Header separately. + """ + self._check_disposed() + values = np.asarray(points, dtype=np.float32) + if values.ndim != 2 or values.shape[1] != 3: + raise ValueError("Voxel points must have shape Nx3") + if not np.isfinite(values).all(): + raise ValueError("Voxel points must be finite") + if timestamp is not None: + self._latest_frame_ts = timestamp + self._impl.add_points(values, timestamp) + self.get_global_pointcloud.invalidate_cache(self) + self.get_global_pointcloud2.invalidate_cache(self) + + def reproject(self, place: Callable[[CdrCloud], CdrCloud]) -> None: """Move every voxel to where `place` puts its stamped original position.""" self._check_disposed() if not isinstance(self._impl, PackedVoxels): @@ -113,7 +135,7 @@ def get_global_pointcloud(self) -> o3d.t.geometry.PointCloud: self._check_disposed() cpu = o3c.Device("CPU:0") out = o3d.t.geometry.PointCloud(device=cpu) - out.point["positions"] = o3c.Tensor(self._impl.points(), device=cpu) + out.point["positions"] = o3c.Tensor(self._impl.points(), dtype=o3c.float32, device=cpu) return out def size(self) -> int: diff --git a/dimos/mapping/voxels/impl/o3d.py b/dimos/mapping/voxels/impl/o3d.py index 82e80982b97..7a88ac16601 100644 --- a/dimos/mapping/voxels/impl/o3d.py +++ b/dimos/mapping/voxels/impl/o3d.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING import numpy as np +from numpy.typing import NDArray from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 @@ -74,7 +75,18 @@ def add_frame(self, frame: PointCloud2) -> None: if pcd.is_empty(): return - pts = pcd.point["positions"].to(self._dev, o3c.float32) + self._insert_points(pcd.point["positions"].to(self._dev, o3c.float32)) + + def add_points(self, points: NDArray[np.float32], _timestamp: float | None = None) -> None: + import open3d.core as o3c # type: ignore[import-untyped] + + if not len(points): + return + self._insert_points(o3c.Tensor(points, dtype=o3c.float32, device=self._dev)) + + def _insert_points(self, pts: o3c.Tensor) -> None: + import open3d.core as o3c # type: ignore[import-untyped] + vox = (pts / self._voxel_size).floor().to(self._key_dtype) keys_Nx3 = vox.contiguous() @@ -108,7 +120,7 @@ def _carve_and_insert(self, new_keys: o3c.Tensor) -> None: value_element_shapes=[o3c.SizeVector([1])], device=self._dev, ) - dummy_vals = o3c.Tensor.zeros((xy_keys.shape[0], 1), o3c.uint8, self._dev) + dummy_vals = o3c.Tensor.zeros(o3c.SizeVector([xy_keys.shape[0], 1]), o3c.uint8, self._dev) xy_hashmap.insert(xy_keys, dummy_vals) active_indices = self._hashmap.active_buf_indices() @@ -142,5 +154,6 @@ def size(self) -> int: def dispose(self) -> None: """Free GPU resources.""" - self.vbg = None - self._hashmap = None + if self.vbg is not None: + self.vbg = None + del self._hashmap diff --git a/dimos/mapping/voxels/impl/packed.py b/dimos/mapping/voxels/impl/packed.py index 83f9d5f2e25..99e96245924 100644 --- a/dimos/mapping/voxels/impl/packed.py +++ b/dimos/mapping/voxels/impl/packed.py @@ -17,8 +17,13 @@ from collections.abc import Callable import numpy as np +from numpy.typing import NDArray from dimos.mapping.voxels.keys import FIELD_BITS as _BITS, FIELD_MASK as _MASK, KEY_OFFSET as _BIAS +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.message_utils.time import time_from_seconds +from dimos.messages.sensor_msgs.msg import PointCloud2 as CdrCloud +from dimos.messages.std_msgs.msg import Header from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 @@ -56,7 +61,9 @@ def _centers(self, k: np.ndarray) -> np.ndarray: return (vox + np.float32(0.5)) * np.float32(self._voxel_size) def add_frame(self, frame: PointCloud2) -> None: - pts = frame.points_f32() + self.add_points(frame.points_f32(), frame.ts) + + def add_points(self, pts: NDArray[np.float32], timestamp: float | None = None) -> None: if not len(pts): return new = np.unique(self._pack(pts)) @@ -84,17 +91,25 @@ def add_frame(self, frame: PointCloud2) -> None: pos, new = pos[fresh], new[fresh] self._keys = np.insert(keys, pos, new) if self.stamped: - self._ts = np.insert(self._ts, pos, frame.ts) + self._ts = np.insert(self._ts, pos, np.nan if timestamp is None else timestamp) self._raw = np.insert(self._raw, pos, new) - def reproject(self, place: Callable[[PointCloud2], PointCloud2]) -> None: + def reproject(self, place: Callable[[CdrCloud], CdrCloud]) -> None: """Move every voxel to where `place` puts its stamped original position.""" if not self.stamped: raise RuntimeError("reproject needs a stamped store") if not len(self._keys): return keys = self._pack( - place(PointCloud2.from_numpy(self._centers(self._raw), stamps=self._ts)).points_f32() + pointcloud_xyz( + place( + pointcloud_from_xyz( + self._centers(self._raw), + header=Header(stamp=time_from_seconds(0), frame_id=""), + stamps=self._ts, + ) + ) + ) ) order = np.lexsort((-self._ts, keys)) # by key, latest stamp first keys, ts, raw = keys[order], self._ts[order], self._raw[order] diff --git a/dimos/mapping/voxels/module.py b/dimos/mapping/voxels/module.py index b273889a4c0..9248dbf2f8b 100644 --- a/dimos/mapping/voxels/module.py +++ b/dimos/mapping/voxels/module.py @@ -14,8 +14,11 @@ from __future__ import annotations +from copy import deepcopy from typing import TYPE_CHECKING, Any +import numpy as np + from dimos.core.core import rpc from dimos.core.module import ModuleConfig from dimos.core.stream import In, Out @@ -23,7 +26,9 @@ from dimos.memory.module import StreamModule from dimos.memory.stream import Stream from dimos.memory.transform import Transformer -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.message_utils.time import to_seconds +from dimos.messages.sensor_msgs.msg import PointCloud2 if TYPE_CHECKING: from collections.abc import Iterator @@ -53,8 +58,11 @@ def _make_obs( self, grid: VoxelGrid, last_obs: Observation[PointCloud2], count: int ) -> Observation[PointCloud2]: # pose=None: the global map is in world frame, per-observation pose is meaningless + header = deepcopy(last_obs.data.header) + header.frame_id = self._grid_kwargs.get("frame_id", "world") + points = grid.get_global_pointcloud().point["positions"].numpy() return last_obs.derive( - data=grid.get_global_pointcloud2(), + data=pointcloud_from_xyz(points, header=header), pose=None, tags={**last_obs.tags, "frame_count": count}, ) @@ -68,7 +76,9 @@ def __call__( count = 0 for obs in upstream: - grid.add_frame(obs.data) + coordinates = pointcloud_xyz(obs.data) + points = coordinates[np.isfinite(coordinates).all(axis=1)].astype(np.float32) + grid.add_points(points, timestamp=to_seconds(obs.data.header.stamp)) last_obs = obs count += 1 diff --git a/dimos/mapping/voxels/test_generated_stream.py b/dimos/mapping/voxels/test_generated_stream.py new file mode 100644 index 00000000000..f3672b26b28 --- /dev/null +++ b/dimos/mapping/voxels/test_generated_stream.py @@ -0,0 +1,228 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from contextlib import ExitStack +from copy import deepcopy +from threading import Event + +from dimos_cdrgen.registry import decode, encode +import numpy as np +import open3d.core as o3c +import pytest +from reactivex.scheduler import ThreadPoolScheduler + +from dimos.core.transport import LCMTransport +from dimos.mapping.costmapper import CostMapper +from dimos.mapping.pointclouds.occupancy import SimpleOccupancyConfig +from dimos.mapping.voxels.grid import VoxelGrid +from dimos.mapping.voxels.impl.o3d import O3dVoxels +from dimos.mapping.voxels.impl.packed import PackedVoxels +from dimos.mapping.voxels.module import VoxelGridMapper, VoxelMapTransformer +from dimos.memory.type.observation import Observation +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.message_utils.time import time_from_nanoseconds, to_nanoseconds +from dimos.messages.nav_msgs.msg import OccupancyGrid +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 as LegacyCloud +from dimos.protocol.rpc.pubsubrpc import LCMRPC +from dimos.utils.threadpool import get_max_workers + + +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +def keys(points): + return {tuple(row) for row in np.floor(np.asarray(points) / np.float32(0.05)).astype(np.int64)} + + +@pytest.mark.parametrize("carve", [False, True]) +def test_array_input_matches_legacy_voxel_ingestion(carve): + with ExitStack() as cleanup: + old = VoxelGrid(device="CPU:0", carve_columns=carve, show_startup_log=False) + new = VoxelGrid(device="CPU:0", carve_columns=carve, show_startup_log=False) + cleanup.callback(old.dispose) + cleanup.callback(new.dispose) + for index, points in enumerate( + ([[0.025, 0.025, 0.025], [0.075, 0.025, 0.025]], [[0.025, 0.025, 0.125]]) + ): + values = np.array(points, dtype=np.float32) + old.add_frame(LegacyCloud.from_numpy(values, timestamp=index)) + new.add_points(values, timestamp=index) + assert keys(old.get_global_pointcloud().point["positions"].numpy()) == keys( + new.get_global_pointcloud().point["positions"].numpy() + ) + assert len(old) == len(new) == (2 if carve else 3) + + +@pytest.mark.parametrize("carve", [False, True]) +def test_actual_open3d_cpu_array_entry_matches_packed_store(carve): + with ExitStack() as cleanup: + packed = PackedVoxels(0.05, carve) + tensor = O3dVoxels(0.05, 128, carve, o3c.Device("CPU:0")) + cleanup.callback(packed.dispose) + cleanup.callback(tensor.dispose) + for points in ([[0.025, 0.025, 0.025], [0.075, 0.025, 0.025]], [[0.025, 0.025, 0.125]]): + values = np.array(points, dtype=np.float32) + packed.add_points(values, 0) + tensor.add_points(values, 0) + assert keys(packed.points()) == keys(tensor.points()) + assert packed.size() == tensor.size() == (2 if carve else 3) + + +@pytest.mark.parametrize("points", [np.zeros((3,)), np.zeros((1, 2)), np.array([[np.nan, 0, 0]])]) +def test_array_entry_rejects_invalid_coordinates(points): + grid = VoxelGrid(device="CPU:0", show_startup_log=False) + try: + with pytest.raises(ValueError, match="Voxel points"): + grid.add_points(points) + finally: + grid.dispose() + + +@pytest.mark.parametrize("base", [0, -1, 1700000000123456789]) +@pytest.mark.parametrize("emit_every,indices", [(0, [2]), (1, [0, 1, 2]), (2, [1, 2])]) +def test_generated_stream_retains_exact_headers_batch_tail_and_observation_metadata( + base, emit_every, indices +): + frames = [] + for index in range(3): + cloud = pointcloud_from_xyz( + np.array([[0.025 + 0.05 * index, 0.025, 0.025]], dtype=np.float32), + header=Header(frame_id="world", stamp=time_from_nanoseconds(base + index)), + ) + frames.append( + Observation(ts=42 + index, data_type=PointCloud2, _data=cloud, tags={"source": "lidar"}) + ) + original = deepcopy(frames[-1].data.header) + transform = VoxelMapTransformer( + emit_every=emit_every, + device="CPU:0", + voxel_size=0.05, + frame_id="map", + show_startup_log=False, + ) + + result = list(transform(iter(frames))) + + assert [to_nanoseconds(item.data.header.stamp) for item in result] == [ + base + index for index in indices + ] + assert [item.tags["frame_count"] for item in result] == [index + 1 for index in indices] + assert [item.ts for item in result] == [42 + index for index in indices] + assert all(item.tags["source"] == "lidar" and item.pose is None for item in result) + assert all(item.data.header.frame_id == "map" for item in result) + received = decode(encode(result[-1].data), PointCloud2) + assert len(pointcloud_xyz(received)) == 3 + assert frames[-1].data.header == original + + +def test_generated_empty_and_missing_points_keep_zero_source_stamp(): + header = Header(frame_id="world", stamp=time_from_nanoseconds(0)) + cloud = pointcloud_from_xyz(np.array([[np.nan, 0, 0]], dtype=np.float32), header=header) + transform = VoxelMapTransformer(device="CPU:0", show_startup_log=False) + + result = list(transform(iter([Observation(ts=123, _data=cloud, data_type=PointCloud2)]))) + + assert pointcloud_xyz(result[0].data).shape == (0, 3) + assert to_nanoseconds(result[0].data.header.stamp) == 0 + assert result[0].ts == 123 + + +def test_voxel_mapper_real_generated_ports_accumulate_over_cdr_memory_transport(owned_scheduler): + with ExitStack() as cleanup: + incoming = LCMTransport("voxel_input", PointCloud2, url="memq://") + outgoing = LCMTransport("voxel_map", PointCloud2, url="memq://") + cleanup.callback(incoming.stop) + cleanup.callback(outgoing.stop) + module = VoxelGridMapper(rpc_transport=MemoryRPC, device="CPU:0", voxel_size=0.05) + cleanup.callback(module.stop) + assert module.lidar.type is PointCloud2 + assert module.global_map.type is PointCloud2 + module.lidar.transport = incoming + module.global_map.transport = outgoing + received = [] + ready = Event() + + def receive(value): + received.append(value) + ready.set() + + cleanup.callback(outgoing.subscribe(receive)) + module.start() + source = pointcloud_from_xyz( + np.array([[0.025, 0.025, 0.025]], dtype=np.float32), + header=Header(frame_id="world", stamp=time_from_nanoseconds(1700000000123456789)), + ) + + incoming.broadcast(None, source) + + assert ready.wait(5) + assert len(received) == 1 + assert received[0].header == source.header + assert keys(pointcloud_xyz(received[0])) == {(0, 0, 0)} + + +def test_actual_voxel_to_costmapper_chain_preserves_source_header_over_cdr(owned_scheduler): + with ExitStack() as cleanup: + source_wire = LCMTransport("voxel_chain_source", PointCloud2, url="memq://") + map_wire = LCMTransport("voxel_chain_map", PointCloud2, url="memq://") + merged_wire = LCMTransport("voxel_chain_merged", PointCloud2, url="memq://") + cost_wire = LCMTransport("voxel_chain_cost", OccupancyGrid, url="memq://") + for wire in (source_wire, map_wire, merged_wire, cost_wire): + cleanup.callback(wire.stop) + voxel = VoxelGridMapper(rpc_transport=MemoryRPC, device="CPU:0", voxel_size=0.05) + cost = CostMapper( + rpc_transport=MemoryRPC, algo="simple", config=SimpleOccupancyConfig(resolution=0.25) + ) + cleanup.callback(voxel.stop) + cleanup.callback(cost.stop) + voxel.lidar.transport = source_wire + voxel.global_map.transport = map_wire + cost.global_map.transport = map_wire + cost.merged_map.transport = merged_wire + cost.global_costmap.transport = cost_wire + received = [] + ready = Event() + + def receive(value): + received.append(value) + ready.set() + + cleanup.callback(cost_wire.subscribe(receive)) + cost.start() + voxel.start() + cloud = pointcloud_from_xyz( + np.array([[0.025, 0.025, 0.225], [1.025, 1.025, 0.025]], dtype=np.float32), + header=Header(frame_id="world", stamp=time_from_nanoseconds(1700000000123456789)), + ) + + source_wire.broadcast(None, cloud) + + assert ready.wait(5) + assert len(received) == 1 + assert received[0].header == cloud.header + assert received[0].info.map_load_time == cloud.header.stamp + assert 100 in received[0].data + + +@pytest.fixture +def owned_scheduler(monkeypatch): + scheduler = ThreadPoolScheduler(max_workers=get_max_workers()) + monkeypatch.setattr("dimos.utils.reactive.get_scheduler", lambda: scheduler) + monkeypatch.setattr("dimos.utils.threadpool.get_scheduler", lambda: scheduler) + yield scheduler + scheduler.executor.shutdown(wait=True) diff --git a/dimos/mapping/voxels/test_stamped.py b/dimos/mapping/voxels/test_stamped.py index 8fc126fef7d..b3367b21fb2 100644 --- a/dimos/mapping/voxels/test_stamped.py +++ b/dimos/mapping/voxels/test_stamped.py @@ -16,6 +16,7 @@ import pytest from dimos.mapping.voxels.grid import VoxelGrid +from dimos.message_utils.pointcloud import pointcloud_stamps, pointcloud_with_xyz, pointcloud_xyz from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 VOXEL = 0.05 @@ -41,8 +42,8 @@ def _drifted_walls(): def _undrift(cloud): - shift = np.where(cloud.stamps_f64()[:, None] == 1.0, [0.5, 0.0, 0.0], 0.0) - return PointCloud2.from_numpy(cloud.points_f32() + shift) + shift = np.where(pointcloud_stamps(cloud)[:, None] == 1.0, [0.5, 0.0, 0.0], 0.0) + return pointcloud_with_xyz(cloud, pointcloud_xyz(cloud) + shift) def test_stamped_holds_the_same_voxels(): diff --git a/dimos/memory/type/observation.py b/dimos/memory/type/observation.py index 3a42b4c80d0..543ed3ac17b 100644 --- a/dimos/memory/type/observation.py +++ b/dimos/memory/type/observation.py @@ -65,9 +65,8 @@ def _to_tuple(p: Any) -> PoseTuple | None: Accepts: `None`, a 3- or 7-tuple, or any object with `translation`+`rotation` (Transform) or `position`+`orientation` - (Pose / PoseStamped) attributes. 3-tuples are padded with the identity - quaternion. Duck-typed to avoid importing `dimos.msgs.geometry_msgs` - at module load. + (Pose) attributes, including the standard generated stamped/covariance/ + odometry nesting. 3-tuples are padded with the identity quaternion. """ if p is None: return None @@ -78,6 +77,15 @@ def _to_tuple(p: Any) -> PoseTuple | None: x, y, z = p return (float(x), float(y), float(z), 0.0, 0.0, 0.0, 1.0) raise TypeError(f"Pose tuple must have length 3 or 7, got {len(p)}") + # Generated ROS messages retain their declared nesting; only unwrap the + # pose-bearing shapes, leaving the indexed storage tuple contract unchanged. + name = getattr(type(p), "__msgtype__", None) + if name in ("nav_msgs/msg/Odometry", "geometry_msgs/msg/PoseWithCovarianceStamped"): + p = p.pose.pose + elif name in ("geometry_msgs/msg/PoseStamped", "geometry_msgs/msg/PoseWithCovariance"): + p = p.pose + elif name == "geometry_msgs/msg/TransformStamped": + p = p.transform # Use ``is not None`` rather than ``or`` — a Vector3 at the origin is # falsy but valid, and falling through to ``position`` on a Transform # would crash. diff --git a/dimos/memory/type/test_observation.py b/dimos/memory/type/test_observation.py new file mode 100644 index 00000000000..5105690aa1b --- /dev/null +++ b/dimos/memory/type/test_observation.py @@ -0,0 +1,92 @@ +# 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 copy import deepcopy + +from dimos_cdrgen.registry import encode +import numpy as np +import pytest + +from dimos.memory.type.observation import Observation +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + PoseWithCovariance, + Quaternion, + Transform, + TransformStamped, + Twist, + TwistWithCovariance, + Vector3, +) +from dimos.messages.nav_msgs.msg import Odometry +from dimos.messages.std_msgs.msg import Header + + +def canonical_pose_values(): + header = Header(frame_id="odom", stamp=time_from_nanoseconds(1700000000123456789)) + pose = Pose( + position=Point(x=1.5, y=-2, z=0.7), orientation=Quaternion(x=0.1, y=0.2, z=0.3, w=0.927) + ) + transform = Transform( + translation=Vector3(x=1.5, y=-2, z=0.7), rotation=deepcopy(pose.orientation) + ) + covariance = PoseWithCovariance(pose=deepcopy(pose), covariance=np.zeros(36, dtype=np.float64)) + return [ + pose, + PoseStamped(header=deepcopy(header), pose=deepcopy(pose)), + transform, + TransformStamped( + header=deepcopy(header), child_frame_id="base", transform=deepcopy(transform) + ), + covariance, + Odometry( + header=deepcopy(header), + child_frame_id="base", + pose=deepcopy(covariance), + twist=TwistWithCovariance( + twist=Twist(linear=Vector3(x=0, y=0, z=0), angular=Vector3(x=0, y=0, z=0)), + covariance=np.zeros(36, dtype=np.float64), + ), + ), + ] + + +@pytest.mark.parametrize("index", range(6)) +def test_canonical_pose_payloads_normalize_metadata_without_loading_or_mutation(index): + value = canonical_pose_values()[index] + before = encode(value) + loads = [] + + def load(): + loads.append("loaded") + return "payload" + + observation = Observation(id=7, ts=4.5, _loader=load) + posed = observation.with_pose(value) + assert posed.pose_tuple == (1.5, -2.0, 0.7, 0.1, 0.2, 0.3, 0.927) + assert observation.pose_tuple is None and posed.id == 7 and posed.ts == 4.5 + assert loads == [] + assert posed.data == observation.data == "payload" + assert loads == ["loaded"] + assert encode(value) == before + assert observation.with_pose(None).pose_tuple is None + + +def test_unsupported_message_is_not_silently_treated_as_a_pose(): + header = Header(frame_id="odom", stamp=time_from_nanoseconds(0)) + with pytest.raises(TypeError, match="Cannot coerce"): + Observation(pose=header, _data="unchanged") diff --git a/dimos/message_utils/occupancy.py b/dimos/message_utils/occupancy.py new file mode 100644 index 00000000000..e6a59187f15 --- /dev/null +++ b/dimos/message_utils/occupancy.py @@ -0,0 +1,167 @@ +# 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. + +"""Array operations on generated occupancy grids.""" + +from copy import deepcopy +from enum import IntEnum +import math +from pathlib import Path +from typing import Any + +import numpy as np +from numpy.typing import NDArray +from PIL import Image + +from dimos.message_utils.geometry import pose_matrix +from dimos.messages.geometry_msgs.msg import Point, Pose, Quaternion +from dimos.messages.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos.messages.std_msgs.msg import Header + + +class CostValues(IntEnum): + """Standard cost values for occupancy grid cells. + + These values follow the ROS nav_msgs/OccupancyGrid convention: + - 0: Free space + - 1-99: Occupied space with varying cost levels + - 100: Lethal obstacle (definitely occupied) + - -1: Unknown space + """ + + UNKNOWN = -1 # Unknown space + FREE = 0 # Free space + OCCUPIED = 100 # Occupied/lethal space + + +def occupancy_view(msg: OccupancyGrid) -> NDArray[np.int8]: + """Borrow a read-only row-major occupancy grid, retaining its owner.""" + if len(msg.data) != msg.info.height * msg.info.width: + raise ValueError("occupancy data length does not match dimensions") + view = msg.data.view().reshape(msg.info.height, msg.info.width) + view.setflags(write=False) + return view + + +def block_max_reduce(cells: NDArray[np.int8], factor: int) -> NDArray[np.int8]: + """Coarsen an occupancy grid by taking the max of factor x factor blocks. + + Block maximum (not mean) so coarsening never erases an obstacle; a block + is unknown (-1) only when every cell in it is unknown. Trailing remainder + rows/cols are trimmed, keeping row 0/col 0 - the origin corner - so the + grid origin stays valid. Grids thinner than `factor` pass through. + """ + h, w = cells.shape[:2] + new_h, new_w = h // factor, w // factor + if new_h == 0 or new_w == 0: + return cells + trimmed = cells[: new_h * factor, : new_w * factor] + blocks = trimmed.reshape(new_h, factor, new_w, factor) + # Sink unknown below every known value for the max, then map it back. + as_int = blocks.astype(np.int16) + known = np.where(as_int < 0, -1000, as_int) + reduced = known.max(axis=(1, 3)) + reduced[reduced == -1000] = -1 + result: NDArray[np.int8] = reduced.astype(np.int8) + return result + + +def occupancy_extent(message: OccupancyGrid) -> tuple[float, float]: + """Return physical width and height, rejecting an invalid grid resolution.""" + info = message.info + if not math.isfinite(info.resolution) or info.resolution <= 0: + raise ValueError("OccupancyGrid resolution must be finite and positive") + return info.width * info.resolution, info.height * info.resolution + + +def world_to_grid(message: OccupancyGrid, point: Point) -> tuple[float, float]: + """Convert a world point to continuous grid coordinates, including origin rotation.""" + occupancy_extent(message) + matrix = pose_matrix(message.info.origin) + coordinates = np.array([point.x, point.y, point.z]) + if not np.isfinite(coordinates).all(): + raise ValueError("grid coordinates must be finite") + local = matrix[:3, :3].T @ (coordinates - matrix[:3, 3]) + if not np.isfinite(local).all(): + raise ValueError("grid coordinates must be finite") + return float(local[0] / message.info.resolution), float(local[1] / message.info.resolution) + + +def grid_to_world(message: OccupancyGrid, coordinates: tuple[float, float]) -> Point: + """Convert continuous grid coordinates to the grid plane in world coordinates.""" + occupancy_extent(message) + x, y = coordinates + if not math.isfinite(x) or not math.isfinite(y): + raise ValueError("grid coordinates must be finite") + result = pose_matrix(message.info.origin) @ np.array( + [x * message.info.resolution, y * message.info.resolution, 0.0, 1.0] + ) + return Point(x=float(result[0]), y=float(result[1]), z=float(result[2])) + + +def occupancy_from_file(path: Path, *, header: Header, resolution: float = 0.05) -> OccupancyGrid: + """Load a static .npy or grayscale PNG grid with explicit frame/time metadata. + + Cells must be ROS occupancy values (-1 unknown, 0..100 known); PNG 255 is + unknown. NumPy object/pickle loading is disabled. Row order is unchanged. + """ + if not math.isfinite(resolution) or resolution <= 0: + raise ValueError("OccupancyGrid resolution must be finite and positive") + cells: NDArray[Any] + if path.suffix == ".npy": + cells = np.load(path, allow_pickle=False) + elif path.suffix == ".png": + with Image.open(path) as image: + cells = np.asarray(image.convert("L")).astype(np.int16) + cells[cells == 255] = -1 + else: + raise ValueError(f"Unsupported occupancy file format: {path.suffix}") + return occupancy_from_array(cells, header=header, resolution=resolution) + + +def occupancy_from_array( + cells: NDArray[Any], *, header: Header, resolution: float, origin: Pose | None = None +) -> OccupancyGrid: + """Copy a 2D ROS occupancy array with explicit map metadata.""" + if not math.isfinite(resolution) or resolution <= 0: + raise ValueError("OccupancyGrid resolution must be finite and positive") + if cells.ndim != 2 or cells.dtype.kind not in "iu": + raise ValueError("Occupancy cells must be a 2D integer array") + if np.any((cells < -1) | (cells > 100)): + raise ValueError("Occupancy cells must be -1 or 0..100") + return OccupancyGrid( + header=deepcopy(header), + info=MapMetaData( + map_load_time=deepcopy(header.stamp), + resolution=resolution, + height=cells.shape[0], + width=cells.shape[1], + origin=deepcopy(origin) + if origin is not None + else Pose( + orientation=Quaternion(w=1, x=0.0, y=0.0, z=0.0), + position=Point(x=0.0, y=0.0, z=0.0), + ), + ), + data=np.asarray(cells.astype(np.int8).reshape(-1), dtype=np.int8), + ) + + +def occupancy_cell_value(message: OccupancyGrid, point: Point) -> int: + """Read a world cell, retaining the planner's unknown result outside the map.""" + x, y = world_to_grid(message, point) + column, row = int(x), int(y) + if not (0 <= column < message.info.width and 0 <= row < message.info.height): + return CostValues.UNKNOWN + return int(occupancy_view(message)[row, column]) diff --git a/dimos/message_utils/pointcloud.py b/dimos/message_utils/pointcloud.py index fed75827913..1d7985d120d 100644 --- a/dimos/message_utils/pointcloud.py +++ b/dimos/message_utils/pointcloud.py @@ -235,13 +235,27 @@ def transform_cloud(message: PointCloud2, transform: TransformStamped) -> PointC """ if message.header.frame_id != transform.child_frame_id: raise ValueError("point cloud frame does not match transform child frame") + matrix = transform_matrix(transform.transform) + xyz = pointcloud_xyz(message) + transformed = xyz @ matrix[:3, :3].T + matrix[:3, 3] + output = pointcloud_with_xyz(message, transformed) + output.header.frame_id = transform.header.frame_id + return output + + +def pointcloud_with_xyz(message: PointCloud2, coordinates: NDArray[Any]) -> PointCloud2: + """Replace XYZ on a copy, preserving all other fields, padding and exact Header. + + Accept one coordinate triple per point in row order. This supports per-point + motion correction without losing capture times, intensity or organized layout. + """ view = pointcloud_view(message) + transformed = np.asarray(coordinates) + if transformed.shape != (message.height * message.width, 3): + raise ValueError("replacement coordinates must have shape (height * width, 3)") for name in ("x", "y", "z"): if name not in (view.dtype.names or ()) or view.dtype[name].kind != "f": raise ValueError("transform requires scalar floating-point XYZ fields") - matrix = transform_matrix(transform.transform) - xyz = pointcloud_xyz(message) - transformed = xyz @ matrix[:3, :3].T + matrix[:3, 3] for index, name in enumerate(("x", "y", "z")): values = transformed[:, index] if np.any(np.abs(values[np.isfinite(values)]) > np.finfo(view.dtype[name]).max): @@ -257,7 +271,7 @@ def transform_cloud(message: PointCloud2, transform: TransformStamped) -> PointC writable[name] = transformed[:, index].reshape(message.height, message.width) output = replace(message, header=deepcopy(message.header), fields=deepcopy(message.fields)) output.data = np.frombuffer(payload, dtype=np.uint8).copy() - output.header = Header(stamp=deepcopy(message.header.stamp), frame_id=transform.header.frame_id) + output.is_dense = message.is_dense and bool(np.isfinite(transformed).all()) return output diff --git a/dimos/message_utils/test_occupancy.py b/dimos/message_utils/test_occupancy.py new file mode 100644 index 00000000000..9df64f26118 --- /dev/null +++ b/dimos/message_utils/test_occupancy.py @@ -0,0 +1,208 @@ +# 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. + +import math + +from dimos_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +import numpy as np +from PIL import Image +import pytest + +from dimos.message_utils.geometry import quaternion_from_euler +from dimos.message_utils.occupancy import ( + grid_to_world, + occupancy_extent, + occupancy_from_array, + occupancy_from_file, + occupancy_view, + world_to_grid, +) +from dimos.messages.builtin_interfaces.msg import Time +from dimos.messages.geometry_msgs.msg import Point, Pose, Quaternion +from dimos.messages.nav_msgs.msg import MapMetaData, OccupancyGrid +from dimos.messages.std_msgs.msg import Header + + +def test_occupancy_cells_keep_ros_row_order_and_do_not_alias_message(): + message = OccupancyGrid( + info=MapMetaData( + width=3, + height=2, + resolution=0.1, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([-1, 0, 100, 25, 50, 75], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + decoded = cdr_decode(cdr_encode(message), OccupancyGrid) + values = occupancy_view(decoded) + np.testing.assert_array_equal(values, [[-1, 0, 100], [25, 50, 75]]) + with pytest.raises(ValueError, match="read-only"): + values[0, 0] = 0 + copied = values.copy() + copied[0, 0] = 0 + assert decoded.data[0] == -1 + + +def test_occupancy_rejects_dimensions_that_do_not_match_payload(): + message = OccupancyGrid( + info=MapMetaData( + width=2, + height=2, + resolution=1.0, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([0], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + with pytest.raises(ValueError, match="data length"): + occupancy_view(message) + + +@pytest.mark.parametrize("resolution", [0.0, -0.1, math.nan, math.inf]) +def test_occupancy_rejects_invalid_resolution_for_nonempty_grid(resolution): + message = OccupancyGrid( + info=MapMetaData( + width=1, + height=1, + resolution=resolution, + map_load_time=Time(sec=0, nanosec=0), + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([0], dtype=np.int8), + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + ) + with pytest.raises(ValueError, match="resolution"): + occupancy_extent(message) + + +def test_empty_default_occupancy_grid_has_no_cells(): + assert occupancy_view( + OccupancyGrid( + header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), + info=MapMetaData( + map_load_time=Time(sec=0, nanosec=0), + resolution=0.0, + width=0, + height=0, + origin=Pose( + position=Point(x=0.0, y=0.0, z=0.0), + orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0), + ), + ), + data=np.array([], dtype=np.int8), + ) + ).shape == (0, 0) + + +@pytest.mark.parametrize("suffix", [".npy", ".png"]) +def test_static_occupancy_loader_preserves_cells_header_and_identity(tmp_path, suffix): + cells = np.array([[-1, 0, 100], [25, 50, 75]], dtype=np.int8) + path = tmp_path / ("grid" + suffix) + if suffix == ".npy": + np.save(path, cells) + else: + Image.fromarray(cells.astype(np.uint8)).save(path) + header = Header(frame_id="map", stamp=Time(sec=123, nanosec=456)) + output = cdr_decode( + cdr_encode(occupancy_from_file(path, header=header, resolution=0.2)), + OccupancyGrid, + ) + np.testing.assert_array_equal(occupancy_view(output), cells) + assert output.header == header + assert output.info.map_load_time == header.stamp + assert output.info.origin.orientation.w == 1 + assert output.info.resolution == float(np.float32(0.2)) + + +@pytest.mark.parametrize( + "cells", + [np.zeros((2, 2), dtype=float), np.zeros(3, dtype=int), np.array([[101]]), np.array([[-2]])], +) +def test_static_occupancy_loader_rejects_invalid_cells(tmp_path, cells): + path = tmp_path / "grid.npy" + np.save(path, cells) + with pytest.raises(ValueError): + occupancy_from_file(path, header=Header(stamp=Time(sec=0, nanosec=0), frame_id="")) + + +@pytest.mark.parametrize("ns", [0, -1, 1700000000123456789]) +def test_array_constructor_copies_data_and_metadata_without_rounding_source_time(ns): + seconds, nanoseconds = divmod(ns, 1_000_000_000) + cells = np.array([[-1, 0], [100, 50]], dtype=np.int8) + header = Header(frame_id="map", stamp=Time(sec=seconds, nanosec=nanoseconds)) + origin = Pose( + position=Point(x=10, y=20, z=3), orientation=quaternion_from_euler(0, 0, math.pi / 2) + ) + + message = occupancy_from_array(cells, header=header, resolution=0.5, origin=origin) + header.stamp.sec = 77 + cells[0, 0] = 44 + origin.position.x = 99 + actual = cdr_decode(cdr_encode(message), OccupancyGrid) + + assert actual.header.stamp.sec == seconds + assert actual.header.stamp.nanosec == nanoseconds + assert actual.info.map_load_time == actual.header.stamp + assert actual.info.origin.position.x == 10 + np.testing.assert_array_equal(occupancy_view(actual), [[-1, 0], [100, 50]]) + + +def test_coordinate_helpers_account_for_rotated_origin_and_grid_plane(): + origin = Pose( + position=Point(x=10, y=20, z=3), orientation=quaternion_from_euler(0, 0, math.pi / 2) + ) + message = occupancy_from_array( + np.zeros((4, 4), dtype=np.int8), + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), + resolution=0.5, + origin=origin, + ) + + point = grid_to_world(message, (2, 3)) + + assert (point.x, point.y, point.z) == pytest.approx((8.5, 21, 3)) + assert world_to_grid(message, point) == pytest.approx((2, 3)) + assert occupancy_extent(message) == (2, 2) + + +@pytest.mark.parametrize("value", [math.nan, math.inf, -math.inf]) +def test_coordinate_helpers_reject_nonfinite_coordinates(value): + message = occupancy_from_array( + np.zeros((1, 1), dtype=np.int8), + header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0)), + resolution=1, + ) + with pytest.raises(ValueError, match="finite"): + grid_to_world(message, (value, 0)) + with pytest.raises(ValueError, match="finite"): + world_to_grid(message, Point(x=value, y=0, z=0)) + + +def test_numpy_loader_never_unpickles_objects(tmp_path): + path = tmp_path / "objects.npy" + np.save(path, np.array([[{}]], dtype=object), allow_pickle=True) + with pytest.raises(ValueError, match="allow_pickle=False"): + occupancy_from_file(path, header=Header(frame_id="map", stamp=Time(sec=0, nanosec=0))) diff --git a/dimos/message_utils/test_pointcloud.py b/dimos/message_utils/test_pointcloud.py index 48411b23d56..2b0959ecc96 100644 --- a/dimos/message_utils/test_pointcloud.py +++ b/dimos/message_utils/test_pointcloud.py @@ -32,6 +32,7 @@ pointcloud_rgb, pointcloud_stamps, pointcloud_view, + pointcloud_with_xyz, pointcloud_xyz, select_points, transform_cloud, @@ -460,3 +461,29 @@ def test_per_point_stamps_reject_wrong_count_or_nonfinite_values(stamps): header=Header(stamp=Time(sec=0, nanosec=0), frame_id=""), stamps=stamps, ) + + +def test_per_point_coordinate_replacement_keeps_organized_records(padded_cloud): + padded_cloud.header.stamp = Time(sec=1700000000, nanosec=123456789) + before = cdr_encode(padded_cloud) + result = pointcloud_with_xyz( + padded_cloud, pointcloud_xyz(padded_cloud) + np.array([0.5, -1, 2]) + ) + np.testing.assert_array_equal( + pointcloud_view(result)["tags"], pointcloud_view(padded_cloud)["tags"] + ) + assert result.header == padded_cloud.header + assert (result.height, result.width, result.row_step) == (2, 2, 64) + assert result.fields == padded_cloud.fields and result.is_bigendian == padded_cloud.is_bigendian + np.testing.assert_allclose( + pointcloud_xyz(result), [[1.5, 1, 5], [2.5, 3, 8], [3.5, 5, 11], [4.5, 7, 14]] + ) + assert bytes(result.data)[48:64] == bytes(padded_cloud.data)[48:64] + assert bytes(result.data)[112:128] == bytes(padded_cloud.data)[112:128] + assert cdr_encode(padded_cloud) == before + + +@pytest.mark.parametrize("coordinates", [np.zeros((3, 3)), np.zeros((2, 2, 3))]) +def test_per_point_coordinate_replacement_rejects_wrong_count(padded_cloud, coordinates): + with pytest.raises(ValueError, match="shape"): + pointcloud_with_xyz(padded_cloud, coordinates) diff --git a/dimos/message_utils/test_regions.py b/dimos/message_utils/test_regions.py new file mode 100644 index 00000000000..fb320fb93e0 --- /dev/null +++ b/dimos/message_utils/test_regions.py @@ -0,0 +1,129 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from copy import deepcopy +import subprocess +import sys + +from dimos_cdrgen.registry import decode, encode, schema +from mcap.reader import make_reader +from mcap.writer import Writer +import numpy as np +import pytest + +from dimos.message_utils.pointcloud import pointcloud_from_xyz +from dimos.message_utils.time import time_from_nanoseconds, to_nanoseconds +from dimos.messages.dimos_msgs.msg import ( + LineSegment3D, + LineSegments3D, + LineSegments3DRegion, + PointCloudRegion, + RegionBounds, +) +from dimos.messages.geometry_msgs.msg import Point +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header + + +def region_values(region_id): + header = Header(frame_id="map", stamp=time_from_nanoseconds(1700000000123456789)) + cloud = pointcloud_from_xyz(np.array([[1, 2, 3]], dtype=np.float32), header=deepcopy(header)) + lines = LineSegments3D( + header=deepcopy(header), + segments=[LineSegment3D(start=Point(x=1, y=2, z=3), end=Point(x=4, y=5, z=6), weight=0.75)], + ) + return [ + PointCloudRegion(header=deepcopy(header), region_id=region_id, cloud=cloud), + RegionBounds( + header=deepcopy(header), + region_id=region_id, + center=Point(x=1, y=2, z=0), + radius=3.5, + z_min=-1.25, + z_max=4.5, + ), + LineSegments3DRegion(header=deepcopy(header), region_id=region_id, lines=lines), + ] + + +@pytest.mark.parametrize("region_id", [-(2**31), 0, 2**31 - 1]) +@pytest.mark.parametrize("index", range(3)) +def test_region_id_and_canonical_dependency_types_survive_cdr(region_id, index): + value = region_values(region_id)[index] + result = decode(encode(value), type(value)) + assert encode(result) == encode(value) + assert result.region_id == region_id + assert type(result.header) is Header + assert to_nanoseconds(result.header.stamp) == 1700000000123456789 + if index == 0: + assert type(result.cloud) is PointCloud2 and type(result.cloud.header) is Header + elif index == 1: + assert type(result.center) is Point + else: + assert type(result.lines) is LineSegments3D + assert type(result.lines.segments[0].start) is Point + + +def test_region_mcap_embeds_all_dependency_schemas_for_an_independent_reader(tmp_path): + path = tmp_path / "regions.mcap" + values = region_values(-2147483648) + with path.open("wb") as stream: + writer = Writer(stream) + writer.start(profile="ros2") + for value in values: + name = value.__msgtype__ + schema_id = writer.register_schema(name, "ros2msg", schema(name).encode()) + channel_id = writer.register_channel("/" + name, "cdr", schema_id) + writer.add_message(channel_id, 1700000000123456790, encode(value), 1700000000123456789) + writer.finish() + with path.open("rb") as stream: + messages = list(make_reader(stream, validate_crcs=True).iter_messages()) + assert len(messages) == 3 + for stored_schema, channel, message in messages: + assert stored_schema.encoding == "ros2msg" and channel.message_encoding == "cdr" + assert message.publish_time == 1700000000123456789 + assert "MSG: std_msgs/Header" in stored_schema.data.decode() + # The independent type store owns a global module; use a separate interpreter + # so it cannot replace canonical Python message class identities in this test run. + subprocess.run( + [ + sys.executable, + "-I", + "-c", + """ +from mcap.reader import make_reader +from rosbags.typesys import Stores, get_typestore, get_types_from_msg +import sys +reference = get_typestore(Stores.EMPTY) +with open(sys.argv[1], "rb") as stream: + for schema, _, message in make_reader(stream).iter_messages(): + reference.register(get_types_from_msg(schema.data.decode(), schema.name)) + value = reference.deserialize_cdr(message.data, schema.name) + assert value.region_id == -2147483648 + assert value.header.stamp.sec == 1700000000 + assert value.header.stamp.nanosec == 123456789 + if schema.name.endswith("PointCloudRegion"): + assert value.cloud.width == 1 and value.cloud.point_step == 12 + elif schema.name.endswith("RegionBounds"): + assert (value.center.x,value.center.y,value.radius,value.z_min,value.z_max) == (1,2,3.5,-1.25,4.5) + else: + assert value.lines.segments[0].weight == 0.75 + assert value.lines.segments[0].end.z == 6 +""", + str(path), + ], + check=True, + capture_output=True, + text=True, + ) diff --git a/dimos/navigation/base.py b/dimos/navigation/base.py index dc9cbbe9c27..4cd8ceca025 100644 --- a/dimos/navigation/base.py +++ b/dimos/navigation/base.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 # Copyright 2025-2026 Dimensional Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -16,7 +15,7 @@ from abc import ABC, abstractmethod from enum import Enum -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.messages.geometry_msgs.msg import PoseStamped class NavigationState(Enum): @@ -37,7 +36,6 @@ def set_goal(self, goal: PoseStamped) -> bool: Returns: True if goal was accepted, False otherwise """ - pass @abstractmethod def get_state(self) -> NavigationState: @@ -47,7 +45,6 @@ def get_state(self) -> NavigationState: Returns: Current navigation state """ - pass @abstractmethod def is_goal_reached(self) -> bool: @@ -57,7 +54,6 @@ def is_goal_reached(self) -> bool: Returns: True if goal was reached, False otherwise """ - pass @abstractmethod def cancel_goal(self) -> bool: @@ -67,4 +63,3 @@ def cancel_goal(self) -> bool: Returns: True if goal was cancelled, False if no goal was active """ - pass diff --git a/dimos/navigation/global_planner/mls_planner/mls_planner_native.py b/dimos/navigation/global_planner/mls_planner/mls_planner_native.py index f194af722ed..d05a9cb3113 100644 --- a/dimos/navigation/global_planner/mls_planner/mls_planner_native.py +++ b/dimos/navigation/global_planner/mls_planner/mls_planner_native.py @@ -17,13 +17,12 @@ from __future__ import annotations from dimos.core.native_module import NativeModule, NativeModuleConfig -from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.core.stream import IO, In, Out +from dimos.messages.dimos_msgs.msg import LineSegments3DRegion, PointCloudRegion, RegionBounds +from dimos.messages.geometry_msgs.msg import PointStamped +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.tf2_msgs.msg import TFMessage from dimos.navigation import spec @@ -71,15 +70,16 @@ class MLSPlannerNative(NativeModule, spec.GlobalPlanner): global_map: In[PointCloud2] local_map: In[PointCloud2] - region_bounds: In[PoseStamped] + region_bounds: In[RegionBounds] # A seeded map's regions as RayTracingVoxelMap lands them, applied through # the region pipeline between live updates. Live updates keep priority. - seed_map: In[PointCloud2] - seed_bounds: In[PoseStamped] + seed_map: In[PointCloudRegion] + seed_bounds: In[RegionBounds] goal: In[PointStamped] - tf: In[TFMessage] + tf: IO[TFMessage] + tf_static: In[TFMessage] path: Out[Path] - surface_map: Out[PointCloud2] + surface_map: Out[PointCloudRegion] nodes: Out[PointCloud2] - node_edges: Out[LineSegments3D] + node_edges: Out[LineSegments3DRegion] diff --git a/dimos/navigation/global_planner/mls_planner/rust/Cargo.lock b/dimos/navigation/global_planner/mls_planner/rust/Cargo.lock index 8af4f7f7a1e..46652094f36 100644 --- a/dimos/navigation/global_planner/mls_planner/rust/Cargo.lock +++ b/dimos/navigation/global_planner/mls_planner/rust/Cargo.lock @@ -249,6 +249,15 @@ dependencies = [ "shlex", ] +[[package]] +name = "cdr" +version = "0.2.4" +source = "git+https://github.com/TomCC7/cdr-rs?rev=38a4368255d408faf374726653b39d4a4e9735cb#38a4368255d408faf374726653b39d4a4e9735cb" +dependencies = [ + "byteorder", + "serde", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -579,6 +588,17 @@ dependencies = [ "subtle", ] +[[package]] +name = "dimos-generated-messages" +version = "0.0.14" +dependencies = [ + "cdr", + "heapless", + "ros2msg", + "serde", + "serde-big-array", +] + [[package]] name = "dimos-lcm" version = "0.1.0" @@ -594,8 +614,8 @@ name = "dimos-mls-planner" version = "0.1.0" dependencies = [ "ahash", + "dimos-generated-messages", "dimos-module", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen)", "rayon", "serde", "smallvec", @@ -618,12 +638,13 @@ dependencies = [ [[package]] name = "dimos-module" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#fc6a4e58da89bb073a38dce5ac63728c6559d3a6" +version = "0.2.0" dependencies = [ + "cdr", + "dimos-generated-messages", "dimos-lcm", "dimos-module-macros", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549)", + "lcm-msgs", "nalgebra", "rayon", "serde", @@ -631,6 +652,7 @@ dependencies = [ "tokio", "tracing", "tracing-subscriber", + "transforms", "url", "validator", "zenoh", @@ -639,7 +661,6 @@ dependencies = [ [[package]] name = "dimos-module-macros" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#fc6a4e58da89bb073a38dce5ac63728c6559d3a6" dependencies = [ "proc-macro2", "quote", @@ -978,6 +999,15 @@ version = "0.33.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e762b9b188634fc508d4ec54b62ea3d352c43fd431d18d05d46f559f217f4725" +[[package]] +name = "hash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" +dependencies = [ + "byteorder", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -1019,6 +1049,22 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] + +[[package]] +name = "heapless" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" +dependencies = [ + "hash32", + "serde", + "stable_deref_trait", +] [[package]] name = "heck" @@ -1425,14 +1471,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "lcm-msgs" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" -dependencies = [ - "byteorder", -] - [[package]] name = "lcm-msgs" version = "0.1.0" @@ -2017,6 +2055,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -2433,6 +2481,23 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "ros2msg" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fe291be5cb6d4b8e6bf9c9de80c40bd1cb09ecac50deafa2fdcae0027e04988" +dependencies = [ + "heck", + "pest", + "pest_derive", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "syn 2.0.119", + "thiserror", +] + [[package]] name = "rustc-hash" version = "2.1.3" @@ -2673,6 +2738,15 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-big-array" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11fc7cc2c76d73e0f27ee52abbd64eec84d46f370c88371120433196934e4b7f" +dependencies = [ + "serde", +] + [[package]] name = "serde_core" version = "1.0.229" @@ -3354,6 +3428,18 @@ dependencies = [ "tracing-serde", ] +[[package]] +name = "transforms" +version = "2.2.0" +source = "git+https://github.com/TomCC7/transforms?rev=5ae9563de0c02638c04e393aadfad28444a4398e#5ae9563de0c02638c04e393aadfad28444a4398e" +dependencies = [ + "approx", + "hashbrown 0.17.1", + "libm", + "serde", + "thiserror", +] + [[package]] name = "twox-hash" version = "1.6.3" diff --git a/dimos/navigation/global_planner/mls_planner/rust/Cargo.toml b/dimos/navigation/global_planner/mls_planner/rust/Cargo.toml index da594faefc7..1c49d50c281 100644 --- a/dimos/navigation/global_planner/mls_planner/rust/Cargo.toml +++ b/dimos/navigation/global_planner/mls_planner/rust/Cargo.toml @@ -22,23 +22,24 @@ python = "dimos.navigation.global_planner.mls_planner.mls_planner_native:MLSPlan threads = 2 [package.metadata.dimos.module.mls_planner.inputs] -global_map = "sensor_msgs.PointCloud2" -local_map = "sensor_msgs.PointCloud2" -region_bounds = "geometry_msgs.PoseStamped" -seed_map = "sensor_msgs.PointCloud2" -seed_bounds = "geometry_msgs.PoseStamped" -goal = "geometry_msgs.PointStamped" -tf = "tf2_msgs.TFMessage" +global_map = "sensor_msgs/msg/PointCloud2" +local_map = "sensor_msgs/msg/PointCloud2" +region_bounds = "dimos_msgs/msg/RegionBounds" +seed_map = "dimos_msgs/msg/PointCloudRegion" +seed_bounds = "dimos_msgs/msg/RegionBounds" +goal = "geometry_msgs/msg/PointStamped" +tf = "tf2_msgs/msg/TFMessage" +tf_static = "tf2_msgs/msg/TFMessage" [package.metadata.dimos.module.mls_planner.outputs] -surface_map = "sensor_msgs.PointCloud2" -nodes = "sensor_msgs.PointCloud2" -node_edges = "nav_msgs.LineSegments3D" -path = "nav_msgs.Path" +surface_map = "dimos_msgs/msg/PointCloudRegion" +nodes = "sensor_msgs/msg/PointCloud2" +node_edges = "dimos_msgs/msg/LineSegments3DRegion" +path = "nav_msgs/msg/Path" [dependencies] -dimos-module = { git = "https://github.com/dimensionalOS/dimos", rev = "fc6a4e58da89bb073a38dce5ac63728c6559d3a6", version = "=0.1.0" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +dimos-module = { path = "../../../../../native/rust/dimos-module" } +dimos-generated-messages = { path = "../../../../../build/message-source/rust" } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "time"] } serde = { version = "1", features = ["derive"] } ahash = "0.8" diff --git a/dimos/navigation/global_planner/mls_planner/rust/crate-hashes.json b/dimos/navigation/global_planner/mls_planner/rust/crate-hashes.json index 187d28513c6..3ee4f926e23 100644 --- a/dimos/navigation/global_planner/mls_planner/rust/crate-hashes.json +++ b/dimos/navigation/global_planner/mls_planner/rust/crate-hashes.json @@ -1,7 +1,6 @@ { - "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#lcm-msgs@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=", + "git+https://github.com/TomCC7/cdr-rs?rev=38a4368255d408faf374726653b39d4a4e9735cb#cdr@0.2.4": "sha256-akg/bLBY3wQ1YRr2bijE5ZqB211Cst5lKjELLrgiG6o=", + "git+https://github.com/TomCC7/transforms?rev=5ae9563de0c02638c04e393aadfad28444a4398e#transforms@2.2.0": "sha256-Kh8aVjj7zQ+Xd1QHOcAhSJbubg7azaWFSMpaI/ZAgM4=", "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549#dimos-lcm@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=", - "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549#lcm-msgs@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=", - "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#dimos-module-macros@0.1.0": "sha256-ECU2xYh/j+a3XwyhO3UR/GPruo5jNzM3n/9S5BIvVrA=", - "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#dimos-module@0.1.0": "sha256-ECU2xYh/j+a3XwyhO3UR/GPruo5jNzM3n/9S5BIvVrA=" + "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549#lcm-msgs@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=" } \ No newline at end of file diff --git a/dimos/navigation/global_planner/mls_planner/rust/flake.nix b/dimos/navigation/global_planner/mls_planner/rust/flake.nix index 6ae4c57eb16..0be11d2e566 100644 --- a/dimos/navigation/global_planner/mls_planner/rust/flake.nix +++ b/dimos/navigation/global_planner/mls_planner/rust/flake.nix @@ -15,7 +15,32 @@ pkgs = nixpkgs.legacyPackages.${system}; name = "dimos-mls-planner"; - src = nix-filter.lib { root = ./.; exclude = [ "target" "build" "result" "__pycache__" ]; }; + moduleSource = nix-filter.lib { root = ./.; exclude = [ "target" "build" "result" "__pycache__" ]; }; + sharedArchive = pkgs.fetchurl { + url = "https://codeload.github.com/dimensionalOS/dimos/tar.gz/a9365487b4323151ea85041c6a944b509c0c3268"; + sha256 = "e339487dd2281980fcde7ab03c46dfd2e274491e0f73e66d9c6cbcaeea8902b2"; + }; + sharedSource = pkgs.runCommand "dimos-mls-cdr-source-input" { + nativeBuildInputs = [ pkgs.gnutar pkgs.gzip ]; + } '' + mkdir -p $out + tar -xzf ${sharedArchive} --strip-components=1 -C $out + ''; + messageSource = (import "${sharedSource}/native/message-source.nix" { inherit pkgs; }).source; + src = pkgs.runCommand "mls-cdr-source" { } '' + cp -r ${moduleSource} $out + chmod -R u+w $out + mkdir -p $out/sdk + cp -r ${sharedSource}/native/rust/dimos-module ${sharedSource}/native/rust/dimos-module-macros $out/sdk/ + cp -r ${messageSource}/rust $out/messages + substituteInPlace $out/sdk/dimos-module/Cargo.toml \ + --replace-fail '../../../build/message-source/rust' '../../messages' + substituteInPlace $out/Cargo.toml \ + --replace-fail '../../../../../native/rust/dimos-module' 'sdk/dimos-module' \ + --replace-fail '../../../../../build/message-source/rust' 'messages' \ + --replace-fail '[workspace]' '[workspace] + exclude = ["sdk/dimos-module", "sdk/dimos-module-macros", "messages"]' + ''; generated = crate2nix.tools.${system}.generatedCargoNix { inherit name src; }; diff --git a/dimos/navigation/global_planner/mls_planner/rust/src/module.rs b/dimos/navigation/global_planner/mls_planner/rust/src/module.rs index 3913a0b17da..252dd1812aa 100644 --- a/dimos/navigation/global_planner/mls_planner/rust/src/module.rs +++ b/dimos/navigation/global_planner/mls_planner/rust/src/module.rs @@ -17,15 +17,31 @@ use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; -use crate::mls_planner::{ms_since, Config, Planner, RegionBounds, DEBUG_LOG_PERIOD}; +use crate::mls_planner::{ + ms_since, Config, Planner, RegionBounds as SurfaceBounds, DEBUG_LOG_PERIOD, +}; use crate::region_viz::{pack_cell, Cell, RegionUpdate, RegionViz, Segment}; use crate::voxel::{surface_point_xyz, voxelize, VoxelKey}; -use dimos_module::time::now; +use dimos_generated_messages::{ + builtin_interfaces::msg::time::Time, + dimos_msgs::msg::{ + line_segment3_d::LineSegment3D, line_segments3_d::LineSegments3D, + line_segments3_d_region::LineSegments3DRegion, point_cloud_region::PointCloudRegion, + region_bounds::RegionBounds, + }, + geometry_msgs::msg::{ + point::Point, point_stamped::PointStamped, pose::Pose, pose_stamped::PoseStamped, + quaternion::Quaternion, + }, + nav_msgs::msg::path::Path, + sensor_msgs::msg::{point_cloud2::PointCloud2, point_field::PointField}, + std_msgs::msg::header::Header, +}; + +use dimos_module::cdr; +use dimos_module::pointcloud::{extract_xyz_rows, ExtractError}; use dimos_module::{debug_throttled, error_throttled, warn_throttled, Input, Module, Output, Tf}; -use lcm_msgs::geometry_msgs::{Point, PointStamped, Pose, PoseStamped, Quaternion}; -use lcm_msgs::nav_msgs::Path; -use lcm_msgs::sensor_msgs::{PointCloud2, PointField}; -use lcm_msgs::std_msgs::{Header, Time}; +use std::time::{SystemTime, UNIX_EPOCH}; use tokio::sync::mpsc; use tokio::sync::Notify; use tracing::{debug, info, warn}; @@ -43,7 +59,7 @@ type Shared = Arc>>; enum MapUpdate { Region { points: Vec, - bounds: PoseStamped, + bounds: RegionBounds, received: Instant, }, Global { @@ -53,15 +69,15 @@ enum MapUpdate { /// One region of a seeded map, as the ray tracer hands them on. struct SeedRegion { - cloud: PointCloud2, - bounds: PoseStamped, + cloud: PointCloudRegion, + bounds: RegionBounds, } /// A seed region with its points extracted, queued for the worker. They /// queue, since every one must land. struct QueuedSeed { points: Vec, - bounds: PoseStamped, + bounds: RegionBounds, } /// One viz tick's content, read by the worker and published off it. @@ -175,31 +191,31 @@ impl SeedProgress { } /// Seed clouds and bounds waiting for their counterpart, keyed by the region -/// number in their header seq, since every region of a seed shares one stamp. +/// number in their explicit region_id, since every region of a seed shares one stamp. #[derive(Default)] struct SeedPairs { - clouds: HashMap, - bounds: HashMap, + clouds: HashMap, + bounds: HashMap, } impl SeedPairs { - fn pair_cloud(&mut self, msg: PointCloud2) -> Option { + fn pair_cloud(&mut self, msg: PointCloudRegion) -> Option { self.expire_before(Instant::now()); - match self.bounds.remove(&msg.header.seq) { + match self.bounds.remove(&msg.region_id) { Some((bounds, _)) => Some(SeedRegion { cloud: msg, bounds }), None => { - self.clouds.insert(msg.header.seq, (msg, Instant::now())); + self.clouds.insert(msg.region_id, (msg, Instant::now())); None } } } - fn pair_bounds(&mut self, msg: PoseStamped) -> Option { + fn pair_bounds(&mut self, msg: RegionBounds) -> Option { self.expire_before(Instant::now()); - match self.clouds.remove(&msg.header.seq) { + match self.clouds.remove(&msg.region_id) { Some((cloud, _)) => Some(SeedRegion { cloud, bounds: msg }), None => { - self.bounds.insert(msg.header.seq, (msg, Instant::now())); + self.bounds.insert(msg.region_id, (msg, Instant::now())); None } } @@ -225,40 +241,40 @@ impl SeedPairs { #[derive(Module)] #[module(name = "mls_planner", setup = spawn_worker, teardown = stop_worker)] pub struct MlsPlanner { - #[input(decode = PointCloud2::decode, handler = on_global_map)] + #[input(decode = dimos_module::cdr::decode, handler = on_global_map)] global_map: Input, - #[input(decode = PointCloud2::decode, handler = on_local_map)] + #[input(decode = dimos_module::cdr::decode, handler = on_local_map)] local_map: Input, - #[input(decode = PoseStamped::decode, handler = on_region_bounds)] - region_bounds: Input, + #[input(decode = dimos_module::cdr::decode, handler = on_region_bounds)] + region_bounds: Input, // A seeded map's regions as the ray tracer lands them, applied through // the region pipeline between live updates. Live updates keep priority. - #[input(decode = PointCloud2::decode, handler = on_seed_map)] - seed_map: Input, + #[input(decode = dimos_module::cdr::decode, handler = on_seed_map)] + seed_map: Input, - #[input(decode = PoseStamped::decode, handler = on_seed_bounds)] - seed_bounds: Input, + #[input(decode = dimos_module::cdr::decode, handler = on_seed_bounds)] + seed_bounds: Input, - #[input(decode = PointStamped::decode, handler = on_goal)] + #[input(decode = dimos_module::cdr::decode, handler = on_goal)] goal: Input, #[tf] tf: Tf, - #[output(encode = PointCloud2::encode)] - surface_map: Output, + #[output(encode = cdr::encode)] + surface_map: Output, - #[output(encode = PointCloud2::encode)] + #[output(encode = cdr::encode)] nodes: Output, - // The wire payload is a Path. dimos names the channel LineSegments3D. - #[output(encode = Path::encode, msg = "LineSegments3D")] - node_edges: Output, + // Canonical weighted segments with an explicit viewer-region key. + #[output(encode = cdr::encode)] + node_edges: Output, - #[output(encode = Path::encode)] + #[output(encode = cdr::encode)] path: Output, #[config] @@ -266,7 +282,7 @@ pub struct MlsPlanner { // Held on the handle loop until stamps match, then handed off paired. pending_local: Option, - pending_bounds: Option, + pending_bounds: Option, pending_seeds: SeedPairs, // Written by the handle loop, read by the worker, so the loop never blocks @@ -327,18 +343,18 @@ impl MlsPlanner { self.try_pair(); } - async fn on_region_bounds(&mut self, msg: PoseStamped) { + async fn on_region_bounds(&mut self, msg: RegionBounds) { self.pending_bounds = Some(msg); self.try_pair(); } - async fn on_seed_map(&mut self, msg: PointCloud2) { + async fn on_seed_map(&mut self, msg: PointCloudRegion) { if let Some(region) = self.pending_seeds.pair_cloud(msg) { self.queue_seed(region); } } - async fn on_seed_bounds(&mut self, msg: PoseStamped) { + async fn on_seed_bounds(&mut self, msg: RegionBounds) { if let Some(region) = self.pending_seeds.pair_bounds(msg) { self.queue_seed(region); } @@ -362,7 +378,7 @@ impl MlsPlanner { } fn queue_seed(&self, region: SeedRegion) { - let Some(points) = cloud_points(®ion.cloud, "seed_map") else { + let Some(points) = cloud_points(®ion.cloud.cloud, "seed_map") else { return; }; self.seed_regions @@ -399,7 +415,7 @@ impl MlsPlanner { } /// True when bounds and a local cloud are both present with matching stamps. -fn stamps_paired(bounds: Option<&PoseStamped>, cloud: Option<&PointCloud2>) -> bool { +fn stamps_paired(bounds: Option<&RegionBounds>, cloud: Option<&PointCloud2>) -> bool { match (bounds, cloud) { (Some(b), Some(c)) => same_stamp(&b.header.stamp, &c.header.stamp), _ => false, @@ -434,9 +450,9 @@ struct VizPublisher { batches: mpsc::Receiver, voxel_size: f32, frame: String, - surface_map: Output, + surface_map: Output, nodes: Output, - node_edges: Output, + node_edges: Output, } impl VizPublisher { @@ -661,7 +677,8 @@ impl Worker { let t = self .tf .get_latest(&self.config.world_frame, &self.config.base_frame)? - .translation(); + .transform + .translation; Some((t.x as f32, t.y as f32, t.z as f32)) } @@ -761,26 +778,25 @@ fn replan_due(goal_changed: bool, live_update_applied: bool) -> bool { goal_changed || live_update_applied } -/// The region a bounds message describes: position is the center, orientation -/// carries radius, z_min and z_max. -fn region_bounds(msg: &PoseStamped) -> RegionBounds { - RegionBounds { - origin_x: msg.pose.position.x as f32, - origin_y: msg.pose.position.y as f32, - radius: msg.pose.orientation.x as f32, - z_min: msg.pose.orientation.y as f32, - z_max: msg.pose.orientation.z as f32, +/// The explicit cylinder a bounds message describes. +fn region_bounds(msg: &RegionBounds) -> SurfaceBounds { + SurfaceBounds { + origin_x: msg.center.x as f32, + origin_y: msg.center.y as f32, + radius: msg.radius as f32, + z_min: msg.z_min as f32, + z_max: msg.z_max as f32, } } -/// One cell's surface as a message, the cell in its header seq. +/// One cell's surface as a message, the cell in its explicit region_id. fn surface_message( cell: Cell, surface: Vec<(VoxelKey, f32)>, voxel_size: f32, frame: &str, stamp: Time, -) -> PointCloud2 { +) -> PointCloudRegion { let points: Vec = surface .into_iter() .map(|((ix, iy, iz), clearance)| { @@ -788,22 +804,28 @@ fn surface_message( (x, y, z, clearance) }) .collect(); - let mut msg = build_pc2_xyzi(&points, frame, stamp); - msg.header.seq = pack_cell(cell); - msg + let cloud = build_pc2_xyzi(&points, frame, stamp); + PointCloudRegion { + header: cloud.header.clone(), + region_id: pack_cell(cell), + cloud, + } } -/// One cell's edge corridors as a message, the cell in its header seq. +/// One cell's edge corridors as a message, the cell in its explicit region_id. fn edges_message( cell: Cell, segments: Vec, voxel_size: f32, frame: &str, stamp: Time, -) -> Path { - let mut msg = build_segments_path(segments, voxel_size, frame, stamp); - msg.header.seq = pack_cell(cell); - msg +) -> LineSegments3DRegion { + let lines = build_segments(segments, voxel_size, frame, stamp); + LineSegments3DRegion { + header: lines.header.clone(), + region_id: pack_cell(cell), + lines, + } } /// True if within tolerance of the goal on the ground plane. @@ -812,10 +834,10 @@ fn is_at_goal(start: Xyz, goal: Xyz, tol: f32) -> bool { } fn same_stamp(a: &Time, b: &Time) -> bool { - a.sec == b.sec && a.nsec == b.nsec + a.sec == b.sec && a.nanosec == b.nanosec } -async fn publish_cloud(out: &Output, cloud: &PointCloud2) { +async fn publish_cloud(out: &Output, cloud: &T) { if let Err(e) = out.publish(cloud).await { error_throttled!( Duration::from_secs(1), @@ -826,7 +848,7 @@ async fn publish_cloud(out: &Output, cloud: &PointCloud2) { } } -async fn publish_path(out: &Output, msg: &Path) { +async fn publish_path(out: &Output, msg: &T) { if let Err(e) = out.publish(msg).await { error_throttled!( Duration::from_secs(1), @@ -839,7 +861,6 @@ async fn publish_path(out: &Output, msg: &Path) { fn header(frame_id: &str, stamp: Time) -> Header { Header { - seq: 0, stamp, frame_id: frame_id.into(), } @@ -886,30 +907,37 @@ fn build_path_from_waypoints(waypoints: &[(f32, f32, f32)], frame_id: &str, stam } } -/// Emit edges as alternating PoseStamped pairs with orientation.w carrying -/// the per-edge cost. -fn build_segments_path( +/// Emit actual weighted line segments rather than a Path-shaped legacy encoding. +fn build_segments( segments: Vec<(VoxelKey, VoxelKey, f32)>, voxel_size: f32, frame_id: &str, stamp: Time, -) -> Path { - let mut poses: Vec = Vec::with_capacity(segments.len() * 2); - for (a, b, cost) in segments { - let pa = surface_point_xyz(a.0, a.1, a.2, voxel_size); - let pb = surface_point_xyz(b.0, b.1, b.2, voxel_size); - poses.push(pose_stamped(pa, cost as f64, frame_id, stamp.clone())); - poses.push(pose_stamped(pb, cost as f64, frame_id, stamp.clone())); - } - Path { +) -> LineSegments3D { + let point = |key: VoxelKey| { + let p = surface_point_xyz(key.0, key.1, key.2, voxel_size); + Point { + x: p.0 as f64, + y: p.1 as f64, + z: p.2 as f64, + } + }; + LineSegments3D { header: header(frame_id, stamp), - poses, + segments: segments + .into_iter() + .map(|(a, b, weight)| LineSegment3D { + start: point(a), + end: point(b), + weight: weight as f64, + }) + .collect(), } } /// Like `build_pc2_xyz` plus an `intensity` float carrying the cell's wall clearance. fn build_pc2_xyzi(points: &[Xyzi], frame_id: &str, stamp: Time) -> PointCloud2 { - let n = points.len() as i32; + let n = points.len() as u32; let mut data = Vec::with_capacity(points.len() * 16); for &(x, y, z, i) in points { data.extend_from_slice(&x.to_le_bytes()); @@ -917,10 +945,10 @@ fn build_pc2_xyzi(points: &[Xyzi], frame_id: &str, stamp: Time) -> PointCloud2 { data.extend_from_slice(&z.to_le_bytes()); data.extend_from_slice(&i.to_le_bytes()); } - let make_field = |name: &str, off: i32| PointField { + let make_field = |name: &str, off: u32| PointField { name: name.into(), offset: off, - datatype: PointField::FLOAT32 as u8, + datatype: PointField::FLOAT32, count: 1, }; PointCloud2 { @@ -942,17 +970,17 @@ fn build_pc2_xyzi(points: &[Xyzi], frame_id: &str, stamp: Time) -> PointCloud2 { } fn build_pc2_xyz(points: &[(f32, f32, f32)], frame_id: &str, stamp: Time) -> PointCloud2 { - let n = points.len() as i32; + let n = points.len() as u32; let mut data = Vec::with_capacity(points.len() * 12); for &(x, y, z) in points { data.extend_from_slice(&x.to_le_bytes()); data.extend_from_slice(&y.to_le_bytes()); data.extend_from_slice(&z.to_le_bytes()); } - let make_field = |name: &str, off: i32| PointField { + let make_field = |name: &str, off: u32| PointField { name: name.into(), offset: off, - datatype: PointField::FLOAT32 as u8, + datatype: PointField::FLOAT32, count: 1, }; PointCloud2 { @@ -968,13 +996,6 @@ fn build_pc2_xyz(points: &[(f32, f32, f32)], frame_id: &str, stamp: Time) -> Poi } } -struct ExtractError(&'static str); -impl std::fmt::Display for ExtractError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.write_str(self.0) - } -} - /// The xyz points of a cloud, or None with a warning when it cannot be read. fn cloud_points(msg: &PointCloud2, topic: &'static str) -> Option> { match extract_xyz(msg) { @@ -991,63 +1012,34 @@ fn cloud_points(msg: &PointCloud2, topic: &'static str) -> Option> { } } -fn extract_xyz(msg: &PointCloud2) -> Result, ExtractError> { - let mut x_off: Option = None; - let mut y_off: Option = None; - let mut z_off: Option = None; - for f in &msg.fields { - if f.datatype != PointField::FLOAT32 as u8 { - continue; - } - match f.name.as_str() { - "x" => x_off = Some(f.offset as usize), - "y" => y_off = Some(f.offset as usize), - "z" => z_off = Some(f.offset as usize), - _ => {} - } - } - let xo = x_off.ok_or(ExtractError("missing float32 x field"))?; - let yo = y_off.ok_or(ExtractError("missing float32 y field"))?; - let zo = z_off.ok_or(ExtractError("missing float32 z field"))?; - - let n = (msg.width as usize) * (msg.height as usize); - let step = msg.point_step as usize; - if step == 0 { - return Err(ExtractError("point_step is 0")); - } - if msg.data.len() < n * step { - return Err(ExtractError( - "data buffer shorter than width*height*point_step", - )); - } - if xo + 4 > step || yo + 4 > step || zo + 4 > step { - return Err(ExtractError( - "xyz field offsets do not fit within point_step", - )); - } - if msg.is_bigendian { - return Err(ExtractError("big-endian point data not supported")); - } - - let mut out = Vec::with_capacity(n); - for i in 0..n { - let base = i * step; - let x = read_f32_le(&msg.data, base + xo); - let y = read_f32_le(&msg.data, base + yo); - let z = read_f32_le(&msg.data, base + zo); - if x.is_finite() && y.is_finite() && z.is_finite() { - out.push((x, y, z)); - } - } - Ok(out) +fn extract_xyz(msg: &PointCloud2) -> Result, ExtractError> { + let offset = |name: &str| { + msg.fields + .iter() + .find(|f| f.name == name && f.datatype == PointField::FLOAT32 && f.count == 1) + .map(|f| f.offset as usize) + .ok_or(ExtractError("missing scalar float32 XYZ field")) + }; + Ok(extract_xyz_rows( + &msg.data, + [msg.width as usize, msg.height as usize], + [msg.point_step as usize, msg.row_step as usize], + [offset("x")?, offset("y")?, offset("z")?], + msg.is_bigendian, + )? + .into_iter() + .map(|[x, y, z]| (x, y, z)) + .collect()) } -#[inline] -fn read_f32_le(buf: &[u8], off: usize) -> f32 { - let bytes: [u8; 4] = buf[off..off + 4] - .try_into() - .expect("bounds checked by caller"); - f32::from_le_bytes(bytes) +fn now() -> Time { + let duration = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default(); + Time { + sec: duration.as_secs().min(i32::MAX as u64) as i32, + nanosec: duration.subsec_nanos(), + } } #[cfg(test)] @@ -1063,32 +1055,43 @@ mod tests { fn stamped(stamp: Time) -> Header { Header { stamp, - ..Default::default() + frame_id: String::new(), } } - - fn bounds_at(stamp: Time) -> PoseStamped { - PoseStamped { + fn bounds_at(stamp: Time) -> RegionBounds { + RegionBounds { header: stamped(stamp), - ..Default::default() + region_id: 0, + center: Point { + x: 0., + y: 0., + z: 0., + }, + radius: 0., + z_min: 0., + z_max: 0., } } - fn cloud_at(stamp: Time) -> PointCloud2 { - PointCloud2 { - header: stamped(stamp), - ..Default::default() + build_pc2_xyz(&[], "", stamp) + } + fn region_at(stamp: Time) -> PointCloudRegion { + let cloud = cloud_at(stamp); + PointCloudRegion { + header: cloud.header.clone(), + region_id: 0, + cloud, } } #[test] fn stamps_paired_only_when_both_present_and_stamps_match() { - let s = Time { sec: 2, nsec: 3 }; + let s = Time { sec: 2, nanosec: 3 }; let b = bounds_at(s.clone()); let c = cloud_at(s); assert!(stamps_paired(Some(&b), Some(&c))); - let other = cloud_at(Time { sec: 2, nsec: 4 }); + let other = cloud_at(Time { sec: 2, nanosec: 4 }); assert!(!stamps_paired(Some(&b), Some(&other))); assert!(!stamps_paired(Some(&b), None)); @@ -1097,102 +1100,158 @@ mod tests { } #[test] - fn seed_pairs_match_on_seq_whichever_side_lands_first() { - let stamp = Time { sec: 2, nsec: 3 }; + fn seed_pairs_match_on_region_id_whichever_side_lands_first() { + let stamp = Time { sec: 2, nanosec: 3 }; let mut pairs = SeedPairs::default(); let mut b1 = bounds_at(stamp.clone()); - b1.header.seq = 1; + b1.region_id = 1; let mut b2 = bounds_at(stamp.clone()); - b2.header.seq = 2; - let mut c1 = cloud_at(stamp.clone()); - c1.header.seq = 1; - let mut c2 = cloud_at(stamp); - c2.header.seq = 2; + b2.region_id = 2; + let mut c1 = region_at(stamp.clone()); + c1.region_id = 1; + let mut c2 = region_at(stamp); + c2.region_id = 2; assert!(pairs.pair_bounds(b1).is_none()); assert!(pairs.pair_bounds(b2).is_none()); let first = pairs .pair_cloud(c1) .expect("region 1 pairs with its own bounds"); - assert_eq!(first.bounds.header.seq, 1); + assert_eq!(first.bounds.region_id, 1); let second = pairs .pair_cloud(c2) .expect("region 2 pairs with its own bounds"); - assert_eq!(second.bounds.header.seq, 2); + assert_eq!(second.bounds.region_id, 2); assert!(pairs.clouds.is_empty() && pairs.bounds.is_empty()); } #[test] fn seed_pairs_match_cloud_first_and_interleaved() { - let stamp = Time { sec: 2, nsec: 3 }; + let stamp = Time { sec: 2, nanosec: 3 }; let mut pairs = SeedPairs::default(); - let mut c1 = cloud_at(stamp.clone()); - c1.header.seq = 1; + let mut c1 = region_at(stamp.clone()); + c1.region_id = 1; let mut b1 = bounds_at(stamp.clone()); - b1.header.seq = 1; + b1.region_id = 1; assert!(pairs.pair_cloud(c1).is_none()); assert_eq!( pairs .pair_bounds(b1) .expect("cloud first pairs") .cloud - .header - .seq, + .region_id, 1 ); // Interleaved: cloud 2, bounds 3, bounds 2, cloud 3. - let mut c2 = cloud_at(stamp.clone()); - c2.header.seq = 2; + let mut c2 = region_at(stamp.clone()); + c2.region_id = 2; let mut b3 = bounds_at(stamp.clone()); - b3.header.seq = 3; + b3.region_id = 3; let mut b2 = bounds_at(stamp.clone()); - b2.header.seq = 2; - let mut c3 = cloud_at(stamp); - c3.header.seq = 3; + b2.region_id = 2; + let mut c3 = region_at(stamp); + c3.region_id = 3; assert!(pairs.pair_cloud(c2).is_none()); assert!(pairs.pair_bounds(b3).is_none()); - assert_eq!(pairs.pair_bounds(b2).expect("region 2").cloud.header.seq, 2); - assert_eq!(pairs.pair_cloud(c3).expect("region 3").bounds.header.seq, 3); + assert_eq!(pairs.pair_bounds(b2).expect("region 2").cloud.region_id, 2); + assert_eq!(pairs.pair_cloud(c3).expect("region 3").bounds.region_id, 3); assert!(pairs.clouds.is_empty() && pairs.bounds.is_empty()); } #[test] - fn region_messages_carry_the_cell_in_their_headers_even_when_empty() { + fn region_messages_carry_the_cell_explicitly_even_when_empty() { let cell = (1, -2); - let surface = surface_message(cell, vec![((1, 2, 0), 0.5)], 0.1, "odom", Time::default()); + let surface = surface_message( + cell, + vec![((1, 2, 0), 0.5)], + 0.1, + "odom", + Time { sec: 0, nanosec: 0 }, + ); let segments = vec![((1, 2, 0), (2, 2, 0), 1.5)]; - let edges = edges_message(cell, segments, 0.1, "odom", Time::default()); - assert_eq!(surface.header.seq, pack_cell(cell)); - assert_eq!(edges.header.seq, pack_cell(cell)); - assert_eq!((surface.width, edges.poses.len()), (1, 2)); + let edges = edges_message(cell, segments, 0.1, "odom", Time { sec: 0, nanosec: 0 }); + assert_eq!(surface.region_id, pack_cell(cell)); + assert_eq!(edges.region_id, pack_cell(cell)); + assert_eq!((surface.cloud.width, edges.lines.segments.len()), (1, 1)); - let surface = surface_message(cell, Vec::new(), 0.1, "odom", Time::default()); - let edges = edges_message(cell, Vec::new(), 0.1, "odom", Time::default()); - assert_eq!(surface.header.seq, pack_cell(cell)); - assert_eq!(edges.header.seq, pack_cell(cell)); - assert_eq!((surface.width, edges.poses.len()), (0, 0)); + let surface = surface_message(cell, Vec::new(), 0.1, "odom", Time { sec: 0, nanosec: 0 }); + let edges = edges_message(cell, Vec::new(), 0.1, "odom", Time { sec: 0, nanosec: 0 }); + assert_eq!(surface.region_id, pack_cell(cell)); + assert_eq!(edges.region_id, pack_cell(cell)); + assert_eq!((surface.cloud.width, edges.lines.segments.len()), (0, 0)); } #[test] fn an_unpaired_seed_half_expires_after_the_timeout() { - let stamp = Time { sec: 2, nsec: 3 }; + let stamp = Time { sec: 2, nanosec: 3 }; let mut pairs = SeedPairs::default(); let mut b1 = bounds_at(stamp.clone()); - b1.header.seq = 1; + b1.region_id = 1; assert!(pairs.pair_bounds(b1).is_none()); pairs.expire_before(Instant::now() + SEED_PAIR_TIMEOUT / 2); assert_eq!(pairs.bounds.len(), 1); pairs.expire_before(Instant::now() + SEED_PAIR_TIMEOUT * 2); assert!(pairs.bounds.is_empty()); - let mut c1 = cloud_at(stamp); - c1.header.seq = 1; + let mut c1 = region_at(stamp); + c1.region_id = 1; assert!( pairs.pair_cloud(c1).is_none(), "the expired bounds no longer pair" ); } + #[test] + fn canonical_edges_keep_geometry_weight_and_region_id_in_cdr() { + let stamp = Time { + sec: 1_700_000_000, + nanosec: 123_456_789, + }; + let message = edges_message( + (-1, 2), + vec![((1, 2, 0), (2, 2, 0), 1.5)], + 0.1, + "odom", + stamp.clone(), + ); + for little in [false, true] { + let bytes = dimos_module::cdr::encode_endian(&message, little).unwrap(); + let received: LineSegments3DRegion = dimos_module::cdr::decode(&bytes).unwrap(); + assert_eq!(received.region_id, pack_cell((-1, 2))); + assert_eq!(received.header.stamp, stamp); + assert_eq!(received.lines.header, received.header); + let edge = &received.lines.segments[0]; + assert!((edge.start.x - 0.15).abs() < 1e-6 && (edge.end.x - 0.25).abs() < 1e-6); + assert_eq!(edge.weight, 1.5); + } + } + + #[test] + fn organized_big_endian_clouds_use_declared_row_stride() { + let mut cloud = build_pc2_xyz( + &[(1., 2., 3.), (4., 5., 6.)], + "map", + Time { sec: 0, nanosec: 0 }, + ); + cloud.width = 1; + cloud.height = 2; + cloud.row_step = 16; + cloud.is_bigendian = true; + cloud.data.clear(); + for point in [[1_f32, 2., 3.], [4., 5., 6.]] { + for value in point { + cloud.data.extend(value.to_be_bytes()); + } + cloud.data.extend([0; 4]); + } + assert_eq!( + extract_xyz(&cloud).unwrap(), + vec![(1., 2., 3.), (4., 5., 6.)] + ); + cloud.data.pop(); + assert!(extract_xyz(&cloud).is_err()); + } + #[test] fn seed_progress_counts_applied_regions() { const REGIONS: usize = 20; diff --git a/dimos/navigation/global_planner/mls_planner/test_generated_native.py b/dimos/navigation/global_planner/mls_planner/test_generated_native.py new file mode 100644 index 00000000000..c64aff1bc9d --- /dev/null +++ b/dimos/navigation/global_planner/mls_planner/test_generated_native.py @@ -0,0 +1,259 @@ +# 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 contextlib import ExitStack +from copy import deepcopy +import json +import os +from pathlib import Path as FilePath +import socket +from threading import Event +import time + +import numpy as np +import pytest +import zenoh + +from dimos.core.global_config import global_config +from dimos.core.transport import ZenohTransport +from dimos.mapping.ray_tracing.module import RayTracingVoxelMap +from dimos.message_utils.pointcloud import pointcloud_from_xyz +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.dimos_msgs.msg import LineSegments3DRegion, PointCloudRegion, RegionBounds +from dimos.messages.geometry_msgs.msg import ( + Point, + PointStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage +from dimos.navigation.global_planner.mls_planner.mls_planner_native import MLSPlannerNative +from dimos.protocol.pubsub.impl.zenohpubsub import Topic as ZenohTopic +from dimos.protocol.rpc.pubsubrpc import LCMRPC +from dimos.protocol.service.zenohservice import ZenohConfig + + +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +def test_mls_native_ports_reuse_canonical_region_and_geometry_types(): + module = MLSPlannerNative(rpc_transport=MemoryRPC) + try: + assert module.region_bounds.type is RegionBounds and module.seed_bounds.type is RegionBounds + assert ( + module.seed_map.type is PointCloudRegion and module.surface_map.type is PointCloudRegion + ) + assert module.node_edges.type is LineSegments3DRegion + assert module.goal.type is PointStamped and module.path.type is Path + finally: + module.stop() + + +@pytest.mark.native_e2e +def test_actual_native_mapper_to_mls_cdr_route_and_weighted_viewer_regions(monkeypatch): + paths = [ + os.environ.get("DIMOS_NATIVE_RAY_TRACING_BINARY"), + os.environ.get("DIMOS_NATIVE_MLS_BINARY"), + ] + assert all(paths), "Set both DIMOS_NATIVE_RAY_TRACING_BINARY and DIMOS_NATIVE_MLS_BINARY" + mapper_binary, planner_binary = [str(FilePath(path).resolve()) for path in paths] + assert all(FilePath(path).is_file() for path in [mapper_binary, planner_binary]) + monkeypatch.setattr(global_config, "transport", "zenoh") + with socket.socket() as reservation: + reservation.bind(("127.0.0.1", 0)) + endpoint = f"tcp/127.0.0.1:{reservation.getsockname()[1]}" + config = zenoh.Config() + for key, value in { + "mode": "router", + "listen/endpoints": [endpoint], + "scouting/multicast/enabled": False, + "scouting/gossip/enabled": False, + }.items(): + config.insert_json5(key, json.dumps(value)) + with ExitStack() as cleanup: + cleanup.enter_context(zenoh.open(config)) + session = ZenohConfig( + mode="client", + connect=[endpoint], + listen=[], + scouting=False, + multicast=False, + gossip=False, + connect_timeout=5, + ) + mapper = RayTracingVoxelMap( + rpc_transport=MemoryRPC, + executable=mapper_binary, + source_dir=None, + build_command=None, + session=session, + voxel_size=0.2, + support_min=0, + min_health=0, + max_health=1, + emit_every=1, + global_emit_every=0, + viz_emit_every=0, + world_frame="odom", + ) + planner = MLSPlannerNative( + rpc_transport=MemoryRPC, + executable=planner_binary, + source_dir=None, + build_command=None, + session=session, + voxel_size=0.2, + robot_height=1.0, + start_z_offset_m=1.0, + world_frame="odom", + base_frame="base_link", + viz_publish_hz=10.0, + ) + wires = {} + wire_types = {} + shared = {"local_map", "region_bounds", "seed_map", "seed_bounds", "tf"} + for owner, name in [ + (mapper, name) + for name in [ + "lidar", + "loaded_map", + "tf", + "tf_static", + "global_map", + "local_map", + "local_map_fine", + "region_bounds", + "map_regions", + "seed_map", + "seed_bounds", + ] + ] + [ + (planner, name) + for name in [ + "global_map", + "local_map", + "region_bounds", + "seed_map", + "seed_bounds", + "goal", + "tf", + "tf_static", + "path", + "surface_map", + "nodes", + "node_edges", + ] + ]: + key = name if name in shared else ("map/" if owner is mapper else "plan/") + name + port = getattr(owner, name) + if key not in wires: + wire_types[key] = port.type + wires[key] = ZenohTransport( + ZenohTopic("native/mls-chain/" + key, port.type), **session.model_dump() + ) + cleanup.callback(wires[key].stop) + assert wire_types[key] is port.type + port.transport = wires[key] + cleanup.callback(mapper.stop) + cleanup.callback(planner.stop) + ready, surface_ready, edges_ready = Event(), Event(), Event() + all_ready = Event() + routes, surfaces, edges = [], [], [] + + def mark_ready(): + if ready.is_set() and surface_ready.is_set() and edges_ready.is_set(): + all_ready.set() + + def route(value): + if len(value.poses) >= 2: + routes.append(value) + ready.set() + mark_ready() + + def surface(value): + if value.cloud.width: + surfaces.append(value) + surface_ready.set() + mark_ready() + + def edge(value): + if value.lines.segments: + edges.append(value) + edges_ready.set() + mark_ready() + + cleanup.callback(wires["plan/path"].subscribe(route)) + cleanup.callback(wires["plan/surface_map"].subscribe(surface)) + cleanup.callback(wires["plan/node_edges"].subscribe(edge)) + planner.start() + mapper.start() + coords = np.arange(-3.0, 3.0, 0.1, dtype=np.float32) + xx, yy = np.meshgrid(coords, coords) + # Sensor pose is (-2,-2,1.1); register these points onto the world floor. + points = np.stack( + [xx.ravel() + 2, yy.ravel() + 2, np.full(xx.size, -1.1, dtype=np.float32)], axis=-1 + ) + deadline = time.monotonic() + 20 + sequence = 0 + while not all_ready.is_set() and time.monotonic() < deadline: + header = Header( + frame_id="odom", stamp=time_from_nanoseconds(1700000000123456789 + sequence) + ) + sequence += 1 + wires["tf"].publish( + TFMessage( + transforms=[ + TransformStamped( + header=deepcopy(header), + child_frame_id=child, + transform=Transform( + translation=Vector3(x=-2, y=-2, z=1.1), + rotation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) + for child in ["lidar", "base_link"] + ] + ) + ) + sensor_header = deepcopy(header) + sensor_header.frame_id = "lidar" + wires["map/lidar"].publish(pointcloud_from_xyz(points, header=sensor_header)) + wires["plan/goal"].publish(PointStamped(header=header, point=Point(x=2, y=2, z=0))) + all_ready.wait(0.15) + for module in [mapper, planner]: + assert module._process is not None and module._process.poll() is None + assert ready.is_set() and surface_ready.is_set() and edges_ready.is_set(), ( + "No native CDR route and weighted regions before deadline" + ) + path = routes[0] + assert type(path) is Path and path.header.frame_id == "odom" + assert all(p.header == path.header and type(p.pose.position) is Point for p in path.poses) + assert ( + np.linalg.norm([path.poses[-1].pose.position.x - 2, path.poses[-1].pose.position.y - 2]) + < 0.4 + ) + assert ( + type(surfaces[0]) is PointCloudRegion and surfaces[0].header == surfaces[0].cloud.header + ) + assert type(edges[0]) is LineSegments3DRegion and edges[0].header == edges[0].lines.header + assert all( + np.isfinite(segment.weight) and segment.weight >= 0 + for segment in edges[0].lines.segments + ) diff --git a/dimos/navigation/global_planner/mls_planner/test_transformer.py b/dimos/navigation/global_planner/mls_planner/test_transformer.py index 875e62ae31e..9e2118adf24 100644 --- a/dimos/navigation/global_planner/mls_planner/test_transformer.py +++ b/dimos/navigation/global_planner/mls_planner/test_transformer.py @@ -20,8 +20,13 @@ pytest.importorskip("dimos_mls_planner") +from dimos_cdrgen.registry import encode + from dimos.memory.type.observation import Observation -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import pointcloud_from_xyz +from dimos.message_utils.time import time_from_nanoseconds, to_nanoseconds +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header from dimos.navigation.global_planner.mls_planner.transformer import MLSPlan @@ -35,7 +40,9 @@ def _obs( ts=0.0, pose=pose, tags={"region_bounds": region_bounds}, - _data=PointCloud2.from_numpy(points), + _data=pointcloud_from_xyz( + points, header=Header(frame_id="world", stamp=time_from_nanoseconds(0)) + ), ) @@ -70,7 +77,9 @@ def test_poseless_obs_is_skipped() -> None: ts=0.0, pose=None, tags={"region_bounds": (0.0, 0.0, 5.0, -1.0, 2.0)}, - _data=PointCloud2.from_numpy(points), + _data=pointcloud_from_xyz( + points, header=Header(frame_id="world", stamp=time_from_nanoseconds(0)) + ), ) posed = _obs(points, pose=(-2.0, -2.0, 1.0), region_bounds=(0.0, 0.0, 5.0, -1.0, 2.0)) @@ -106,3 +115,18 @@ def test_no_route_yields_empty_path_with_planned_false() -> None: assert out.tags["planned"] is False assert out.data.poses == [] + + +@pytest.mark.parametrize("stamp_ns", [-1, 0, 1700000000123456789]) +def test_generated_path_preserves_input_header_nanoseconds_without_aliasing(stamp_ns): + obs = _obs(_flat_floor(), pose=(-2.0, -2.0, 1.0), region_bounds=(0.0, 0.0, 5.0, -1.0, 2.0)) + obs.data.header.stamp = time_from_nanoseconds(stamp_ns) + before = encode(obs.data) + [result] = list(MLSPlan(goal=(2.0, 2.0, 0.0), voxel_size=0.2, robot_height=1.0)(iter([obs]))) + assert len(result.data.poses) >= 2 + assert to_nanoseconds(result.data.header.stamp) == stamp_ns + assert all(p.header == result.data.header for p in result.data.poses) + result.data.header.frame_id = "changed" + assert obs.data.header.frame_id == "world" + assert result.data.poses[0].header.frame_id == "world" + assert encode(obs.data) == before diff --git a/dimos/navigation/global_planner/mls_planner/test_viz.py b/dimos/navigation/global_planner/mls_planner/test_viz.py index 01457cf2c7d..116f55156f0 100644 --- a/dimos/navigation/global_planner/mls_planner/test_viz.py +++ b/dimos/navigation/global_planner/mls_planner/test_viz.py @@ -14,10 +14,28 @@ from functools import partial +from dimos_cdrgen.registry import decode, encode import numpy as np -from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import pointcloud_from_xyz +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.dimos_msgs.msg import ( + LineSegment3D, + LineSegments3D, + LineSegments3DRegion, + PointCloudRegion, +) +from dimos.messages.geometry_msgs.msg import Point +from dimos.messages.std_msgs.msg import Header + + +def _region(points): + header = Header(frame_id="odom", stamp=time_from_nanoseconds(0)) + return PointCloudRegion( + header=header, region_id=0, cloud=pointcloud_from_xyz(points, header=header) + ) + + from dimos.navigation.global_planner.mls_planner import viz @@ -62,19 +80,22 @@ def test_graph_nodes_are_lifted_off_the_surface() -> None: def test_region_renders_are_static_and_an_empty_cell_still_lands() -> None: - cell = PointCloud2.from_numpy(np.array([[1.0, 1.0, 0.0]], dtype=np.float32)) - cell.seq = 1 << 16 + cell = _region(np.array([[1.0, 1.0, 0.0]], dtype=np.float32)) + cell.region_id = 1 << 16 (path, arch, static) = viz.render_surface_region(cell, 0.1, 0.1, 1.0)[0] assert path == "world/surface_map/1_0" and static assert len(arch.positions.as_arrow_array()) == 1 - emptied = PointCloud2.from_numpy(np.zeros((0, 3), dtype=np.float32)) - emptied.seq = 1 << 16 + emptied = _region(np.zeros((0, 3), dtype=np.float32)) + emptied.region_id = 1 << 16 (path, arch, static) = viz.render_surface_region(emptied, 0.1, 0.1, 1.0)[0] assert path == "world/surface_map/1_0" and static assert len(arch.positions.as_arrow_array()) == 0 - edges = LineSegments3D(segments=np.zeros((0, 2, 3)), seq=(2 << 16) | 3) + header = Header(frame_id="odom", stamp=time_from_nanoseconds(0)) + edges = LineSegments3DRegion( + header=header, region_id=(2 << 16) | 3, lines=LineSegments3D(header=header, segments=[]) + ) (path, arch, static) = viz.render_edge_region(edges)[0] assert path == "world/node_edges/2_3" and static assert arch.strips.as_arrow_array().to_pylist() == [] @@ -93,3 +114,27 @@ def test_overrides_follow_the_planner_publish_rate() -> None: } assert on["world/nodes"] is viz.render_nodes assert on["world/node_edges"] is viz.render_edge_region + + +def test_cdr_weighted_region_renders_geometry_and_keeps_source_immutable(): + header = Header(frame_id="odom", stamp=time_from_nanoseconds(1700000000123456789)) + value = LineSegments3DRegion( + header=header, + region_id=(-2 << 16) | 3, + lines=LineSegments3D( + header=header, + segments=[ + LineSegment3D(start=Point(x=1, y=2, z=0), end=Point(x=2, y=3, z=0), weight=0.75) + ], + ), + ) + payload = encode(value) + received = decode(payload, LineSegments3DRegion) + [(entity, arch, static)] = viz.render_edge_region(received) + assert entity == "world/node_edges/-2_3" and static + np.testing.assert_allclose( + arch.strips.as_arrow_array().to_pylist(), + [[[1, 2, viz.GRAPH_Z_LIFT], [2, 3, viz.GRAPH_Z_LIFT]]], + atol=1e-6, + ) + assert encode(received) == payload diff --git a/dimos/navigation/global_planner/mls_planner/transformer.py b/dimos/navigation/global_planner/mls_planner/transformer.py index 9937c260006..f78323ed794 100644 --- a/dimos/navigation/global_planner/mls_planner/transformer.py +++ b/dimos/navigation/global_planner/mls_planner/transformer.py @@ -14,20 +14,24 @@ from __future__ import annotations +from copy import deepcopy import time from typing import TYPE_CHECKING, Any +import numpy as np + from dimos.memory.transform import Transformer -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import pointcloud_xyz +from dimos.messages.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header from dimos.navigation.global_planner.mls_planner.mls_planner import MLSPlanner from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from collections.abc import Iterator - import numpy as np from numpy.typing import NDArray from dimos.memory.type.observation import Observation @@ -51,19 +55,22 @@ def __init__( self.robot_height = robot_height self._planner_kwargs = planner_kwargs - def _path_from_waypoints(self, waypoints: NDArray[np.float32] | None, ts: float) -> Path: - poses: list[PoseStamped] = [] + def _path_from_waypoints(self, waypoints: NDArray[np.float32] | None, source: Header) -> Path: + header = deepcopy(source) + header.frame_id = "world" + poses = [] if waypoints is not None: for x, y, z in waypoints: poses.append( PoseStamped( - ts=ts, - frame_id="world", - position=(float(x), float(y), float(z)), - orientation=(0.0, 0.0, 0.0, 1.0), + header=deepcopy(header), + pose=Pose( + position=Point(x=float(x), y=float(y), z=float(z)), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), ) ) - return Path(ts=ts, frame_id="world", poses=poses) + return Path(header=header, poses=poses) def __call__( self, @@ -83,11 +90,18 @@ def __call__( ox, oy, radius, z_min, z_max = obs.tags["region_bounds"] t_update = time.perf_counter() - planner.update_region(obs.data.points_f32(), (ox, oy), radius, z_min, z_max, float(z)) + planner.update_region( + np.asarray(pointcloud_xyz(obs.data), dtype=np.float32), + (ox, oy), + radius, + z_min, + z_max, + float(z), + ) t_plan = time.perf_counter() waypoints = planner.plan(start, self.goal) t_done = time.perf_counter() - path = self._path_from_waypoints(waypoints, obs.ts) + path = self._path_from_waypoints(waypoints, obs.data.header) timings = { "update_ms": (t_plan - t_update) * 1000, diff --git a/dimos/navigation/global_planner/mls_planner/viz.py b/dimos/navigation/global_planner/mls_planner/viz.py index 5174a621fbf..27ca067f700 100644 --- a/dimos/navigation/global_planner/mls_planner/viz.py +++ b/dimos/navigation/global_planner/mls_planner/viz.py @@ -21,9 +21,11 @@ from typing import TYPE_CHECKING import numpy as np +import rerun as rr -from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import pointcloud_view, pointcloud_xyz +from dimos.messages.dimos_msgs.msg import LineSegments3D, LineSegments3DRegion, PointCloudRegion +from dimos.messages.sensor_msgs.msg import PointCloud2 from dimos.visualization.rerun.bridge import RerunEntry, keyed_by_seq, region_entity if TYPE_CHECKING: @@ -32,8 +34,6 @@ from dimos.visualization.rerun.bridge import RerunMulti, VisualOverride -# rerun is imported inside the renderers: it is heavy and only viewers load it. - SURFACE_MAP_ENTITY = "world/surface_map" NODES_ENTITY = "world/nodes" NODE_EDGES_ENTITY = "world/node_edges" @@ -61,8 +61,6 @@ def surface_points( clearance_clamp_m: float, ) -> Archetype: """Floor cells colored by wall clearance, dropping the ones the robot cannot fit on.""" - import rerun as rr - passable = clearance >= wall_clearance_m pts, clearance = pts[passable], clearance[passable] return rr.Points3D( @@ -74,7 +72,7 @@ def surface_points( @keyed_by_seq def render_surface_region( - msg: PointCloud2, + msg: PointCloudRegion, voxel_size: float, wall_clearance_m: float, clearance_clamp_m: float, @@ -85,24 +83,32 @@ def render_surface_region( Clearance rides the intensity channel, flat blue without it. Points outside z_band (world frame) are dropped, to show one storey of a multi-storey map. """ - seq = msg.seq - pts = msg.points_f32() - clearance = msg.intensities_f32() + cloud = msg.cloud + pts = np.asarray(pointcloud_xyz(cloud), dtype=np.float32) + records = pointcloud_view(cloud) + clearance = ( + np.asarray(records["intensity"].ravel(), dtype=np.float32) + if "intensity" in (records.dtype.names or ()) + else None + ) + cell: Archetype if clearance is None or len(clearance) != len(pts): if z_band is not None: - msg = msg.filter_by_height(min_height=z_band[0], max_height=z_band[1]) - cell = msg.to_rerun(voxel_size=voxel_size, colors=[40, 75, 130]) + pts = pts[(pts[:, 2] >= z_band[0]) & (pts[:, 2] <= z_band[1])] + cell = ( + rr.Points3D(pts, colors=[40, 75, 130], radii=voxel_size / 2) + if len(pts) + else rr.Points3D([]) + ) else: if z_band is not None: keep = (pts[:, 2] >= z_band[0]) & (pts[:, 2] <= z_band[1]) pts, clearance = pts[keep], clearance[keep] cell = surface_points(pts, clearance, voxel_size, wall_clearance_m, clearance_clamp_m) - return [RerunEntry(region_entity(SURFACE_MAP_ENTITY, seq), cell, static=True)] + return [RerunEntry(region_entity(SURFACE_MAP_ENTITY, msg.region_id), cell, static=True)] def graph_nodes(pts: NDArray[np.float32]) -> Archetype: - import rerun as rr - if len(pts) == 0: return rr.Points3D([]) lifted = pts.copy() @@ -111,27 +117,51 @@ def graph_nodes(pts: NDArray[np.float32]) -> Archetype: def render_nodes(msg: PointCloud2) -> Archetype: - return graph_nodes(msg.points_f32()) + return graph_nodes(np.asarray(pointcloud_xyz(msg), dtype=np.float32)) + + +def _weighted_lines(strips: NDArray[np.float32], weights: NDArray[np.float64]) -> Archetype: + if len(strips) == 0: + return rr.LineStrips3D([]) + positions = strips.astype(np.float32, copy=True) + positions[:, :, 2] += GRAPH_Z_LIFT + log_weights = np.log10(np.maximum(weights, 1e-6)) + low, high = float(log_weights.min()), float(log_weights.max()) + norm = (log_weights - low) / (high - low) if high > low else np.zeros_like(log_weights) + red = (255 * norm).astype(np.uint8) + green = (255 * (1 - norm)).astype(np.uint8) + colors = np.column_stack([red, green, np.full_like(red, 60), np.full_like(red, 220)]) + return rr.LineStrips3D(positions, colors=colors, radii=0.01) def graph_edges(edges: NDArray[np.float32]) -> Archetype: - """Edges as ``[x0, y0, z0, x1, y1, z1, cost]`` rows, colored green to red by cost.""" - segments = LineSegments3D( - segments=edges[:, :6].reshape(-1, 2, 3) if len(edges) else None, - weights=edges[:, 6] if len(edges) else None, + """Edges as [x0,y0,z0,x1,y1,z1,cost] rows, colored green to red by cost.""" + return _weighted_lines( + edges[:, :6].reshape(-1, 2, 3), np.asarray(edges[:, 6], dtype=np.float64) ) - return render_node_edges(segments) def render_node_edges(msg: LineSegments3D) -> Archetype: - return msg.to_rerun(z_offset=GRAPH_Z_LIFT, radii=0.01) + points = np.array( + [ + [[edge.start.x, edge.start.y, edge.start.z], [edge.end.x, edge.end.y, edge.end.z]] + for edge in msg.segments + ], + dtype=np.float32, + ).reshape(-1, 2, 3) + weights = np.array([edge.weight for edge in msg.segments], dtype=np.float64) + return _weighted_lines(points, weights) @keyed_by_seq -def render_edge_region(msg: LineSegments3D) -> RerunMulti: - """One cell of the edge corridors on its own static entity.""" +def render_edge_region(msg: LineSegments3DRegion) -> RerunMulti: + """One keyed cell of weighted edge corridors, including empty clears.""" return [ - RerunEntry(region_entity(NODE_EDGES_ENTITY, msg.seq), render_node_edges(msg), static=True) + RerunEntry( + region_entity(NODE_EDGES_ENTITY, msg.region_id), + render_node_edges(msg.lines), + static=True, + ) ] diff --git a/dimos/navigation/global_planner/test_viz.py b/dimos/navigation/global_planner/test_viz.py index 1deb9ba4dcb..555df3bb3af 100644 --- a/dimos/navigation/global_planner/test_viz.py +++ b/dimos/navigation/global_planner/test_viz.py @@ -19,25 +19,59 @@ import rerun as rr from dimos.mapping.ray_tracing.viz import render_map_region -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.geometry import quaternion_from_euler +from dimos.message_utils.pointcloud import pointcloud_from_xyz +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.dimos_msgs.msg import PointCloudRegion +from dimos.messages.geometry_msgs.msg import Point, PointStamped, Pose, PoseStamped +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Header from dimos.navigation.global_planner import viz from dimos.navigation.global_planner.mls_planner.viz import render_surface_region from dimos.visualization.rerun.bridge import is_keyed_by_seq def test_goal_placeholder_is_not_drawn() -> None: - assert viz.render_goal(PointStamped(x=float("nan"), y=0.0, z=0.0)) is None - goal = viz.render_goal(PointStamped(x=1.0, y=2.0, z=0.0)) + assert ( + viz.render_goal( + PointStamped( + header=Header(frame_id="odom", stamp=time_from_nanoseconds(0)), + point=Point(x=float("nan"), y=0.0, z=0.0), + ) + ) + is None + ) + goal = viz.render_goal( + PointStamped( + header=Header(frame_id="odom", stamp=time_from_nanoseconds(0)), + point=Point(x=1.0, y=2.0, z=0.0), + ) + ) assert goal is not None assert goal.positions.as_arrow_array().to_pylist() == [[1.0, 2.0, 0.0]] def test_empty_path_keeps_the_last_one_drawn() -> None: - assert viz.render_path(Path(poses=[])) is None - path = viz.render_path(Path(poses=[PoseStamped(1.0, 0.0, 0.0), PoseStamped(2.0, 0.0, 0.0)])) + assert ( + viz.render_path( + Path(header=Header(frame_id="odom", stamp=time_from_nanoseconds(0)), poses=[]) + ) + is None + ) + path = viz.render_path( + Path( + header=Header(frame_id="odom", stamp=time_from_nanoseconds(0)), + poses=[ + PoseStamped( + header=Header(frame_id="odom", stamp=time_from_nanoseconds(0)), + pose=Pose( + position=Point(x=x, y=0, z=0), orientation=quaternion_from_euler(0, 0, 0) + ), + ) + for x in [1.0, 2.0] + ], + ) + ) assert path is not None lift = viz.PATH_Z_LIFT np.testing.assert_allclose( @@ -63,7 +97,11 @@ def test_bridge_config_pickles_for_the_workers() -> None: regions = overrides["world/map_regions"] assert regions.func is render_map_region assert regions.keywords == {"voxel_size": 0.08, "height_range": viz.HEIGHT_RANGE} - region = PointCloud2.from_numpy(np.array([[1.0, 1.0, 0.0]], dtype=np.float32)) - region.seq = (4 << 16) | 7 + header = Header(frame_id="map", stamp=time_from_nanoseconds(0)) + region = PointCloudRegion( + header=header, + region_id=(4 << 16) | 7, + cloud=pointcloud_from_xyz(np.array([[1.0, 1.0, 0.0]], dtype=np.float32), header=header), + ) (path, _, static_entry) = regions(region)[0] assert path == "world/map_regions/4_7" and static_entry diff --git a/dimos/navigation/global_planner/viz.py b/dimos/navigation/global_planner/viz.py index 21f720bbd31..db509d0bde0 100644 --- a/dimos/navigation/global_planner/viz.py +++ b/dimos/navigation/global_planner/viz.py @@ -32,8 +32,8 @@ from numpy.typing import NDArray from rerun._baseclasses import Archetype - from dimos.msgs.geometry_msgs.PointStamped import PointStamped - from dimos.msgs.nav_msgs.Path import Path + from dimos.messages.geometry_msgs.msg import PointStamped + from dimos.messages.nav_msgs.msg import Path from dimos.visualization.rerun.bridge import VisualOverride # rerun is imported inside the renderers: it is heavy and only viewers load it. @@ -66,7 +66,12 @@ def render_path(msg: Path) -> Archetype | None: """The planned route. An empty path means no route, so the last one stays drawn.""" if len(msg.poses) == 0: return None - return path_strip(np.array([[p.x, p.y, p.z] for p in msg.poses], dtype=np.float32)) + return path_strip( + np.array( + [[p.pose.position.x, p.pose.position.y, p.pose.position.z] for p in msg.poses], + dtype=np.float32, + ) + ) def goal_point(xyz: tuple[float, float, float]) -> Archetype: @@ -77,9 +82,9 @@ def goal_point(xyz: tuple[float, float, float]) -> Archetype: def render_goal(msg: PointStamped) -> Archetype | None: """The active goal. NaN is the movement manager's placeholder for none.""" - if any(math.isnan(v) for v in (msg.x, msg.y, msg.z)): + if any(math.isnan(v) for v in (msg.point.x, msg.point.y, msg.point.z)): return None - return goal_point((msg.x, msg.y, msg.z)) + return goal_point((msg.point.x, msg.point.y, msg.point.z)) def robot_body_box(length: float, width: float, height: float) -> Archetype: diff --git a/dimos/navigation/go2/loop_closure/module.py b/dimos/navigation/go2/loop_closure/module.py index fc72ecd34a1..56d73a41059 100644 --- a/dimos/navigation/go2/loop_closure/module.py +++ b/dimos/navigation/go2/loop_closure/module.py @@ -14,6 +14,7 @@ from __future__ import annotations +from copy import deepcopy import time from reactivex.disposable import Disposable @@ -21,13 +22,10 @@ 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.PoseArray import PoseArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.LineSegments3D import LineSegments3D -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.std_msgs.Float32 import Float32 -from dimos.msgs.std_msgs.Header import Header +from dimos.messages.dimos_msgs.msg import LineSegment3D, LineSegments3D +from dimos.messages.geometry_msgs.msg import Point, Pose, PoseArray, PoseStamped, Quaternion +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Float32, Header from dimos.navigation.go2.loop_closure.pgo import PGOConfig from dimos.navigation.go2.loop_closure.pgo_map import LIVE_PGO, PGOMap from dimos.robot.unitree.type.lidar import repair_stale_ts @@ -95,7 +93,7 @@ def _on_odom(self, pose: PoseStamped) -> None: def _on_lidar(self, cloud: PointCloud2) -> None: assert self._map is not None t0, loops = time.perf_counter(), self._map.n_loops - rebuilt = self._map.add(cloud, self._pose) + rebuilt = self._map.add(cloud, None if self._pose is None else self._pose.pose) self._frames += 1 emit = rebuilt or self._frames % self.config.emit_every == 0 global_map = self._map.global_map() if emit else None @@ -105,28 +103,35 @@ def _on_lidar(self, cloud: PointCloud2) -> None: self.global_map.publish(global_map) if self._map.n_keyframes != self._graph_keyframes: self._graph_keyframes = self._map.n_keyframes - self._publish_graph(cloud.ts) + self._publish_graph(cloud.header) if self.config.telemetry: - self.mapper_frame_ms.publish(Float32(frame_ms)) + self.mapper_frame_ms.publish(Float32(data=frame_ms)) if self._map.n_loops != loops: - self.pgo_loop_ms.publish(Float32(self._map.loop_ms)) + self.pgo_loop_ms.publish(Float32(data=self._map.loop_ms)) if rebuilt: - self.pgo_rebuild_ms.publish(Float32(self._map.rebuild_ms)) + self.pgo_rebuild_ms.publish(Float32(data=self._map.rebuild_ms)) - def _publish_graph(self, ts: float) -> None: + def _publish_graph(self, source: Header) -> None: assert self._map is not None - frame_id = self.config.frame_id + header = deepcopy(source) + header.frame_id = self.config.frame_id positions, quats = self._map.keyframe_poses() self.pgo_keyframes.publish( PoseArray( - Header(ts, frame_id), - [ - Pose(*xyz, *quat) + header=deepcopy(header), + poses=[ + Pose(position=Point(*xyz), orientation=Quaternion(*quat)) for xyz, quat in zip(positions.tolist(), quats.tolist(), strict=True) ], ) ) segments, scores = self._map.loop_segments() self.pgo_loops.publish( - LineSegments3D(ts=ts, frame_id=frame_id, segments=segments, weights=scores) + LineSegments3D( + header=header, + segments=[ + LineSegment3D(start=Point(*ends[0]), end=Point(*ends[1]), weight=float(score)) + for ends, score in zip(segments.tolist(), scores.tolist(), strict=True) + ], + ) ) diff --git a/dimos/navigation/go2/loop_closure/pgo.py b/dimos/navigation/go2/loop_closure/pgo.py index b50fd60e4d0..abd921269cf 100644 --- a/dimos/navigation/go2/loop_closure/pgo.py +++ b/dimos/navigation/go2/loop_closure/pgo.py @@ -45,6 +45,7 @@ from __future__ import annotations from collections.abc import Iterator +from copy import deepcopy from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, TypedDict, TypeVar, cast @@ -53,11 +54,29 @@ from dimos.memory.transform import Transformer from dimos.memory.type.observation import Observation -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.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.geometry import ( + compose_transform_values, + inverse_transform, + quaternion_array, + quaternion_from_matrix, + transform_from_matrix, + transform_matrix, + vector_array, +) +from dimos.message_utils.pointcloud import ( + concatenate_clouds, + pointcloud_from_xyz, + pointcloud_stamps, + pointcloud_to_open3d, + pointcloud_with_xyz, + pointcloud_xyz, + transform_cloud, + voxel_downsample_cloud, +) +from dimos.message_utils.time import time_from_seconds, to_seconds +from dimos.messages.geometry_msgs.msg import Pose, Quaternion, Transform, TransformStamped, Vector3 +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger @@ -137,8 +156,8 @@ class Keyframe: """ ts: float - local: Transform - optimized: Transform + local: TransformStamped + optimized: TransformStamped @dataclass(frozen=True) @@ -151,8 +170,8 @@ class LoopClosure: tighter match. """ - source: Transform - target: Transform + source: TransformStamped + target: TransformStamped score: float @@ -168,58 +187,65 @@ class PoseGraph(Transformer[Any, Any]): keyframes: tuple[Keyframe, ...] = () loops: tuple[LoopClosure, ...] = () - def correct(self, pose: Transform) -> Transform: - """Map a raw-world pose into the drift-corrected frame at ``pose.ts``.""" - # Transform.__add__ composes: (T_corr + T_raw) applies T_corr after T_raw. - return self.correction_at(pose.ts) + pose + def correct(self, pose: TransformStamped) -> TransformStamped: + """Map a raw-world pose into the drift-corrected frame at its Header time.""" + correction = self.correction_at(to_seconds(pose.header.stamp)) + return TransformStamped( + header=Header(stamp=deepcopy(pose.header.stamp), frame_id=FRAME_WORLD_CORRECTED), + child_frame_id=pose.child_frame_id, + transform=compose_transform_values(correction.transform, pose.transform), + ) - def correction_at(self, ts: float) -> Transform: + def correction_at(self, ts: float) -> TransformStamped: """The raw ``world_corrected <- world_raw`` transform at ``ts``. Useful when applying the correction to non-pose data (e.g. point - clouds): ``pointcloud.transform(graph.correction_at(obs.ts))``. + clouds): ``transform_cloud(cloud, graph.correction_at(obs.ts))``. """ R, t = self._corrections(np.array([ts])) - return Transform( - translation=Vector3(t[0]), - rotation=Quaternion.from_rotation_matrix(R[0].as_matrix()), - frame_id=FRAME_WORLD_CORRECTED, + return TransformStamped( + header=Header(stamp=time_from_seconds(float(ts)), frame_id=FRAME_WORLD_CORRECTED), child_frame_id=FRAME_WORLD_RAW, - ts=float(ts), + transform=Transform( + translation=Vector3(*t[0]), + rotation=quaternion_from_matrix(R[0].as_matrix()), + ), ) def place(self, cloud: PointCloud2) -> PointCloud2: """Move each point by the correction for its capture time, anchored at the latest keyframe. Capture time is the cloud's per-point ``stamps`` when present, else - ``cloud.ts``. The latest keyframe's correction maps to identity, so + the cloud Header time. The latest keyframe's correction maps to identity, so current data stays where raw odometry puts it and older data moves. """ - pts = cloud.points_f32() + pts = pointcloud_xyz(cloud) if not len(pts): return cloud - stamps = cloud.stamps_f64() + stamps = pointcloud_stamps(cloud) if stamps is None: - stamps = np.full(len(pts), cloud.ts) + stamps = np.full(len(pts), to_seconds(cloud.header.stamp)) unique, inverse = np.unique(stamps, return_inverse=True) R, t = self._corrections(unique) now_R, now_t = self._corrections(np.array([np.inf])) anchor = now_R[0].inv() R, t = anchor * R, anchor.apply(t - now_t[0]) - return PointCloud2.from_numpy( - R[inverse].apply(pts.astype(np.float64)) + t[inverse], - frame_id=cloud.frame_id, - timestamp=cloud.ts, - ) + return pointcloud_with_xyz(cloud, R[inverse].apply(pts) + t[inverse]) def __call__(self, upstream: Iterator[Observation[Any]]) -> Iterator[Observation[Any]]: """Rewrite obs.pose via :meth:`correct`; pass through pose-less obs unchanged.""" for obs in upstream: - ps = obs.pose_stamped - if ps is None: + values = obs.pose_tuple + if values is None: yield obs continue - raw_tf = Transform.from_pose(FRAME_BODY, ps) + raw_tf = TransformStamped( + header=Header(stamp=time_from_seconds(obs.ts), frame_id=FRAME_WORLD_RAW), + child_frame_id=FRAME_BODY, + transform=Transform( + translation=Vector3(*values[:3]), rotation=Quaternion(*values[3:]) + ), + ) yield obs.derive(data=obs.data, pose=self.correct(raw_tf)) def _corrections(self, ts: np.ndarray) -> tuple[Rotation, np.ndarray]: @@ -240,9 +266,12 @@ def _drifts(self) -> tuple[np.ndarray, Slerp, np.ndarray]: # Per-keyframe drift: world_corrected <- body <- world_raw. ts_list = [kf.ts for kf in self.keyframes] - drifts = [(kf.optimized + kf.local.inverse()) for kf in self.keyframes] - quat_list = [tf.rotation.to_numpy() for tf in drifts] - t_list = [tf.translation.to_numpy() for tf in drifts] + drifts = [ + compose_transform_values(kf.optimized.transform, inverse_transform(kf.local).transform) + for kf in self.keyframes + ] + quat_list = [quaternion_array(tf.rotation) for tf in drifts] + t_list = [vector_array(tf.translation) for tf in drifts] # Slerp needs ≥2 keyframes; pad a len==1 list with a near-duplicate. if len(ts_list) == 1: @@ -285,12 +314,12 @@ def __call__( last_kf_ts: float | None = None for obs in upstream: - pose = obs.pose + pose = obs.pose_tuple if pose is None: continue # Placeholder filter: zero translation OR uninitialized (all-zero) # quaternion. Identity rotation (qw=1) is valid and stays. - if pose.position.is_zero() or pose.orientation.is_zero(): + if not any(pose[:3]) or not any(pose[3:]): continue pgo.process(_obs_to_pose3(obs), obs.ts, obs.data) @@ -316,10 +345,13 @@ def __call__( def _obs_to_pose3(obs: Observation[Any]) -> gtsam.Pose3: """Convert an observation's pose to a `gtsam.Pose3`.""" - pose = obs.pose + import gtsam # type: ignore[import-not-found,import-untyped] + + pose = obs.pose_tuple if pose is None: raise LookupError("No pose set on this observation") - return _pose_to_pose3(pose) + x, y, z, qx, qy, qz, qw = pose + return gtsam.Pose3(gtsam.Rot3.Quaternion(qw, qx, qy, qz), gtsam.Point3(x, y, z)) def _pose_to_pose3(pose: Pose) -> gtsam.Pose3: @@ -404,20 +436,26 @@ def process( ts: float, world_cloud: PointCloud2, ) -> None: - if len(world_cloud) == 0: + if world_cloud.width * world_cloud.height == 0: return if not self._is_keyframe(local_pose): return # Unregister: lift world-frame scan back into body frame using its # odom pose, so PGO can re-project it via the optimized pose later. - body_cloud = world_cloud.transform( - _pose3_to_transform( - local_pose.inverse(), - ts=ts, - frame_id=FRAME_BODY, - child_frame_id=FRAME_WORLD_RAW, - ) - ).voxel_downsample(self._cfg.submap_resolution) + # ICP consumes geometry only; the incoming cloud remains unchanged. + geometry = pointcloud_from_xyz(pointcloud_xyz(world_cloud), header=world_cloud.header) + body_cloud = voxel_downsample_cloud( + transform_cloud( + geometry, + _pose3_to_transform( + local_pose.inverse(), + ts=ts, + frame_id=FRAME_BODY, + child_frame_id=world_cloud.header.frame_id, + ), + ), + self._cfg.submap_resolution, + ) self._add_keyframe(local_pose, ts, body_cloud) self._search_for_loops() self._smooth_and_update() @@ -539,26 +577,34 @@ def _get_submap(self, idx: int, half_range: int) -> PointCloud2: lo = max(0, idx - half_range) hi = min(len(self._key_poses) - 1, idx + half_range) if lo > hi: - return PointCloud2() - cloud = self._key_poses[lo].body_cloud.transform( + return pointcloud_from_xyz( + np.empty((0, 3)), + header=Header(stamp=time_from_seconds(0), frame_id=FRAME_WORLD_CORRECTED), + ) + cloud = transform_cloud( + self._key_poses[lo].body_cloud, _pose3_to_transform( self._key_poses[lo].optimized, ts=self._key_poses[lo].timestamp, frame_id=FRAME_WORLD_CORRECTED, child_frame_id=FRAME_BODY, - ) + ), ) for i in range(lo + 1, hi + 1): kp = self._key_poses[i] - cloud = cloud + kp.body_cloud.transform( - _pose3_to_transform( - kp.optimized, - ts=kp.timestamp, - frame_id=FRAME_WORLD_CORRECTED, - child_frame_id=FRAME_BODY, - ) + cloud = concatenate_clouds( + cloud, + transform_cloud( + kp.body_cloud, + _pose3_to_transform( + kp.optimized, + ts=kp.timestamp, + frame_id=FRAME_WORLD_CORRECTED, + child_frame_id=FRAME_BODY, + ), + ), ) - return cloud.voxel_downsample(self._cfg.submap_resolution) + return voxel_downsample_cloud(cloud, self._cfg.submap_resolution) def _search_for_loops(self) -> None: if len(self._key_poses) < self._cfg.min_keyframes_for_loop_search: @@ -682,23 +728,20 @@ def _pose3_to_transform( ts: float, frame_id: str, child_frame_id: str, -) -> Transform: +) -> TransformStamped: """PGO-internal: build a Transform from a Pose3.""" - t = np.asarray(pose.translation()) - return Transform( - translation=Vector3(float(t[0]), float(t[1]), float(t[2])), - rotation=Quaternion.from_rotation_matrix(pose.rotation().matrix()), - frame_id=frame_id, + return TransformStamped( + header=Header(stamp=time_from_seconds(ts), frame_id=frame_id), child_frame_id=child_frame_id, - ts=ts, + transform=transform_from_matrix(np.asarray(pose.matrix())), ) -def _transform_to_pose3(tf: Transform) -> gtsam.Pose3: +def _transform_to_pose3(tf: TransformStamped) -> gtsam.Pose3: """PGO-internal: build a Pose3 from a Transform.""" import gtsam # type: ignore[import-not-found,import-untyped] - return gtsam.Pose3(tf.to_matrix()) + return gtsam.Pose3(transform_matrix(tf.transform)) def _icp( @@ -708,8 +751,8 @@ def _icp( max_dist: float = 1.0, tol: float = 1e-6, min_inliers: int = 10, - init: Transform | None = None, -) -> tuple[Transform, float]: + init: TransformStamped | None = None, +) -> tuple[TransformStamped, float]: """Point-to-plane ICP using Open3D's tensor pipeline. Returns ``(tf, fitness)`` where ``fitness`` is mean squared inlier @@ -720,11 +763,15 @@ def _icp( import open3d as o3d # type: ignore[import-untyped] import open3d.core as o3c # type: ignore[import-untyped] - if len(source) < min_inliers or len(target) < min_inliers: - return Transform.identity(), float("inf") + if source.width * source.height < min_inliers or target.width * target.height < min_inliers: + return TransformStamped( + header=deepcopy(source.header), + child_frame_id=source.header.frame_id, + transform=Transform(translation=Vector3(0, 0, 0), rotation=Quaternion(0, 0, 0, 1)), + ), float("inf") - src_pcd = source.pointcloud_tensor - tgt_pcd = target.pointcloud_tensor + src_pcd = o3d.t.geometry.PointCloud.from_legacy(pointcloud_to_open3d(source)) + tgt_pcd = o3d.t.geometry.PointCloud.from_legacy(pointcloud_to_open3d(target)) # Normals on the target enable point-to-plane ICP — converges tighter # than point-to-point on indoor scenes (walls give unambiguous normals @@ -733,7 +780,7 @@ def _icp( device = src_pcd.device init_T = o3c.Tensor( - init.to_matrix() if init is not None else np.eye(4), + transform_matrix(init.transform) if init is not None else np.eye(4), dtype=o3c.float64, device=device, ) @@ -755,9 +802,17 @@ def _icp( ) if float(result.fitness) == 0.0: - return Transform.identity(), float("inf") + return TransformStamped( + header=deepcopy(source.header), + child_frame_id=source.header.frame_id, + transform=Transform(translation=Vector3(0, 0, 0), rotation=Quaternion(0, 0, 0, 1)), + ), float("inf") # Frames intentionally unlabeled: caller only reads .to_matrix(). - tf = Transform.from_matrix(result.transformation.numpy(), ts=source.ts) + tf = TransformStamped( + header=deepcopy(source.header), + child_frame_id=source.header.frame_id, + transform=transform_from_matrix(result.transformation.numpy()), + ) rmse = float(result.inlier_rmse) return tf, rmse * rmse diff --git a/dimos/navigation/go2/loop_closure/pgo_auto.py b/dimos/navigation/go2/loop_closure/pgo_auto.py index 1a5fe756d9f..cf06335d0f4 100644 --- a/dimos/navigation/go2/loop_closure/pgo_auto.py +++ b/dimos/navigation/go2/loop_closure/pgo_auto.py @@ -25,7 +25,7 @@ lidar: Stream[PointCloud2] # pose-stamped scans -> pgo_keyframes(...) -> Stream[Keyframe] # one per kf - -> keyframes_to_corrections(...) -> Stream[Transform] # world_corrected <- world_raw + -> keyframes_to_corrections(...) -> Stream[TransformStamped] # world_corrected <- world_raw -> apply_corrections(any_stream, corrections) -> Stream[T] # obs.pose shuffled `pgo_keyframes` runs ISAM2 + ICP loop closure across the input, emitting @@ -55,6 +55,7 @@ from __future__ import annotations from collections.abc import Callable, Iterable, Iterator +from copy import deepcopy from dataclasses import dataclass from typing import TYPE_CHECKING, Any, TypedDict, TypeVar @@ -64,10 +65,26 @@ from dimos.memory.store.memory import MemoryStore from dimos.memory.stream import Stream from dimos.memory.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.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.geometry import ( + compose_transform_values, + compose_transforms, + inverse_transform, + quaternion_from_matrix, + transform_matrix, +) +from dimos.message_utils.pointcloud import ( + concatenate_clouds, + pointcloud_from_xyz, + pointcloud_to_open3d, + pointcloud_xyz, + transform_cloud, + voxel_downsample_cloud, +) +from dimos.message_utils.time import time_from_seconds, to_seconds +from dimos.messages.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.navigation.go2.loop_closure.pgo import _pose3_to_transform, _transform_to_pose3 from dimos.protocol.service.spec import BaseConfig from dimos.utils.logging_config import setup_logger @@ -142,8 +159,8 @@ class Keyframe: """ ts: float - local: Transform - optimized: Transform + local: TransformStamped + optimized: TransformStamped @dataclass(frozen=True) @@ -156,8 +173,8 @@ class LoopClosure: tighter match. """ - source: Transform - target: Transform + source: TransformStamped + target: TransformStamped score: float @@ -199,27 +216,26 @@ def pgo_keyframes( return out -def keyframes_to_corrections(keyframes: Stream[Keyframe]) -> Stream[Transform]: +def keyframes_to_corrections(keyframes: Stream[Keyframe]) -> Stream[TransformStamped]: """Per-keyframe drift correction as Transform(world_corrected <- world_raw).""" mem = MemoryStore() - out: Stream[Transform] = mem.stream("corrections", Transform) + out: Stream[TransformStamped] = mem.stream("corrections", TransformStamped) for obs in keyframes: kf = obs.data # world_corrected <- body <- world_raw composes to world_corrected <- world_raw. - drift = kf.optimized + kf.local.inverse() + drift = compose_transforms(kf.optimized, inverse_transform(kf.local)) out.append(drift, ts=kf.ts) return out -def make_interpolator(corrections: Stream[Transform]) -> Callable[[float], Transform]: +def make_interpolator(corrections: Stream[TransformStamped]) -> Callable[[float], TransformStamped]: """Materialize corrections once; return a fast ts -> Transform lookup.""" ts_list: list[float] = [] R_list: list[np.ndarray] = [] t_list: list[np.ndarray] = [] for obs in corrections: R, t = _r_t_from_transform(obs.data) - # obs.ts is authoritative; obs.data.ts can be mutated by Transform's - # ts=0.0 -> time.time() fallback in its constructor. + # Indexed stream time remains authoritative for interpolation. ts_list.append(obs.ts) R_list.append(R) t_list.append(t) @@ -238,7 +254,7 @@ def make_interpolator(corrections: Stream[Transform]) -> Callable[[float], Trans R_stack = np.stack(R_list) t_stack = np.stack(t_list) - def interp(ts: float) -> Transform: + def interp(ts: float) -> TransformStamped: ts_clip = float(np.clip(ts, ts_arr[0], ts_arr[-1])) idx = int(np.searchsorted(ts_arr, ts_clip)) if idx == 0: @@ -261,7 +277,7 @@ def interp(ts: float) -> Transform: def apply_corrections( stream: Stream[T], - corrections: Stream[Transform], + corrections: Stream[TransformStamped], ) -> Stream[T]: """Shuffle obs.pose on `stream` by the interpolated correction at each obs.ts. @@ -272,13 +288,14 @@ def apply_corrections( def xf(upstream: Iterator[Observation[T]]) -> Iterator[Observation[T]]: for obs in upstream: - if obs.pose is None: + values = obs.pose_tuple + if values is None: yield obs continue - raw_tf = Transform.from_pose(FRAME_BODY, obs.pose_stamped) - # Transform.__add__ composes: (T_corr + T_raw) applies T_corr after T_raw. - # Observation normalizes Transform back to 7-tuple via __post_init__. - corrected = interp(obs.ts) + raw_tf + corrected = compose_transform_values( + interp(obs.ts).transform, + Transform(translation=Vector3(*values[:3]), rotation=Quaternion(*values[3:])), + ) yield obs.derive(data=obs.data, pose=corrected) return stream.transform(xf) @@ -297,19 +314,25 @@ def __init__( self, keyframes: tuple[Keyframe, ...], loops: tuple[LoopClosure, ...], - correct: Callable[[float], Transform], + correct: Callable[[float], TransformStamped], ) -> None: self.keyframes = keyframes self.loops = loops self._correct = correct - def correction_at(self, ts: float) -> Transform: + def correction_at(self, ts: float) -> TransformStamped: """Drift correction (world_corrected <- world_raw) at `ts`.""" return self._correct(ts) - def correct(self, pose: Transform) -> Transform: + def correct(self, pose: TransformStamped) -> TransformStamped: """Lift a world_raw pose into world_corrected.""" - return self._correct(pose.ts) + pose + return TransformStamped( + header=Header(stamp=deepcopy(pose.header.stamp), frame_id=FRAME_WORLD_CORRECTED), + child_frame_id=pose.child_frame_id, + transform=compose_transform_values( + self._correct(to_seconds(pose.header.stamp)).transform, pose.transform + ), + ) class PGO: @@ -333,27 +356,25 @@ def __call__( last_ts = keyframes[-1].ts else: # No valid keyframes (e.g. pose-less input): correction is identity. - correct = lambda ts: Transform( # noqa: E731 - frame_id=FRAME_WORLD_CORRECTED, child_frame_id=FRAME_WORLD_RAW, ts=ts - ) + correct = lambda ts: _transform_from_r_t(np.eye(3), np.zeros(3), ts=ts) # noqa: E731 last_ts = 0.0 yield Observation(ts=last_ts, _data=PoseGraph(keyframes, tuple(loops), correct)) -def _r_t_from_transform(tf: Transform) -> tuple[np.ndarray, np.ndarray]: - q = tf.rotation +def _r_t_from_transform(tf: TransformStamped) -> tuple[np.ndarray, np.ndarray]: + q = tf.transform.rotation R = Rotation.from_quat([q.x, q.y, q.z, q.w]).as_matrix() - t = np.array([tf.translation.x, tf.translation.y, tf.translation.z]) + t = np.array( + [tf.transform.translation.x, tf.transform.translation.y, tf.transform.translation.z] + ) return R, t -def _transform_from_r_t(R: np.ndarray, t: np.ndarray, *, ts: float) -> Transform: - return Transform( - translation=Vector3(float(t[0]), float(t[1]), float(t[2])), - rotation=Quaternion.from_rotation_matrix(R), - frame_id=FRAME_WORLD_CORRECTED, +def _transform_from_r_t(R: np.ndarray, t: np.ndarray, *, ts: float) -> TransformStamped: + return TransformStamped( + header=Header(stamp=time_from_seconds(ts), frame_id=FRAME_WORLD_CORRECTED), child_frame_id=FRAME_WORLD_RAW, - ts=ts, + transform=Transform(translation=Vector3(*t), rotation=quaternion_from_matrix(R)), ) @@ -419,15 +440,25 @@ def process( ts: float, world_cloud: PointCloud2, ) -> None: - if len(world_cloud) == 0: + if world_cloud.width * world_cloud.height == 0: return if not self._is_keyframe(local_pose): return # Unregister: lift world-frame scan back into body frame using its # odom pose, so PGO can re-project it via the optimized pose later. - body_cloud = world_cloud.transform( - _pose3_to_transform(local_pose.inverse(), ts=ts) - ).voxel_downsample(self._cfg.submap_resolution) + geometry = pointcloud_from_xyz(pointcloud_xyz(world_cloud), header=world_cloud.header) + body_cloud = voxel_downsample_cloud( + transform_cloud( + geometry, + _pose3_to_transform( + local_pose.inverse(), + ts=ts, + frame_id=FRAME_BODY, + child_frame_id=world_cloud.header.frame_id, + ), + ), + self._cfg.submap_resolution, + ) self._add_keyframe(local_pose, ts, body_cloud) self._search_for_loops() self._smooth_and_update() @@ -558,16 +589,34 @@ def _get_submap(self, idx: int, half_range: int) -> PointCloud2: lo = max(0, idx - half_range) hi = min(len(self._key_poses) - 1, idx + half_range) if lo > hi: - return PointCloud2() - cloud = self._key_poses[lo].body_cloud.transform( - _pose3_to_transform(self._key_poses[lo].optimized, ts=self._key_poses[lo].timestamp) + return pointcloud_from_xyz( + np.empty((0, 3)), + header=Header(stamp=time_from_seconds(0), frame_id=FRAME_WORLD_CORRECTED), + ) + cloud = transform_cloud( + self._key_poses[lo].body_cloud, + _pose3_to_transform( + self._key_poses[lo].optimized, + ts=self._key_poses[lo].timestamp, + frame_id=FRAME_WORLD_CORRECTED, + child_frame_id=FRAME_BODY, + ), ) for i in range(lo + 1, hi + 1): kp = self._key_poses[i] - cloud = cloud + kp.body_cloud.transform( - _pose3_to_transform(kp.optimized, ts=kp.timestamp) + cloud = concatenate_clouds( + cloud, + transform_cloud( + kp.body_cloud, + _pose3_to_transform( + kp.optimized, + ts=kp.timestamp, + frame_id=FRAME_WORLD_CORRECTED, + child_frame_id=FRAME_BODY, + ), + ), ) - return cloud.voxel_downsample(self._cfg.submap_resolution) + return voxel_downsample_cloud(cloud, self._cfg.submap_resolution) def _search_for_loops( self, @@ -668,7 +717,7 @@ def _collect(local_idxs: set[int], time_thresh: float) -> list[tuple[float, int] # already and extra loops add noise. Early-exit on a very tight # match to save redundant ICP runs in the common case. source = self._get_submap(cur_idx, source_half_range) - accepted: list[tuple[int, Transform, float]] = [] + accepted: list[tuple[int, TransformStamped, float]] = [] k = ( candidates_per_iter if candidates_per_iter is not None @@ -765,31 +814,6 @@ def _smooth_and_update(self) -> None: self._world_correction = last.optimized.compose(last.local.inverse()) -def _pose3_to_transform( - pose: gtsam.Pose3, - *, - ts: float, - frame_id: str = "", - child_frame_id: str = "", -) -> Transform: - """PGO-internal: build a Transform from a Pose3.""" - t = np.asarray(pose.translation()) - return Transform( - translation=Vector3(float(t[0]), float(t[1]), float(t[2])), - rotation=Quaternion.from_rotation_matrix(pose.rotation().matrix()), - frame_id=frame_id, - child_frame_id=child_frame_id, - ts=ts, - ) - - -def _transform_to_pose3(tf: Transform) -> gtsam.Pose3: - """PGO-internal: build a Pose3 from a Transform.""" - import gtsam # type: ignore[import-not-found,import-untyped] - - return gtsam.Pose3(tf.to_matrix()) - - def _icp( source: PointCloud2, target: PointCloud2, @@ -797,8 +821,8 @@ def _icp( max_dist: float = 1.0, tol: float = 1e-6, min_inliers: int = 10, - init: Transform | None = None, -) -> tuple[Transform, float]: + init: TransformStamped | None = None, +) -> tuple[TransformStamped, float]: """Point-to-plane ICP using Open3D's tensor pipeline. Returns ``(tf, fitness)`` where ``fitness`` is mean squared inlier @@ -809,11 +833,13 @@ def _icp( import open3d as o3d # type: ignore[import-untyped] import open3d.core as o3c # type: ignore[import-untyped] - if len(source) < min_inliers or len(target) < min_inliers: - return Transform.identity(), float("inf") + if source.width * source.height < min_inliers or target.width * target.height < min_inliers: + return _transform_from_r_t( + np.eye(3), np.zeros(3), ts=to_seconds(source.header.stamp) + ), float("inf") - src_pcd = source.pointcloud_tensor - tgt_pcd = target.pointcloud_tensor + src_pcd = o3d.t.geometry.PointCloud.from_legacy(pointcloud_to_open3d(source)) + tgt_pcd = o3d.t.geometry.PointCloud.from_legacy(pointcloud_to_open3d(target)) # Normals on the target enable point-to-plane ICP — converges tighter # than point-to-point on indoor scenes (walls give unambiguous normals @@ -822,7 +848,7 @@ def _icp( device = src_pcd.device init_T = o3c.Tensor( - init.to_matrix() if init is not None else np.eye(4), + transform_matrix(init.transform) if init is not None else np.eye(4), dtype=o3c.float64, device=device, ) @@ -845,11 +871,17 @@ def _icp( ), ) except RuntimeError: - return Transform.identity(), float("inf") + return _transform_from_r_t( + np.eye(3), np.zeros(3), ts=to_seconds(source.header.stamp) + ), float("inf") if float(result.fitness) == 0.0: - return Transform.identity(), float("inf") + return _transform_from_r_t( + np.eye(3), np.zeros(3), ts=to_seconds(source.header.stamp) + ), float("inf") T_mat = result.transformation.numpy() rmse = float(result.inlier_rmse) - return _transform_from_r_t(T_mat[:3, :3], T_mat[:3, 3], ts=source.ts), rmse * rmse + return _transform_from_r_t( + T_mat[:3, :3], T_mat[:3, 3], ts=to_seconds(source.header.stamp) + ), rmse * rmse diff --git a/dimos/navigation/go2/loop_closure/pgo_map.py b/dimos/navigation/go2/loop_closure/pgo_map.py index 3ab5abcd876..8371e0418a7 100644 --- a/dimos/navigation/go2/loop_closure/pgo_map.py +++ b/dimos/navigation/go2/loop_closure/pgo_map.py @@ -27,6 +27,7 @@ from __future__ import annotations from collections.abc import Iterator +from copy import deepcopy import math import sys import time @@ -34,8 +35,11 @@ import numpy as np from dimos.mapping.voxels.grid import VoxelGrid -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.message_utils.time import time_from_seconds, to_seconds +from dimos.messages.geometry_msgs.msg import Point, Pose, Quaternion +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header from dimos.navigation.go2.loop_closure.pgo import PGOConfig, PoseGraph, _PGOState, _pose_to_pose3 from dimos.utils.logging_config import setup_logger @@ -61,6 +65,7 @@ def __init__( self._grid = VoxelGrid( voxel_size=voxel_size, frame_id=frame_id, stamped=True, show_startup_log=False ) + self._header = Header(stamp=time_from_seconds(0), frame_id=frame_id) self._cooldown = rebuild_cooldown_s self._placed_loops = 0 self._last_rebuild_ts = -math.inf @@ -78,19 +83,29 @@ def n_loops(self) -> int: def add(self, cloud: PointCloud2, pose: Pose | None) -> bool: """Insert a world-frame lidar frame taken at odom `pose`. True if the map was rebuilt.""" - self._grid.add_frame(cloud) - if pose is not None and not (pose.position.is_zero() or pose.orientation.is_zero()): + timestamp = to_seconds(cloud.header.stamp) + self._header.stamp = deepcopy(cloud.header.stamp) + self._grid.add_points( + np.asarray(pointcloud_xyz(cloud), dtype=np.float32), timestamp=timestamp + ) + if ( + pose is not None + and any((pose.position.x, pose.position.y, pose.position.z)) + and any( + (pose.orientation.x, pose.orientation.y, pose.orientation.z, pose.orientation.w) + ) + ): t0, loops = time.perf_counter(), self.n_loops - self._pgo.process(_pose_to_pose3(pose), cloud.ts, cloud) + self._pgo.process(_pose_to_pose3(pose), timestamp, cloud) if self.n_loops != loops: self.loop_ms = (time.perf_counter() - t0) * 1e3 # frame time, not wall time, so replay and tests pace the same as the robot - if self.n_loops == self._placed_loops or cloud.ts - self._last_rebuild_ts < self._cooldown: + if self.n_loops == self._placed_loops or timestamp - self._last_rebuild_ts < self._cooldown: return False t0 = time.perf_counter() self._grid.reproject(self.graph().place) - self._placed_loops, self._last_rebuild_ts = self.n_loops, cloud.ts + self._placed_loops, self._last_rebuild_ts = self.n_loops, timestamp self.rebuild_ms = (time.perf_counter() - t0) * 1e3 logger.info( "PGO map rebuilt", @@ -113,7 +128,9 @@ def loop_segments(self) -> tuple[np.ndarray, np.ndarray]: return self._pgo.loop_segments() def global_map(self) -> PointCloud2: - return self._grid.get_global_pointcloud2() + return pointcloud_from_xyz( + self._grid.get_global_pointcloud().point["positions"].numpy(), header=self._header + ) def dispose(self) -> None: self._grid.dispose() @@ -144,8 +161,15 @@ def main(dataset: str) -> None: def frames() -> Iterator[tuple[PointCloud2, Pose | None]]: for obs in lidar: cloud = obs.data - cloud.ts = obs.ts # old recordings carry Unitree's stale lidar stamps - yield cloud, obs.pose + # Explicit offline replay time override for recordings with stale sensor stamps. + cloud.header.stamp = time_from_seconds(obs.ts) + values = obs.pose_tuple + pose = ( + None + if values is None + else Pose(position=Point(*values[:3]), orientation=Quaternion(*values[3:])) + ) + yield cloud, pose live = PGOMap() step_ms, rebuilds = [], 0 @@ -161,9 +185,12 @@ def frames() -> Iterator[tuple[PointCloud2, Pose | None]]: live._grid.reproject(graph.place) gold = VoxelGrid(device="CPU:0", show_startup_log=False) for cloud, _ in frames(): - gold.add_frame(graph.place(cloud)) + gold.add_points( + np.asarray(pointcloud_xyz(graph.place(cloud)), dtype=np.float32), + timestamp=to_seconds(cloud.header.stamp), + ) - got, want = live.global_map().points_f32(), gold.get_global_pointcloud2().points_f32() + got, want = pointcloud_xyz(live.global_map()), gold.get_global_pointcloud2().points_f32() missing = 1 - _within_one_voxel(want, got, 0.05) extra = 1 - _within_one_voxel(got, want, 0.05) steps = np.array(step_ms) diff --git a/dimos/navigation/go2/loop_closure/test_generated_module.py b/dimos/navigation/go2/loop_closure/test_generated_module.py new file mode 100644 index 00000000000..1ef4411b89f --- /dev/null +++ b/dimos/navigation/go2/loop_closure/test_generated_module.py @@ -0,0 +1,109 @@ +# 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 contextlib import ExitStack +from threading import Event + +from dimos_cdrgen.registry import encode +import numpy as np +import pytest +from reactivex.scheduler import ThreadPoolScheduler + +from dimos.core.transport import LCMTransport +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.dimos_msgs.msg import LineSegments3D +from dimos.messages.geometry_msgs.msg import Point, Pose, PoseArray, PoseStamped, Quaternion +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.navigation.go2.loop_closure.module import PGOVoxelMapper +from dimos.protocol.rpc.pubsubrpc import LCMRPC +from dimos.utils.threadpool import get_max_workers + + +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +def test_actual_pgo_module_ports_publish_generated_map_graph_and_telemetry(owned_scheduler): + with ExitStack() as cleanup: + module = PGOVoxelMapper(rpc_transport=MemoryRPC, telemetry=True) + cleanup.callback(module.stop) + received = {} + events = {} + transports = {} + + def collector(name): + def receive(value): + received[name] = value + events[name].set() + + return receive + + for name in ( + "lidar", + "odom", + "global_map", + "pgo_keyframes", + "pgo_loops", + "mapper_frame_ms", + "pgo_loop_ms", + "pgo_rebuild_ms", + ): + port = getattr(module, name) + transport = LCMTransport("generated_pgo_" + name, port.type, url="memq://") + transports[name] = transport + cleanup.callback(transport.stop) + port.transport = transport + if name not in ("lidar", "odom"): + events[name] = Event() + cleanup.callback(transport.subscribe(collector(name))) + assert module.lidar.type is PointCloud2 and module.odom.type is PoseStamped + assert module.pgo_keyframes.type is PoseArray and module.pgo_loops.type is LineSegments3D + module.start() + pose_ready = Event() + cleanup.callback(module.odom.subscribe(lambda _: pose_ready.set())) + stamp = time_from_nanoseconds(1700000000123456789) + header = Header(frame_id="world", stamp=stamp) + pose = PoseStamped( + header=header, pose=Pose(position=Point(1, 0, 0), orientation=Quaternion(0, 0, 0, 1)) + ) + transports["odom"].broadcast(None, pose) + assert pose_ready.wait(5) + cloud = pointcloud_from_xyz( + np.array([[1.025, 0.025, 0.025]], dtype=np.float32), header=header + ) + before = encode(cloud) + transports["lidar"].broadcast(None, cloud) + for name in ("global_map", "pgo_keyframes", "pgo_loops", "mapper_frame_ms"): + assert events[name].wait(5), name + output = received["global_map"] + assert output.header == header + np.testing.assert_allclose(pointcloud_xyz(output), [[1.025, 0.025, 0.025]], atol=1e-6) + graph = received["pgo_keyframes"] + assert graph.header == header and len(graph.poses) == 1 + assert graph.poses[0].position.x == 1 and graph.poses[0].orientation.w == 1 + assert received["pgo_loops"].header == header and received["pgo_loops"].segments == [] + assert received["mapper_frame_ms"].data >= 0 + assert encode(cloud) == before + + +@pytest.fixture +def owned_scheduler(monkeypatch): + scheduler = ThreadPoolScheduler(max_workers=get_max_workers()) + monkeypatch.setattr("dimos.utils.reactive.get_scheduler", lambda: scheduler) + monkeypatch.setattr("dimos.utils.threadpool.get_scheduler", lambda: scheduler) + yield scheduler + scheduler.executor.shutdown(wait=True) diff --git a/dimos/navigation/go2/loop_closure/test_pgo.py b/dimos/navigation/go2/loop_closure/test_pgo.py index 9f70d4625e3..f5dec0210cd 100644 --- a/dimos/navigation/go2/loop_closure/test_pgo.py +++ b/dimos/navigation/go2/loop_closure/test_pgo.py @@ -14,21 +14,33 @@ from __future__ import annotations +from dimos_cdrgen.registry import decode, encode import numpy as np import pytest from scipy.spatial.transform import Rotation from dimos.memory.store.memory import MemoryStore from dimos.memory.stream import Stream -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.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.geometry import quaternion_from_matrix, transform_matrix, vector_array +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz, transform_cloud +from dimos.message_utils.time import time_from_nanoseconds, time_from_seconds, to_seconds +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header from dimos.navigation.go2.loop_closure.pgo import ( PGO, Keyframe, PGOConfig, PoseGraph, + _icp, _obs_to_pose3, _PGOState, _pose3_to_transform, @@ -39,6 +51,25 @@ pytest.importorskip("gtsam") +def _transform(*, translation=None, rotation=None, ts=0.0): + return TransformStamped( + header=Header(stamp=time_from_seconds(ts), frame_id="world_raw"), + child_frame_id="body", + transform=Transform( + translation=translation if translation is not None else Vector3(0, 0, 0), + rotation=rotation if rotation is not None else Quaternion(0, 0, 0, 1), + ), + ) + + +def _cloud(points, *, timestamp=0.0, stamps=None): + return pointcloud_from_xyz( + points, + header=Header(stamp=time_from_seconds(timestamp), frame_id="world_raw"), + stamps=stamps, + ) + + def _random_R(rng: np.random.Generator) -> np.ndarray: """Random uniform rotation matrix via random quaternion.""" q = rng.standard_normal(4) @@ -76,7 +107,7 @@ def test_observation_normalizes_transform_pose(self) -> None: """Constructing/deriving with pose=Transform should coerce to 7-tuple.""" from dimos.memory.type.observation import Observation - tf = Transform( + tf = _transform( translation=Vector3(1.5, -2.0, 0.7), rotation=Quaternion(0.1, 0.2, 0.3, 0.927), ts=1.0, @@ -92,9 +123,11 @@ def test_observation_normalizes_transform_pose(self) -> None: def test_observation_normalizes_posestamped(self) -> None: from dimos.memory.type.observation import Observation - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped - ps = PoseStamped(ts=1.0, position=(1.0, 2.0, 3.0), orientation=(0.0, 0.0, 0.0, 1.0)) + ps = PoseStamped( + header=Header(stamp=time_from_seconds(1.0), frame_id="world_raw"), + pose=Pose(position=Point(1.0, 2.0, 3.0), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), + ) obs: Observation[int] = Observation(id=0, ts=1.0, pose=ps, _data=0) assert obs.pose_tuple == (1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0) @@ -104,7 +137,7 @@ def test_obs_to_pose3_roundtrip(self) -> None: rng = np.random.default_rng(4) R = _random_R(rng) t = rng.uniform(-3, 3, size=3) - tf = Transform(translation=Vector3(t), rotation=Quaternion.from_rotation_matrix(R), ts=1.0) + tf = _transform(translation=Vector3(*t), rotation=quaternion_from_matrix(R), ts=1.0) obs: Observation[int] = Observation(id=0, ts=1.0, pose=tf, _data=0) p = _obs_to_pose3(obs) np.testing.assert_allclose(p.rotation().matrix(), R, atol=1e-9) @@ -118,8 +151,8 @@ def test_pose3_to_transform(self) -> None: t = rng.uniform(-3, 3, size=3) p = gtsam.Pose3(gtsam.Rot3(R), gtsam.Point3(t)) tf = _pose3_to_transform(p, ts=7.89, frame_id="world", child_frame_id="body") - np.testing.assert_allclose(tf.rotation.to_rotation_matrix(), R, atol=1e-10) - np.testing.assert_allclose(tf.translation.to_numpy(), t, atol=1e-10) + np.testing.assert_allclose(transform_matrix(tf.transform)[:3, :3], R, atol=1e-10) + np.testing.assert_allclose(vector_array(tf.transform.translation), t, atol=1e-10) def test_pose3_to_transform_with_frames(self) -> None: import gtsam @@ -129,10 +162,10 @@ def test_pose3_to_transform_with_frames(self) -> None: t = rng.uniform(-3, 3, size=3) p = gtsam.Pose3(gtsam.Rot3(R), gtsam.Point3(t)) tf = _pose3_to_transform(p, ts=1.0, frame_id="world_corrected", child_frame_id="body") - assert tf.frame_id == "world_corrected" + assert tf.header.frame_id == "world_corrected" assert tf.child_frame_id == "body" - np.testing.assert_allclose(tf.rotation.to_rotation_matrix(), R, atol=1e-10) - np.testing.assert_allclose(tf.translation.to_numpy(), t, atol=1e-10) + np.testing.assert_allclose(transform_matrix(tf.transform)[:3, :3], R, atol=1e-10) + np.testing.assert_allclose(vector_array(tf.transform.translation), t, atol=1e-10) def _make_lidar_stream(n_frames: int = 12, points_per_frame: int = 500) -> Stream[PointCloud2]: @@ -153,7 +186,7 @@ def _make_lidar_stream(n_frames: int = 12, points_per_frame: int = 500) -> Strea body = rng.uniform(-1, 1, size=(points_per_frame, 3)).astype(np.float32) world = (R_world @ body.T).T + np.array([i, 0, 0], dtype=np.float32) lidar.append( - PointCloud2.from_numpy(world.astype(np.float32)), + _cloud(world.astype(np.float32)), ts=float(i), pose=(float(i), 0.0, 0.0, qx, qy, qz, qw), ) @@ -210,9 +243,15 @@ def _graph_with_drift_at(drifts: list[Transform]) -> PoseGraph: return PoseGraph( keyframes=tuple( Keyframe( - ts=d.ts, - local=Transform(translation=identity, rotation=identity_rot, ts=d.ts), - optimized=Transform(translation=d.translation, rotation=d.rotation, ts=d.ts), + ts=to_seconds(d.header.stamp), + local=_transform( + translation=identity, rotation=identity_rot, ts=to_seconds(d.header.stamp) + ), + optimized=_transform( + translation=d.transform.translation, + rotation=d.transform.rotation, + ts=to_seconds(d.header.stamp), + ), ) for d in drifts ) @@ -226,67 +265,68 @@ def test_empty_raises(self) -> None: def test_single_keyframe_returns_constant(self) -> None: R = Rotation.from_euler("z", np.pi / 4).as_matrix() - only = Transform( + only = _transform( translation=Vector3(1.0, 2.0, 3.0), - rotation=Quaternion.from_rotation_matrix(R), + rotation=quaternion_from_matrix(R), ts=10.0, ) graph = _graph_with_drift_at([only]) for query_ts in (0.0, 10.0, 100.0): out = graph.correction_at(query_ts) - assert out.translation.x == pytest.approx(1.0, abs=1e-10) - assert out.translation.y == pytest.approx(2.0, abs=1e-10) - assert out.translation.z == pytest.approx(3.0, abs=1e-10) + assert out.transform.translation.x == pytest.approx(1.0, abs=1e-10) + assert out.transform.translation.y == pytest.approx(2.0, abs=1e-10) + assert out.transform.translation.z == pytest.approx(3.0, abs=1e-10) def test_out_of_range_clips_to_endpoints(self) -> None: - # Transform's ctor maps ts=0.0 -> time.time(); use ts>0 for determinism. - a = Transform(translation=Vector3(0.0, 0.0, 0.0), ts=1.0) - b = Transform(translation=Vector3(10.0, 0.0, 0.0), ts=11.0) + a = _transform(translation=Vector3(0.0, 0.0, 0.0), ts=1.0) + b = _transform(translation=Vector3(10.0, 0.0, 0.0), ts=11.0) graph = _graph_with_drift_at([a, b]) # Below range -> clipped to a - assert graph.correction_at(-5.0).translation.x == pytest.approx(0.0, abs=1e-10) + assert graph.correction_at(-5.0).transform.translation.x == pytest.approx(0.0, abs=1e-10) # Above range -> clipped to b - assert graph.correction_at(100.0).translation.x == pytest.approx(10.0, abs=1e-10) + assert graph.correction_at(100.0).transform.translation.x == pytest.approx(10.0, abs=1e-10) # In-range midpoint - assert graph.correction_at(6.0).translation.x == pytest.approx(5.0, abs=1e-10) + assert graph.correction_at(6.0).transform.translation.x == pytest.approx(5.0, abs=1e-10) def test_frozen(self) -> None: graph = PoseGraph() with pytest.raises(Exception): - graph.keyframes = (Keyframe(ts=0, local=Transform(), optimized=Transform()),) # type: ignore[misc] + graph.keyframes = (Keyframe(ts=0, local=_transform(), optimized=_transform()),) # type: ignore[misc] class TestPlace: def _graph(self) -> PoseGraph: - turn = Quaternion.from_rotation_matrix(Rotation.from_euler("z", np.pi / 2).as_matrix()) + turn = quaternion_from_matrix(Rotation.from_euler("z", np.pi / 2).as_matrix()) return _graph_with_drift_at( [ - Transform(translation=Vector3(0.0, 0.0, 0.0), ts=1.0), - Transform(translation=Vector3(10.0, 0.0, 0.0), rotation=turn, ts=11.0), + _transform(translation=Vector3(0.0, 0.0, 0.0), ts=1.0), + _transform(translation=Vector3(10.0, 0.0, 0.0), rotation=turn, ts=11.0), ] ) def test_latest_data_stays_put(self) -> None: - cloud = PointCloud2.from_numpy(np.array([[1.0, 2.0, 3.0]]), timestamp=11.0) + cloud = _cloud(np.array([[1.0, 2.0, 3.0]]), timestamp=11.0) np.testing.assert_allclose( - self._graph().place(cloud).points_f32(), [[1.0, 2.0, 3.0]], atol=1e-5 + pointcloud_xyz(self._graph().place(cloud)), [[1.0, 2.0, 3.0]], atol=1e-5 ) def test_old_data_moves_relative_to_the_latest_keyframe(self) -> None: # C(now)^-1 . C(1): undo the latest correction (turn 90 deg, +10 x). - cloud = PointCloud2.from_numpy(np.array([[10.0, 1.0, 0.0]]), timestamp=1.0) + cloud = _cloud(np.array([[10.0, 1.0, 0.0]]), timestamp=1.0) np.testing.assert_allclose( - self._graph().place(cloud).points_f32(), [[1.0, 0.0, 0.0]], atol=1e-5 + pointcloud_xyz(self._graph().place(cloud)), [[1.0, 0.0, 0.0]], atol=1e-5 ) def test_per_point_stamps_override_the_cloud_ts(self) -> None: - cloud = PointCloud2.from_numpy( + cloud = _cloud( np.array([[10.0, 1.0, 0.0], [1.0, 2.0, 3.0]]), timestamp=1.0, stamps=np.array([1.0, 11.0]), ) np.testing.assert_allclose( - self._graph().place(cloud).points_f32(), [[1.0, 0.0, 0.0], [1.0, 2.0, 3.0]], atol=1e-5 + pointcloud_xyz(self._graph().place(cloud)), + [[1.0, 0.0, 0.0], [1.0, 2.0, 3.0]], + atol=1e-5, ) @@ -299,14 +339,14 @@ def test_pure_translation_shifts_poses(self) -> None: lidar: Stream[PointCloud2] = mem.stream("lidar", PointCloud2) for i in range(3): lidar.append( - PointCloud2.from_numpy(np.zeros((1, 3), dtype=np.float32)), + _cloud(np.zeros((1, 3), dtype=np.float32)), ts=float(i + 1), pose=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0), ) graph = _graph_with_drift_at( [ - Transform(translation=Vector3(5.0, 0.0, 0.0), ts=1.0), - Transform(translation=Vector3(5.0, 0.0, 0.0), ts=3.0), + _transform(translation=Vector3(5.0, 0.0, 0.0), ts=1.0), + _transform(translation=Vector3(5.0, 0.0, 0.0), ts=3.0), ] ) for obs in lidar.transform(graph): @@ -320,14 +360,14 @@ def test_passes_through_pose_none(self) -> None: mem = MemoryStore() lidar: Stream[PointCloud2] = mem.stream("lidar", PointCloud2) lidar.append( - PointCloud2.from_numpy(np.zeros((1, 3), dtype=np.float32)), + _cloud(np.zeros((1, 3), dtype=np.float32)), ts=1.0, pose=None, ) graph = _graph_with_drift_at( [ - Transform(translation=Vector3(5.0, 0.0, 0.0), ts=1.0), - Transform(translation=Vector3(5.0, 0.0, 0.0), ts=2.0), + _transform(translation=Vector3(5.0, 0.0, 0.0), ts=1.0), + _transform(translation=Vector3(5.0, 0.0, 0.0), ts=2.0), ] ) for obs in lidar.transform(graph): @@ -336,7 +376,7 @@ def test_passes_through_pose_none(self) -> None: class TestKeyframeType: def test_keyframe_is_frozen(self) -> None: - identity = Transform( + identity = _transform( translation=Vector3(0.0, 0.0, 0.0), rotation=Quaternion(0.0, 0.0, 0.0, 1.0), ts=1.0, @@ -344,5 +384,61 @@ def test_keyframe_is_frozen(self) -> None: kf = Keyframe(ts=1.0, local=identity, optimized=identity) with pytest.raises(Exception): kf.ts = 2.0 # type: ignore[misc] - assert isinstance(kf.local, Transform) - assert isinstance(kf.optimized, Transform) + assert isinstance(kf.local, TransformStamped) + assert isinstance(kf.optimized, TransformStamped) + + +def test_place_retains_capture_fields_header_and_input_bytes(): + graph = _graph_with_drift_at( + [_transform(translation=Vector3(5, 0, 0), ts=1), _transform(ts=11)] + ) + cloud = _cloud(np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), stamps=np.array([1.0, 11.0])) + cloud.header.stamp = time_from_nanoseconds(1700000000123456789) + before = encode(cloud) + output = decode(encode(graph.place(cloud)), PointCloud2) + np.testing.assert_allclose(pointcloud_xyz(output), [[6, 2, 3], [4, 5, 6]]) + assert output.header == cloud.header and output.fields == cloud.fields + assert encode(cloud) == before + + +def test_correct_preserves_exact_input_stamp_and_child_frame(): + graph = _graph_with_drift_at([_transform(translation=Vector3(5, 0, 0), ts=1)]) + pose = _transform(translation=Vector3(1, 2, 3), ts=1) + pose.header.stamp = time_from_nanoseconds(1700000000123456789) + before = encode(pose) + corrected = decode(encode(graph.correct(pose)), TransformStamped) + assert corrected.header.stamp == pose.header.stamp + assert corrected.header.frame_id == "world_corrected" and corrected.child_frame_id == "body" + np.testing.assert_allclose(vector_array(corrected.transform.translation), [6, 2, 3]) + assert encode(pose) == before + + +def test_actual_point_to_plane_icp_recovers_rigid_motion(): + axis = np.linspace(-1, 1, 17) + u, v = np.meshgrid(axis, axis) + walls = np.concatenate( + [ + np.column_stack([np.full(u.size, -1), u.ravel(), v.ravel()]), + np.column_stack([u.ravel(), np.full(u.size, -1), v.ravel()]), + np.column_stack([u.ravel(), v.ravel(), np.full(u.size, -1)]), + ] + ) + target = _cloud(walls, timestamp=7.25) + motion = _transform( + translation=Vector3(0.03, -0.02, 0.01), + rotation=quaternion_from_matrix( + Rotation.from_euler("xyz", [0.01, -0.015, 0.02]).as_matrix() + ), + ts=7.25, + ) + motion.header.frame_id = motion.child_frame_id = target.header.frame_id + source = transform_cloud(target, motion) + before = encode(source) + correction, fitness = _icp(source, target, max_dist=0.2) + np.testing.assert_allclose( + transform_matrix(correction.transform), + np.linalg.inv(transform_matrix(motion.transform)), + atol=1e-4, + ) + assert fitness < 1e-8 + assert encode(source) == before diff --git a/dimos/navigation/go2/loop_closure/test_pgo_auto.py b/dimos/navigation/go2/loop_closure/test_pgo_auto.py new file mode 100644 index 00000000000..771d8c6f4f3 --- /dev/null +++ b/dimos/navigation/go2/loop_closure/test_pgo_auto.py @@ -0,0 +1,108 @@ +# 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_cdrgen.registry import decode as cdr_decode, encode as cdr_encode +import numpy as np +import pytest + +from dimos.memory.store.memory import MemoryStore +from dimos.memory.type.observation import Observation +from dimos.message_utils.geometry import quaternion_from_euler +from dimos.message_utils.pointcloud import pointcloud_from_xyz +from dimos.message_utils.time import time_from_seconds, to_seconds +from dimos.messages.geometry_msgs.msg import Transform, TransformStamped, Vector3 +from dimos.messages.std_msgs.msg import Header +from dimos.navigation.go2.loop_closure.pgo_auto import ( + FRAME_BODY, + FRAME_WORLD_CORRECTED, + FRAME_WORLD_RAW, + PGO, + Keyframe, + apply_corrections, + keyframes_to_corrections, + make_interpolator, +) + + +def _pose(x: float, frame: str, ts: float) -> TransformStamped: + return TransformStamped( + header=Header(frame_id=frame, stamp=time_from_seconds(ts)), + child_frame_id=FRAME_BODY, + transform=Transform( + translation=Vector3(x=x, y=0.0, z=0.0), rotation=quaternion_from_euler(0, 0, 0) + ), + ) + + +def test_generated_corrections_preserve_frames_time_and_nearest_neighbor() -> None: + with MemoryStore() as store: + keyframes = store.stream("keyframes", Keyframe) + keyframes.append( + Keyframe(0.0, _pose(1, FRAME_WORLD_RAW, 0), _pose(3, FRAME_WORLD_CORRECTED, 0)), ts=0 + ) + keyframes.append( + Keyframe(2.0, _pose(2, FRAME_WORLD_RAW, 2), _pose(6, FRAME_WORLD_CORRECTED, 2)), ts=2 + ) + corrections = keyframes_to_corrections(keyframes) + lookup = make_interpolator(corrections) + for ts, expected in [(-1, 2), (0.9, 2), (1.1, 4), (3, 4)]: + correction = cdr_decode(cdr_encode(lookup(ts)), TransformStamped) + assert correction.transform.translation.x == expected + assert correction.header.frame_id == FRAME_WORLD_CORRECTED + assert correction.child_frame_id == FRAME_WORLD_RAW + assert to_seconds(correction.header.stamp) == ts + stream = store.stream("values", str) + stream.append("payload", ts=0, pose=(1, 0, 0, 0, 0, 0, 1)) + stream.append("no pose", ts=2) + output = list(apply_corrections(stream, corrections)) + assert output[0].data == "payload" + assert output[0].pose_tuple == (3, 0, 0, 0, 0, 0, 1) + assert output[0].ts == 0 + assert output[1].pose is None + assert stream.first().pose_tuple == (1, 0, 0, 0, 0, 0, 1) + + +def test_empty_pgo_has_generated_identity_correction() -> None: + pytest.importorskip("gtsam") + graph = next(PGO()(iter(()))).data + assert not graph.keyframes and not graph.loops + correction = graph.correction_at(0) + assert to_seconds(correction.header.stamp) == 0 + assert correction.transform.rotation.w == 1 + result = graph.correct(_pose(5, FRAME_WORLD_RAW, 0)) + assert result.transform.translation.x == 5 + assert result.header.frame_id == FRAME_WORLD_CORRECTED + + +def test_pgo_consumes_generated_cloud_without_mutating_input() -> None: + pytest.importorskip("gtsam") + cloud = pointcloud_from_xyz( + np.array([[1.0, 0, 0], [1, 1, 0], [1, 0, 1]]), + header=Header(frame_id=FRAME_WORLD_RAW, stamp=time_from_seconds(12.25)), + ) + before = cdr_encode(cloud) + quaternion = quaternion_from_euler(0, 0, 0.1) + observation = Observation( + ts=12.25, + _data=cloud, + pose=(1, 0, 0, quaternion.x, quaternion.y, quaternion.z, quaternion.w), + ) + graph = next(PGO()(iter([observation]))).data + assert len(graph.keyframes) == 1 + assert graph.keyframes[0].ts == 12.25 + assert graph.keyframes[0].local.header.frame_id == FRAME_WORLD_RAW + correction = graph.correction_at(12.25) + assert correction.transform.translation.x == pytest.approx(0, abs=1e-6) + assert correction.transform.rotation.w == pytest.approx(1) + assert cdr_encode(cloud) == before diff --git a/dimos/navigation/go2/loop_closure/test_pgo_map.py b/dimos/navigation/go2/loop_closure/test_pgo_map.py index 42de2edcefe..09f92fa4fee 100644 --- a/dimos/navigation/go2/loop_closure/test_pgo_map.py +++ b/dimos/navigation/go2/loop_closure/test_pgo_map.py @@ -12,13 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos_cdrgen.registry import encode import numpy as np import pytest -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.message_utils.time import time_from_nanoseconds, time_from_seconds +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos.messages.std_msgs.msg import Header from dimos.navigation.go2.loop_closure.pgo_map import PGOMap # TODO(PY311): drop. Skip before importing test_pgo, which skips mid-import without gtsam. @@ -29,34 +37,33 @@ VOXEL = 0.05 -class _FakePGO: - """Stands in for the optimizer: the test decides when a loop closes.""" - - n_keyframes = 0 - n_loops = 0 - graph = None - - def process(self, *args): - pass - - def snapshot(self): - return self.graph - - def _wall(x, ts): ys, zs = np.meshgrid(np.arange(0.025, 1, VOXEL), np.arange(0.025, 1, VOXEL)) pts = np.stack([np.full(ys.size, x), ys.ravel(), zs.ravel()], axis=1) - return PointCloud2.from_numpy(pts, timestamp=ts) + return pointcloud_from_xyz(pts, header=Header(stamp=time_from_seconds(ts), frame_id="world")) def _xs(world_map): - return {int(x) for x in np.floor(world_map.global_map().points_f32()[:, 0] / VOXEL)} + return {int(x) for x in np.floor(pointcloud_xyz(world_map.global_map())[:, 0] / VOXEL)} -def test_rebuilds_on_loop_closure_and_respects_cooldown(): +@pytest.fixture +def controlled_map(mocker): + pgo = mocker.patch( + "dimos.navigation.go2.loop_closure.pgo_map._PGOState", autospec=True + ).return_value + pgo.n_keyframes = 0 + pgo.n_loops = 0 world_map = PGOMap(rebuild_cooldown_s=10.0) - world_map._pgo = pgo = _FakePGO() - pose = PoseStamped(1.0, 0.0, 0.0) + try: + yield world_map, pgo + finally: + world_map.dispose() + + +def test_rebuilds_on_loop_closure_and_respects_cooldown(controlled_map): + world_map, pgo = controlled_map + pose = Pose(position=Point(1, 0, 0), orientation=Quaternion(0, 0, 0, 1)) # the same wall seen twice, the first time 0.5 m off; no loop yet assert not world_map.add(_wall(1.025, 100.0), pose) @@ -65,10 +72,20 @@ def test_rebuilds_on_loop_closure_and_respects_cooldown(): # loop closes: the frame at t=100 was really 0.5 m further along x pgo.n_loops = 1 - pgo.graph = _graph_with_drift_at( + pgo.snapshot.return_value = _graph_with_drift_at( [ - Transform(translation=Vector3(0.5, 0.0, 0.0), ts=100.0), - Transform(translation=Vector3(0.0, 0.0, 0.0), ts=101.0), + TransformStamped( + header=Header(stamp=time_from_seconds(100), frame_id="world"), + child_frame_id="body", + transform=Transform( + translation=Vector3(0.5, 0, 0), rotation=Quaternion(0, 0, 0, 1) + ), + ), + TransformStamped( + header=Header(stamp=time_from_seconds(101), frame_id="world"), + child_frame_id="body", + transform=Transform(translation=Vector3(0, 0, 0), rotation=Quaternion(0, 0, 0, 1)), + ), ] ) assert world_map.add(_wall(1.525, 102.0), pose) @@ -78,3 +95,25 @@ def test_rebuilds_on_loop_closure_and_respects_cooldown(): pgo.n_loops = 2 assert not world_map.add(_wall(1.525, 105.0), pose) assert world_map.add(_wall(1.525, 112.5), pose) + + +@pytest.fixture +def actual_map(): + value = PGOMap() + try: + yield value + finally: + value.dispose() + + +def test_real_optimizer_map_preserves_sensor_stamp_and_world_frame(actual_map): + cloud = _wall(1.025, 100) + cloud.header.stamp = time_from_nanoseconds(1700000000123456789) + before = encode(cloud) + pose = Pose(position=Point(1, 0, 0), orientation=Quaternion(0, 0, 0, 1)) + assert not actual_map.add(cloud, pose) + assert actual_map.n_keyframes == 1 and actual_map.n_loops == 0 + output = actual_map.global_map() + assert output.header.stamp == cloud.header.stamp and output.header.frame_id == "world" + assert _xs(actual_map) == {20} + assert encode(cloud) == before diff --git a/dimos/navigation/go2/replanning_a_star/controllers.py b/dimos/navigation/go2/replanning_a_star/controllers.py index 6ad0d8f4106..9a006c3e424 100644 --- a/dimos/navigation/go2/replanning_a_star/controllers.py +++ b/dimos/navigation/go2/replanning_a_star/controllers.py @@ -18,12 +18,18 @@ import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.message_utils.geometry import yaw +from dimos.messages.geometry_msgs.msg import PoseStamped, Twist, Vector3 from dimos.utils.trigonometry import angle_diff +def _velocity(linear: float = 0.0, angular: float = 0.0) -> Twist: + return Twist( + linear=Vector3(x=float(linear), y=0.0, z=0.0), + angular=Vector3(x=0.0, y=0.0, z=float(angular)), + ) + + class Controller(Protocol): def advance(self, lookahead_point: NDArray[np.float64], current_odom: PoseStamped) -> Twist: ... @@ -51,15 +57,15 @@ def __init__(self, simulation: bool, speed: float, control_frequency: float) -> self._control_frequency = control_frequency def advance(self, lookahead_point: NDArray[np.float64], current_odom: PoseStamped) -> Twist: - current_pos = np.array([current_odom.position.x, current_odom.position.y]) + current_pos = np.array([current_odom.pose.position.x, current_odom.pose.position.y]) direction = lookahead_point - current_pos distance = np.linalg.norm(direction) if distance < 1e-6: # Robot is coincidentally at the lookahead point; skip this cycle. - return Twist() + return _velocity() - robot_yaw = current_odom.orientation.euler[2] + robot_yaw = yaw(current_odom.pose.orientation) desired_yaw = np.arctan2(direction[1], direction[0]) yaw_error = angle_diff(desired_yaw, robot_yaw) @@ -73,10 +79,7 @@ def advance(self, lookahead_point: NDArray[np.float64], current_odom: PoseStampe linear_velocity = self._speed * (1.0 - abs(yaw_error) / self._rotation_threshold) linear_velocity = self._apply_min_velocity(linear_velocity, self._min_linear_velocity) - return Twist( - linear=Vector3(linear_velocity, 0.0, 0.0), - angular=Vector3(0.0, 0.0, angular_velocity), - ) + return _velocity(linear_velocity, angular_velocity) def rotate(self, yaw_error: float) -> Twist: angular_velocity = self._compute_angular_velocity(yaw_error) @@ -107,10 +110,7 @@ def _angular_twist(self, angular_velocity: float) -> Twist: if self._simulation and abs(angular_velocity) < 0.8: angular_velocity = 0.8 * np.sign(angular_velocity) - return Twist( - linear=Vector3(0.0, 0.0, 0.0), - angular=Vector3(0.0, 0.0, angular_velocity), - ) + return _velocity(angular=angular_velocity) class PdController(PController): diff --git a/dimos/navigation/go2/replanning_a_star/global_planner.py b/dimos/navigation/go2/replanning_a_star/global_planner.py index 2bc41856a9f..e5f4bba0933 100644 --- a/dimos/navigation/go2/replanning_a_star/global_planner.py +++ b/dimos/navigation/go2/replanning_a_star/global_planner.py @@ -12,11 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy import math from threading import Event, RLock, Thread, current_thread import time -from dimos_lcm.std_msgs import Bool import numpy as np from reactivex import Subject from reactivex.disposable import CompositeDisposable @@ -24,11 +24,17 @@ from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT from dimos.core.resource import Resource from dimos.mapping.occupancy.path_resampling import smooth_resample_path -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path +from dimos.message_utils.geometry import point_distance, yaw +from dimos.message_utils.occupancy import ( + CostValues, + occupancy_cell_value, + occupancy_view, + world_to_grid, +) +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import Point, PoseStamped, Twist +from dimos.messages.nav_msgs.msg import OccupancyGrid, Path +from dimos.messages.std_msgs.msg import Bool, Header from dimos.navigation.base import NavigationState from dimos.navigation.go2.replanning_a_star.goal_validator import find_safe_goal from dimos.navigation.go2.replanning_a_star.local_planner import LocalPlanner, StopMessage @@ -172,11 +178,21 @@ def cancel_goal(self, *, but_will_try_again: bool = False, arrived: bool = False self._goal_reached = arrived self._replan_limiter.reset() - self.path.on_next(Path()) + self.path.on_next(self._empty_path()) self._local_planner.stop_planning() if not but_will_try_again: - self.goal_reached.on_next(Bool(arrived)) + self.goal_reached.on_next(Bool(data=arrived)) + + def _empty_path(self) -> Path: + with self._lock: + source = self._current_goal or self._current_odom + header = ( + deepcopy(source.header) + if source is not None + else Header(frame_id="world", stamp=time_from_nanoseconds(0)) + ) + return Path(header=header, poses=[]) def set_replanning_enabled(self, enabled: bool) -> None: with self._lock: @@ -230,9 +246,12 @@ def _thread_entrypoint(self) -> None: continue if ( - current_goal.position.distance(current_odom.position) < self._goal_tolerance + point_distance(current_goal.pose.position, current_odom.pose.position) + < self._goal_tolerance and abs( - angle_diff(current_goal.orientation.euler[2], current_odom.orientation.euler[2]) + angle_diff( + yaw(current_goal.pose.orientation), yaw(current_odom.pose.orientation) + ) ) < self._rotation_tolerance ): @@ -277,7 +296,7 @@ def _on_stopped_navigating(self, stop_message: StopMessage) -> None: def _handle_stop_message(self, stop_message: StopMessage) -> None: # Note, this runs in the monitoring thread. - self.path.on_next(Path()) + self.path.on_next(self._empty_path()) if stop_message == "arrived": logger.info("Arrived at goal.") @@ -302,7 +321,10 @@ def _replan_path(self) -> None: assert current_odom is not None assert current_goal is not None - if current_goal.position.distance(current_odom.position) < self._replan_goal_tolerance: + if ( + point_distance(current_goal.pose.position, current_odom.pose.position) + < self._replan_goal_tolerance + ): self.cancel_goal(arrived=True) return @@ -310,7 +332,7 @@ def _replan_path(self) -> None: self.cancel_goal() return - if not self._replan_limiter.can_retry(current_odom.position): + if not self._replan_limiter.can_retry(current_odom.pose.position): self.cancel_goal() return @@ -331,16 +353,18 @@ def _plan_path(self) -> None: logger.warning("Cannot handle goal request: missing odometry.") return - safe_goal = self._find_safe_goal(current_goal.position) + safe_goal = self._find_safe_goal(current_goal.pose.position) if not safe_goal: logger.warning( - "No safe goal found.", x=round(current_goal.x, 3), y=round(current_goal.y, 3) + "No safe goal found.", + x=round(current_goal.pose.position.x, 3), + y=round(current_goal.pose.position.y, 3), ) self.cancel_goal() return - path = self._find_wide_path(safe_goal, current_odom.position) + path = self._find_wide_path(safe_goal, current_odom.pose.position) if not path: logger.warning( @@ -349,18 +373,18 @@ def _plan_path(self) -> None: self.cancel_goal() return - resampled_path = smooth_resample_path(path, current_goal, 0.1) + resampled_path = smooth_resample_path(path, current_goal.pose, 0.1) self.path.on_next(resampled_path) self._local_planner.start_planning(resampled_path) - def _find_wide_path(self, goal: Vector3, robot_pos: Vector3) -> Path | None: + def _find_wide_path(self, goal: Point, robot_pos: Point) -> Path | None: # sizes_to_try: list[float] = [2.2, 1.7, 1.3, 1] sizes_to_try: list[float] = [1.1] for size in sizes_to_try: - distance = robot_pos.distance(goal) + distance = point_distance(robot_pos, goal) navigation_map = self._navigation_map if distance > 1.5 else self._navigation_map_near costmap = navigation_map.make_gradient_costmap(size) self._clear_robot_footprint(costmap, navigation_map.binary_costmap, robot_pos) @@ -372,7 +396,7 @@ def _find_wide_path(self, goal: Vector3, robot_pos: Vector3) -> Path | None: return None def _clear_robot_footprint( - self, costmap: OccupancyGrid, binary: OccupancyGrid, robot_pos: Vector3 + self, costmap: OccupancyGrid, binary: OccupancyGrid, robot_pos: Point ) -> None: """Make the cells under the robot passable for planning. @@ -385,32 +409,35 @@ def _clear_robot_footprint( the raw binary costmap (actually observed obstacle points) stay blocked. """ - if binary.grid.shape != costmap.grid.shape or binary.origin != costmap.origin: + binary_values = occupancy_view(binary) + cost_values = occupancy_view(costmap).copy() + if binary_values.shape != cost_values.shape or binary.info.origin != costmap.info.origin: # A newer map update raced in between building the two grids; # skip clearing rather than misalign cells. return - center = costmap.world_to_grid(robot_pos) - center_x, center_y = int(center.x), int(center.y) - cells = int(self._robot_rotation_diameter / 2 / costmap.resolution) + 1 + center = world_to_grid(costmap, robot_pos) + center_x, center_y = int(center[0]), int(center[1]) + cells = int(self._robot_rotation_diameter / 2 / costmap.info.resolution) + 1 - height, width = costmap.grid.shape + height, width = cost_values.shape y0, y1 = max(0, center_y - cells), min(height, center_y + cells + 1) x0, x1 = max(0, center_x - cells), min(width, center_x + cells + 1) if y0 >= y1 or x0 >= x1: return - region = costmap.grid[y0:y1, x0:x1] - binary_region = binary.grid[y0:y1, x0:x1] + region = cost_values[y0:y1, x0:x1] + binary_region = binary_values[y0:y1, x0:x1] rows, columns = np.ogrid[y0:y1, x0:x1] disc = (rows - center_y) ** 2 + (columns - center_x) ** 2 <= cells**2 clearable = disc & (region >= CostValues.OCCUPIED) & (binary_region < CostValues.OCCUPIED) region[clearable] = CostValues.OCCUPIED - 1 + costmap.data = cost_values.reshape(-1) - def _find_safe_goal(self, goal: Vector3) -> Vector3 | None: + def _find_safe_goal(self, goal: Point) -> Point | None: costmap = self._navigation_map.binary_costmap - if costmap.cell_value(goal) == CostValues.UNKNOWN: + if occupancy_cell_value(costmap, goal) == CostValues.UNKNOWN: return goal safe_goal = find_safe_goal( @@ -426,7 +453,7 @@ def _find_safe_goal(self, goal: Vector3) -> Vector3 | None: logger.warning("No safe goal found near requested target.") return None - goals_distance = safe_goal.distance(goal) + goals_distance = point_distance(safe_goal, goal) if goals_distance > 0.2: logger.warning(f"Travelling to goal {goals_distance}m away from requested goal.") diff --git a/dimos/navigation/go2/replanning_a_star/goal_validator.py b/dimos/navigation/go2/replanning_a_star/goal_validator.py index b717c762952..bb8852ea59f 100644 --- a/dimos/navigation/go2/replanning_a_star/goal_validator.py +++ b/dimos/navigation/go2/replanning_a_star/goal_validator.py @@ -13,22 +13,59 @@ # limitations under the License. from collections import deque +from typing import overload import numpy as np +from dimos.mapping.occupancy.operations import _grid_coordinates, _grid_values +from dimos.message_utils.occupancy import grid_to_world +from dimos.messages.geometry_msgs.msg import Point +from dimos.messages.nav_msgs.msg import OccupancyGrid as CdrOccupancyGrid from dimos.msgs.geometry_msgs.Vector3 import Vector3, VectorLike from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid +def _world_point( + costmap: OccupancyGrid | CdrOccupancyGrid, coordinates: tuple[int, int] +) -> Point | Vector3: + if isinstance(costmap, CdrOccupancyGrid): + return grid_to_world(costmap, coordinates) + return costmap.grid_to_world(coordinates) + + +@overload +def find_safe_goal( + costmap: CdrOccupancyGrid, + goal: Point | VectorLike, + algorithm: str = "bfs", + cost_threshold: int = 50, + min_clearance: float = 0.3, + max_search_distance: float = 5.0, + connectivity_check_radius: int = 3, +) -> Point | None: ... + + +@overload def find_safe_goal( costmap: OccupancyGrid, - goal: VectorLike, + goal: Point | VectorLike, algorithm: str = "bfs", cost_threshold: int = 50, min_clearance: float = 0.3, max_search_distance: float = 5.0, connectivity_check_radius: int = 3, -) -> Vector3 | None: +) -> Vector3 | None: ... + + +def find_safe_goal( + costmap: OccupancyGrid | CdrOccupancyGrid, + goal: Point | VectorLike, + algorithm: str = "bfs", + cost_threshold: int = 50, + min_clearance: float = 0.3, + max_search_distance: float = 5.0, + connectivity_check_radius: int = 3, +) -> Vector3 | Point | None: """ Find a safe goal position when the original goal is in collision or too close to obstacles. @@ -68,13 +105,13 @@ def find_safe_goal( def _find_safe_goal_bfs( - costmap: OccupancyGrid, - goal: VectorLike, + costmap: OccupancyGrid | CdrOccupancyGrid, + goal: Point | VectorLike, cost_threshold: int, min_clearance: float, max_search_distance: float, connectivity_check_radius: int, -) -> Vector3 | None: +) -> Vector3 | Point | None: """ BFS-based search for nearest safe goal position. This guarantees finding the closest valid position. @@ -90,16 +127,17 @@ def _find_safe_goal_bfs( """ # Convert goal to grid coordinates - goal_grid = costmap.world_to_grid(goal) - gx, gy = int(goal_grid.x), int(goal_grid.y) + point = goal if isinstance(goal, Point) else Vector3(goal) + x, y = _grid_coordinates(costmap, point) + gx, gy = int(x), int(y) # Convert distances to grid cells - clearance_cells = int(np.ceil(min_clearance / costmap.resolution)) - max_search_cells = int(np.ceil(max_search_distance / costmap.resolution)) + clearance_cells = int(np.ceil(min_clearance / costmap.info.resolution)) + max_search_cells = int(np.ceil(max_search_distance / costmap.info.resolution)) # BFS queue and visited set queue = deque([(gx, gy, 0)]) - visited = set([(gx, gy)]) + visited = {(gx, gy)} # 8-connected neighbors neighbors = [(0, 1), (1, 0), (0, -1), (-1, 0), (1, 1), (1, -1), (-1, 1), (-1, -1)] @@ -116,29 +154,32 @@ def _find_safe_goal_bfs( costmap, x, y, cost_threshold, clearance_cells, connectivity_check_radius ): # Convert back to world coordinates - return costmap.grid_to_world((x, y)) + return _world_point(costmap, (x, y)) # Add neighbors to queue for dx, dy in neighbors: nx, ny = x + dx, y + dy # Check bounds - if 0 <= nx < costmap.width and 0 <= ny < costmap.height: - if (nx, ny) not in visited: - visited.add((nx, ny)) - queue.append((nx, ny, dist + 1)) + if ( + 0 <= nx < costmap.info.width + and 0 <= ny < costmap.info.height + and (nx, ny) not in visited + ): + visited.add((nx, ny)) + queue.append((nx, ny, dist + 1)) return None def _find_safe_goal_bfs_contiguous( - costmap: OccupancyGrid, - goal: VectorLike, + costmap: OccupancyGrid | CdrOccupancyGrid, + goal: Point | VectorLike, cost_threshold: int, min_clearance: float, max_search_distance: float, connectivity_check_radius: int, -) -> Vector3 | None: +) -> Vector3 | Point | None: """ BFS-based search for nearest safe goal position, only following passable cells. Unlike regular BFS, this only expands through cells with occupancy < 100, @@ -155,16 +196,18 @@ def _find_safe_goal_bfs_contiguous( """ # Convert goal to grid coordinates - goal_grid = costmap.world_to_grid(goal) - gx, gy = int(goal_grid.x), int(goal_grid.y) + point = goal if isinstance(goal, Point) else Vector3(goal) + x, y = _grid_coordinates(costmap, point) + gx, gy = int(x), int(y) + cells = _grid_values(costmap) # Convert distances to grid cells - clearance_cells = int(np.ceil(min_clearance / costmap.resolution)) - max_search_cells = int(np.ceil(max_search_distance / costmap.resolution)) + clearance_cells = int(np.ceil(min_clearance / costmap.info.resolution)) + max_search_cells = int(np.ceil(max_search_distance / costmap.info.resolution)) # BFS queue and visited set queue = deque([(gx, gy, 0)]) - visited = set([(gx, gy)]) + visited = {(gx, gy)} # 8-connected neighbors neighbors = [(0, 1), (1, 0), (0, -1), (-1, 0), (1, 1), (1, -1), (-1, 1), (-1, -1)] @@ -181,25 +224,28 @@ def _find_safe_goal_bfs_contiguous( costmap, x, y, cost_threshold, clearance_cells, connectivity_check_radius ): # Convert back to world coordinates - return costmap.grid_to_world((x, y)) + return _world_point(costmap, (x, y)) # Add neighbors to queue for dx, dy in neighbors: nx, ny = x + dx, y + dy # Check bounds - if 0 <= nx < costmap.width and 0 <= ny < costmap.height: - if (nx, ny) not in visited: - # Only expand through passable cells (occupancy < 100) - if costmap.grid[ny, nx] < 100: - visited.add((nx, ny)) - queue.append((nx, ny, dist + 1)) + # Only expand through passable cells (occupancy < 100). + if ( + 0 <= nx < costmap.info.width + and 0 <= ny < costmap.info.height + and (nx, ny) not in visited + and cells[ny, nx] < 100 + ): + visited.add((nx, ny)) + queue.append((nx, ny, dist + 1)) return None def _is_position_safe( - costmap: OccupancyGrid, + costmap: OccupancyGrid | CdrOccupancyGrid, x: int, y: int, cost_threshold: int, @@ -220,23 +266,26 @@ def _is_position_safe( True if position is safe, False otherwise """ + cells = _grid_values(costmap) # Check bounds first - if not (0 <= x < costmap.width and 0 <= y < costmap.height): + if not (0 <= x < costmap.info.width and 0 <= y < costmap.info.height): return False # Check if position itself is free - if costmap.grid[y, x] >= cost_threshold or costmap.grid[y, x] == CostValues.UNKNOWN: + if cells[y, x] >= cost_threshold or cells[y, x] == CostValues.UNKNOWN: return False # Check clearance around position for dy in range(-clearance_cells, clearance_cells + 1): for dx in range(-clearance_cells, clearance_cells + 1): nx, ny = x + dx, y + dy - if 0 <= nx < costmap.width and 0 <= ny < costmap.height: - # Check if within circular clearance - if dx * dx + dy * dy <= clearance_cells * clearance_cells: - if costmap.grid[ny, nx] >= cost_threshold: - return False + if ( + 0 <= nx < costmap.info.width + and 0 <= ny < costmap.info.height + and dx * dx + dy * dy <= clearance_cells * clearance_cells + and cells[ny, nx] >= cost_threshold + ): + return False # Check connectivity (not surrounded by obstacles) # Count free neighbors in a larger radius @@ -249,16 +298,10 @@ def _is_position_safe( continue nx, ny = x + dx, y + dy - if 0 <= nx < costmap.width and 0 <= ny < costmap.height: + if 0 <= nx < costmap.info.width and 0 <= ny < costmap.info.height: total_count += 1 - if ( - costmap.grid[ny, nx] < cost_threshold - and costmap.grid[ny, nx] != CostValues.UNKNOWN - ): + if cells[ny, nx] < cost_threshold and cells[ny, nx] != CostValues.UNKNOWN: free_count += 1 # Require at least 50% of neighbors to be free (not surrounded) - if total_count > 0 and free_count < total_count * 0.5: - return False - - return True + return not (total_count > 0 and free_count < total_count * 0.5) diff --git a/dimos/navigation/go2/replanning_a_star/local_planner.py b/dimos/navigation/go2/replanning_a_star/local_planner.py index 3b1a6716ae8..378253a311a 100644 --- a/dimos/navigation/go2/replanning_a_star/local_planner.py +++ b/dimos/navigation/go2/replanning_a_star/local_planner.py @@ -22,12 +22,11 @@ from reactivex import Subject from dimos.core.resource import Resource -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.nav_msgs.Path import Path +from dimos.message_utils.geometry import point_distance, yaw +from dimos.messages.geometry_msgs.msg import PoseStamped, Twist +from dimos.messages.nav_msgs.msg import OccupancyGrid, Path from dimos.navigation.base import NavigationState -from dimos.navigation.go2.replanning_a_star.controllers import Controller, PController +from dimos.navigation.go2.replanning_a_star.controllers import Controller, PController, _velocity from dimos.navigation.go2.replanning_a_star.navigation_map import NavigationMap from dimos.navigation.go2.replanning_a_star.path_clearance import PathClearance from dimos.navigation.go2.replanning_a_star.path_distancer import PathDistancer @@ -124,7 +123,7 @@ def start_planning(self, path: Path) -> None: self._thread.start() def stop_planning(self) -> None: - self.cmd_vel.on_next(Twist()) + self.cmd_vel.on_next(_velocity()) self._stop_planning_event.set() with self._lock: @@ -157,7 +156,7 @@ def _thread_entrypoint(self) -> None: self.stopped_navigating.on_next("error") finally: self._reset_state() - self.cmd_vel.on_next(Twist()) + self.cmd_vel.on_next(_velocity()) def _change_state(self, new_state: PlannerState) -> None: if new_state == self._state: @@ -180,14 +179,14 @@ def _loop(self) -> None: # Determine initial state: skip initial_rotation if already aligned. new_state: PlannerState = "initial_rotation" if current_odom is not None and len(path.poses) > 0: - first_yaw = path.poses[0].orientation.euler[2] - robot_yaw = current_odom.orientation.euler[2] + first_yaw = yaw(path.poses[0].pose.orientation) + robot_yaw = yaw(current_odom.pose.orientation) initial_yaw_error = angle_diff(first_yaw, robot_yaw) self._controller.reset_yaw_error(initial_yaw_error) angle_in_tolerance = abs(initial_yaw_error) < self._orientation_tolerance if angle_in_tolerance: position_in_tolerance = ( - path.poses[0].position.distance(current_odom.position) < 0.01 + point_distance(path.poses[0].pose.position, current_odom.pose.position) < 0.01 ) if position_in_tolerance: new_state = "final_rotation" @@ -245,8 +244,8 @@ def _compute_initial_rotation(self) -> Twist: assert current_odom is not None first_pose = path.poses[0] - first_yaw = first_pose.orientation.euler[2] - robot_yaw = current_odom.orientation.euler[2] + first_yaw = yaw(first_pose.pose.orientation) + robot_yaw = yaw(current_odom.pose.orientation) yaw_error = angle_diff(first_yaw, robot_yaw) if abs(yaw_error) < self._orientation_tolerance: @@ -264,7 +263,7 @@ def get_distance_to_path(self) -> float | None: if path_distancer is None or current_odom is None: return None - current_pos = np.array([current_odom.position.x, current_odom.position.y]) + current_pos = np.array([current_odom.pose.position.x, current_odom.pose.position.y]) return path_distancer.get_distance_to_path(current_pos) @@ -276,7 +275,7 @@ def _compute_path_following(self) -> Twist: assert path_distancer is not None assert current_odom is not None - current_pos = np.array([current_odom.position.x, current_odom.position.y]) + current_pos = np.array([current_odom.pose.position.x, current_odom.pose.position.y]) if path_distancer.distance_to_goal(current_pos) < self._goal_tolerance: logger.info("Reached goal position, starting final rotation") @@ -301,15 +300,15 @@ def _compute_final_rotation(self) -> Twist: assert path is not None assert current_odom is not None - goal_yaw = path.poses[-1].orientation.euler[2] - robot_yaw = current_odom.orientation.euler[2] + goal_yaw = yaw(path.poses[-1].pose.orientation) + robot_yaw = yaw(current_odom.pose.orientation) yaw_error = angle_diff(goal_yaw, robot_yaw) if abs(yaw_error) < self._orientation_tolerance: logger.info("Final rotation complete, goal reached") with self._lock: self._change_state("arrived") - return Twist() + return _velocity() return self._controller.rotate(yaw_error) diff --git a/dimos/navigation/go2/replanning_a_star/min_cost_astar.py b/dimos/navigation/go2/replanning_a_star/min_cost_astar.py index 03311db6897..81dc7c7f81b 100644 --- a/dimos/navigation/go2/replanning_a_star/min_cost_astar.py +++ b/dimos/navigation/go2/replanning_a_star/min_cost_astar.py @@ -12,11 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy import heapq - +from typing import overload + +import numpy as np +from numpy.typing import NDArray + +from dimos.message_utils.occupancy import grid_to_world, occupancy_view, world_to_grid +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped as CdrPoseStamped, + Quaternion as CdrQuaternion, +) +from dimos.messages.nav_msgs.msg import OccupancyGrid as CdrOccupancyGrid, Path as CdrPath from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import VectorLike +from dimos.msgs.geometry_msgs.Vector3 import Vector3, VectorLike from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid from dimos.msgs.nav_msgs.Path import Path from dimos.utils.logging_config import setup_logger @@ -64,89 +77,99 @@ def _heuristic(x1: int, y1: int, x2: int, y2: int) -> float: def _reconstruct_path( parents: dict[tuple[int, int], tuple[int, int]], current: tuple[int, int], - costmap: OccupancyGrid, start_tuple: tuple[int, int], - goal_tuple: tuple[int, int], -) -> Path: - frame_id = costmap.frame_id - waypoints: list[PoseStamped] = [] +) -> list[tuple[int, int]]: + coordinates = [] while current in parents: - world_point = costmap.grid_to_world(current) - pose = PoseStamped( - frame_id=frame_id, - position=[world_point.x, world_point.y, 0.0], - orientation=Quaternion(0, 0, 0, 1), # Identity quaternion - ) - waypoints.append(pose) + coordinates.append(current) current = parents[current] + coordinates.append(start_tuple) + coordinates.reverse() + return coordinates - start_world_point = costmap.grid_to_world(start_tuple) - start_pose = PoseStamped( - frame_id=frame_id, - position=[start_world_point.x, start_world_point.y, 0.0], - orientation=Quaternion(0, 0, 0, 1), - ) - waypoints.append(start_pose) - - waypoints.reverse() - - # Add the goal position if it's not already included - goal_point = costmap.grid_to_world(goal_tuple) - if ( - not waypoints - or (waypoints[-1].x - goal_point.x) ** 2 + (waypoints[-1].y - goal_point.y) ** 2 > 1e-10 - ): - goal_pose = PoseStamped( - frame_id=frame_id, - position=[goal_point.x, goal_point.y, 0.0], - orientation=Quaternion(0, 0, 0, 1), +def _reconstruct_path_from_coords( + path_coords: list[tuple[int, int]], + costmap: OccupancyGrid | CdrOccupancyGrid, +) -> Path | CdrPath: + if isinstance(costmap, CdrOccupancyGrid): + return CdrPath( + header=deepcopy(costmap.header), + poses=[ + CdrPoseStamped( + header=deepcopy(costmap.header), + pose=Pose( + position=grid_to_world(costmap, point), + orientation=CdrQuaternion(x=0, y=0, z=0, w=1), + ), + ) + for point in path_coords + ], + ) + waypoints = [] + for coordinate in path_coords: + point = costmap.grid_to_world(coordinate) + waypoints.append( + PoseStamped( + frame_id=costmap.frame_id, + position=[point.x, point.y, 0.0], + orientation=Quaternion(0, 0, 0, 1), + ) ) - waypoints.append(goal_pose) + return Path(frame_id=costmap.frame_id, poses=waypoints) - return Path(frame_id=frame_id, poses=waypoints) +def _cell(costmap: OccupancyGrid | CdrOccupancyGrid, value: Point | VectorLike) -> tuple[int, int]: + vector = Vector3((value.x, value.y, value.z)) if isinstance(value, Point) else Vector3(value) + if isinstance(costmap, CdrOccupancyGrid): + x, y = world_to_grid(costmap, Point(x=vector.x, y=vector.y, z=vector.z)) + else: + coordinates = costmap.world_to_grid(vector) + x, y = coordinates.x, coordinates.y + return int(x), int(y) -def _reconstruct_path_from_coords( - path_coords: list[tuple[int, int]], - costmap: OccupancyGrid, -) -> Path: - frame_id = costmap.frame_id - waypoints: list[PoseStamped] = [] - - for gx, gy in path_coords: - world_point = costmap.grid_to_world((gx, gy)) - pose = PoseStamped( - frame_id=frame_id, - position=[world_point.x, world_point.y, 0.0], - orientation=Quaternion(0, 0, 0, 1), - ) - waypoints.append(pose) - return Path(frame_id=frame_id, poses=waypoints) +@overload +def min_cost_astar( + costmap: CdrOccupancyGrid, + goal: Point | VectorLike, + start: Point | VectorLike = (0.0, 0.0), + cost_threshold: int = 100, + unknown_penalty: float = 0.8, + use_cpp: bool = True, +) -> CdrPath | None: ... +@overload def min_cost_astar( costmap: OccupancyGrid, - goal: VectorLike, - start: VectorLike = (0.0, 0.0), + goal: Point | VectorLike, + start: Point | VectorLike = (0.0, 0.0), cost_threshold: int = 100, unknown_penalty: float = 0.8, use_cpp: bool = True, -) -> Path | None: - start_vector = costmap.world_to_grid(start) - goal_vector = costmap.world_to_grid(goal) +) -> Path | None: ... - start_tuple = (int(start_vector.x), int(start_vector.y)) - goal_tuple = (int(goal_vector.x), int(goal_vector.y)) - if not (0 <= goal_tuple[0] < costmap.width and 0 <= goal_tuple[1] < costmap.height): +def min_cost_astar( + costmap: OccupancyGrid | CdrOccupancyGrid, + goal: Point | VectorLike, + start: Point | VectorLike = (0.0, 0.0), + cost_threshold: int = 100, + unknown_penalty: float = 0.8, + use_cpp: bool = True, +) -> Path | CdrPath | None: + start_tuple = _cell(costmap, start) + goal_tuple = _cell(costmap, goal) + cells = occupancy_view(costmap) if isinstance(costmap, CdrOccupancyGrid) else costmap.grid + height, width = cells.shape + if not (0 <= goal_tuple[0] < width and 0 <= goal_tuple[1] < height): return None if use_cpp: if _USE_CPP: path_coords = _astar_cpp( - costmap.grid, + cells, start_tuple[0], start_tuple[1], goal_tuple[0], @@ -163,6 +186,20 @@ def min_cost_astar( _CPP_IMPORT_ERROR, ) + coordinates = _find_path(cells, start_tuple, goal_tuple, cost_threshold, unknown_penalty) + if coordinates is None: + return None + return _reconstruct_path_from_coords(coordinates, costmap) + + +def _find_path( + cells: NDArray[np.int8], + start_tuple: tuple[int, int], + goal_tuple: tuple[int, int], + cost_threshold: int, + unknown_penalty: float, +) -> list[tuple[int, int]] | None: + height, width = cells.shape open_set: list[tuple[float, float, tuple[int, int]]] = [] # Priority queue for nodes to explore closed_set: set[tuple[int, int]] = set() # Set of explored nodes @@ -185,7 +222,7 @@ def min_cost_astar( continue if current == goal_tuple: - return _reconstruct_path(parents, current, costmap, start_tuple, goal_tuple) + return _reconstruct_path(parents, current, start_tuple) closed_set.add(current) @@ -193,13 +230,13 @@ def min_cost_astar( neighbor_x, neighbor_y = current_x + dx, current_y + dy neighbor = (neighbor_x, neighbor_y) - if not (0 <= neighbor_x < costmap.width and 0 <= neighbor_y < costmap.height): + if not (0 <= neighbor_x < width and 0 <= neighbor_y < height): continue if neighbor in closed_set: continue - neighbor_val = costmap.grid[neighbor_y, neighbor_x] + neighbor_val = cells[neighbor_y, neighbor_x] if neighbor_val >= cost_threshold: continue diff --git a/dimos/navigation/go2/replanning_a_star/module.py b/dimos/navigation/go2/replanning_a_star/module.py index 58bbc6a2cb4..4a7eac32058 100644 --- a/dimos/navigation/go2/replanning_a_star/module.py +++ b/dimos/navigation/go2/replanning_a_star/module.py @@ -12,22 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +from copy import deepcopy import os from typing import Any -from dimos_lcm.std_msgs import Bool, String 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.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.nav_msgs.Path import Path +from dimos.messages.geometry_msgs.msg import PointStamped, Pose, PoseStamped, Quaternion, Twist +from dimos.messages.nav_msgs.msg import OccupancyGrid, Odometry, Path +from dimos.messages.std_msgs.msg import Bool, String from dimos.navigation.base import NavigationInterface, NavigationState from dimos.navigation.go2.replanning_a_star.global_planner import GlobalPlanner from dimos.utils.logging_config import setup_logger @@ -74,13 +71,7 @@ def start(self) -> None: super().start() self.register_disposable(Disposable(self.odom.subscribe(self._planner.handle_odom))) - self.register_disposable( - Disposable( - self.odometry.subscribe( - lambda msg: self._planner.handle_odom(msg.to_pose_stamped()) - ) - ) - ) + self.register_disposable(Disposable(self.odometry.subscribe(self._handle_odometry))) self.register_disposable( Disposable(self.global_costmap.subscribe(self._planner.handle_global_costmap)) ) @@ -92,11 +83,7 @@ def start(self) -> None: ) self.register_disposable( - Disposable( - self.clicked_point.subscribe( - lambda pt: self._planner.handle_goal_request(pt.to_pose_stamped()) - ) - ) + Disposable(self.clicked_point.subscribe(self._handle_clicked_point)) ) if self.stop_movement.transport is not None: @@ -124,6 +111,21 @@ def stop(self) -> None: super().stop() + def _handle_odometry(self, msg: Odometry) -> None: + self._planner.handle_odom( + PoseStamped(header=deepcopy(msg.header), pose=deepcopy(msg.pose.pose)) + ) + + def _handle_clicked_point(self, msg: PointStamped) -> None: + self._planner.handle_goal_request( + PoseStamped( + header=deepcopy(msg.header), + pose=Pose( + position=deepcopy(msg.point), orientation=Quaternion(x=0.0, y=0.0, z=0.0, w=1.0) + ), + ) + ) + def _on_stop_movement(self, msg: Bool) -> None: if msg.data: self.cancel_goal() diff --git a/dimos/navigation/go2/replanning_a_star/module_spec.py b/dimos/navigation/go2/replanning_a_star/module_spec.py index c9ec73dc473..a4e21ea823a 100644 --- a/dimos/navigation/go2/replanning_a_star/module_spec.py +++ b/dimos/navigation/go2/replanning_a_star/module_spec.py @@ -14,7 +14,7 @@ from typing import Protocol -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.messages.geometry_msgs.msg import PoseStamped from dimos.navigation.base import NavigationState from dimos.spec.utils import Spec diff --git a/dimos/navigation/go2/replanning_a_star/navigation_map.py b/dimos/navigation/go2/replanning_a_star/navigation_map.py index 36ba910d01f..5218c84a1c8 100644 --- a/dimos/navigation/go2/replanning_a_star/navigation_map.py +++ b/dimos/navigation/go2/replanning_a_star/navigation_map.py @@ -16,7 +16,7 @@ from dimos.mapping.occupancy.gradient import GradientStrategy from dimos.mapping.occupancy.path_map import make_navigation_map -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid +from dimos.messages.nav_msgs.msg import OccupancyGrid class NavigationMap: diff --git a/dimos/navigation/go2/replanning_a_star/path_clearance.py b/dimos/navigation/go2/replanning_a_star/path_clearance.py index c7cf276c402..d995c30c6ca 100644 --- a/dimos/navigation/go2/replanning_a_star/path_clearance.py +++ b/dimos/navigation/go2/replanning_a_star/path_clearance.py @@ -17,14 +17,17 @@ import numpy as np from numpy.typing import NDArray +from dimos.mapping.occupancy.operations import _grid_values from dimos.mapping.occupancy.path_mask import make_path_mask +from dimos.message_utils.geometry import point_distance +from dimos.messages.nav_msgs.msg import OccupancyGrid as CdrOccupancyGrid, Path as CdrPath from dimos.msgs.nav_msgs.OccupancyGrid import CostValues, OccupancyGrid from dimos.msgs.nav_msgs.Path import Path class PathClearance: - _costmap: OccupancyGrid | None = None - _last_costmap: OccupancyGrid | None = None + _costmap: OccupancyGrid | CdrOccupancyGrid | None = None + _last_costmap: OccupancyGrid | CdrOccupancyGrid | None = None _path_lookup_distance: float = 3.0 _max_distance_cache: float = 1.0 _last_used_shape: tuple[int, ...] | None = None @@ -32,16 +35,16 @@ class PathClearance: _last_used_pose: int | None = None _robot_width: float _lock: RLock - _path: Path + _path: Path | CdrPath _pose_index: int - def __init__(self, robot_width: float, path: Path) -> None: + def __init__(self, robot_width: float, path: Path | CdrPath) -> None: self._robot_width = robot_width self._path = path self._pose_index = 0 self._lock = RLock() - def update_costmap(self, costmap: OccupancyGrid) -> None: + def update_costmap(self, costmap: OccupancyGrid | CdrOccupancyGrid) -> None: with self._lock: self._costmap = costmap @@ -60,7 +63,7 @@ def mask(self) -> NDArray[np.bool_]: if ( self._last_mask is not None and self._last_used_pose is not None - and costmap.grid.shape == self._last_used_shape + and _grid_values(costmap).shape == self._last_used_shape and self._pose_distance(self._last_used_pose, pose_index) < self._max_distance_cache ): return self._last_mask @@ -73,7 +76,7 @@ def mask(self) -> NDArray[np.bool_]: max_length=self._path_lookup_distance, ) - self._last_used_shape = costmap.grid.shape + self._last_used_shape = _grid_values(costmap).shape self._last_used_pose = pose_index return self._last_mask @@ -85,9 +88,13 @@ def is_obstacle_ahead(self) -> bool: if costmap is None: return True - return bool(np.any(costmap.grid[self.mask] == CostValues.OCCUPIED)) + return bool(np.any(_grid_values(costmap)[self.mask] == CostValues.OCCUPIED)) def _pose_distance(self, index1: int, index2: int) -> float: + if isinstance(self._path, CdrPath): + return point_distance( + self._path.poses[index1].pose.position, self._path.poses[index2].pose.position + ) p1 = self._path.poses[index1].position p2 = self._path.poses[index2].position return p1.distance(p2) diff --git a/dimos/navigation/go2/replanning_a_star/path_distancer.py b/dimos/navigation/go2/replanning_a_star/path_distancer.py index 9a2c28dad78..533f7b7ea82 100644 --- a/dimos/navigation/go2/replanning_a_star/path_distancer.py +++ b/dimos/navigation/go2/replanning_a_star/path_distancer.py @@ -17,6 +17,8 @@ import numpy as np from numpy.typing import NDArray +from dimos.messages.geometry_msgs.msg import PoseStamped as CdrPoseStamped +from dimos.messages.nav_msgs.msg import Path as CdrPath from dimos.msgs.nav_msgs.Path import Path @@ -25,9 +27,16 @@ class PathDistancer: _path: NDArray[np.float64] _cumulative_dists: NDArray[np.float64] - def __init__(self, path: Path, lookahead_dist: float = 0.5) -> None: + def __init__(self, path: Path | CdrPath, lookahead_dist: float = 0.5) -> None: self._lookahead_dist = lookahead_dist - self._path = np.array([[p.position.x, p.position.y] for p in path.poses]) + self._path = np.array( + [ + [p.pose.position.x, p.pose.position.y] + if isinstance(p, CdrPoseStamped) + else [p.position.x, p.position.y] + for p in path.poses + ] + ) self._cumulative_dists = _make_cumulative_distance_array(self._path) def find_lookahead_point(self, start_idx: int) -> NDArray[np.float64]: diff --git a/dimos/navigation/go2/replanning_a_star/position_tracker.py b/dimos/navigation/go2/replanning_a_star/position_tracker.py index 7d8249b562b..739da3271c1 100644 --- a/dimos/navigation/go2/replanning_a_star/position_tracker.py +++ b/dimos/navigation/go2/replanning_a_star/position_tracker.py @@ -19,7 +19,7 @@ import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.messages.geometry_msgs.msg import PoseStamped _max_points_per_second = 1000 @@ -51,7 +51,7 @@ def reset_data(self) -> None: def add_position(self, pose: PoseStamped) -> None: with self._lock: self._timestamps[self._index] = time.time() - self._positions[self._index] = (pose.position.x, pose.position.y) + self._positions[self._index] = (pose.pose.position.x, pose.pose.position.y) self._index = (self._index + 1) % self._max_points self._size = min(self._size + 1, self._max_points) diff --git a/dimos/navigation/go2/replanning_a_star/replan_limiter.py b/dimos/navigation/go2/replanning_a_star/replan_limiter.py index 8cc630f3dff..56b2e74e572 100644 --- a/dimos/navigation/go2/replanning_a_star/replan_limiter.py +++ b/dimos/navigation/go2/replanning_a_star/replan_limiter.py @@ -14,7 +14,8 @@ from threading import RLock -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.message_utils.geometry import point_distance +from dimos.messages.geometry_msgs.msg import Point from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -28,7 +29,7 @@ class ReplanLimiter: _max_attempts: int = 6 _reset_distance: float = 2.0 - _attempt_pos: Vector3 | None = None + _attempt_pos: Point | None = None _lock: RLock _attempt: int @@ -37,13 +38,13 @@ def __init__(self) -> None: self._lock = RLock() self._attempt = 0 - def can_retry(self, position: Vector3) -> bool: + def can_retry(self, position: Point) -> bool: with self._lock: if self._attempt == 0: self._attempt_pos = position - if self._attempt >= 1 and self._attempt_pos: - distance = self._attempt_pos.distance(position) + if self._attempt >= 1 and self._attempt_pos is not None: + distance = point_distance(self._attempt_pos, position) if distance >= self._reset_distance: logger.info( "Traveled enough to reset attempts", diff --git a/dimos/navigation/go2/replanning_a_star/spec.py b/dimos/navigation/go2/replanning_a_star/spec.py index d7aaab0ba7c..aed7a90bea8 100644 --- a/dimos/navigation/go2/replanning_a_star/spec.py +++ b/dimos/navigation/go2/replanning_a_star/spec.py @@ -14,7 +14,7 @@ from typing import Protocol -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.messages.geometry_msgs.msg import PoseStamped from dimos.navigation.base import NavigationState from dimos.spec.utils import Spec diff --git a/dimos/navigation/go2/replanning_a_star/test_generated_astar.py b/dimos/navigation/go2/replanning_a_star/test_generated_astar.py new file mode 100644 index 00000000000..66c2cd73750 --- /dev/null +++ b/dimos/navigation/go2/replanning_a_star/test_generated_astar.py @@ -0,0 +1,134 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from copy import deepcopy +from itertools import pairwise +import math + +from dimos_cdrgen.registry import decode, encode +import numpy as np +import pytest + +from dimos.mapping.occupancy.path_map import make_navigation_map +from dimos.message_utils.occupancy import ( + grid_to_world, + occupancy_from_array, + occupancy_view, + world_to_grid, +) +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import Point, Pose, Quaternion +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid as LegacyGrid +from dimos.navigation.go2.replanning_a_star.min_cost_astar import _USE_CPP, min_cost_astar + + +@pytest.mark.parametrize("use_cpp", [False, True]) +@pytest.mark.parametrize("stamp", [0, -1, 1700000000123456789]) +@pytest.mark.parametrize("rotated", [False, True]) +def test_generated_astar_routes_around_obstacles_and_retains_exact_metadata( + use_cpp, stamp, rotated +): + if use_cpp: + assert _USE_CPP, "Acceptance requires the actual compiled C++ extension" + cells = np.zeros((6, 7), dtype=np.int8) + cells[:5, 3] = 100 + origin = Pose( + position=Point(x=10, y=-20, z=0), + orientation=Quaternion( + x=0, + y=0, + z=math.sin(math.pi / 4) if rotated else 0, + w=math.cos(math.pi / 4) if rotated else 1, + ), + ) + grid = occupancy_from_array( + cells, + header=Header(frame_id="map", stamp=time_from_nanoseconds(stamp)), + resolution=0.5, + origin=origin, + ) + grid = decode(encode(grid), type(grid)) + before = deepcopy(grid) + start = grid_to_world(grid, (1.25, 1.25)) + goal = grid_to_world(grid, (5.25, 1.25)) + + result = min_cost_astar(grid, goal, start, use_cpp=use_cpp) + + assert type(result) is Path + assert result.header == grid.header + assert all(p.header == grid.header and p.pose.orientation.w == 1 for p in result.poses) + coordinates = [ + tuple(round(v) for v in world_to_grid(grid, p.pose.position)) for p in result.poses + ] + assert coordinates[0] == (1, 1) + assert coordinates[-1] == (5, 1) + assert (3, 5) in coordinates + assert all(cells[y, x] != 100 for x, y in coordinates) + assert all(max(abs(a[0] - b[0]), abs(a[1] - b[1])) == 1 for a, b in pairwise(coordinates)) + assert decode(encode(result), Path) == result + result.header.frame_id = "changed" + result.poses[0].header.frame_id = "changed" + assert result.poses[-1].header == before.header + assert grid.header == before.header and grid.info == before.info + np.testing.assert_array_equal(grid.data, before.data) + + +@pytest.mark.parametrize("use_cpp", [False, True]) +def test_generated_astar_unknown_penalty_and_unreachable_outside_goal(use_cpp): + cells = np.array([[0, -1, 0]], dtype=np.int8) + grid = occupancy_from_array( + cells, header=Header(frame_id="map", stamp=time_from_nanoseconds(0)), resolution=1 + ) + assert min_cost_astar(grid, (2, 0), unknown_penalty=1, use_cpp=use_cpp) is None + path = min_cost_astar(grid, (2, 0), unknown_penalty=0.1, use_cpp=use_cpp) + assert path is not None + assert [p.pose.position.x for p in path.poses] == [0, 1, 2] + assert min_cost_astar(grid, (3, 0), use_cpp=use_cpp) is None + + +@pytest.mark.parametrize("use_cpp", [False, True]) +def test_generated_and_existing_astar_share_identical_cell_route(use_cpp): + cells = np.zeros((7, 7), dtype=np.int8) + cells[:5, 3] = 100 + generated = occupancy_from_array( + cells, header=Header(frame_id="map", stamp=time_from_nanoseconds(-1)), resolution=0.5 + ) + legacy = LegacyGrid(cells, resolution=0.5, frame_id="map", ts=0) + old = min_cost_astar(legacy, (2.5, 0.5), (0.5, 0.5), use_cpp=use_cpp) + new = min_cost_astar(generated, (2.5, 0.5), (0.5, 0.5), use_cpp=use_cpp) + assert old is not None and new is not None + assert [(p.x, p.y, p.z) for p in old.poses] == [ + (p.pose.position.x, p.pose.position.y, p.pose.position.z) for p in new.poses + ] + + +@pytest.mark.parametrize("strategy", ["simple", "mixed"]) +@pytest.mark.parametrize("gradient_strategy", ["gradient", "voronoi"]) +def test_navigation_costmap_generation_matches_existing_math_and_preserves_metadata( + strategy, gradient_strategy +): + cells = np.zeros((10, 10), dtype=np.int8) + cells[4:7, 4:7] = 100 + cells[0, :] = -1 + generated = occupancy_from_array( + cells, header=Header(frame_id="map", stamp=time_from_nanoseconds(-1)), resolution=0.5 + ) + generated.info.map_load_time = time_from_nanoseconds(17) + old = make_navigation_map(LegacyGrid(cells, resolution=0.5), 0.5, strategy, gradient_strategy) + new = make_navigation_map(generated, 0.5, strategy, gradient_strategy) + assert new.header == generated.header and new.info == generated.info + np.testing.assert_array_equal(occupancy_view(new), old.grid) + np.testing.assert_array_equal(occupancy_view(generated), cells) diff --git a/dimos/navigation/go2/replanning_a_star/test_generated_planner.py b/dimos/navigation/go2/replanning_a_star/test_generated_planner.py new file mode 100644 index 00000000000..6679aa5149e --- /dev/null +++ b/dimos/navigation/go2/replanning_a_star/test_generated_planner.py @@ -0,0 +1,266 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from contextlib import ExitStack +from copy import deepcopy +import math +from threading import Event + +from dimos_cdrgen.registry import decode, encode +import numpy as np +import pytest + +from dimos.core.transport import LCMTransport +from dimos.message_utils.geometry import quaternion_from_euler, yaw +from dimos.message_utils.occupancy import occupancy_cell_value, occupancy_from_array +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import ( + Point, + PointStamped, + Pose, + PoseStamped, + PoseWithCovariance, + Quaternion, + TwistWithCovariance, +) +from dimos.messages.nav_msgs.msg import Odometry +from dimos.messages.std_msgs.msg import Bool, Header +from dimos.navigation.go2.replanning_a_star.controllers import PController, PdController, _velocity +from dimos.navigation.go2.replanning_a_star.module import ReplanningAStarPlanner +from dimos.navigation.go2.replanning_a_star.position_tracker import PositionTracker +from dimos.navigation.go2.replanning_a_star.replan_limiter import ReplanLimiter +from dimos.protocol.rpc.pubsubrpc import LCMRPC + + +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +def pose_at(x=0.0, y=0.0, angle=0.0, stamp=1700000000123456789): + return PoseStamped( + header=Header(frame_id="world", stamp=time_from_nanoseconds(stamp)), + pose=Pose(position=Point(x=x, y=y, z=0), orientation=quaternion_from_euler(0, 0, angle)), + ) + + +def odometry_at(pose): + return Odometry( + header=deepcopy(pose.header), + child_frame_id="base_link", + pose=PoseWithCovariance(pose=deepcopy(pose.pose), covariance=np.zeros(36)), + twist=TwistWithCovariance(twist=_velocity(), covariance=np.zeros(36)), + ) + + +@pytest.mark.parametrize( + "target,expected", + [ + ((1.0, 0.0), (0.55, 0.0)), + ((0.0, 0.0), (0.0, 0.0)), + ((math.cos(2), math.sin(2)), (0.0, 0.55)), + ], +) +def test_generated_p_controller_retains_advance_and_rotate_then_drive_commands(target, expected): + controller = PController(False, 0.55, 10) + command = controller.advance(np.array(target), pose_at()) + assert (command.linear.x, command.angular.z) == expected + assert (command.linear.y, command.linear.z, command.angular.x, command.angular.y) == ( + 0, + 0, + 0, + 0, + ) + assert decode(encode(command), type(command)) == command + + +def test_generated_pd_controller_retains_derivative_and_acceleration_state(): + controller = PdController(False, 0.55, 10) + assert controller.rotate(1).angular.z == -0.2 + assert controller.rotate(1).angular.z == 0 + controller.reset_errors() + controller.reset_yaw_error(1) + assert controller.rotate(1).angular.z == 0.2 + + +def test_generated_position_tracker_uses_existing_wall_clock_window(monkeypatch): + now = [0.0] + monkeypatch.setattr( + "dimos.navigation.go2.replanning_a_star.position_tracker.time.time", lambda: now[0] + ) + tracker = PositionTracker(2, 0.4) + assert not tracker.is_stuck() + tracker.add_position(pose_at(stamp=-1)) + assert tracker.is_stuck() + tracker.add_position(pose_at(x=1, stamp=0)) + assert not tracker.is_stuck() + now[0] = 3.0 + assert not tracker.is_stuck() + tracker.reset_data() + tracker.add_position(pose_at()) + assert tracker.is_stuck() + + +def test_generated_replanning_attempts_reset_only_after_existing_distance_limit(): + limiter = ReplanLimiter() + for _ in range(6): + assert limiter.can_retry(Point(x=0, y=0, z=0)) + limiter.will_retry() + assert not limiter.can_retry(Point(x=1, y=0, z=0)) + assert limiter.can_retry(Point(x=2, y=0, z=0)) + assert limiter.get_attempt() == 0 + + +@pytest.mark.parametrize( + "position,expected", [((0, 0), 0), ((1, 0), 100), ((2, 0), -1), ((3, 0), -1), ((-1, 0), -1)] +) +def test_generated_world_cell_lookup_preserves_unknown_outside_map(position, expected): + grid = occupancy_from_array( + np.array([[0, 100, -1]], dtype=np.int8), + header=Header(frame_id="world", stamp=time_from_nanoseconds(0)), + resolution=1, + ) + assert occupancy_cell_value(grid, Point(x=position[0], y=position[1], z=0)) == expected + + +@pytest.fixture +def planner_module(): + module = ReplanningAStarPlanner(rpc_transport=MemoryRPC) + yield module + module.stop() + + +@pytest.mark.parametrize("stamp", [0, -1, 1700000000123456789]) +def test_odometry_conversion_retains_owned_header_and_nested_pose(planner_module, stamp): + source = odometry_at(pose_at(x=1, y=2, stamp=stamp)) + planner_module._handle_odometry(source) + result = planner_module._planner._current_odom + assert result.header == source.header and result.pose == source.pose.pose + source.header.frame_id = "changed" + source.pose.pose.position.x = 100 + assert result.header.frame_id == "world" and result.pose.position.x == 1 + + +@pytest.mark.parametrize("stamp", [0, -1, 1700000000123456789]) +def test_clicked_point_conversion_retains_owned_header_and_identity_orientation( + planner_module, stamp +): + source = PointStamped( + header=Header(frame_id="world", stamp=time_from_nanoseconds(stamp)), + point=Point(x=1, y=2, z=0), + ) + planner_module._handle_clicked_point(source) + result = planner_module._planner._current_goal + assert result.header == source.header and result.pose.position == source.point + assert result.pose.orientation == Quaternion(x=0, y=0, z=0, w=1) + source.header.frame_id = "changed" + source.point.x = 100 + assert result.header.frame_id == "world" and result.pose.position.x == 1 + + +def test_actual_planner_module_emits_path_velocity_and_stop_over_cdr_memory_ports(): + with ExitStack() as cleanup: + module = ReplanningAStarPlanner(rpc_transport=MemoryRPC) + wires = {} + for name in [ + "odom", + "odometry", + "global_costmap", + "goal_request", + "clicked_point", + "target", + "stop_movement", + "goal_reached", + "navigation_state", + "nav_cmd_vel", + "path", + "navigation_costmap", + ]: + port = getattr(module, name) + wire = LCMTransport("planner_" + name, port.type, url="memq://") + port.transport = wire + wires[name] = wire + cleanup.callback(wire.stop) + threads = [] + + def shutdown(): + current = module._planner._local_planner._thread + if current is not None: + threads.append(current) + module.stop() + for thread in threads: + thread.join(5) + assert not thread.is_alive() + + cleanup.callback(shutdown) + paths = [] + commands = [] + path_ready = Event() + velocity_ready = Event() + motion_ready = Event() + stopped = Event() + + def receive_path(value): + paths.append(value) + if value.poses: + path_ready.set() + + def receive_command(value): + commands.append(value) + if value.linear.x != 0 or value.angular.z != 0: + motion_ready.set() + if value.linear.x > 0: + velocity_ready.set() + elif velocity_ready.is_set() and value == _velocity(): + stopped.set() + + cleanup.callback(wires["path"].subscribe(receive_path)) + cleanup.callback(wires["nav_cmd_vel"].subscribe(receive_command)) + module.start() + pose = pose_at(x=0.125, y=0.125) + costmap = occupancy_from_array( + np.zeros((20, 20), dtype=np.int8), header=deepcopy(pose.header), resolution=0.25 + ) + map_ready = Event() + odom_ready = Event() + cleanup.callback(wires["global_costmap"].subscribe(lambda _: map_ready.set())) + cleanup.callback(wires["odometry"].subscribe(lambda _: odom_ready.set())) + wires["global_costmap"].publish(costmap) + wires["odometry"].publish(odometry_at(pose)) + assert map_ready.wait(5) and odom_ready.wait(5) + wires["goal_request"].publish(pose_at(x=1.5, y=0.125)) + assert path_ready.wait(5) + assert motion_ready.wait(5) + route = next(p for p in paths if p.poses) + # Supply the same orientation feedback a moving robot would publish. + wires["odometry"].publish( + odometry_at( + pose_at( + x=0.125, + y=0.125, + angle=yaw(route.poses[0].pose.orientation), + stamp=1700000000123456790, + ) + ) + ) + assert velocity_ready.wait(5) + assert route.header == costmap.header + assert all(p.header == costmap.header for p in route.poses) + assert commands[-1].linear.x > 0 + current = module._planner._local_planner._thread + if current is not None: + threads.append(current) + wires["stop_movement"].publish(Bool(data=True)) + assert stopped.wait(5) + assert module._planner._current_goal is None diff --git a/dimos/navigation/go2/replanning_a_star/test_generated_safety.py b/dimos/navigation/go2/replanning_a_star/test_generated_safety.py new file mode 100644 index 00000000000..e69eb40b10d --- /dev/null +++ b/dimos/navigation/go2/replanning_a_star/test_generated_safety.py @@ -0,0 +1,168 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from copy import deepcopy +import math + +from dimos_cdrgen.registry import decode, encode +import numpy as np +import pytest + +from dimos.message_utils.occupancy import grid_to_world, occupancy_from_array, world_to_grid +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import Point, Pose, PoseStamped, Quaternion +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Header +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped as LegacyStamped +from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid as LegacyGrid +from dimos.msgs.nav_msgs.Path import Path as LegacyPath +from dimos.navigation.go2.replanning_a_star.goal_validator import find_safe_goal +from dimos.navigation.go2.replanning_a_star.path_clearance import PathClearance +from dimos.navigation.go2.replanning_a_star.path_distancer import PathDistancer + + +def grid_from_cells(cells, rotated=False): + return occupancy_from_array( + cells, + header=Header(frame_id="map", stamp=time_from_nanoseconds(-1)), + resolution=0.5, + origin=Pose( + position=Point(x=0, y=0, z=0), + orientation=Quaternion( + x=0, + y=0, + z=math.sin(math.pi / 4) if rotated else 0, + w=math.cos(math.pi / 4) if rotated else 1, + ), + ), + ) + + +@pytest.mark.parametrize("algorithm", ["bfs", "bfs_contiguous"]) +@pytest.mark.parametrize("rotated", [False, True]) +@pytest.mark.parametrize("clearance", [0, 0.5]) +def test_generated_safe_goal_preserves_existing_search_choices_and_rotated_coordinates( + algorithm, rotated, clearance +): + cells = np.zeros((12, 12), dtype=np.int8) + cells[5:7, 5:7] = 100 + grid = grid_from_cells(cells, rotated) + grid = decode(encode(grid), type(grid)) + before = encode(grid) + point = grid_to_world(grid, (5.25, 5.25)) + expected = find_safe_goal( + LegacyGrid(cells, resolution=0.5), + (5.25 * 0.5, 5.25 * 0.5), + algorithm=algorithm, + min_clearance=clearance, + connectivity_check_radius=0, + ) + + result = find_safe_goal( + grid, point, algorithm=algorithm, min_clearance=clearance, connectivity_check_radius=0 + ) + + assert type(result) is Point + assert expected is not None + coordinate = world_to_grid(grid, result) + np.testing.assert_allclose(coordinate, (expected.x / 0.5, expected.y / 0.5), atol=1e-12, rtol=0) + assert encode(grid) == before + assert decode(encode(result), Point) == result + + +@pytest.mark.parametrize("algorithm", ["bfs", "bfs_contiguous"]) +def test_generated_safe_goal_returns_none_when_no_safe_cell_exists(algorithm): + grid = grid_from_cells(np.full((7, 7), 100, dtype=np.int8)) + assert ( + find_safe_goal(grid, Point(x=1, y=1, z=0), algorithm=algorithm, connectivity_check_radius=0) + is None + ) + + +@pytest.mark.parametrize("algorithm", ["bfs", "bfs_contiguous"]) +def test_generated_safe_goal_rejects_unknown_cells(algorithm): + grid = grid_from_cells(np.full((7, 7), -1, dtype=np.int8)) + assert find_safe_goal(grid, (1, 1), algorithm=algorithm, connectivity_check_radius=0) is None + + +def make_path(): + header = Header(frame_id="map", stamp=time_from_nanoseconds(1700000000123456789)) + return Path( + header=deepcopy(header), + poses=[ + PoseStamped( + header=deepcopy(header), + pose=Pose( + position=Point(x=x, y=1, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1) + ), + ) + for x in [1.0, 2.0, 3.0, 4.0, 5.0] + ], + ) + + +@pytest.mark.parametrize("index,expected", [(0, (1.5, 1)), (1, (2.5, 1)), (4, (5.0, 1))]) +def test_generated_path_distances_and_lookahead_match_existing_geometry(index, expected): + path = make_path() + original = deepcopy(path) + helper = PathDistancer(path) + old = PathDistancer( + LegacyPath( + poses=[ + LegacyStamped(position=[p.pose.position.x, p.pose.position.y, 0]) + for p in path.poses + ] + ) + ) + np.testing.assert_array_equal(helper.find_lookahead_point(index), expected) + np.testing.assert_array_equal( + helper.find_lookahead_point(index), old.find_lookahead_point(index) + ) + assert helper.find_closest_point_index(np.array([2.1, 1.0])) == 1 + assert helper.get_distance_to_path(np.array([2.0, 2.0])) == 1 + assert helper.distance_to_goal(np.array([2.0, 1.0])) == 3 + assert path == original + + +def test_generated_path_clearance_detects_new_obstacle_using_retained_path_mask(): + path = make_path() + helper = PathClearance(0.2, path) + assert helper.is_obstacle_ahead() + cells = np.zeros((15, 15), dtype=np.int8) + clear = grid_from_cells(cells) + helper.update_costmap(clear) + assert not helper.is_obstacle_ahead() + cells[2, 4] = 100 + helper.update_costmap(grid_from_cells(cells)) + assert helper.is_obstacle_ahead() + helper.update_pose_index(3) + assert not helper.is_obstacle_ahead() + assert clear.header.stamp == time_from_nanoseconds(-1) + + +def test_contiguous_generated_search_cannot_cross_an_enclosing_obstacle(): + cells = np.zeros((12, 12), dtype=np.int8) + cells[4:8, 4:8] = 100 + grid = grid_from_cells(cells) + goal = grid_to_world(grid, (5.25, 5.25)) + assert ( + find_safe_goal( + grid, goal, algorithm="bfs_contiguous", min_clearance=0, connectivity_check_radius=0 + ) + is None + ) + assert ( + find_safe_goal(grid, goal, algorithm="bfs", min_clearance=0, connectivity_check_radius=0) + is not None + ) diff --git a/dimos/navigation/go2/replanning_a_star/test_global_planner.py b/dimos/navigation/go2/replanning_a_star/test_global_planner.py index 55428a8a625..d89d73d83f1 100644 --- a/dimos/navigation/go2/replanning_a_star/test_global_planner.py +++ b/dimos/navigation/go2/replanning_a_star/test_global_planner.py @@ -13,14 +13,23 @@ # limitations under the License. import numpy as np +import pytest -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid +from dimos.message_utils.occupancy import occupancy_from_array +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import Point +from dimos.messages.std_msgs.msg import Header from dimos.navigation.go2.replanning_a_star.global_planner import GlobalPlanner -def test_find_wide_path_with_start_inside_inflation() -> None: +@pytest.fixture +def planner(): + value = GlobalPlanner(robot_width=0.3, robot_rotation_diameter=0.6) + yield value + value.stop() + + +def test_find_wide_path_with_start_inside_inflation(planner) -> None: """A wall observed at the last moment can be so close that its inflation covers the robot's own cell (the robot drove there before the costmap caught up). Planning must still find a way out instead of failing.""" @@ -28,18 +37,23 @@ def test_find_wide_path_with_start_inside_inflation() -> None: resolution = 0.05 grid = np.zeros((60, 60), dtype=np.int8) grid[20:40, 30] = 100 # wall at x=1.5m spanning y=1.0..2.0m - costmap = OccupancyGrid(grid=grid, resolution=resolution, origin=Pose(), frame_id="world") + costmap = occupancy_from_array( + grid, + resolution=resolution, + header=Header(frame_id="world", stamp=time_from_nanoseconds(1700000000123456789)), + ) - planner = GlobalPlanner(robot_width=0.3, robot_rotation_diameter=0.6) planner.handle_global_costmap(costmap) # 7 cm in front of the wall: within the inflation radius # (robot_width * 1.1 / 2 = 0.165m), so the start cell is engulfed. - robot = Vector3(1.43, 1.5, 0) + robot = Point(x=1.43, y=1.5, z=0) # On the other side of the wall; the path must round a wall end. - goal = Vector3(2.75, 1.5, 0) + goal = Point(x=2.75, y=1.5, z=0) path = planner._find_wide_path(goal, robot) assert path is not None assert len(path.poses) > 0 + + assert path.header == costmap.header diff --git a/dimos/navigation/local_planner/native.py b/dimos/navigation/local_planner/native.py index 41d59844e30..a93df090624 100644 --- a/dimos/navigation/local_planner/native.py +++ b/dimos/navigation/local_planner/native.py @@ -20,9 +20,9 @@ from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import IO, In, Out -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.tf2_msgs.msg import TFMessage from dimos.navigation import spec from dimos.navigation.embodiment.base import Embodiment @@ -72,6 +72,7 @@ class LocalPlannerNative(NativeModule, spec.MapLocalPlanner): planner_path: In[Path] # IO, not In: rust `#[tf]` subscribes and publishes, and refuses to start on a mismatch. tf: IO[TFMessage] + tf_static: In[TFMessage] path: Out[Path] # The floor voxels the ground plane was fit to, moved onto it. Empty while the body height stands in. diff --git a/dimos/navigation/local_planner/profile.py b/dimos/navigation/local_planner/profile.py index e2a2a0a442c..29983b1fdec 100644 --- a/dimos/navigation/local_planner/profile.py +++ b/dimos/navigation/local_planner/profile.py @@ -25,13 +25,15 @@ import numpy as np from numpy.typing import NDArray +from dimos.message_utils.time import NANOSECONDS_PER_SECOND, to_nanoseconds +from dimos.messages.nav_msgs.msg import Path as CdrPath from dimos.msgs.nav_msgs.Path import Path from dimos.navigation.embodiment.base import Embodiment _FAN_EPS = 1e-6 # segment shorter than this is a rotation, not a move -def decode_ceilings(path: Path, lo: float, hi: float) -> NDArray[np.float64] | None: +def decode_ceilings(path: Path | CdrPath, lo: float, hi: float) -> NDArray[np.float64] | None: """Per-waypoint speed ceiling (m/s) from the stamps; None if unstamped. Unstamped = flat or non-monotone ts. Fan segments inherit the previous @@ -40,11 +42,15 @@ def decode_ceilings(path: Path, lo: float, hi: float) -> NDArray[np.float64] | N n = len(path.poses) if n < 2: return None - ts = np.array([p.ts for p in path.poses]) - dt = np.diff(ts) + if isinstance(path, CdrPath): + stamps = np.asarray([to_nanoseconds(p.header.stamp) for p in path.poses], dtype=np.int64) + dt = np.diff(stamps).astype(np.float64) / NANOSECONDS_PER_SECOND + xy = np.asarray([[p.pose.position.x, p.pose.position.y] for p in path.poses]) + else: + dt = np.diff(np.asarray([p.ts for p in path.poses])) + xy = np.asarray([[p.position.x, p.position.y] for p in path.poses]) if np.any(dt < 0) or not np.any(dt > 0): return None - xy = np.array([[p.position.x, p.position.y] for p in path.poses]).reshape(-1, 2) ds = np.linalg.norm(np.diff(xy, axis=0), axis=1) # a config is free to set min above max; ordering here rather than # trusting the caller keeps the clip a clip diff --git a/dimos/navigation/local_planner/rust/Cargo.lock b/dimos/navigation/local_planner/rust/Cargo.lock index 24e92d272ff..58922f8f8fe 100644 --- a/dimos/navigation/local_planner/rust/Cargo.lock +++ b/dimos/navigation/local_planner/rust/Cargo.lock @@ -249,6 +249,15 @@ dependencies = [ "shlex", ] +[[package]] +name = "cdr" +version = "0.2.4" +source = "git+https://github.com/TomCC7/cdr-rs?rev=38a4368255d408faf374726653b39d4a4e9735cb#38a4368255d408faf374726653b39d4a4e9735cb" +dependencies = [ + "byteorder", + "serde", +] + [[package]] name = "cfg-if" version = "1.0.5" @@ -579,6 +588,17 @@ dependencies = [ "subtle", ] +[[package]] +name = "dimos-generated-messages" +version = "0.0.14" +dependencies = [ + "cdr", + "heapless", + "ros2msg", + "serde", + "serde-big-array", +] + [[package]] name = "dimos-lcm" version = "0.1.0" @@ -593,8 +613,8 @@ dependencies = [ name = "dimos-local-planner" version = "0.1.0" dependencies = [ + "dimos-generated-messages", "dimos-module", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen)", "serde", "serde_json", "tokio", @@ -605,12 +625,13 @@ dependencies = [ [[package]] name = "dimos-module" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#fc6a4e58da89bb073a38dce5ac63728c6559d3a6" +version = "0.2.0" dependencies = [ + "cdr", + "dimos-generated-messages", "dimos-lcm", "dimos-module-macros", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549)", + "lcm-msgs", "nalgebra", "rayon", "serde", @@ -618,6 +639,7 @@ dependencies = [ "tokio", "tracing", "tracing-subscriber", + "transforms", "url", "validator", "zenoh", @@ -626,7 +648,6 @@ dependencies = [ [[package]] name = "dimos-module-macros" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#fc6a4e58da89bb073a38dce5ac63728c6559d3a6" dependencies = [ "proc-macro2", "quote", @@ -965,6 +986,15 @@ version = "0.33.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e762b9b188634fc508d4ec54b62ea3d352c43fd431d18d05d46f559f217f4725" +[[package]] +name = "hash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" +dependencies = [ + "byteorder", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -1006,6 +1036,28 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] + +[[package]] +name = "heapless" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" +dependencies = [ + "hash32", + "serde", + "stable_deref_trait", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "hermit-abi" @@ -1397,14 +1449,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "lcm-msgs" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" -dependencies = [ - "byteorder", -] - [[package]] name = "lcm-msgs" version = "0.1.0" @@ -1949,6 +1993,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -2302,6 +2356,23 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "ros2msg" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fe291be5cb6d4b8e6bf9c9de80c40bd1cb09ecac50deafa2fdcae0027e04988" +dependencies = [ + "heck", + "pest", + "pest_derive", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "syn 2.0.119", + "thiserror", +] + [[package]] name = "rustc-hash" version = "2.1.3" @@ -2542,6 +2613,15 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-big-array" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11fc7cc2c76d73e0f27ee52abbd64eec84d46f370c88371120433196934e4b7f" +dependencies = [ + "serde", +] + [[package]] name = "serde_core" version = "1.0.229" @@ -3217,6 +3297,18 @@ dependencies = [ "tracing-serde", ] +[[package]] +name = "transforms" +version = "2.2.0" +source = "git+https://github.com/TomCC7/transforms?rev=5ae9563de0c02638c04e393aadfad28444a4398e#5ae9563de0c02638c04e393aadfad28444a4398e" +dependencies = [ + "approx", + "hashbrown 0.17.1", + "libm", + "serde", + "thiserror", +] + [[package]] name = "twox-hash" version = "1.6.3" diff --git a/dimos/navigation/local_planner/rust/Cargo.toml b/dimos/navigation/local_planner/rust/Cargo.toml index 1d90441cb59..6557c489fc2 100644 --- a/dimos/navigation/local_planner/rust/Cargo.toml +++ b/dimos/navigation/local_planner/rust/Cargo.toml @@ -17,8 +17,8 @@ required-features = ["module"] [dependencies] serde = { version = "1", features = ["derive"] } serde_json = { version = "1", features = ["float_roundtrip"] } -dimos-module = { git = "https://github.com/dimensionalOS/dimos", rev = "fc6a4e58da89bb073a38dce5ac63728c6559d3a6", optional = true, version = "=0.1.0" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen", optional = true } +dimos-module = { path = "../../../../native/rust/dimos-module", optional = true } +dimos-generated-messages = { path = "../../../../build/message-source/rust", optional = true } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "time"], optional = true } tracing = { version = "0.1", optional = true } tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true } @@ -27,7 +27,7 @@ validator = { version = "0.21", features = ["derive"], optional = true } [features] # Nothing by default: `module` is the native module that runs on the robot. default = [] -module = ["dep:dimos-module", "dep:lcm-msgs", "dep:tokio", "dep:tracing", "dep:tracing-subscriber", "dep:validator"] +module = ["dep:dimos-module", "dep:dimos-generated-messages", "dep:tokio", "dep:tracing", "dep:tracing-subscriber", "dep:validator"] # `Emb::fixture()`, a test body for this crate and its dependants. test-bodies = [] @@ -42,13 +42,14 @@ python = "dimos.navigation.local_planner.native:LocalPlannerNative" threads = 2 [package.metadata.dimos.module.local_planner.inputs] -local_map = "sensor_msgs.PointCloud2" -planner_path = "nav_msgs.Path" -tf = "tf2_msgs.TFMessage" +local_map = "sensor_msgs/msg/PointCloud2" +planner_path = "nav_msgs/msg/Path" +tf = "tf2_msgs/msg/TFMessage" +tf_static = "tf2_msgs/msg/TFMessage" [package.metadata.dimos.module.local_planner.outputs] -path = "nav_msgs.Path" -ground_support = "sensor_msgs.PointCloud2" +path = "nav_msgs/msg/Path" +ground_support = "sensor_msgs/msg/PointCloud2" [profile.release] lto = "thin" diff --git a/dimos/navigation/local_planner/rust/crate-hashes.json b/dimos/navigation/local_planner/rust/crate-hashes.json index 187d28513c6..3ee4f926e23 100644 --- a/dimos/navigation/local_planner/rust/crate-hashes.json +++ b/dimos/navigation/local_planner/rust/crate-hashes.json @@ -1,7 +1,6 @@ { - "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#lcm-msgs@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=", + "git+https://github.com/TomCC7/cdr-rs?rev=38a4368255d408faf374726653b39d4a4e9735cb#cdr@0.2.4": "sha256-akg/bLBY3wQ1YRr2bijE5ZqB211Cst5lKjELLrgiG6o=", + "git+https://github.com/TomCC7/transforms?rev=5ae9563de0c02638c04e393aadfad28444a4398e#transforms@2.2.0": "sha256-Kh8aVjj7zQ+Xd1QHOcAhSJbubg7azaWFSMpaI/ZAgM4=", "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549#dimos-lcm@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=", - "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549#lcm-msgs@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=", - "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#dimos-module-macros@0.1.0": "sha256-ECU2xYh/j+a3XwyhO3UR/GPruo5jNzM3n/9S5BIvVrA=", - "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#dimos-module@0.1.0": "sha256-ECU2xYh/j+a3XwyhO3UR/GPruo5jNzM3n/9S5BIvVrA=" + "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549#lcm-msgs@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=" } \ No newline at end of file diff --git a/dimos/navigation/local_planner/rust/flake.nix b/dimos/navigation/local_planner/rust/flake.nix index 4e8e518e592..de999f620a3 100644 --- a/dimos/navigation/local_planner/rust/flake.nix +++ b/dimos/navigation/local_planner/rust/flake.nix @@ -15,7 +15,31 @@ pkgs = nixpkgs.legacyPackages.${system}; name = "dimos-local-planner"; - src = nix-filter.lib { root = ./.; exclude = [ "target" "build" "result" "__pycache__" ]; }; + moduleSource = nix-filter.lib { root = ./.; exclude = [ "target" "build" "result" "__pycache__" ]; }; + sdkArchive = pkgs.fetchurl { + url = "https://codeload.github.com/dimensionalOS/dimos/tar.gz/a9365487b4323151ea85041c6a944b509c0c3268"; + sha256 = "e339487dd2281980fcde7ab03c46dfd2e274491e0f73e66d9c6cbcaeea8902b2"; + }; + sdkSource = pkgs.runCommand "dimos-local-planner-sdk-source" { + nativeBuildInputs = [ pkgs.gnutar pkgs.gzip ]; + } '' + mkdir -p $out + tar -xzf ${sdkArchive} --strip-components=1 -C $out + ''; + messageSource = (import "${sdkSource}/native/message-source.nix" { inherit pkgs; }).source; + # Standalone flake builds resolve explicit source artifacts, not checkout parents. + src = pkgs.runCommand "local-planner-cdr-source" { } '' + cp -r ${moduleSource} $out + chmod -R u+w $out + mkdir -p $out/sdk + cp -r ${sdkSource}/native/rust/dimos-module ${sdkSource}/native/rust/dimos-module-macros $out/sdk/ + cp -r ${messageSource}/rust $out/messages + substituteInPlace $out/sdk/dimos-module/Cargo.toml \ + --replace-fail '../../../build/message-source/rust' '../../messages' + substituteInPlace $out/Cargo.toml \ + --replace-fail '../../../../native/rust/dimos-module' 'sdk/dimos-module' \ + --replace-fail '../../../../build/message-source/rust' 'messages' + ''; generated = crate2nix.tools.${system}.generatedCargoNix { inherit name src; }; diff --git a/dimos/navigation/local_planner/rust/src/module/msg.rs b/dimos/navigation/local_planner/rust/src/module/msg.rs index cd8fcdf0018..26705c59dfa 100644 --- a/dimos/navigation/local_planner/rust/src/module/msg.rs +++ b/dimos/navigation/local_planner/rust/src/module/msg.rs @@ -17,11 +17,17 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use dimos_generated_messages::builtin_interfaces::msg::time::Time; +use dimos_generated_messages::geometry_msgs::msg::{ + point::Point, pose::Pose, pose_stamped::PoseStamped, quaternion::Quaternion, twist::Twist, + vector3::Vector3, +}; +use dimos_generated_messages::nav_msgs::msg::path::Path; +use dimos_generated_messages::sensor_msgs::msg::{ + point_cloud2::PointCloud2, point_field::PointField, +}; +use dimos_generated_messages::std_msgs::msg::header::Header; use dimos_module::{error_throttled, Output}; -use lcm_msgs::geometry_msgs::{Point, Pose, PoseStamped, Quaternion, Vector3}; -use lcm_msgs::nav_msgs::Path; -use lcm_msgs::sensor_msgs::{PointCloud2, PointField}; -use lcm_msgs::std_msgs::{Header, Time}; /// A plan waypoint as the laws and the planner see it: `(x, y, yaw)`. pub type State = [f64; 3]; @@ -46,19 +52,25 @@ pub fn quat_of_yaw(yaw: f64) -> Quaternion { /// Seconds since the unix epoch as one float, the python `PoseStamped.ts`. pub fn secs_of(t: &Time) -> f64 { - t.sec as f64 + t.nsec as f64 * 1e-9 + t.sec as f64 + t.nanosec as f64 * 1e-9 } -/// A float second count as a `Time`, matching `Header.__init__`'s truncation -/// (`sec = int(ts)`, `nsec = int((ts - sec) * 1e9)`). +pub fn nanoseconds(time: &Time) -> i64 { + time.sec as i64 * 1_000_000_000 + time.nanosec as i64 +} + +/// Normalize a floating second count into the canonical ROS Time fields. pub fn time_of_secs(ts: f64) -> Time { if !ts.is_finite() { - return Time { sec: 0, nsec: 0 }; + return Time { sec: 0, nanosec: 0 }; } - let sec = ts.trunc(); + let ns = (ts * 1_000_000_000.).round() as i64; + let sec = ns + .div_euclid(1_000_000_000) + .clamp(i32::MIN as i64, i32::MAX as i64) as i32; Time { - sec: sec.clamp(i32::MIN as f64, i32::MAX as f64) as i32, - nsec: ((ts - sec) * 1e9) as i32, + sec, + nanosec: ns.rem_euclid(1_000_000_000) as u32, } } @@ -72,7 +84,6 @@ pub fn now_secs() -> f64 { pub fn header(frame_id: &str, ts: f64) -> Header { Header { - seq: 0, stamp: time_of_secs(ts), frame_id: frame_id.into(), } @@ -122,20 +133,20 @@ pub fn build_cloud(points: &[[f32; 3]], frame_id: &str, ts: f64) -> PointCloud2 data.extend_from_slice(&v.to_le_bytes()); } } - let field = |name: &str, offset: i32| PointField { + let field = |name: &str, offset: u32| PointField { name: name.into(), offset, - datatype: PointField::FLOAT32 as u8, + datatype: PointField::FLOAT32, count: 1, }; PointCloud2 { header: header(frame_id, ts), height: 1, - width: points.len() as i32, + width: points.len() as u32, fields: vec![field("x", 0), field("y", 4), field("z", 8)], is_bigendian: false, point_step: 12, - row_step: 12 * points.len() as i32, + row_step: 12 * points.len() as u32, data, is_dense: true, } @@ -163,9 +174,21 @@ pub fn path_stamps(path: &Path) -> Vec { .collect() } +/// Relative seconds for laws that consume only profile deltas. +pub fn relative_path_stamps(path: &Path) -> Vec { + let Some(first) = path.poses.first() else { + return Vec::new(); + }; + let base = nanoseconds(&first.header.stamp); + path.poses + .iter() + .map(|pose| (nanoseconds(&pose.header.stamp) - base) as f64 / 1_000_000_000.) + .collect() +} + /// A body-frame twist, the only shape either module publishes. -pub fn twist(vx: f64, vy: f64, wz: f64) -> lcm_msgs::geometry_msgs::Twist { - lcm_msgs::geometry_msgs::Twist { +pub fn twist(vx: f64, vy: f64, wz: f64) -> Twist { + Twist { linear: Vector3 { x: vx, y: vy, @@ -192,11 +215,34 @@ pub async fn publish(out: &Output, msg: &T) { } // Reading a cloud off the wire is dimos-module's; the path stays `msg::extract_xyz`. -pub use dimos_module::pointcloud::{extract_xyz, ExtractError}; +pub use dimos_module::pointcloud::ExtractError; + +pub fn extract_xyz( + cloud: &dimos_generated_messages::sensor_msgs::msg::point_cloud2::PointCloud2, +) -> Result, ExtractError> { + use dimos_generated_messages::sensor_msgs::msg::point_field::PointField; + let offset = |name: &str| { + cloud + .fields + .iter() + .find(|f| f.name == name && f.datatype == PointField::FLOAT32 && f.count == 1) + .map(|f| f.offset as usize) + .ok_or(ExtractError("missing scalar float32 xyz field")) + }; + dimos_module::pointcloud::extract_xyz_rows( + &cloud.data, + [cloud.width as usize, cloud.height as usize], + [cloud.point_step as usize, cloud.row_step as usize], + [offset("x")?, offset("y")?, offset("z")?], + cloud.is_bigendian, + ) +} #[cfg(test)] mod tests { use super::*; + + use dimos_module::cdr; use std::f64::consts::PI; #[test] @@ -219,8 +265,8 @@ mod tests { #[test] fn a_non_finite_stamp_does_not_wrap_around() { // a saturating `as i32` would give i32::MAX seconds; zero is the honest "no time" - assert_eq!(time_of_secs(f64::NAN), Time { sec: 0, nsec: 0 }); - assert_eq!(time_of_secs(f64::INFINITY), Time { sec: 0, nsec: 0 }); + assert_eq!(time_of_secs(f64::NAN), Time { sec: 0, nanosec: 0 }); + assert_eq!(time_of_secs(f64::INFINITY), Time { sec: 0, nanosec: 0 }); } #[test] @@ -245,4 +291,83 @@ mod tests { let path = build_path(&states, &[], 5.0, "odom", 0.0); assert_eq!(path_stamps(&path), vec![5.0, 5.0]); } + + #[test] + fn generated_path_codec_and_relative_profile_keep_nanosecond_differences() { + let mut path = build_path( + &[[0., 0., 0.], [1., 0., 0.], [2., 0., 0.]], + &[], + 0., + "map", + 0., + ); + let base = 1_700_000_000_123_456_789i64; + for (pose, delta) in path.poses.iter_mut().zip([0, 1, 3]) { + let ns = base + delta; + pose.header.stamp = Time { + sec: ns.div_euclid(1_000_000_000) as i32, + nanosec: ns.rem_euclid(1_000_000_000) as u32, + }; + } + for little in [false, true] { + let payload = cdr::encode_endian(&path, little).unwrap(); + let decoded: Path = cdr::decode(&payload).unwrap(); + assert_eq!(relative_path_stamps(&decoded), vec![0., 1e-9, 3e-9]); + assert_eq!(decoded.header.frame_id, "map"); + } + assert_eq!( + time_of_secs(-0.25), + Time { + sec: -1, + nanosec: 750_000_000 + } + ); + } + + #[test] + fn generated_cloud_layout_handles_padding_endian_and_cdr() { + use dimos_generated_messages::sensor_msgs::msg::{ + point_cloud2::PointCloud2, point_field::PointField, + }; + let values = [[1f32, 2., 3.], [4., 5., 6.], [7., 8., 9.], [10., 11., 12.]]; + for big in [false, true] { + let mut data = vec![0u8; 80]; + for (index, point) in values.iter().enumerate() { + for (axis, value) in point.iter().enumerate() { + let offset = (index / 2) * 40 + (index % 2) * 16 + axis * 4; + let bytes = if big { + value.to_be_bytes() + } else { + value.to_le_bytes() + }; + data[offset..offset + 4].copy_from_slice(&bytes); + } + } + let fields = ["x", "y", "z"] + .iter() + .enumerate() + .map(|(i, name)| PointField { + name: (*name).into(), + offset: (i * 4) as u32, + datatype: PointField::FLOAT32, + count: 1, + }) + .collect(); + let cloud = PointCloud2 { + header: header("map", 0.), + height: 2, + width: 2, + fields, + is_bigendian: big, + point_step: 16, + row_step: 40, + data, + is_dense: true, + }; + let mut decoded: PointCloud2 = cdr::decode(&cdr::encode(&cloud).unwrap()).unwrap(); + assert_eq!(extract_xyz(&decoded).unwrap(), values); + decoded.row_step = 7; + assert!(extract_xyz(&decoded).is_err()); + } + } } diff --git a/dimos/navigation/local_planner/rust/src/module/planner.rs b/dimos/navigation/local_planner/rust/src/module/planner.rs index e3167c90745..85d3923d5fc 100644 --- a/dimos/navigation/local_planner/rust/src/module/planner.rs +++ b/dimos/navigation/local_planner/rust/src/module/planner.rs @@ -22,9 +22,9 @@ use std::time::{Duration, Instant}; use crate::planner::{plan_explored, step_reach, Emb, Terrain, COMMIT_MARGIN}; use crate::{clearance, stamps}; +use dimos_generated_messages::nav_msgs::msg::path::Path; +use dimos_generated_messages::sensor_msgs::msg::point_cloud2::PointCloud2; use dimos_module::{debug_throttled, native_config, warn_throttled, Input, Module, Output, Tf}; -use lcm_msgs::nav_msgs::Path; -use lcm_msgs::sensor_msgs::PointCloud2; use tracing::{debug, info, warn}; use validator::ValidationError; @@ -102,17 +102,17 @@ struct Shared { #[derive(Module)] #[module(name = "local_planner", setup = spawn_worker, teardown = stop_worker)] pub struct LocalPlanner { - #[input(decode = PointCloud2::decode, handler = on_local_map)] + #[input(decode = dimos_module::cdr::decode, handler = on_local_map)] local_map: Input, - #[input(decode = Path::decode, handler = on_planner_path)] + #[input(decode = dimos_module::cdr::decode, handler = on_planner_path)] planner_path: Input, - #[output(encode = Path::encode)] + #[output(encode = dimos_module::cdr::encode)] path: Output, /// The floor voxels the ground plane was fit to, moved onto it. Empty while the body height stands in. - #[output(encode = PointCloud2::encode)] + #[output(encode = dimos_module::cdr::encode)] ground_support: Output, #[config] diff --git a/dimos/navigation/local_planner/rust/src/module/tf_pose.rs b/dimos/navigation/local_planner/rust/src/module/tf_pose.rs index 6351ee970b6..b8ae3537237 100644 --- a/dimos/navigation/local_planner/rust/src/module/tf_pose.rs +++ b/dimos/navigation/local_planner/rust/src/module/tf_pose.rs @@ -18,8 +18,8 @@ use std::time::Instant; -use dimos_module::nalgebra::{Isometry3, Translation3}; -use dimos_module::{Tf, Transform}; +use dimos_module::nalgebra::Isometry3; +use dimos_module::{Tf, TransformStamped}; use crate::module::msg::{yaw_of, State}; @@ -35,7 +35,7 @@ pub struct BasePose { pub struct PoseWatch { max_age_s: f64, /// The stamp last seen, and when it was first seen. - seen: Option<(f64, Instant)>, + seen: Option<(i64, Instant)>, } impl PoseWatch { @@ -47,7 +47,7 @@ impl PoseWatch { } pub fn get(&mut self, tf: &Tf, world: &str, base: &str, now: Instant) -> Option { - self.observe(tf.get_latest(world, base).map(edge_of), now) + self.observe_nanoseconds(tf.get_latest(world, base).and_then(edge_of), now) } /// The tick's pose off what tf handed back, as `(stamp, transform)`. @@ -55,6 +55,19 @@ impl PoseWatch { &mut self, edge: Option<(f64, Isometry3)>, now: Instant, + ) -> Option { + let edge = edge.and_then(|(seconds, iso)| { + let ns = (seconds * 1_000_000_000.).round(); + (ns.is_finite() && ns >= i64::MIN as f64 && ns < i64::MAX as f64) + .then_some((ns as i64, iso)) + }); + self.observe_nanoseconds(edge, now) + } + + pub fn observe_nanoseconds( + &mut self, + edge: Option<(i64, Isometry3)>, + now: Instant, ) -> Option { let Some((ts, iso)) = edge else { self.seen = None; @@ -77,28 +90,30 @@ impl PoseWatch { } } -fn edge_of(t: Transform) -> (f64, Isometry3) { - ( - t.ts, - Isometry3::from_parts(Translation3::from(t.translation()), t.rotation()), - ) +fn edge_of(t: TransformStamped) -> Option<(i64, Isometry3)> { + Some(( + dimos_module::tf::nanoseconds(&t.header.stamp).ok()?, + dimos_module::tf::isometry(&t).ok()?, + )) } pub fn state_of(iso: &Isometry3) -> State { let q = iso.rotation.quaternion(); - let yaw = yaw_of(&lcm_msgs::geometry_msgs::Quaternion { - x: q.i, - y: q.j, - z: q.k, - w: q.w, - }); + let yaw = yaw_of( + &dimos_generated_messages::geometry_msgs::msg::quaternion::Quaternion { + x: q.i, + y: q.j, + z: q.k, + w: q.w, + }, + ); [iso.translation.x, iso.translation.y, yaw] } #[cfg(test)] mod tests { use super::*; - use dimos_module::nalgebra::UnitQuaternion; + use dimos_module::nalgebra::{Translation3, UnitQuaternion}; use std::time::Duration; fn body_at(x: f64, yaw: f64, ts: f64) -> Option<(f64, Isometry3)> { @@ -167,7 +182,7 @@ mod tests { } #[test] - fn the_yaw_round_trips_through_the_lcm_quaternion() { + fn the_yaw_round_trips_through_the_generated_quaternion() { for yaw in [-3.0, -1.0, 0.0, 0.7, 3.0] { let iso = Isometry3::from_parts( Translation3::new(0.0, 0.0, 0.0), @@ -177,4 +192,31 @@ mod tests { assert!((got[2] - yaw).abs() < 1e-12, "yaw {yaw} -> {got:?}"); } } + #[test] + fn epoch_nanosecond_progress_refreshes_the_existing_deadman() { + let mut watch = PoseWatch::new(2.5); + let now = Instant::now(); + let ns = 1_700_000_000_123_456_789; + let first = dimos_module::tf::from_isometry( + "world", + "base", + dimos_module::tf::time_from_nanoseconds(ns).unwrap(), + Isometry3::identity(), + ); + assert!(watch + .observe_nanoseconds(edge_of(first.clone()), now) + .is_some()); + assert!(watch + .observe_nanoseconds(edge_of(first), now + Duration::from_secs(3)) + .is_none()); + let fresh = dimos_module::tf::from_isometry( + "world", + "base", + dimos_module::tf::time_from_nanoseconds(ns + 1).unwrap(), + Isometry3::identity(), + ); + assert!(watch + .observe_nanoseconds(edge_of(fresh), now + Duration::from_secs(3)) + .is_some()); + } } diff --git a/dimos/navigation/local_planner/test_generated_native.py b/dimos/navigation/local_planner/test_generated_native.py new file mode 100644 index 00000000000..922f23c82ec --- /dev/null +++ b/dimos/navigation/local_planner/test_generated_native.py @@ -0,0 +1,174 @@ +# 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 contextlib import ExitStack +from copy import deepcopy +from itertools import pairwise +import json +import os +from pathlib import Path as FilePath +import socket +from threading import Event +import time + +import numpy as np +import pytest +import zenoh + +from dimos.core.global_config import global_config +from dimos.core.transport import ZenohTransport +from dimos.message_utils.pointcloud import pointcloud_from_xyz +from dimos.message_utils.time import time_from_nanoseconds, to_nanoseconds +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.std_msgs.msg import Header +from dimos.messages.tf2_msgs.msg import TFMessage +from dimos.navigation.embodiment.go2 import GO2 +from dimos.navigation.local_planner.native import LocalPlannerNative +from dimos.protocol.pubsub.impl.zenohpubsub import Topic as ZenohTopic +from dimos.protocol.rpc.pubsubrpc import LCMRPC +from dimos.protocol.service.zenohservice import ZenohConfig + + +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +def test_native_planner_declares_canonical_generated_port_identities(): + module = LocalPlannerNative(rpc_transport=MemoryRPC, embodiment=GO2) + try: + assert module.local_map.type is PointCloud2 + assert module.planner_path.type is Path and module.path.type is Path + assert module.tf.type is TFMessage + finally: + module.stop() + + +@pytest.mark.native_e2e +def test_actual_native_planner_consumes_python_cdr_and_returns_canonical_path(monkeypatch): + executable = os.environ.get("DIMOS_NATIVE_LOCAL_PLANNER_BINARY") + assert executable, "Build and set the native binary before running native_e2e" + binary = FilePath(executable).resolve() + assert binary.is_file() + monkeypatch.setattr(global_config, "transport", "zenoh") + with socket.socket() as reservation: + reservation.bind(("127.0.0.1", 0)) + endpoint = f"tcp/127.0.0.1:{reservation.getsockname()[1]}" + config = zenoh.Config() + for key, value in { + "mode": "router", + "listen/endpoints": [endpoint], + "scouting/multicast/enabled": False, + "scouting/gossip/enabled": False, + }.items(): + config.insert_json5(key, json.dumps(value)) + with ExitStack() as cleanup: + cleanup.enter_context(zenoh.open(config)) + session = ZenohConfig( + mode="client", + connect=[endpoint], + listen=[], + scouting=False, + multicast=False, + gossip=False, + connect_timeout=5, + ) + module = LocalPlannerNative( + embodiment=GO2, + rpc_transport=MemoryRPC, + executable=str(binary), + source_dir=None, + build_command=None, + session=session, + ) + wires = {} + for name in ["local_map", "planner_path", "tf", "tf_static", "path", "ground_support"]: + port = getattr(module, name) + wire = ZenohTransport( + ZenohTopic("native/planner/" + name, port.type), **session.model_dump() + ) + port.transport = wire + wires[name] = wire + cleanup.callback(wire.stop) + cleanup.callback(module.stop) + received = [] + ready = Event() + + def receive(value): + if len(value.poses) >= 2: + received.append(value) + ready.set() + + cleanup.callback(wires["path"].subscribe(receive)) + module.start() + points = np.array( + [[x, y, 0.0] for x in np.arange(-1.0, 4.1, 0.25) for y in np.arange(-2.0, 2.1, 0.25)], + dtype=np.float32, + ) + deadline = time.monotonic() + 15 + index = 0 + while not ready.is_set() and time.monotonic() < deadline: + header = Header( + frame_id="odom", stamp=time_from_nanoseconds(1700000000123456789 + index) + ) + wires["tf"].publish( + TFMessage( + transforms=[ + TransformStamped( + header=deepcopy(header), + child_frame_id="base_link", + transform=Transform( + translation=Vector3(x=0, y=0, z=GO2.base_height), + rotation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) + ] + ) + ) + wires["local_map"].publish(pointcloud_from_xyz(points, header=deepcopy(header))) + route = Path( + header=deepcopy(header), + poses=[ + PoseStamped( + header=deepcopy(header), + pose=Pose( + position=Point(x=x, y=0, z=0), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) + for x in [0.0, 1.0, 2.0] + ], + ) + wires["planner_path"].publish(route) + index += 1 + ready.wait(0.2) + assert ready.is_set(), "No native CDR route before the bounded deadline" + result = received[0] + assert type(result) is Path and type(result.header) is Header + assert result.header.frame_id == "odom" + assert result.poses[-1].pose.position.x > 1.5 + assert abs(result.poses[-1].pose.position.y) < 0.2 + stamps = [to_nanoseconds(p.header.stamp) for p in result.poses] + assert all(b >= a for a, b in pairwise(stamps)) + assert module._process is not None and module._process.poll() is None diff --git a/dimos/navigation/local_planner/test_profile.py b/dimos/navigation/local_planner/test_profile.py index a3c19e354eb..5bc25a3355c 100644 --- a/dimos/navigation/local_planner/test_profile.py +++ b/dimos/navigation/local_planner/test_profile.py @@ -14,10 +14,11 @@ import numpy as np -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.message_utils.geometry import quaternion_from_euler +from dimos.message_utils.time import time_from_nanoseconds, time_from_seconds +from dimos.messages.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Header from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.local_planner.profile import decode_ceilings @@ -26,29 +27,39 @@ def _pose(x: float, y: float, yaw: float = 0.0, ts: float = 0.0) -> PoseStamped: return PoseStamped( - ts=ts, - frame_id="odom", - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0, 0, yaw)), + header=Header(frame_id="odom", stamp=time_from_seconds(ts)), + pose=Pose(position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0, 0, yaw)), ) def _path(n: int = 21, step: float = 0.1) -> Path: - return Path(frame_id="odom", poses=[_pose(i * step, 0.0, ts=1.0) for i in range(n)]) + return Path( + header=Header(frame_id="odom", stamp=time_from_seconds(0.0)), + poses=[_pose(i * step, 0.0, ts=1.0) for i in range(n)], + ) def _stamp(path: Path, dts: list[float], t0: float = 0.0) -> Path: t = t0 - path.poses[0].ts = t + path.poses[0].header.stamp = time_from_seconds(t) for pose, dt in zip(path.poses[1:], dts, strict=True): t += dt - pose.ts = t + pose.header.stamp = time_from_seconds(t) return path def test_ceiling_is_segment_length_over_stamp_delta() -> None: speeds = np.linspace(MIN_SPEED, MAX_SPEED, 20) - path = _stamp(_path(), [0.1 / v for v in speeds]) + # Keep a varied speed sweep while using exactly representable ROS time deltas. + # The old 0.1 / v fixture required sub-nanosecond timestamp precision. + positions = np.concatenate([[0.0], np.cumsum(speeds * 0.5)]) + path = _stamp( + Path( + header=Header(frame_id="odom", stamp=time_from_seconds(0.0)), + poses=[_pose(float(x), 0.0) for x in positions], + ), + [0.5] * 20, + ) ceilings = decode_ceilings(path, MIN_SPEED, MAX_SPEED) assert ceilings is not None # waypoint i carries segment i-1 @@ -62,16 +73,19 @@ def test_unstamped_path_decodes_none() -> None: def test_non_monotone_decodes_none() -> None: path = _stamp(_path(), [0.1 / MAX_SPEED] * 20) - path.poses[5].ts = 999.0 + path.poses[5].header.stamp = time_from_seconds(999.0) assert decode_ceilings(path, MIN_SPEED, MAX_SPEED) is None def test_fan_inherits_ceiling() -> None: - poses = [_pose(0.0, 0.0, 0.0), _pose(0.1, 0.0, 0.0)] - poses += [_pose(0.1, 0.0, yaw) for yaw in (0.3, 0.6, 0.9)] # rotation in place - poses += [_pose(0.2, 0.0, 0.9)] + poses = [_pose(0.0, 0.0, 0.0), _pose(MAX_SPEED * 0.1, 0.0, 0.0)] + poses += [_pose(MAX_SPEED * 0.1, 0.0, yaw) for yaw in (0.3, 0.6, 0.9)] # rotation in place + poses += [_pose(MAX_SPEED * 0.2, 0.0, 0.9)] fan_dt = 0.3 / GO2.max_yaw_rate - path = _stamp(Path(frame_id="odom", poses=poses), [0.1 / MAX_SPEED] + [fan_dt] * 3 + [0.5]) + path = _stamp( + Path(header=Header(frame_id="odom", stamp=time_from_seconds(0.0)), poses=poses), + [0.1] + [fan_dt] * 3 + [0.5], + ) ceilings = decode_ceilings(path, MIN_SPEED, MAX_SPEED) assert ceilings is not None assert ceilings[3] == MAX_SPEED # fan inherits, does not read as slow motion @@ -86,8 +100,22 @@ def test_decode_clips_to_safe_band() -> None: def test_empty_and_single_pose() -> None: - assert decode_ceilings(Path(frame_id="odom", poses=[]), MIN_SPEED, MAX_SPEED) is None - assert decode_ceilings(Path(frame_id="odom", poses=[_pose(0, 0)]), MIN_SPEED, MAX_SPEED) is None + assert ( + decode_ceilings( + Path(header=Header(frame_id="odom", stamp=time_from_seconds(0.0)), poses=[]), + MIN_SPEED, + MAX_SPEED, + ) + is None + ) + assert ( + decode_ceilings( + Path(header=Header(frame_id="odom", stamp=time_from_seconds(0.0)), poses=[_pose(0, 0)]), + MIN_SPEED, + MAX_SPEED, + ) + is None + ) def test_stall_never_reads_as_schedule() -> None: @@ -97,3 +125,12 @@ def test_stall_never_reads_as_schedule() -> None: ca, cb = decode_ceilings(a, MIN_SPEED, MAX_SPEED), decode_ceilings(b, MIN_SPEED, MAX_SPEED) assert ca is not None and cb is not None np.testing.assert_allclose(ca, cb) + + +def test_exact_epoch_nanosecond_deltas_drive_the_profile_without_float_loss(): + path = _path(n=3, step=0.1) + base = 1700000000123456789 + for pose, delta in zip(path.poses, [0, 200000000, 700000000], strict=True): + pose.header.stamp = time_from_nanoseconds(base + delta) + ceilings = decode_ceilings(path, MIN_SPEED, MAX_SPEED) + np.testing.assert_allclose(ceilings, [0.5, 0.5, 0.2], rtol=0, atol=1e-15) diff --git a/dimos/navigation/local_planner/test_viz.py b/dimos/navigation/local_planner/test_viz.py index 361bd8b41e4..a4ce76efb22 100644 --- a/dimos/navigation/local_planner/test_viz.py +++ b/dimos/navigation/local_planner/test_viz.py @@ -17,13 +17,16 @@ from dataclasses import replace import math +from dimos_cdrgen.registry import decode, encode import numpy as np from numpy.typing import NDArray +import rerun as rr -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.message_utils.geometry import quaternion_from_euler +from dimos.message_utils.time import time_from_nanoseconds, time_from_seconds, to_nanoseconds +from dimos.messages.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Header from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.local_planner.viz import ( motion_visual_override, @@ -35,21 +38,26 @@ def _pose(x: float, y: float, yaw: float = 0.0) -> PoseStamped: return PoseStamped( - frame_id="odom", - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0, 0, yaw)), + header=Header( + frame_id="odom", + stamp=time_from_nanoseconds(1_700_000_000_000_000_000 + round(x * 10000)), + ), + pose=Pose(position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0, 0, yaw)), ) def _plan(n: int = 40, yaw: float = 0.0) -> Path: - return Path(frame_id="odom", poses=[_pose(k * 0.1, 0.0, yaw) for k in range(n)]) + return Path( + header=Header(frame_id="odom", stamp=time_from_nanoseconds(0)), + poses=[_pose(k * 0.1, 0.0, yaw) for k in range(n)], + ) def _stamp(plan: Path, speeds: NDArray[np.float64]) -> Path: """Per-segment speeds as stamp deltas over the 0.1 m spacing, the planner's dialect.""" ts = np.concatenate([[0.0], np.cumsum(0.1 / speeds)]) for pose, t in zip(plan.poses, ts, strict=True): - pose.ts = float(t) + pose.header.stamp = time_from_seconds(float(t)) return plan @@ -61,6 +69,19 @@ def _rgba(boxes) -> NDArray[np.float64]: ) +def test_cdr_plan_draws_without_mutating_exact_source_values() -> None: + plan = _plan(n=2, yaw=math.pi / 2) + plan.header.stamp = time_from_nanoseconds(1_700_000_000_123_456_789) + payload = encode(plan) + received = decode(payload, Path) + drawn = render_plan(received, GO2) + assert drawn is not None + assert drawn[0] == ("world/path", rr.Transform3D(parent_frame="tf#/odom")) + assert len(drawn[1][1].centers.pa_array) == 2 + assert to_nanoseconds(received.header.stamp) == 1_700_000_000_123_456_789 + assert encode(received) == payload + + def test_boxes_are_subsampled_by_arc_not_by_index() -> None: """Every waypoint would be an opaque wall hiding the geometry it shows.""" plan = _plan() # 3.9 m of plan discretised at 0.1 m @@ -92,7 +113,7 @@ def test_colour_tracks_the_precision_the_planner_stamped() -> None: def test_an_unstamped_plan_draws_one_flat_colour() -> None: """No pinch is claimed where no precision was stamped. - Freshly built poses carry wall-clock stamps microseconds apart, so the + The fixture carries deterministic stamps microseconds apart, so the decoder clips the absurd speed to cruise rather than returning None. It must not manufacture a varying profile out of noise: every box comes out the same colour. @@ -106,11 +127,22 @@ def test_an_unstamped_plan_draws_one_flat_colour() -> None: def test_the_veto_stub_is_drawn_rather_than_blanked() -> None: """An empty viewport reads as a dead module; a refusal must look refused.""" - veto = render_body(Path(frame_id="odom", poses=[_pose(1.0, 2.0, 0.5)]), GO2) + veto = render_body( + Path( + header=Header(frame_id="odom", stamp=time_from_nanoseconds(0)), + poses=[_pose(1.0, 2.0, 0.5)], + ), + GO2, + ) assert veto is not None assert len(veto.centers.pa_array) == 1 # an EMPTY path is different: hold the last good picture - assert render_body(Path(frame_id="odom", poses=[]), GO2) is None + assert ( + render_body( + Path(header=Header(frame_id="odom", stamp=time_from_nanoseconds(0)), poses=[]), GO2 + ) + is None + ) def test_the_box_sits_on_the_body_not_on_the_pose_point() -> None: @@ -123,18 +155,26 @@ def test_the_box_sits_on_the_body_not_on_the_pose_point() -> None: def test_the_plan_pins_its_frame_and_draws_its_body_off_the_one_path_topic() -> None: - import rerun as rr - drawn = render_plan(_plan(), GO2) assert drawn is not None by_entity = {e: type(a) for e, a in drawn} assert by_entity == {"world/path": rr.Transform3D, "world/path/body": rr.Boxes3D} assert drawn[0] == ("world/path", rr.Transform3D(parent_frame="tf#/odom")) # a refusal stub is drawn too, so it is not mistaken for a dead module - veto = render_plan(Path(frame_id="odom", poses=[_pose(1.0, 2.0)]), GO2) + veto = render_plan( + Path( + header=Header(frame_id="odom", stamp=time_from_nanoseconds(0)), poses=[_pose(1.0, 2.0)] + ), + GO2, + ) assert veto is not None and veto[-1][0] == "world/path/body" # an empty path is "no route": hold the last picture rather than blank it - assert render_plan(Path(frame_id="odom", poses=[]), GO2) is None + assert ( + render_plan( + Path(header=Header(frame_id="odom", stamp=time_from_nanoseconds(0)), poses=[]), GO2 + ) + is None + ) assert set(motion_visual_override(GO2, body_dilate_m=-0.03)) == { "world/path", "world/ground_support", diff --git a/dimos/navigation/local_planner/viz.py b/dimos/navigation/local_planner/viz.py index b6fea546b86..fe525d4c2ec 100644 --- a/dimos/navigation/local_planner/viz.py +++ b/dimos/navigation/local_planner/viz.py @@ -26,8 +26,9 @@ import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.nav_msgs.Path import Path +from dimos.message_utils.geometry import yaw +from dimos.messages.geometry_msgs.msg import PoseStamped +from dimos.messages.nav_msgs.msg import Path from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.local_planner.profile import ( @@ -48,10 +49,10 @@ def _body_centre(pose: PoseStamped, center_off: float) -> tuple[float, float]: """The body centre for a pose point, offset along the pose's own heading.""" - yaw = pose.yaw + heading = yaw(pose.pose.orientation) return ( - pose.position.x + center_off * np.cos(yaw), - pose.position.y + center_off * np.sin(yaw), + pose.pose.position.x + center_off * np.cos(heading), + pose.pose.position.y + center_off * np.sin(heading), ) @@ -84,16 +85,18 @@ def render_body( p = msg.poses[0] cx, cy = _body_centre(p, center_off) return rr.Boxes3D( - centers=[[cx, cy, p.position.z + height / 2.0 + _Z_LIFT]], + centers=[[cx, cy, p.pose.position.z + height / 2.0 + _Z_LIFT]], half_sizes=[[length / 2.0, width / 2.0, height / 2.0]], - rotation_axis_angles=[rr.RotationAxisAngle([0.0, 0.0, 1.0], rr.Angle(rad=p.yaw))], + rotation_axis_angles=[ + rr.RotationAxisAngle([0.0, 0.0, 1.0], rr.Angle(rad=yaw(p.pose.orientation))) + ], colors=[[255, 60, 60, 200]], radii=[line_radius], fill_mode="majorwireframe", labels=["VETO: no safe route"], ) - xy = np.array([[p.position.x, p.position.y] for p in msg.poses]).reshape(-1, 2) + xy = np.array([[p.pose.position.x, p.pose.position.y] for p in msg.poses]).reshape(-1, 2) seg = np.linalg.norm(np.diff(xy, axis=0), axis=1) arcs = np.concatenate([[0.0], np.cumsum(seg)]) # one box per stride_m of arc, plus the last pose so the goal end is drawn @@ -105,8 +108,8 @@ def render_body( for i in picks: p = msg.poses[int(i)] cx, cy = _body_centre(p, center_off) - centers.append([cx, cy, p.position.z + height / 2.0 + _Z_LIFT]) - angles.append(rr.RotationAxisAngle([0.0, 0.0, 1.0], rr.Angle(rad=p.yaw))) + centers.append([cx, cy, p.pose.position.z + height / 2.0 + _Z_LIFT]) + angles.append(rr.RotationAxisAngle([0.0, 0.0, 1.0], rr.Angle(rad=yaw(p.pose.orientation)))) if clear is None: colors.append([100, 160, 255, 160]) # unstamped: the plan line's own blue else: @@ -157,7 +160,7 @@ def render_plan( return None out: RerunMulti = [ # the bridge only pins single-archetype entities to their tf frame; the boxes are the plan - (entity, rr.Transform3D(parent_frame=f"tf#/{msg.frame_id}")), + (entity, rr.Transform3D(parent_frame=f"tf#/{msg.header.frame_id}")), ] boxes = render_body(msg, emb, stride_m, line_radius) if boxes is not None: diff --git a/dimos/navigation/movement_manager/movement_manager.py b/dimos/navigation/movement_manager/movement_manager.py index 4f83428f9a7..61b5d68963d 100644 --- a/dimos/navigation/movement_manager/movement_manager.py +++ b/dimos/navigation/movement_manager/movement_manager.py @@ -22,18 +22,17 @@ import math import threading -import time +from time import monotonic, time_ns from typing import Any -from dimos_lcm.std_msgs import Bool 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.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import Point, PointStamped, Twist, Vector3 +from dimos.messages.std_msgs.msg import Bool, Header from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -82,18 +81,22 @@ def stop(self) -> None: super().stop() def _on_click(self, msg: PointStamped) -> None: - if not all(math.isfinite(v) for v in (msg.x, msg.y, msg.z)): - logger.warning("Ignored invalid click", x=msg.x, y=msg.y, z=msg.z) + if not all(math.isfinite(v) for v in (msg.point.x, msg.point.y, msg.point.z)): + logger.warning("Ignored invalid click", x=msg.point.x, y=msg.point.y, z=msg.point.z) return if ( - abs(msg.x) > MAX_CLICK_HORIZONTAL_M - or abs(msg.y) > MAX_CLICK_HORIZONTAL_M - or abs(msg.z) > MAX_CLICK_VERTICAL_M + abs(msg.point.x) > MAX_CLICK_HORIZONTAL_M + or abs(msg.point.y) > MAX_CLICK_HORIZONTAL_M + or abs(msg.point.z) > MAX_CLICK_VERTICAL_M ): - logger.warning("Ignored out-of-range click", x=msg.x, y=msg.y, z=msg.z) + logger.warning( + "Ignored out-of-range click", x=msg.point.x, y=msg.point.y, z=msg.point.z + ) return - logger.debug("Goal", x=round(msg.x, 1), y=round(msg.y, 1), z=round(msg.z, 1)) + logger.debug( + "Goal", x=round(msg.point.x, 1), y=round(msg.point.y, 1), z=round(msg.point.z, 1) + ) self.way_point.publish(msg) self.goal.publish(msg) @@ -103,7 +106,8 @@ def _cancel_goal(self) -> None: # It can be REALLY bad if a robot is supposed to stop moving but wont # we should probably think a more robust/strict requirement on planners cancel = PointStamped( - ts=time.time(), frame_id="map", x=float("nan"), y=float("nan"), z=float("nan") + header=Header(frame_id="map", stamp=time_from_nanoseconds(time_ns())), + point=Point(x=float("nan"), y=float("nan"), z=float("nan")), ) self.way_point.publish(cancel) self.goal.publish(cancel) @@ -113,7 +117,7 @@ def _on_nav(self, msg: Twist) -> None: with self._lock: if self._teleop_active: # check if cooldown has expired - elapsed = time.monotonic() - self._last_teleop_time + elapsed = monotonic() - self._last_teleop_time if elapsed < self.config.tele_cooldown_sec: return self._teleop_active = False @@ -122,7 +126,7 @@ def _on_nav(self, msg: Twist) -> None: def _on_teleop(self, msg: Twist) -> None: with self._lock: self._teleop_active = True - self._last_teleop_time = time.monotonic() + self._last_teleop_time = monotonic() self._cancel_goal() diff --git a/dimos/navigation/movement_manager/test_movement_manager.py b/dimos/navigation/movement_manager/test_movement_manager.py index e266ad3e983..7aa3928773a 100644 --- a/dimos/navigation/movement_manager/test_movement_manager.py +++ b/dimos/navigation/movement_manager/test_movement_manager.py @@ -15,19 +15,28 @@ from __future__ import annotations from collections.abc import Generator +from contextlib import ExitStack from dataclasses import dataclass, field import math -import time +from threading import Event +from dimos_cdrgen.registry import decode, encode import pytest -from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.core.transport import LCMTransport +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import Point, PointStamped, Twist, Vector3 +from dimos.messages.std_msgs.msg import Bool, Header +from dimos.navigation.movement_manager import movement_manager as movement_module from dimos.navigation.movement_manager.movement_manager import ( MovementManager, ) +from dimos.protocol.rpc.pubsubrpc import LCMRPC + + +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) @dataclass @@ -54,14 +63,14 @@ def _attach(module): @pytest.fixture() def manager_and_captured() -> Generator[tuple[MovementManager, Captured], None, None]: - module = MovementManager(tele_cooldown_sec=0.1) + module = MovementManager(tele_cooldown_sec=0.1, rpc_transport=MemoryRPC) captured, unsubs = _attach(module) try: yield module, captured finally: for unsub in unsubs: unsub() - module._close_module() + module.stop() def _twist(lx=0.0): @@ -69,7 +78,10 @@ def _twist(lx=0.0): def _click(x=1.0, y=2.0, z=0.0): - return PointStamped(ts=time.time(), frame_id="map", x=x, y=y, z=z) + return PointStamped( + header=Header(frame_id="map", stamp=time_from_nanoseconds(1700000000123456789)), + point=Point(x=x, y=y, z=z), + ) def test_teleop_suppresses_nav_and_cancels_goal(manager_and_captured): @@ -88,15 +100,17 @@ def test_teleop_suppresses_nav_and_cancels_goal(manager_and_captured): # Goal cancelled with NaN assert len(captured.goal) == 1 - assert math.isnan(captured.goal[0].x) + assert math.isnan(captured.goal[0].point.x) -def test_nav_resumes_after_cooldown(manager_and_captured): +def test_nav_resumes_after_cooldown(manager_and_captured, monkeypatch): """After the cooldown expires, nav commands pass through again.""" manager, captured = manager_and_captured + now = [0.0] + monkeypatch.setattr(movement_module, "monotonic", lambda: now[0]) manager.config.tele_cooldown_sec = 0.05 manager._on_teleop(_twist(lx=0.3)) - time.sleep(DEFAULT_THREAD_JOIN_TIMEOUT) + now[0] = 0.05 cmd_count_before = len(captured.cmd_vel) manager._on_nav(_twist(lx=0.9)) @@ -139,3 +153,111 @@ def test_tele_cmd_vel_scaling(manager_and_captured): assert published.linear.y == pytest.approx(2.0) assert published.linear.z == pytest.approx(0.0) assert published.angular.z == pytest.approx(0.25) + + +@pytest.mark.parametrize("stamp", [0, -1, 1700000000123456789]) +def test_cancel_goal_preserves_exact_wall_time_and_existing_nan_fallback( + manager_and_captured, monkeypatch, stamp +): + manager, captured = manager_and_captured + monkeypatch.setattr(movement_module, "time_ns", lambda: stamp) + manager._on_teleop(_twist(0.3)) + assert captured.stop_movement == [Bool(data=True)] + cancel = decode(encode(captured.goal[0]), PointStamped) + assert cancel.header == Header(frame_id="map", stamp=time_from_nanoseconds(stamp)) + assert all(math.isnan(value) for value in (cancel.point.x, cancel.point.y, cancel.point.z)) + + +@pytest.mark.parametrize("position", [(500.0, -500.0, 50.0), (-500.0, 500.0, -50.0)]) +def test_click_limit_boundaries_keep_exact_source_header(manager_and_captured, position): + manager, captured = manager_and_captured + click = _click(*position) + manager._on_click(click) + assert captured.goal == [click] and captured.way_point == [click] + assert decode(encode(captured.goal[0]), PointStamped) == click + + +@pytest.mark.parametrize( + "position", + [ + (0.0, 501.0, 0.0), + (0.0, 0.0, 51.0), + (0.0, 0.0, -51.0), + (0.0, float("inf"), 0.0), + (0.0, 0.0, float("nan")), + ], +) +def test_other_invalid_click_components_are_rejected(manager_and_captured, position): + manager, captured = manager_and_captured + manager._on_click(_click(*position)) + assert captured.goal == [] and captured.way_point == [] + + +def test_actual_cdr_movement_ports_keep_teleop_priority_and_resume_nav_at_cooldown(monkeypatch): + now = [0.0] + monkeypatch.setattr(movement_module, "monotonic", lambda: now[0]) + monkeypatch.setattr(movement_module, "time_ns", lambda: 1700000000123456789) + with ExitStack() as cleanup: + manager = MovementManager(rpc_transport=MemoryRPC, tele_cooldown_sec=1.0) + wires = {} + for name in [ + "clicked_point", + "nav_cmd_vel", + "tele_cmd_vel", + "goal", + "way_point", + "cmd_vel", + "stop_movement", + ]: + port = getattr(manager, name) + wire = LCMTransport("movement_" + name, port.type, url="memq://") + port.transport = wire + wires[name] = wire + cleanup.callback(wire.stop) + cleanup.callback(manager.stop) + commands = [] + cancellations = [] + stops = [] + received = Event() + nav_ack = Event() + cancelled = Event() + stopped = Event() + on_nav = manager._on_nav + + def handle_nav(value): + on_nav(value) + nav_ack.set() + + def cancel(value): + cancellations.append(value) + cancelled.set() + + def stop(value): + stops.append(value) + stopped.set() + + monkeypatch.setattr(manager, "_on_nav", handle_nav) + + def receive(value): + commands.append(value) + received.set() + + cleanup.callback(wires["cmd_vel"].subscribe(receive)) + cleanup.callback(wires["goal"].subscribe(cancel)) + cleanup.callback(wires["stop_movement"].subscribe(stop)) + manager.start() + wires["tele_cmd_vel"].publish(_twist(0.3)) + assert received.wait(5) + assert cancelled.wait(5) and stopped.wait(5) + assert commands == [_twist(0.3)] + assert stops == [Bool(data=True)] + assert cancellations[0].header.stamp == time_from_nanoseconds(1700000000123456789) + now[0] = 0.5 + wires["nav_cmd_vel"].publish(_twist(0.9)) + assert nav_ack.wait(5) + assert commands == [_twist(0.3)] + received.clear() + now[0] = 1.0 + wires["nav_cmd_vel"].publish(_twist(0.9)) + assert received.wait(5) + assert commands == [_twist(0.3), _twist(0.9)] diff --git a/dimos/navigation/spec.py b/dimos/navigation/spec.py index 6f16a7c5964..0a6ae622c1c 100644 --- a/dimos/navigation/spec.py +++ b/dimos/navigation/spec.py @@ -19,22 +19,21 @@ A ``path`` with one pose means "hold, no safe route"; with none, "stop". tf is ``IO`` where the implementers are dimos-module natives (``#[tf]`` both -subscribes and publishes) and their python twins; ``In`` on the global planner, -whose native only reads it. +subscribes and publishes) and their python twins; ``In`` on the Python-only +blind local planner. """ from typing import Protocol from dimos.core.stream import IO, In, Out -from dimos.msgs.geometry_msgs.PointStamped import PointStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.messages.geometry_msgs.msg import PointStamped, Twist +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.sensor_msgs.msg import PointCloud2 +from dimos.messages.tf2_msgs.msg import TFMessage class GlobalPlanner(Protocol): - tf: In[TFMessage] + tf: IO[TFMessage] goal: In[PointStamped] path: Out[Path] diff --git a/dimos/navigation/test_tf_pose.py b/dimos/navigation/test_tf_pose.py index 97a2bca2055..39ab71c6c05 100644 --- a/dimos/navigation/test_tf_pose.py +++ b/dimos/navigation/test_tf_pose.py @@ -12,9 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +import pytest + +from dimos.message_utils.time import time_from_nanoseconds, time_from_seconds, to_seconds +from dimos.messages.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos.messages.std_msgs.msg import Header from dimos.navigation.tf_pose import TfPose from dimos.protocol.tf.tf import MultiTBuffer @@ -34,7 +36,7 @@ def get( time_tolerance: float | None = None, *, forward_tolerance: float = 0.0, - ) -> Transform | None: + ) -> TransformStamped | None: self.gets += 1 return super().get( parent_frame, @@ -53,13 +55,11 @@ def __call__(self) -> float: return self.t -def _body_at(x: float, ts: float) -> Transform: - return Transform( - translation=Vector3(x, 0.0, 0.3), - rotation=IDENTITY, - frame_id="odom", +def _body_at(x: float, ts: float) -> TransformStamped: + return TransformStamped( + header=Header(frame_id="odom", stamp=time_from_seconds(ts)), child_frame_id="base_link", - ts=ts, + transform=Transform(translation=Vector3(x=x, y=0.0, z=0.3), rotation=IDENTITY), ) @@ -68,7 +68,12 @@ def test_tf_pose_is_the_world_to_base_edge(): tf.receive_transform(_body_at(1.0, ts=5.0)) pose = TfPose(tf, "base_link", max_age_s=2.5, clock=clock).get("odom") assert pose is not None - assert (pose.position.x, pose.position.z, pose.ts, pose.frame_id) == (1.0, 0.3, 5.0, "odom") + assert ( + pose.pose.position.x, + pose.pose.position.z, + to_seconds(pose.header.stamp), + pose.header.frame_id, + ) == (1.0, 0.3, 5.0, "odom") def test_tf_pose_goes_stale_when_the_stamp_stops_advancing(): @@ -101,3 +106,22 @@ def test_tf_pose_retries_a_missing_edge_once_per_period(): clock.t += TfPose.RETRY_PERIOD_S assert src.get("odom") is None assert tf.gets == 2 + + +@pytest.mark.parametrize("stamp", [0, -1, 1700000000123456789]) +def test_tf_pose_observes_single_nanosecond_progress_without_float_conversion(stamp): + tf, clock = MultiTBuffer(), _Clock() + source = TfPose(tf, "base_link", max_age_s=2.5, clock=clock) + transform = _body_at(1.0, 0.0) + transform.header.stamp = time_from_nanoseconds(stamp) + tf.receive_transform(transform) + first = source.get("odom") + assert first.header == transform.header + first.pose.position.x = 100 + assert source.get("odom").pose.position.x == 1 + clock.t += 2.6 + assert source.get("odom") is None + transform.header.stamp = time_from_nanoseconds(stamp + 1) + tf.receive_transform(transform) + fresh = source.get("odom") + assert fresh.header == transform.header and fresh.pose.position.x == 1 diff --git a/dimos/navigation/tf_pose.py b/dimos/navigation/tf_pose.py index a1a2f08e45d..2ed5337c106 100644 --- a/dimos/navigation/tf_pose.py +++ b/dimos/navigation/tf_pose.py @@ -19,10 +19,13 @@ import time from typing import TYPE_CHECKING +from dimos.message_utils.geometry import pose_from_transform +from dimos.message_utils.time import to_nanoseconds + if TYPE_CHECKING: from collections.abc import Callable - from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + from dimos.messages.geometry_msgs.msg import PoseStamped from dimos.protocol.tf.tf import TFLookup @@ -48,7 +51,7 @@ def __init__( # The stamp last seen and when it was first seen, on OUR clock: the # edge's stamp is the robot's clock, which need not be ours, and what # the deadman asks is how long since tf last moved. - self._seen: tuple[float, float] | None = None + self._seen: tuple[int, float] | None = None def get(self, world_frame: str) -> PoseStamped | None: now = self._clock() @@ -59,8 +62,9 @@ def get(self, world_frame: str) -> PoseStamped | None: self._next_lookup = now + self.RETRY_PERIOD_S self._seen = None return None - if self._seen is None or self._seen[0] != tf.ts: - self._seen = (tf.ts, now) + stamp = to_nanoseconds(tf.header.stamp) + if self._seen is None or self._seen[0] != stamp: + self._seen = (stamp, now) if now - self._seen[1] > self.max_age_s: return None - return tf.to_pose(ts=tf.ts) + return pose_from_transform(tf) diff --git a/dimos/navigation/trajectory_follower/basic/module.py b/dimos/navigation/trajectory_follower/basic/module.py index 9b8ed5dd5c0..66cae094b1f 100644 --- a/dimos/navigation/trajectory_follower/basic/module.py +++ b/dimos/navigation/trajectory_follower/basic/module.py @@ -19,7 +19,6 @@ import time from typing import Any -from dimos_lcm.std_msgs import Bool import numpy as np from numpy.typing import NDArray from reactivex.disposable import Disposable @@ -28,12 +27,13 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import IO, In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.message_utils.geometry import yaw +from dimos.messages.geometry_msgs.msg import PoseStamped, Twist, Vector3 +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Bool +from dimos.messages.tf2_msgs.msg import TFMessage from dimos.navigation import spec +from dimos.navigation.go2.replanning_a_star.controllers import _velocity from dimos.utils.logging_config import setup_logger from dimos.utils.trigonometry import angle_diff @@ -103,7 +103,7 @@ def stop(self) -> None: self._stop_event.set() if self._thread is not None: self._thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish(_velocity()) super().stop() def _lookup_pose(self) -> PoseStamped | None: @@ -120,9 +120,11 @@ def _on_path(self, path: Path) -> None: if len(path.poses) == 0: with self._lock: self._waypoints = None - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish(_velocity()) return - waypoints = np.array([[p.position.x, p.position.y] for p in path.poses], dtype=np.float32) + waypoints = np.array( + [[p.pose.position.x, p.pose.position.y] for p in path.poses], dtype=np.float32 + ) with self._lock: self._waypoints = waypoints @@ -130,7 +132,7 @@ def _on_stop(self, msg: Bool) -> None: if msg.data: with self._lock: self._waypoints = None - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish(_velocity()) def _follow(self) -> None: period = 1.0 / self.config.control_frequency @@ -145,20 +147,20 @@ def _follow(self) -> None: self._stop_event.wait(max(0.0, period - elapsed)) def _step(self, pose: PoseStamped, waypoints: NDArray[np.float32]) -> None: - position = np.array([pose.position.x, pose.position.y], dtype=np.float32) + position = np.array([pose.pose.position.x, pose.pose.position.y], dtype=np.float32) if float(np.linalg.norm(waypoints[-1] - position)) < self.config.goal_tolerance: - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish(_velocity()) with self._lock: if self._waypoints is waypoints: self._waypoints = None - self.goal_reached.publish(Bool(True)) + self.goal_reached.publish(Bool(data=True)) logger.info("Goal reached") return target = self._lookahead_point(waypoints, position) yaw_error = angle_diff( math.atan2(target[1] - position[1], target[0] - position[0]), - pose.orientation.euler[2], + yaw(pose.pose.orientation), ) angular = max( diff --git a/dimos/navigation/trajectory_follower/basic/test_module.py b/dimos/navigation/trajectory_follower/basic/test_module.py index 9cf403aba33..40c05e1da3b 100644 --- a/dimos/navigation/trajectory_follower/basic/test_module.py +++ b/dimos/navigation/trajectory_follower/basic/test_module.py @@ -12,103 +12,78 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Vector3 import Vector3 +import pytest + +from dimos.core.transport import LCMTransport +from dimos.message_utils.time import time_from_seconds +from dimos.messages.geometry_msgs.msg import Quaternion, Transform, TransformStamped, Vector3 +from dimos.messages.std_msgs.msg import Header from dimos.navigation.trajectory_follower.basic.module import BasicPathFollower, lookahead_distance -from dimos.protocol.tf.tf import MultiTBuffer +from dimos.protocol.rpc.pubsubrpc import LCMRPC MOUNT_Z = 0.163 -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", +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +@pytest.fixture +def module_and_tf(): + module = BasicPathFollower(rpc_transport=MemoryRPC) + wire = LCMTransport("basic_tf", module.tf.type, url="memq://") + module.tf.transport = wire + try: + yield module, module.tfbuffer + finally: + module.stop() + wire.stop() + + +def _mount(): + return TransformStamped( + header=Header(frame_id="base_link", stamp=time_from_seconds(1.0)), child_frame_id="mid360_link", - ts=1.0, + transform=Transform( + translation=Vector3(0.0, 0.0, MOUNT_Z), rotation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), ) -def _odom_edge() -> Transform: - return Transform( - translation=Vector3(1.0, 2.0, 3.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="odom", +def _odom_edge(): + return TransformStamped( + header=Header(frame_id="odom", stamp=time_from_seconds(1.0)), child_frame_id="mid360_link", - ts=1.0, + transform=Transform( + translation=Vector3(1.0, 2.0, 3.0), rotation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), ) -def test_lookup_pose_steers_from_the_tf_base_pose() -> None: - tf = FakeTF() - tf.receive_transform(_mount()) - tf.receive_transform(_odom_edge()) - module = BasicPathFollower() - module._tf = tf - try: - pose = module._lookup_pose() - assert pose is not None - assert abs(pose.position.z - (3.0 - MOUNT_Z)) < 1e-9 - finally: - module.stop() +def test_lookup_pose_steers_from_the_tf_base_pose(module_and_tf): + module, tf = module_and_tf + tf.receive_transform(_mount(), _odom_edge()) + pose = module._lookup_pose() + assert pose is not None + assert abs(pose.pose.position.z - (3.0 - MOUNT_Z)) < 1e-9 -def test_lookup_pose_is_none_without_the_mount_tf() -> None: - tf = FakeTF() +def test_lookup_pose_is_none_without_the_mount_tf(module_and_tf): + module, tf = module_and_tf tf.receive_transform(_odom_edge()) - module = BasicPathFollower() - module._tf = tf - try: - assert module._lookup_pose() is None - finally: - module.stop() - - -def test_lookup_retries_are_throttled_during_an_outage() -> None: - tf = FakeTF() - module = BasicPathFollower() - module._tf = tf - try: - assert module._lookup_pose() is None - assert module._lookup_pose() is None - assert tf.gets == 1 - tf.receive_transform(_mount()) - tf.receive_transform(_odom_edge()) - module._next_lookup = 0.0 - assert module._lookup_pose() is not None - finally: - module.stop() + assert module._lookup_pose() is None + + +def test_lookup_retries_are_throttled_during_an_outage(module_and_tf, mocker): + module, tf = module_and_tf + lookup = mocker.spy(tf, "get_pose") + assert module._lookup_pose() is None + assert module._lookup_pose() is None + assert lookup.call_count == 1 + tf.receive_transform(_mount(), _odom_edge()) + module._next_lookup = 0.0 + assert module._lookup_pose() is not None def test_lookahead_floor_at_low_speed() -> None: diff --git a/dimos/navigation/trajectory_follower/fancy/controller.py b/dimos/navigation/trajectory_follower/fancy/controller.py index eef34feb27f..9125ea0fa4e 100644 --- a/dimos/navigation/trajectory_follower/fancy/controller.py +++ b/dimos/navigation/trajectory_follower/fancy/controller.py @@ -30,9 +30,9 @@ import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path +from dimos.message_utils.geometry import yaw +from dimos.messages.geometry_msgs.msg import PoseStamped, Twist +from dimos.messages.nav_msgs.msg import Path from dimos.protocol.service.spec import BaseConfig @@ -87,6 +87,7 @@ def path_xy_yaw(path: Path) -> NDArray[np.float64]: """The plan as a contiguous (N, 3) float64 array, the rust marshalling.""" return np.ascontiguousarray( np.array( - [[p.position.x, p.position.y, p.yaw] for p in path.poses], dtype=np.float64 + [[p.pose.position.x, p.pose.position.y, yaw(p.pose.orientation)] for p in path.poses], + dtype=np.float64, ).reshape(-1, 3) ) diff --git a/dimos/navigation/trajectory_follower/fancy/laws/hinted.py b/dimos/navigation/trajectory_follower/fancy/laws/hinted.py index f7d02853bcf..4126aef532a 100644 --- a/dimos/navigation/trajectory_follower/fancy/laws/hinted.py +++ b/dimos/navigation/trajectory_follower/fancy/laws/hinted.py @@ -25,10 +25,10 @@ import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.message_utils.geometry import yaw +from dimos.message_utils.time import to_nanoseconds +from dimos.messages.geometry_msgs.msg import PoseStamped, Twist, Vector3 +from dimos.messages.nav_msgs.msg import Path from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.local_planner.profile import decode_ceilings @@ -76,12 +76,12 @@ def update( self, pose: PoseStamped, path: Path, t: float, clearance: NDArray[np.float64] | None = None ) -> Twist: cfg, emb = self.config, self.emb - if len(path) < 2: + if len(path.poses) < 2: # empty path or single-pose veto stub: the planner says stop return Twist(Vector3(0, 0, 0), Vector3(0, 0, 0)) - xy = np.array([[p.position.x, p.position.y] for p in path.poses]) - yaws = np.array([p.yaw for p in path.poses]) - px, py, pyaw = pose.position.x, pose.position.y, pose.yaw + xy = np.array([[p.pose.position.x, p.pose.position.y] for p in path.poses]) + yaws = np.array([yaw(p.pose.orientation) for p in path.poses]) + px, py, pyaw = pose.pose.position.x, pose.pose.position.y, yaw(pose.pose.orientation) n = len(xy) seg = np.linalg.norm(np.diff(xy, axis=0), axis=1) if n > 1 else np.zeros(1) @@ -198,10 +198,21 @@ def update( clr = None if clearance is None else np.ascontiguousarray(clearance, dtype=np.float64) # the law reads only stamp deltas: a precision profile, not a schedule ts = np.ascontiguousarray( - np.array([p.ts for p in path.poses], dtype=np.float64).reshape(-1) + np.array( + [ + to_nanoseconds(p.header.stamp) - to_nanoseconds(path.poses[0].header.stamp) + for p in path.poses + ], + dtype=np.float64, + ).reshape(-1) + / 1_000_000_000 ) vx, vy, wz = self._mod.update_hinted( - (float(pose.position.x), float(pose.position.y), float(pose.yaw)), + ( + float(pose.pose.position.x), + float(pose.pose.position.y), + float(yaw(pose.pose.orientation)), + ), path_xy_yaw(path), clr, ts, diff --git a/dimos/navigation/trajectory_follower/fancy/laws/seed.py b/dimos/navigation/trajectory_follower/fancy/laws/seed.py index 14072fe6311..ce7ae85eefc 100644 --- a/dimos/navigation/trajectory_follower/fancy/laws/seed.py +++ b/dimos/navigation/trajectory_follower/fancy/laws/seed.py @@ -25,10 +25,9 @@ import numpy as np from numpy.typing import NDArray -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.message_utils.geometry import yaw +from dimos.messages.geometry_msgs.msg import PoseStamped, Twist, Vector3 +from dimos.messages.nav_msgs.msg import Path from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.trajectory_follower.fancy.controller import ( @@ -67,12 +66,12 @@ def update( self, pose: PoseStamped, path: Path, t: float, clearance: NDArray[np.float64] | None = None ) -> Twist: cfg, emb = self.config, self.emb - if len(path) < 2: + if len(path.poses) < 2: # empty path or a single-pose veto stub: hold return Twist(Vector3(0, 0, 0), Vector3(0, 0, 0)) - xy = np.array([[p.position.x, p.position.y] for p in path.poses]) - yaws = np.array([p.yaw for p in path.poses]) - px, py, pyaw = pose.position.x, pose.position.y, pose.yaw + xy = np.array([[p.pose.position.x, p.pose.position.y] for p in path.poses]) + yaws = np.array([yaw(p.pose.orientation) for p in path.poses]) + px, py, pyaw = pose.pose.position.x, pose.pose.position.y, yaw(pose.pose.orientation) seg = np.linalg.norm(np.diff(xy, axis=0), axis=1) if len(xy) > 1 else np.zeros(1) arcs = np.concatenate([[0.0], np.cumsum(seg)]) @@ -141,7 +140,11 @@ def update( ) -> Twist: clr = None if clearance is None else np.ascontiguousarray(clearance, dtype=np.float64) vx, vy, wz = self._mod.update_seed( - (float(pose.position.x), float(pose.position.y), float(pose.yaw)), + ( + float(pose.pose.position.x), + float(pose.pose.position.y), + float(yaw(pose.pose.orientation)), + ), path_xy_yaw(path), clr, self._emb, diff --git a/dimos/navigation/trajectory_follower/fancy/module.py b/dimos/navigation/trajectory_follower/fancy/module.py index dedf3f80d05..3073dc2ff3b 100644 --- a/dimos/navigation/trajectory_follower/fancy/module.py +++ b/dimos/navigation/trajectory_follower/fancy/module.py @@ -33,13 +33,13 @@ from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.core.stream import IO, In, Out -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.std_msgs.Bool import Bool -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.messages.geometry_msgs.msg import PoseStamped, Twist +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Bool +from dimos.messages.tf2_msgs.msg import TFMessage from dimos.navigation import spec from dimos.navigation.embodiment.base import Embodiment +from dimos.navigation.go2.replanning_a_star.controllers import _velocity from dimos.navigation.tf_pose import TfPose from dimos.navigation.trajectory_follower.fancy.controller import TrajectoryController from dimos.navigation.trajectory_follower.fancy.laws import hinted @@ -83,7 +83,7 @@ class TrajectoryFollowerConfig(ModuleConfig): goal_tolerance: float = 0.20 # planar distance that counts as arrival (m) # The planner's own body: the law decodes the path stamps with its governor band. embodiment: Embodiment - # The pose is the `path.frame_id -> base_frame` edge on tf, read each tick. + # The pose is the `path.header.frame_id -> base_frame` edge on tf, read each tick. base_frame: str = "base_link" # Deadman: zero the twist once the held path is this old, measured from arrival; must clear the ~1 Hz replan cadence. max_path_age_s: float = 2.5 @@ -128,7 +128,7 @@ def stop(self) -> None: self._stop_event.set() if self._thread is not None: self._thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish(_velocity()) super().stop() def _on_path(self, msg: Path) -> None: @@ -143,10 +143,12 @@ def _on_path(self, msg: Path) -> None: self._path_at = time.monotonic() if len(msg.poses) >= 2: # a single-pose stub is a refusal, not an arrival target - self._latch.set_goal((msg.poses[-1].position.x, msg.poses[-1].position.y)) + self._latch.set_goal( + (msg.poses[-1].pose.position.x, msg.poses[-1].pose.position.y) + ) if not msg.poses: # now, not a control period later - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish(_velocity()) def _control_loop(self) -> None: period = 1.0 / self.config.control_frequency @@ -164,29 +166,29 @@ def tick(self) -> None: age = None if self._path_at is None else now - self._path_at pose = None if path is not None and self._pose_src is not None: - pose = self._pose_src.get(path.frame_id) + pose = self._pose_src.get(path.header.frame_id) if path is not None and pose is not None: assert age is not None self.step(pose, path, age) elif path is not None: # a plan with no live pose under it: the deadman on the pose - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish(_velocity()) def step(self, pose: PoseStamped, path: Path, age: float) -> None: """One control tick against a path that arrived `age` seconds ago.""" # the deadman outranks arrival: a goal reached against an unrefreshed plan is a coincidence if age > self.config.max_path_age_s: - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish(_velocity()) return - xy = (pose.position.x, pose.position.y) + xy = (pose.pose.position.x, pose.pose.position.y) with self._lock: arrived, reached = self._latch.arrive(xy), self._latch.reached if arrived: - self.nav_cmd_vel.publish(Twist()) - self.goal_reached.publish(Bool(True)) + self.nav_cmd_vel.publish(_velocity()) + self.goal_reached.publish(Bool(data=True)) logger.info("goal reached", x=round(xy[0], 2), y=round(xy[1], 2)) return if reached: - self.nav_cmd_vel.publish(Twist()) + self.nav_cmd_vel.publish(_velocity()) return self.nav_cmd_vel.publish(self._controller.update(pose, path, time.monotonic())) diff --git a/dimos/navigation/trajectory_follower/fancy/native.py b/dimos/navigation/trajectory_follower/fancy/native.py index 950784be08c..ab4eb90d985 100644 --- a/dimos/navigation/trajectory_follower/fancy/native.py +++ b/dimos/navigation/trajectory_follower/fancy/native.py @@ -29,10 +29,10 @@ from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import IO, In, Out -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.nav_msgs.Path import Path -from dimos.msgs.std_msgs.Bool import Bool -from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.messages.geometry_msgs.msg import Twist +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Bool +from dimos.messages.tf2_msgs.msg import TFMessage from dimos.navigation import spec from dimos.navigation.embodiment.base import Embodiment from dimos.navigation.trajectory_follower.fancy.module import TrajectoryFollowerConfig @@ -70,6 +70,7 @@ class TrajectoryFollowerNative(NativeModule, spec.TrajectoryFollower): path: In[Path] tf: IO[TFMessage] # IO, not In: see planner_native.py + tf_static: In[TFMessage] nav_cmd_vel: Out[Twist] goal_reached: Out[Bool] diff --git a/dimos/navigation/trajectory_follower/fancy/rust/Cargo.lock b/dimos/navigation/trajectory_follower/fancy/rust/Cargo.lock index 9bce5ee4a4b..4629ec6922e 100644 --- a/dimos/navigation/trajectory_follower/fancy/rust/Cargo.lock +++ b/dimos/navigation/trajectory_follower/fancy/rust/Cargo.lock @@ -249,6 +249,15 @@ dependencies = [ "shlex", ] +[[package]] +name = "cdr" +version = "0.2.4" +source = "git+https://github.com/TomCC7/cdr-rs?rev=38a4368255d408faf374726653b39d4a4e9735cb#38a4368255d408faf374726653b39d4a4e9735cb" +dependencies = [ + "byteorder", + "serde", +] + [[package]] name = "cfg-if" version = "1.0.5" @@ -579,6 +588,17 @@ dependencies = [ "subtle", ] +[[package]] +name = "dimos-generated-messages" +version = "0.0.14" +dependencies = [ + "cdr", + "heapless", + "ros2msg", + "serde", + "serde-big-array", +] + [[package]] name = "dimos-lcm" version = "0.1.0" @@ -593,8 +613,8 @@ dependencies = [ name = "dimos-local-planner" version = "0.1.0" dependencies = [ + "dimos-generated-messages", "dimos-module", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen)", "serde", "serde_json", "tokio", @@ -605,12 +625,13 @@ dependencies = [ [[package]] name = "dimos-module" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#fc6a4e58da89bb073a38dce5ac63728c6559d3a6" +version = "0.2.0" dependencies = [ + "cdr", + "dimos-generated-messages", "dimos-lcm", "dimos-module-macros", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549)", + "lcm-msgs", "nalgebra", "rayon", "serde", @@ -618,6 +639,7 @@ dependencies = [ "tokio", "tracing", "tracing-subscriber", + "transforms", "url", "validator", "zenoh", @@ -626,7 +648,6 @@ dependencies = [ [[package]] name = "dimos-module-macros" version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#fc6a4e58da89bb073a38dce5ac63728c6559d3a6" dependencies = [ "proc-macro2", "quote", @@ -638,9 +659,9 @@ dependencies = [ name = "dimos-trajectory-follower" version = "0.1.0" dependencies = [ + "dimos-generated-messages", "dimos-local-planner", "dimos-module", - "lcm-msgs 0.1.0 (git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen)", "numpy", "pyo3", "serde", @@ -982,6 +1003,15 @@ version = "0.33.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e762b9b188634fc508d4ec54b62ea3d352c43fd431d18d05d46f559f217f4725" +[[package]] +name = "hash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" +dependencies = [ + "byteorder", +] + [[package]] name = "hashbrown" version = "0.12.3" @@ -1023,6 +1053,22 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] + +[[package]] +name = "heapless" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" +dependencies = [ + "hash32", + "serde", + "stable_deref_trait", +] [[package]] name = "heck" @@ -1429,14 +1475,6 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" -[[package]] -name = "lcm-msgs" -version = "0.1.0" -source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#04d78e8622500244123ba9cefa4c51b4cb454549" -dependencies = [ - "byteorder", -] - [[package]] name = "lcm-msgs" version = "0.1.0" @@ -2021,6 +2059,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn 2.0.119", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -2436,6 +2484,23 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "ros2msg" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fe291be5cb6d4b8e6bf9c9de80c40bd1cb09ecac50deafa2fdcae0027e04988" +dependencies = [ + "heck", + "pest", + "pest_derive", + "prettyplease", + "proc-macro2", + "quote", + "regex", + "syn 2.0.119", + "thiserror", +] + [[package]] name = "rustc-hash" version = "2.1.3" @@ -2676,6 +2741,15 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-big-array" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11fc7cc2c76d73e0f27ee52abbd64eec84d46f370c88371120433196934e4b7f" +dependencies = [ + "serde", +] + [[package]] name = "serde_core" version = "1.0.229" @@ -3357,6 +3431,18 @@ dependencies = [ "tracing-serde", ] +[[package]] +name = "transforms" +version = "2.2.0" +source = "git+https://github.com/TomCC7/transforms?rev=5ae9563de0c02638c04e393aadfad28444a4398e#5ae9563de0c02638c04e393aadfad28444a4398e" +dependencies = [ + "approx", + "hashbrown 0.17.1", + "libm", + "serde", + "thiserror", +] + [[package]] name = "twox-hash" version = "1.6.3" diff --git a/dimos/navigation/trajectory_follower/fancy/rust/Cargo.toml b/dimos/navigation/trajectory_follower/fancy/rust/Cargo.toml index fd46937e34f..1c77a3415d4 100644 --- a/dimos/navigation/trajectory_follower/fancy/rust/Cargo.toml +++ b/dimos/navigation/trajectory_follower/fancy/rust/Cargo.toml @@ -20,8 +20,8 @@ required-features = ["module"] # The body (`Emb`) the laws are parameterised from; python feature off, see adapter. dimos-local-planner = { path = "../../../local_planner/rust", default-features = false } serde = { version = "1", features = ["derive"], optional = true } -dimos-module = { git = "https://github.com/dimensionalOS/dimos", rev = "fc6a4e58da89bb073a38dce5ac63728c6559d3a6", optional = true, version = "=0.1.0" } -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen", optional = true } +dimos-module = { path = "../../../../../native/rust/dimos-module", optional = true } +dimos-generated-messages = { path = "../../../../../build/message-source/rust", optional = true } tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal", "time"], optional = true } tracing = { version = "0.1", optional = true } tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true } @@ -35,7 +35,7 @@ serde_json = { version = "1", features = ["float_roundtrip"], optional = true } # Python; maturin builds with defaults and gets the extension module. default = [] python = ["dep:pyo3", "dep:numpy", "dep:serde_json"] -module = ["dimos-local-planner/module", "dep:serde", "dep:dimos-module", "dep:lcm-msgs", "dep:tokio", "dep:tracing", "dep:tracing-subscriber", "dep:validator"] +module = ["dimos-local-planner/module", "dep:serde", "dep:dimos-module", "dep:dimos-generated-messages", "dep:tokio", "dep:tracing", "dep:tracing-subscriber", "dep:validator"] # `cargo test --features module` also runs the module's own tests. [dev-dependencies] @@ -48,12 +48,13 @@ python = "dimos.navigation.trajectory_follower.fancy.native:TrajectoryFollowerNa threads = 2 [package.metadata.dimos.module.trajectory_follower.inputs] -path = "nav_msgs.Path" -tf = "tf2_msgs.TFMessage" +path = "nav_msgs/msg/Path" +tf = "tf2_msgs/msg/TFMessage" +tf_static = "tf2_msgs/msg/TFMessage" [package.metadata.dimos.module.trajectory_follower.outputs] -nav_cmd_vel = "geometry_msgs.Twist" -goal_reached = "std_msgs.Bool" +nav_cmd_vel = "geometry_msgs/msg/Twist" +goal_reached = "std_msgs/msg/Bool" [profile.release] lto = "thin" diff --git a/dimos/navigation/trajectory_follower/fancy/rust/crate-hashes.json b/dimos/navigation/trajectory_follower/fancy/rust/crate-hashes.json index 187d28513c6..3ee4f926e23 100644 --- a/dimos/navigation/trajectory_follower/fancy/rust/crate-hashes.json +++ b/dimos/navigation/trajectory_follower/fancy/rust/crate-hashes.json @@ -1,7 +1,6 @@ { - "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#lcm-msgs@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=", + "git+https://github.com/TomCC7/cdr-rs?rev=38a4368255d408faf374726653b39d4a4e9735cb#cdr@0.2.4": "sha256-akg/bLBY3wQ1YRr2bijE5ZqB211Cst5lKjELLrgiG6o=", + "git+https://github.com/TomCC7/transforms?rev=5ae9563de0c02638c04e393aadfad28444a4398e#transforms@2.2.0": "sha256-Kh8aVjj7zQ+Xd1QHOcAhSJbubg7azaWFSMpaI/ZAgM4=", "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549#dimos-lcm@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=", - "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549#lcm-msgs@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=", - "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#dimos-module-macros@0.1.0": "sha256-ECU2xYh/j+a3XwyhO3UR/GPruo5jNzM3n/9S5BIvVrA=", - "git+https://github.com/dimensionalOS/dimos?rev=fc6a4e58da89bb073a38dce5ac63728c6559d3a6#dimos-module@0.1.0": "sha256-ECU2xYh/j+a3XwyhO3UR/GPruo5jNzM3n/9S5BIvVrA=" + "git+https://github.com/dimensionalOS/dimos-lcm.git?rev=04d78e8622500244123ba9cefa4c51b4cb454549#lcm-msgs@0.1.0": "sha256-GGkx4Mn6NYP6KZecmoRLKGWIih/+y8OgNn12DeXX6n8=" } \ No newline at end of file diff --git a/dimos/navigation/trajectory_follower/fancy/rust/flake.lock b/dimos/navigation/trajectory_follower/fancy/rust/flake.lock index b4e7367f4ca..12d4afa66f6 100644 --- a/dimos/navigation/trajectory_follower/fancy/rust/flake.lock +++ b/dimos/navigation/trajectory_follower/fancy/rust/flake.lock @@ -73,37 +73,6 @@ "type": "github" } }, - "dimos-local-planner": { - "inputs": { - "crate2nix": [ - "crate2nix" - ], - "flake-utils": [ - "flake-utils" - ], - "nix-filter": [ - "nix-filter" - ], - "nixpkgs": [ - "nixpkgs" - ] - }, - "locked": { - "dir": "dimos/navigation/local_planner/rust", - "lastModified": 1792746978, - "narHash": "sha256-AYFe5iA3tnfHiZA8vVpX43dAk8Wl380NKTs1IH2OiGs=", - "owner": "dimensionalOS", - "repo": "dimos", - "rev": "77e57872657a4094dff31cdc91cabdd132cbe8a1", - "type": "github" - }, - "original": { - "dir": "dimos/navigation/local_planner/rust", - "owner": "dimensionalOS", - "repo": "dimos", - "type": "github" - } - }, "flake-compat": { "locked": { "lastModified": 1733328505, @@ -322,7 +291,6 @@ "root": { "inputs": { "crate2nix": "crate2nix", - "dimos-local-planner": "dimos-local-planner", "flake-utils": "flake-utils", "nix-filter": "nix-filter", "nixpkgs": "nixpkgs_2" diff --git a/dimos/navigation/trajectory_follower/fancy/rust/flake.nix b/dimos/navigation/trajectory_follower/fancy/rust/flake.nix index 45147d44a36..19686579a50 100644 --- a/dimos/navigation/trajectory_follower/fancy/rust/flake.nix +++ b/dimos/navigation/trajectory_follower/fancy/rust/flake.nix @@ -7,32 +7,47 @@ flake-utils.url = "github:numtide/flake-utils"; crate2nix.url = "github:nix-community/crate2nix"; crate2nix.inputs.nixpkgs.follows = "nixpkgs"; - # The sibling local planner crate it builds on (a path dependency in Cargo.toml), pinned like any shared code. - dimos-local-planner.url = "github:dimensionalOS/dimos?dir=dimos/navigation/local_planner/rust"; - dimos-local-planner.inputs.nixpkgs.follows = "nixpkgs"; - dimos-local-planner.inputs.nix-filter.follows = "nix-filter"; - dimos-local-planner.inputs.flake-utils.follows = "flake-utils"; - dimos-local-planner.inputs.crate2nix.follows = "crate2nix"; }; - outputs = { self, nix-filter, nixpkgs, flake-utils, crate2nix, dimos-local-planner }: + outputs = { self, nix-filter, nixpkgs, flake-utils, crate2nix }: flake-utils.lib.eachSystem [ "x86_64-linux" "aarch64-linux" "aarch64-darwin" ] (system: let pkgs = nixpkgs.legacyPackages.${system}; name = "dimos-trajectory-follower"; - here = "navigation/trajectory_follower/fancy/rust"; own = nix-filter.lib { root = ./.; exclude = [ "target" "build" "result" "__pycache__" ]; }; - # The two crates at their repo-relative places, so the path dependency resolves. - src = pkgs.runCommand "${name}-src" { } '' - mkdir -p $out/${here} $out/navigation/local_planner - cp -r ${own}/. $out/${here} - cp -r ${dimos-local-planner} $out/navigation/local_planner/rust - ''; - generated = crate2nix.tools.${system}.generatedCargoNix { - inherit name src; - cargoToml = "${here}/Cargo.toml"; + plannerArchive = pkgs.fetchurl { + url = "https://codeload.github.com/dimensionalOS/dimos/tar.gz/a9365487b4323151ea85041c6a944b509c0c3268"; + sha256 = "e339487dd2281980fcde7ab03c46dfd2e274491e0f73e66d9c6cbcaeea8902b2"; }; + plannerSource = pkgs.runCommand "dimos-follower-planner-source" { + nativeBuildInputs = [ pkgs.gnutar pkgs.gzip ]; + } '' + mkdir -p $out + tar -xzf ${plannerArchive} --strip-components=1 -C $out + ''; + messageSource = (import "${plannerSource}/native/message-source.nix" { inherit pkgs; }).source; + src = pkgs.runCommand "${name}-cdr-source" { } '' + cp -r ${own} $out + chmod -R u+w $out + mkdir -p $out/sdk + cp -r ${plannerSource}/dimos/navigation/local_planner/rust $out/local-planner + cp -r ${plannerSource}/native/rust/dimos-module ${plannerSource}/native/rust/dimos-module-macros $out/sdk/ + cp -r ${messageSource}/rust $out/messages + substituteInPlace $out/sdk/dimos-module/Cargo.toml \ + --replace-fail '../../../build/message-source/rust' '../../messages' + # The source artifacts own workspace roots; do not make them automatic members. + substituteInPlace $out/Cargo.toml --replace-fail '[workspace]' '[workspace] + exclude = ["local-planner", "sdk/dimos-module", "sdk/dimos-module-macros", "messages"]' + substituteInPlace $out/Cargo.toml \ + --replace-fail '../../../local_planner/rust' 'local-planner' \ + --replace-fail '../../../../../native/rust/dimos-module' 'sdk/dimos-module' \ + --replace-fail '../../../../../build/message-source/rust' 'messages' + substituteInPlace $out/local-planner/Cargo.toml \ + --replace-fail '../../../../native/rust/dimos-module' '../sdk/dimos-module' \ + --replace-fail '../../../../build/message-source/rust' '../messages' + ''; + generated = crate2nix.tools.${system}.generatedCargoNix { inherit name src; }; ours = [ name "dimos-local-planner" "dimos-module" "dimos-module-macros" ]; callWith = mode: import generated { diff --git a/dimos/navigation/trajectory_follower/fancy/rust/src/module/follower.rs b/dimos/navigation/trajectory_follower/fancy/rust/src/module/follower.rs index 2a3effa75ae..d39c65ba8c4 100644 --- a/dimos/navigation/trajectory_follower/fancy/rust/src/module/follower.rs +++ b/dimos/navigation/trajectory_follower/fancy/rust/src/module/follower.rs @@ -20,11 +20,11 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use crate::laws::hinted::update as hinted_update; +use dimos_generated_messages::geometry_msgs::msg::twist::Twist; +use dimos_generated_messages::nav_msgs::msg::path::Path; +use dimos_generated_messages::std_msgs::msg::bool::Bool; use dimos_local_planner::planner::Emb; use dimos_module::{native_config, Input, Module, Output, Tf}; -use lcm_msgs::geometry_msgs::Twist; -use lcm_msgs::nav_msgs::Path; -use lcm_msgs::std_msgs::Bool; // serde derives come in through #[native_config] use tracing::info; @@ -118,13 +118,13 @@ struct Shared { #[derive(Module)] #[module(name = "trajectory_follower", setup = spawn_worker, teardown = stop_worker)] pub struct TrajectoryFollower { - #[input(decode = Path::decode, handler = on_path)] + #[input(decode = dimos_module::cdr::decode, handler = on_path)] path: Input, - #[output(encode = Twist::encode)] + #[output(encode = dimos_module::cdr::encode)] nav_cmd_vel: Output, - #[output(encode = Bool::encode)] + #[output(encode = dimos_module::cdr::encode)] goal_reached: Output, #[config] @@ -307,7 +307,7 @@ impl Worker { let pose = pose.expect("Drive implies a pose"); let path = snap.path.clone().expect("Drive implies a path"); let states = msg::path_states(&path); - let ts = msg::path_stamps(&path); + let ts = msg::relative_path_stamps(&path); // no clearance array: the law decodes the stamps itself let (vx, vy, wz) = hinted_update(pose, &states, None, Some(&ts), &hinted); self.publish_twist(vx, vy, wz).await; @@ -334,6 +334,42 @@ impl Worker { mod tests { use super::*; + #[test] + fn cdr_paths_retain_epoch_nanoseconds_and_arrival_goal_in_both_byte_orders() { + let mut path = path_of(&[[0.0, 0.0, 0.0], [1.0, 2.0, 0.0]]); + for (index, pose) in path.poses.iter_mut().enumerate() { + pose.header.stamp.sec = 1_700_000_000; + pose.header.stamp.nanosec = 123_456_789 + index as u32; + } + for little_endian in [false, true] { + let bytes = dimos_module::cdr::encode_endian(&path, little_endian).unwrap(); + let decoded: Path = dimos_module::cdr::decode(&bytes).unwrap(); + assert_eq!(msg::relative_path_stamps(&decoded), vec![0.0, 1e-9]); + assert_eq!(decoded.poses[0].header.stamp.nanosec, 123_456_789); + let mut latch = latch(); + latch.set_goal(goal_of(&decoded).unwrap()); + assert_eq!( + decide(Some((1.0, 2.0, 0.0)), Some(0.0), 1.0, &mut latch), + Tick::Arrived + ); + assert_eq!( + decide(Some((1.0, 2.0, 0.0)), Some(0.0), 1.0, &mut latch), + Tick::Holding + ); + } + } + + #[test] + fn arrival_boolean_is_plain_cdr_and_empty_path_has_no_goal() { + assert_eq!( + dimos_module::cdr::encode(&Bool { data: true }).unwrap(), + vec![0, 1, 0, 0, 1] + ); + let empty: Path = + dimos_module::cdr::decode(&dimos_module::cdr::encode(&path_of(&[])).unwrap()).unwrap(); + assert_eq!(goal_of(&empty), None); + } + #[test] fn a_hold_stub_is_a_stop_to_every_law() { // the refusal is one pose, under the laws' `path.len() < 2` veto diff --git a/dimos/navigation/trajectory_follower/fancy/test_controller.py b/dimos/navigation/trajectory_follower/fancy/test_controller.py index 167029c2c8f..aa17b3494bd 100644 --- a/dimos/navigation/trajectory_follower/fancy/test_controller.py +++ b/dimos/navigation/trajectory_follower/fancy/test_controller.py @@ -19,10 +19,11 @@ from pydantic import ValidationError import pytest -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.message_utils.geometry import quaternion_from_euler +from dimos.message_utils.time import time_from_seconds +from dimos.messages.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Header from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.trajectory_follower.fancy.controller import ControllerConfig from dimos.navigation.trajectory_follower.fancy.laws.seed import PursuitController @@ -30,21 +31,23 @@ def _pose(x: float, y: float, yaw: float = 0.0) -> PoseStamped: return PoseStamped( - frame_id="world", - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0, 0, yaw)), + header=Header(frame_id="world", stamp=time_from_seconds(0.0)), + pose=Pose(position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0, 0, yaw)), ) def _straight_path(n: int = 40, step: float = 0.1) -> Path: - return Path(frame_id="world", poses=[_pose(i * step, 0.0) for i in range(n)]) + return Path( + header=Header(frame_id="world", stamp=time_from_seconds(0.0)), + poses=[_pose(i * step, 0.0) for i in range(n)], + ) def _fan_path() -> Path: """Rotate in place at the origin: yaw steps with zero displacement.""" yaws = [0.0, 0.3, 0.6, 0.9, 1.2, 1.5] poses = [_pose(0.0, 0.0, y) for y in yaws] + [_pose(0.1, 0.0, 1.5), _pose(1.0, 0.1, 1.5)] - return Path(frame_id="world", poses=poses) + return Path(header=Header(frame_id="world", stamp=time_from_seconds(0.0)), poses=poses) def test_on_path_drives_forward() -> None: @@ -84,7 +87,11 @@ def test_fan_done_resumes_translation() -> None: def test_empty_path_stops() -> None: - tw = PursuitController().update(_pose(0.0, 0.0), Path(frame_id="world", poses=[]), 0.0) + tw = PursuitController().update( + _pose(0.0, 0.0), + Path(header=Header(frame_id="world", stamp=time_from_seconds(0.0)), poses=[]), + 0.0, + ) assert tw.linear.x == 0.0 and tw.angular.z == 0.0 @@ -105,7 +112,7 @@ def test_config_rejects_unknown_fields() -> None: def test_governor_creeps_in_tight_room() -> None: path = _straight_path() - tight = np.full(len(path), 0.06) # barely above the precision floor + tight = np.full(len(path.poses), 0.06) # barely above the precision floor tw = PursuitController().update(_pose(0.0, 0.0), path, 0.0, clearance=tight) # within the bottom twentieth of the governor's band, whatever the ceiling is creep = GO2.min_speed + 0.05 * (GO2.max_speed - GO2.min_speed) @@ -114,7 +121,7 @@ def test_governor_creeps_in_tight_room() -> None: def test_governor_full_speed_in_open_room() -> None: path = _straight_path() - wide = np.full(len(path), 1.0) + wide = np.full(len(path.poses), 1.0) tw_open = PursuitController().update(_pose(-1.0, 0.0), path, 0.0, clearance=wide) tw_bare = PursuitController().update(_pose(-1.0, 0.0), path, 0.0) assert math.hypot(tw_open.linear.x, tw_open.linear.y) == pytest.approx( @@ -124,7 +131,7 @@ def test_governor_full_speed_in_open_room() -> None: def test_governor_reads_room_ahead_not_behind() -> None: path = _straight_path() # 4 m of path - clear = np.full(len(path), 1.0) + clear = np.full(len(path.poses), 1.0) clear[:5] = 0.06 # tight patch already behind the lookahead window tw = PursuitController().update(_pose(1.5, 0.0), path, 0.0, clearance=clear) assert math.hypot(tw.linear.x, tw.linear.y) > GO2.min_speed + 0.1 diff --git a/dimos/navigation/trajectory_follower/fancy/test_generated_native.py b/dimos/navigation/trajectory_follower/fancy/test_generated_native.py new file mode 100644 index 00000000000..169296db8d2 --- /dev/null +++ b/dimos/navigation/trajectory_follower/fancy/test_generated_native.py @@ -0,0 +1,188 @@ +# 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 contextlib import ExitStack +from copy import deepcopy +import json +import os +from pathlib import Path as FilePath +import socket +from threading import Event +import time + +import pytest +import zenoh + +from dimos.core.global_config import global_config +from dimos.core.transport import ZenohTransport +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Twist, + Vector3, +) +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Bool, Header +from dimos.messages.tf2_msgs.msg import TFMessage +from dimos.navigation.embodiment.go2 import GO2 +from dimos.navigation.trajectory_follower.fancy.native import TrajectoryFollowerNative +from dimos.protocol.pubsub.impl.zenohpubsub import Topic as ZenohTopic +from dimos.protocol.rpc.pubsubrpc import LCMRPC +from dimos.protocol.service.zenohservice import ZenohConfig + + +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +def test_native_follower_declares_canonical_generated_port_identities(): + module = TrajectoryFollowerNative(rpc_transport=MemoryRPC, embodiment=GO2) + try: + assert module.path.type is Path and module.tf.type is TFMessage + assert module.tf_static.type is TFMessage + assert module.nav_cmd_vel.type is Twist and module.goal_reached.type is Bool + finally: + module.stop() + + +@pytest.mark.native_e2e +def test_actual_native_follower_cdr_drive_arrival_new_goal_and_empty_stop(monkeypatch): + executable = os.environ.get("DIMOS_NATIVE_FOLLOWER_BINARY") + assert executable, "Build and set the native follower binary before running native_e2e" + binary = FilePath(executable).resolve() + assert binary.is_file() + monkeypatch.setattr(global_config, "transport", "zenoh") + with socket.socket() as reservation: + reservation.bind(("127.0.0.1", 0)) + endpoint = f"tcp/127.0.0.1:{reservation.getsockname()[1]}" + config = zenoh.Config() + for key, value in { + "mode": "router", + "listen/endpoints": [endpoint], + "scouting/multicast/enabled": False, + "scouting/gossip/enabled": False, + }.items(): + config.insert_json5(key, json.dumps(value)) + with ExitStack() as cleanup: + cleanup.enter_context(zenoh.open(config)) + session = ZenohConfig( + mode="client", + connect=[endpoint], + listen=[], + scouting=False, + multicast=False, + gossip=False, + connect_timeout=5, + ) + module = TrajectoryFollowerNative( + rpc_transport=MemoryRPC, + embodiment=GO2, + executable=str(binary), + source_dir=None, + build_command=None, + session=session, + ) + wires = {} + for name in ["path", "tf", "tf_static", "nav_cmd_vel", "goal_reached"]: + port = getattr(module, name) + wire = ZenohTransport( + ZenohTopic("native/follower/" + name, port.type), **session.model_dump() + ) + port.transport = wire + wires[name] = wire + cleanup.callback(wire.stop) + cleanup.callback(module.stop) + moving, arrived, stopped = Event(), Event(), Event() + commands, arrivals = [], [] + + def command(value): + commands.append(value) + if value.linear.x > 0: + moving.set() + if value.linear.x == value.linear.y == value.angular.z == 0: + stopped.set() + + def arrival(value): + arrivals.append(value) + if value.data: + arrived.set() + + cleanup.callback(wires["nav_cmd_vel"].subscribe(command)) + cleanup.callback(wires["goal_reached"].subscribe(arrival)) + module.start() + epoch = 1700000000123456789 + sequence = 0 + + def feed(x, route): + nonlocal sequence + header = Header(frame_id="odom", stamp=time_from_nanoseconds(epoch + sequence)) + sequence += 1 + wires["tf"].publish( + TFMessage( + transforms=[ + TransformStamped( + header=header, + child_frame_id="base_link", + transform=Transform( + translation=Vector3(x=x, y=0, z=0), + rotation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) + ] + ) + ) + wires["path"].publish(route) + + def route(start): + return Path( + header=Header(frame_id="odom", stamp=time_from_nanoseconds(epoch)), + poses=[ + PoseStamped( + header=Header( + frame_id="odom", stamp=time_from_nanoseconds(epoch + i * 125000000) + ), + pose=Pose( + position=Point(x=start + i * 0.125, y=0, z=0), + orientation=Quaternion(x=0, y=0, z=0, w=1), + ), + ) + for i in range(17) + ], + ) + + def until(event, x, path): + deadline = time.monotonic() + 10 + while not event.is_set() and time.monotonic() < deadline: + feed(x, path) + event.wait(0.15) + assert event.is_set(), "Native follower did not reach the expected state" + + first = route(0) + until(moving, 0, first) + assert type(commands[-1]) is Twist + until(arrived, 2, first) + assert type(arrivals[-1]) is Bool and arrivals[-1].data + moving.clear() + second = route(2) + until(moving, 2, second) + stopped.clear() + until(stopped, 2, Path(header=deepcopy(second.header), poses=[])) + assert commands[-1].linear.x == commands[-1].linear.y == commands[-1].angular.z == 0 + assert module._process is not None and module._process.poll() is None diff --git a/dimos/navigation/trajectory_follower/fancy/test_module.py b/dimos/navigation/trajectory_follower/fancy/test_module.py index 8376014d949..c6c5eb38b1e 100644 --- a/dimos/navigation/trajectory_follower/fancy/test_module.py +++ b/dimos/navigation/trajectory_follower/fancy/test_module.py @@ -16,11 +16,18 @@ import pytest from dimos.core.module import ModuleConfig -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.Path import Path +from dimos.message_utils.time import time_from_seconds +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Header from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.tf_pose import TfPose from dimos.navigation.trajectory_follower.fancy.laws import seed @@ -30,8 +37,15 @@ TrajectoryFollowerConfig, ) from dimos.navigation.trajectory_follower.fancy.native import TrajectoryFollowerNativeConfig +from dimos.protocol.rpc.pubsubrpc import LCMRPC from dimos.protocol.tf.tf import MultiTBuffer + +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + # The real constructor stands up an LCM RPC transport per instance; the fixture stops them. _BUILT: list[TrajectoryFollower] = [] @@ -63,25 +77,24 @@ def test_goal_latch_ignores_sub_tolerance_goal_moves(): def _follower(**config): - follower = TrajectoryFollower(**{"embodiment": GO2, **config}) + follower = TrajectoryFollower(rpc_transport=MemoryRPC, **{"embodiment": GO2, **config}) _BUILT.append(follower) return follower def _base_at(z: float) -> PoseStamped: """The tf-resolved base pose; the ground sits emb.base_height under it.""" - return PoseStamped(ts=0.0, frame_id="odom", position=Vector3(0.0, 0.0, z)) + return PoseStamped( + header=Header(frame_id="odom", stamp=time_from_seconds(0.0)), + pose=Pose(position=Point(x=0.0, y=0.0, z=z), orientation=Quaternion(0.0, 0.0, 0.0, 1.0)), + ) def _straight_path(end_x: float = 0.5) -> Path: - return Path( - ts=0.0, - frame_id="odom", - poses=[ - PoseStamped(ts=0.0, frame_id="odom", position=Vector3(0.0, 0.0, 0.0)), - PoseStamped(ts=0.0, frame_id="odom", position=Vector3(end_x, 0.0, 0.0)), - ], - ) + first = _base_at(0.0) + last = _base_at(0.0) + last.pose.position.x = end_x + return Path(header=Header(frame_id="odom", stamp=time_from_seconds(0.0)), poses=[first, last]) class _Heard(list): @@ -133,7 +146,8 @@ def test_a_stale_path_outranks_an_arrival(): out, reached = _driven(follower), _reached(follower) path = _straight_path(0.5) follower._on_path(path) - at_goal = PoseStamped(ts=0.0, frame_id="odom", position=Vector3(0.5, 0.0, 0.01)) + at_goal = _base_at(0.01) + at_goal.pose.position.x = 0.5 follower.step(at_goal, path, age=9.0) assert out[-1].linear.x == 0.0 assert not reached @@ -154,12 +168,12 @@ def _on_tf(follower: TrajectoryFollower, frame: str = "odom") -> tuple[MultiTBuf tf, clock = MultiTBuffer(), _Clock() follower._pose_src = TfPose(tf, "base_link", follower.config.max_path_age_s, clock=clock) tf.receive_transform( - Transform( - translation=Vector3(0.0, 0.0, 0.01), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id=frame, + TransformStamped( + header=Header(frame_id=frame, stamp=time_from_seconds(5.0)), child_frame_id="base_link", - ts=5.0, + transform=Transform( + translation=Vector3(0.0, 0.0, 0.01), rotation=Quaternion(0.0, 0.0, 0.0, 1.0) + ), ) ) return tf, clock @@ -172,7 +186,12 @@ def test_the_pose_comes_off_tf_in_the_path_frame(heard): follower._on_path(_straight_path(5.0)) # an odom plan; the edge on tf is map -> base_link follower.tick() assert (out[-1].linear.x, out[-1].linear.y, out[-1].angular.z) == (0.0, 0.0, 0.0) - follower._on_path(Path(ts=0.0, frame_id="map", poses=_straight_path(5.0).poses)) + follower._on_path( + Path( + header=Header(frame_id="map", stamp=time_from_seconds(0.0)), + poses=_straight_path(5.0).poses, + ) + ) clock.t += TfPose.RETRY_PERIOD_S # the miss parked lookups for a period follower.tick() assert out[-1].linear.x > 0.0 @@ -197,7 +216,7 @@ def test_an_empty_path_drops_the_plan_and_zeroes_the_twist(): follower._on_path(path) follower.step(pose, path, age=0.0) assert out[-1].linear.x > 0.0 - follower._on_path(Path(ts=0.0, frame_id=path.frame_id, poses=[])) + follower._on_path(Path(header=path.header, poses=[])) assert (out[-1].linear.x, out[-1].linear.y, out[-1].angular.z) == (0.0, 0.0, 0.0) assert follower._path is None diff --git a/dimos/navigation/trajectory_follower/fancy/test_rust_parity.py b/dimos/navigation/trajectory_follower/fancy/test_rust_parity.py index f31922ee309..fb9b040f7b1 100644 --- a/dimos/navigation/trajectory_follower/fancy/test_rust_parity.py +++ b/dimos/navigation/trajectory_follower/fancy/test_rust_parity.py @@ -27,10 +27,11 @@ pytest.importorskip("dimos_trajectory_follower") -from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Path import Path +from dimos.message_utils.geometry import quaternion_from_euler +from dimos.message_utils.time import time_from_seconds +from dimos.messages.geometry_msgs.msg import Point, Pose, PoseStamped +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Header from dimos.navigation.embodiment.go2 import GO2 from dimos.navigation.trajectory_follower.fancy.controller import ( ControllerConfig, @@ -51,14 +52,16 @@ def _pose(x: float, y: float, yaw: float = 0.0) -> PoseStamped: return PoseStamped( - frame_id="world", - position=Vector3(x, y, 0.0), - orientation=Quaternion.from_euler(Vector3(0, 0, yaw)), + header=Header(frame_id="world", stamp=time_from_seconds(0.0)), + pose=Pose(position=Point(x=x, y=y, z=0.0), orientation=quaternion_from_euler(0, 0, yaw)), ) def _path(states: list[tuple[float, float, float]]) -> Path: - return Path(frame_id="world", poses=[_pose(*s) for s in states]) + return Path( + header=Header(frame_id="world", stamp=time_from_seconds(0.0)), + poses=[_pose(*s) for s in states], + ) def _straight(rng: np.random.Generator) -> list[tuple[float, float, float]]: @@ -142,10 +145,12 @@ def _cases(seed: int = 20260802, n: int = CASES): # type: ignore[no-untyped-def path_xy_yaw(path), enc, float(srng.uniform(0.0, 1e9)), GOVERNOR ) for q, t in zip(path.poses, ts, strict=True): - q.ts = t + q.header.stamp = time_from_seconds(t) if k % 8 == 3: for q in path.poses: - q.ts = 5.0 # flat: not the dialect, both must ignore it + q.header.stamp = time_from_seconds( + 5.0 + ) # flat: not the dialect, both must ignore it # non-default gains and plant every fifth case: the params tuple order is load-bearing emb = GO2 @@ -153,14 +158,14 @@ def _cases(seed: int = 20260802, n: int = CASES): # type: ignore[no-untyped-def emb = replace( GO2, control=GO2.control.model_copy( - update=dict( - lookahead=float(rng.uniform(0.1, 1.2)), - k_pos=float(rng.uniform(0.5, 4.0)), - k_yaw=float(rng.uniform(0.5, 4.0)), - fan_yaw_per_m=float(rng.uniform(1.0, 6.0)), - fan_yaw_done=float(rng.uniform(0.05, 0.6)), - speed_lookahead=float(rng.uniform(0.5, 4.0)), - ) + update={ + "lookahead": float(rng.uniform(0.1, 1.2)), + "k_pos": float(rng.uniform(0.5, 4.0)), + "k_yaw": float(rng.uniform(0.5, 4.0)), + "fan_yaw_per_m": float(rng.uniform(1.0, 6.0)), + "fan_yaw_done": float(rng.uniform(0.05, 0.6)), + "speed_lookahead": float(rng.uniform(0.5, 4.0)), + } ), max_speed=float(rng.uniform(0.2, 1.5)), min_speed=float(rng.uniform(0.05, 0.3)), @@ -203,9 +208,9 @@ def test_rust_matches_python(law: str) -> None: seen["held"] += 1 if abs(a[2]) >= emb.max_yaw_rate - 1e-12 or math.hypot(a[0], a[1]) >= top - 1e-12: seen["clamped"] += 1 - if case[3] is not None and len(case[3]) == len(case[2]): + if case[3] is not None and len(case[3]) == len(case[2].poses): seen["governed"] += 1 - if len(case[2]) > 2 and math.hypot(a[0], a[1]) < 1e-3 and abs(a[2]) > 1e-3: + if len(case[2].poses) > 2 and math.hypot(a[0], a[1]) < 1e-3 and abs(a[2]) > 1e-3: seen["fan"] += 1 # the sweep is only worth its tolerance if it actually reached the branches assert all(v > 0 for v in seen.values()), f"unexercised branches: {seen}" diff --git a/dimos/navigation/trajectory_follower/test_generated_ports.py b/dimos/navigation/trajectory_follower/test_generated_ports.py new file mode 100644 index 00000000000..c1ff7bbb847 --- /dev/null +++ b/dimos/navigation/trajectory_follower/test_generated_ports.py @@ -0,0 +1,123 @@ +# 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 contextlib import ExitStack +from copy import deepcopy +from threading import Event + +import pytest + +from dimos.core.transport import LCMTransport +from dimos.message_utils.time import time_from_nanoseconds +from dimos.messages.geometry_msgs.msg import ( + Point, + Pose, + PoseStamped, + Quaternion, + Transform, + TransformStamped, + Vector3, +) +from dimos.messages.nav_msgs.msg import Path +from dimos.messages.std_msgs.msg import Bool, Header +from dimos.messages.tf2_msgs.msg import TFMessage +from dimos.navigation.embodiment.go2 import GO2 +from dimos.navigation.trajectory_follower.basic.module import BasicPathFollower +from dimos.navigation.trajectory_follower.fancy.module import TrajectoryFollower +from dimos.protocol.rpc.pubsubrpc import LCMRPC + + +class MemoryRPC(LCMRPC): + def __init__(self, **kwargs): + super().__init__(url="memq://", **kwargs) + + +def route(header, end): + return Path( + header=deepcopy(header), + poses=[ + PoseStamped( + header=deepcopy(header), + pose=Pose( + position=Point(x=x, y=0, z=0), orientation=Quaternion(x=0, y=0, z=0, w=1) + ), + ) + for x in [end * index / 16 for index in range(17)] + ], + ) + + +def base_transform(header, x): + return TFMessage( + transforms=[ + TransformStamped( + header=deepcopy(header), + child_frame_id="base_link", + transform=Transform( + translation=Vector3(x=x, y=0, z=0.01), rotation=Quaternion(x=0, y=0, z=0, w=1) + ), + ) + ] + ) + + +@pytest.mark.parametrize("follower_type", [BasicPathFollower, TrajectoryFollower]) +def test_actual_generated_follower_ports_drive_arrive_and_preempt_over_cdr(follower_type): + with ExitStack() as cleanup: + config = {"embodiment": GO2} if follower_type is TrajectoryFollower else {} + follower = follower_type(rpc_transport=MemoryRPC, **config) + wires = {} + for name in ["tf", "path", "nav_cmd_vel", "goal_reached"]: + port = getattr(follower, name) + wire = LCMTransport("follower_" + name, port.type, url="memq://") + port.transport = wire + wires[name] = wire + cleanup.callback(wire.stop) + cleanup.callback(follower.stop) + commands = [] + goals = [] + moving = Event() + arrived = Event() + zeroed = Event() + + def receive_command(value): + commands.append(value) + if value.linear.x > 0: + moving.set() + elif value.linear.x == 0 and value.linear.y == 0 and value.angular.z == 0: + zeroed.set() + + def receive_goal(value): + goals.append(value) + if value.data: + arrived.set() + + cleanup.callback(wires["nav_cmd_vel"].subscribe(receive_command)) + cleanup.callback(wires["goal_reached"].subscribe(receive_goal)) + follower.start() + header = Header(frame_id="odom", stamp=time_from_nanoseconds(1700000000123456789)) + wires["tf"].publish(base_transform(header, 0.0)) + wires["path"].publish(route(header, 2.0)) + assert moving.wait(5) + header.stamp = time_from_nanoseconds(1700000000123456790) + wires["tf"].publish(base_transform(header, 2.0)) + assert arrived.wait(5) + assert goals == [Bool(data=True)] + moving.clear() + wires["path"].publish(route(header, 4.0)) + assert moving.wait(5) + zeroed.clear() + wires["path"].publish(Path(header=deepcopy(header), poses=[])) + assert zeroed.wait(5) + assert any(c.linear.x > 0 for c in commands) diff --git a/dimos/robot/unitree/type/lidar.py b/dimos/robot/unitree/type/lidar.py index c8e01962e2d..0e2092a17e5 100644 --- a/dimos/robot/unitree/type/lidar.py +++ b/dimos/robot/unitree/type/lidar.py @@ -15,6 +15,7 @@ """Unitree WebRTC lidar message parsing utilities.""" from collections.abc import Callable +from copy import deepcopy import time from typing import TypedDict, TypeVar @@ -22,6 +23,8 @@ from reactivex import operators as ops from reactivex.observable import Observable +from dimos.message_utils.time import time_from_nanoseconds, time_from_seconds, to_nanoseconds +from dimos.messages.sensor_msgs.msg import PointCloud2 as CdrCloud from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.types.timestamped import Timestamped from dimos.utils.logging_config import setup_logger @@ -83,7 +86,7 @@ def pointcloud2_from_webrtc_lidar(raw_message: RawLidarMsg, ts: float | None = N ) -T = TypeVar("T", bound=Timestamped) +T = TypeVar("T", bound=Timestamped | CdrCloud) def repair_stale_ts( @@ -98,40 +101,52 @@ def repair_stale_ts( On new firmware, occasionally frames will revert back to the initial timestamp. In these cases, we update based on the default period. We calibrate through the first few frames to determine which correction method to use. Once it's been determined, it does not change. + + Generated clouds use integer Header nanoseconds; repaired values own a copy + so other subscribers retain the original capture Header. Legacy Timestamped + callers retain their existing mutation behavior until their consumers migrate. """ - prev_good: float | None = None - prev_raw: float | None = None + prev_good: int | None = None + prev_raw: int | None = None n_seen = 0 calibrated = False use_system_time = False + period_ns = to_nanoseconds(time_from_seconds(default_period)) + def _repair(item: T) -> T: nonlocal prev_good, prev_raw, n_seen, calibrated, use_system_time - + raw = ( + to_nanoseconds(item.header.stamp) + if isinstance(item, CdrCloud) + else to_nanoseconds(time_from_seconds(item.ts)) + ) + repaired = raw if use_system_time: - item.ts = now() - return item - - if not calibrated: - if prev_raw is not None and item.ts != prev_raw: + repaired = to_nanoseconds(time_from_seconds(now())) + else: + if not calibrated: + if prev_raw is not None and raw != prev_raw: + calibrated = True + prev_raw = raw + n_seen += 1 + if prev_good is not None and raw <= prev_good: + repaired = prev_good + period_ns + prev_good = repaired + if not calibrated and n_seen >= calibration_frames: calibrated = True - # lidar stamps advancing — using lidar time", - prev_raw = item.ts - n_seen += 1 - - if prev_good is not None and item.ts <= prev_good: - item.ts = prev_good + default_period - - prev_good = item.ts - - if not calibrated and n_seen >= calibration_frames: - calibrated = True - use_system_time = True - logger.warning( - "repair_stale_ts: lidar timestmaps frozen (%d calibration stamps equal) — using system time, upgrade your GO2 firmware", - calibration_frames, - ) - + use_system_time = True + logger.warning( + "repair_stale_ts: lidar timestamps frozen (%d calibration stamps equal) — using system time, upgrade your GO2 firmware", + calibration_frames, + ) + if repaired == raw: + return item + if isinstance(item, CdrCloud): + result = deepcopy(item) + result.header.stamp = time_from_nanoseconds(repaired) + return result + item.ts = repaired / 1_000_000_000 return item return ops.map(_repair) diff --git a/dimos/robot/unitree/type/test_lidar.py b/dimos/robot/unitree/type/test_lidar.py index 631e7f680ef..11cd8e06ac1 100644 --- a/dimos/robot/unitree/type/test_lidar.py +++ b/dimos/robot/unitree/type/test_lidar.py @@ -16,9 +16,14 @@ import itertools from typing import cast +from dimos_cdrgen.registry import encode +import numpy as np import pytest import reactivex as rx +from dimos.message_utils.pointcloud import pointcloud_from_xyz, pointcloud_xyz +from dimos.message_utils.time import time_from_nanoseconds, to_nanoseconds +from dimos.messages.std_msgs.msg import Header from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.robot.unitree.type.lidar import ( RawLidarMsg, @@ -101,3 +106,54 @@ def test_repair_stale_ts_calibration_boundary_one_differs() -> None: repair_stale_ts(default_period=0.130, calibration_frames=10, now=lambda: 999.0) ).subscribe(on_next=lambda item: out.append(item.ts)) assert 999.0 not in out + + +@pytest.mark.parametrize( + "stamps,expected", + [ + ( + [1700000000123456789, 1700000000123456790, 1700000000123456791], + [1700000000123456789, 1700000000123456790, 1700000000123456791], + ), + ( + [1700000000123456789, 1700000000123456789, 1700000000253456789], + [1700000000123456789, 1700000000253456789, 1700000000383456789], + ), + ], +) +def test_generated_stale_repair_uses_nanoseconds_and_preserves_source(stamps, expected): + clouds = [ + pointcloud_from_xyz( + np.array([[1, 2, 3]], dtype=np.float32), + header=Header(frame_id="world", stamp=time_from_nanoseconds(stamp)), + ) + for stamp in stamps + ] + originals = [encode(cloud) for cloud in clouds] + received = [] + rx.from_iterable(clouds).pipe(repair_stale_ts()).subscribe(received.append) + assert [to_nanoseconds(cloud.header.stamp) for cloud in received] == expected + assert [encode(cloud) for cloud in clouds] == originals + assert [cloud.header.frame_id for cloud in received] == ["world"] * 3 + assert all(pointcloud_xyz(cloud).tolist() == [[1, 2, 3]] for cloud in received) + + +def test_generated_frozen_firmware_uses_clock_after_same_calibration_window(): + stamp = 1700000000123456789 + clouds = [ + pointcloud_from_xyz( + np.array([[1, 2, 3]], dtype=np.float32), + header=Header(frame_id="world", stamp=time_from_nanoseconds(stamp)), + ) + for _ in range(13) + ] + before = [encode(cloud) for cloud in clouds] + received = [] + rx.from_iterable(clouds).pipe( + repair_stale_ts(calibration_frames=10, now=lambda: 999.0) + ).subscribe(received.append) + assert [to_nanoseconds(cloud.header.stamp) for cloud in received[:10]] == [ + stamp + i * 130000000 for i in range(10) + ] + assert [to_nanoseconds(cloud.header.stamp) for cloud in received[10:]] == [999000000000] * 3 + assert [encode(cloud) for cloud in clouds] == before diff --git a/dimos/spec/mapping.py b/dimos/spec/mapping.py index f35778f40b4..57ade6e36ed 100644 --- a/dimos/spec/mapping.py +++ b/dimos/spec/mapping.py @@ -15,8 +15,8 @@ from typing import Protocol from dimos.core.stream import Out -from dimos.msgs.nav_msgs.OccupancyGrid import OccupancyGrid -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.messages.nav_msgs.msg import OccupancyGrid +from dimos.messages.sensor_msgs.msg import PointCloud2 class GlobalPointcloud(Protocol): diff --git a/docs/capabilities/navigation/deep_dive.md b/docs/capabilities/navigation/deep_dive.md index 1b702a98f08..66577d36bd2 100644 --- a/docs/capabilities/navigation/deep_dive.md +++ b/docs/capabilities/navigation/deep_dive.md @@ -80,7 +80,7 @@ Live column-carving has no loop closure. We trust Go2 odometry, which is stable The [`CostMapper`](/dimos/mapping/costmapper.py) converts the 3D voxel map into a 2D occupancy grid. The default algorithm (`height_cost`) maps rate of change of Z, with some smoothing. -Algorithm settings live in [`occupancy.py`](/dimos/mapping/pointclouds/occupancy.py) and can be configured per robot. +Algorithm settings live in [`mapping/pointclouds/occupancy.py`](/dimos/mapping/pointclouds/occupancy.py) and can be configured per robot. #### Configuration diff --git a/docs/capabilities/navigation/relocalization.md b/docs/capabilities/navigation/relocalization.md index 4785038f30f..3e93720e4ea 100644 --- a/docs/capabilities/navigation/relocalization.md +++ b/docs/capabilities/navigation/relocalization.md @@ -13,7 +13,7 @@ This guide takes four steps: 3. Test relocalization in replay, no robot needed 4. Deploy on the live Go2 -Throughout this guide, `{DB_NAME}` is the stem of your recording, for example `recording_go2` for `recording_go2.db`. For `map_file`, pass the same stem and dimOS appends `.pc2.lcm` automatically. +Throughout this guide, `{DB_NAME}` is the stem of your recording, for example `recording_go2` for `recording_go2.db`. For `map_file`, pass the same stem and dimOS appends `.pc2.cdr` automatically. ## 1. Record a run @@ -44,7 +44,7 @@ dimos map replay recording_go2 --duration 60 ## 2. Build the premap -Export a loop-closed global map as `.pc2.lcm`: +Export a loop-closed global map as `.pc2.cdr`: ```bash dimos map global recording_go2 --export @@ -52,7 +52,7 @@ dimos map global recording_go2 --export | Flag | Effect | |------|--------| -| `--export` | Run PGO and write `./{DB_NAME}.pc2.lcm` to the current working directory (implies `--pgo`) | +| `--export` | Run PGO and write `./{DB_NAME}.pc2.cdr` to the current working directory (implies `--pgo`) | | `--no-gui` | Skip launching Rerun for headless servers or CI | | `--pgo-tol 0.3` | Spatial dedup tolerance for keyframes in meters. Use `0` to keep all posed frames | | `--voxel 0.05` | Voxel size in meters (default matches live mapper) | @@ -76,8 +76,8 @@ Sample log: ``` running PGO twopass map... Pass 1: 908 frames, 1 keyframes -exporting PGO twopass map to .../recording_go2.pc2.lcm... -wrote .../recording_go2.pc2.lcm +exporting PGO twopass map to .../recording_go2.pc2.cdr... +wrote .../recording_go2.pc2.cdr ``` Open the companion `{DB_NAME}.rrd` in Rerun to verify loop closure before deploying to hardware. @@ -91,7 +91,7 @@ dimos --replay --replay-db recording_go2 run unitree-go2-relocalization \ --map-file=recording_go2 ``` -`map_file` resolves `{DB_NAME}.pc2.lcm` with the same search order as above (cwd, then project root, then `data/`). +`map_file` resolves `{DB_NAME}.pc2.cdr` with the same search order as above (cwd, then project root, then `data/`). ### Reading the logs @@ -161,7 +161,7 @@ Note that [`CostMapper`](/dimos/mapping/costmapper.py) builds the costmap from t | File | Format | Produced by | Consumed by | |------|--------|-------------|-------------| | `{name}.db` | memory SQLite (`lidar`, `odom`, `color_image`, …) | `unitree-go2-memory` | `dimos map *`, `--replay-db` | -| `{name}.pc2.lcm` | LCM-encoded `PointCloud2` premap | `dimos map global --export` | `RelocalizationModule` (`map_file`) | +| `{name}.pc2.cdr` | CDR-encoded `PointCloud2` premap | `dimos map global --export` | `RelocalizationModule` (`map_file`) | | `{name}.rrd` | Rerun recording (visual QA) | `dimos map global` | Rerun viewer | ## Configuration reference @@ -172,7 +172,7 @@ for example `--relocalizationmodule.map-file=…`. | Field | Default | Description | |-------|---------|-------------| -| `map_file` | `None` (module disabled) | Premap stem or path. dimOS appends `.pc2.lcm` automatically | +| `map_file` | `None` (module disabled) | Premap stem or path. dimOS appends `.pc2.cdr` automatically | | `fitness_threshold` | `0.45` | Minimum ICP fitness to accept a relocalization (0 to 1) | | `tf_interval` | `10.0` | Seconds between tf republishes of the accepted fix (published immediately on every fix) | | `republish_loaded_map` | `0.0` | Seconds between `loaded_map` republishes once placed; `0` publishes once per fix | @@ -199,7 +199,7 @@ dimos run unitree-go2-relocalization \ | Symptom | Likely cause | Fix | |---------|--------------|-----| | `Relocalization module disabled (no map_file configured)` | Missing `--map-file=…` | Set `map_file` to your premap stem | -| File not found for `.pc2.lcm` | Export not run or wrong cwd | Run `dimos map global … --export` and check cwd or `data/` | +| File not found for `.pc2.cdr` | Export not run or wrong cwd | Run `dimos map global … --export` and check cwd or `data/` | | Long stretch of `relocalize skipped` | Map still accumulating points | Wait or drive slowly through mapped geometry | | Repeated `relocalize rejected` | Poor overlap with premap or wrong space | Start in a known area and check premap in `.rrd` | | Nav works but map looks misaligned | Low fitness accepted in debug mode | Raise `fitness_threshold` back to default `0.45` | diff --git a/docs/platforms/quadruped/go2/index.md b/docs/platforms/quadruped/go2/index.md index 59ce88e3ea8..2f45d2f541f 100644 --- a/docs/platforms/quadruped/go2/index.md +++ b/docs/platforms/quadruped/go2/index.md @@ -17,7 +17,7 @@ | `dimos --record run unitree-go2` | Navigation + record every stream to `recordings//memory.db` ([Recording](/docs/usage/recording.md)) | | `dimos run unitree-go2-memory` | Navigation + record `lidar`/`odom`/`color_image` to `.db` | | `dimos --record run unitree-go2-joystick-record` | Teleop from the local keyboard window + record `joystick`, `cmd_vel` and every sensor; no mapper, viewer capped at 2 GB | -| `dimos run unitree-go2-relocalization` | Navigation + align live scans to a saved `.pc2.lcm` premap | +| `dimos run unitree-go2-relocalization` | Navigation + align live scans to a saved `.pc2.cdr` premap | ## Deep Dive diff --git a/packages/dimos-generated/custom/dimos_msgs/msg/LineSegments3DRegion.msg b/packages/dimos-generated/custom/dimos_msgs/msg/LineSegments3DRegion.msg new file mode 100644 index 00000000000..62b2fcb256b --- /dev/null +++ b/packages/dimos-generated/custom/dimos_msgs/msg/LineSegments3DRegion.msg @@ -0,0 +1,7 @@ +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +# A keyed region of weighted graph edges. Header matches lines.header. +std_msgs/Header header +int32 region_id +LineSegments3D lines diff --git a/packages/dimos-generated/custom/dimos_msgs/msg/PointCloudRegion.msg b/packages/dimos-generated/custom/dimos_msgs/msg/PointCloudRegion.msg new file mode 100644 index 00000000000..aabf2d2eaa4 --- /dev/null +++ b/packages/dimos-generated/custom/dimos_msgs/msg/PointCloudRegion.msg @@ -0,0 +1,7 @@ +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +# A keyed region of a point-cloud map. Header matches cloud.header. +std_msgs/Header header +int32 region_id +sensor_msgs/PointCloud2 cloud diff --git a/packages/dimos-generated/custom/dimos_msgs/msg/RegionBounds.msg b/packages/dimos-generated/custom/dimos_msgs/msg/RegionBounds.msg new file mode 100644 index 00000000000..c9ec3113841 --- /dev/null +++ b/packages/dimos-generated/custom/dimos_msgs/msg/RegionBounds.msg @@ -0,0 +1,10 @@ +# Copyright 2026 Dimensional Inc. +# SPDX-License-Identifier: Apache-2.0 + +# Cylindrical map region; region_id pairs seeded regions independently of stamp. +std_msgs/Header header +int32 region_id +geometry_msgs/Point center +float64 radius +float64 z_min +float64 z_max