From 86606d370bc6922971f9695544e0a2843273da82 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 23 May 2026 14:19:58 -0700 Subject: [PATCH 01/71] fixup go2 transform frames --- dimos/robot/unitree/go2/connection.py | 98 +++++++++++++-------------- docs/capabilities/memory/plot.md | 9 +-- 2 files changed, 50 insertions(+), 57 deletions(-) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 50286594ff..dff3a37fdc 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -67,6 +67,10 @@ class Go2Mode(str, Enum): class ConnectionConfig(ModuleConfig): ip: str = Field(default_factory=lambda m: m["g"].robot_ip) mode: Go2Mode = Go2Mode.DEFAULT + frame_id: str | None = "base_link" + camera_link_frame: str = "camera_link" + camera_optical_frame: str = "camera_optical" + static_transform_publish_rate: float = 1.0 class Go2ConnectionProtocol(Protocol): @@ -96,21 +100,6 @@ def _camera_info_static() -> CameraInfo: return CameraInfo.from_yaml(str(yaml_path)) -# Static camera mount chain: base_link -> camera_link -> camera_optical. -# TODO we need a standardized way to specify this for all cameras in dimos -BASE_TO_OPTICAL: Transform = Transform( - translation=Vector3(0.3, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", - child_frame_id="camera_link", -) + Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id="camera_link", - child_frame_id="camera_optical", -) - - def make_connection(ip: str | None, cfg: GlobalConfig) -> Go2ConnectionProtocol: connection_type = cfg.unitree_connection_type @@ -210,7 +199,22 @@ class GO2Connection(Module, Camera, Pointcloud): connection: Go2ConnectionProtocol camera_info_static: CameraInfo = _camera_info_static() - _camera_info_thread: Thread | None = None + static_transforms = dict( + # key=default child transform name + camera_link=Transform( + translation=Vector3(0.3, 0.0, 0.0), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="base_link", + child_frame_id="camera_link", + ), + camera_optical=Transform( + translation=Vector3(0.0, 0.0, 0.0), + rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), + frame_id="camera_link", + child_frame_id="camera_optical", + ), + ) + _static_publish_thread: Thread | None = None _latest_video_frame: Image | None = None @classmethod @@ -246,11 +250,11 @@ def onimage(image: Image) -> None: self.register_disposable(self.connection.video_stream().subscribe(onimage)) self.register_disposable(Disposable(self.cmd_vel.subscribe(self.move))) - self._camera_info_thread = Thread( - target=self.publish_camera_info, + self._static_publish_thread = Thread( + target=self._static_publish, daemon=True, ) - self._camera_info_thread.start() + self._static_publish_thread.start() self.standup() time.sleep(3) @@ -261,8 +265,6 @@ def onimage(image: Image) -> None: self.connection.set_obstacle_avoidance(self.config.g.obstacle_avoidance) - # self.record("go2_bigoffice") - @rpc def stop(self) -> None: self.liedown() @@ -270,45 +272,39 @@ def stop(self) -> None: if self.connection: self.connection.stop() - if self._camera_info_thread and self._camera_info_thread.is_alive(): - self._camera_info_thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + if self._static_publish_thread and self._static_publish_thread.is_alive(): + self._static_publish_thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) super().stop() - @classmethod - def _odom_to_tf(cls, odom: PoseStamped) -> list[Transform]: - camera_link = Transform( - translation=Vector3(0.3, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", - child_frame_id="camera_link", - ts=odom.ts, - ) - - camera_optical = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id="camera_link", - child_frame_id="camera_optical", - ts=odom.ts, - ) - - return [ - Transform.from_pose("base_link", odom), - camera_link, - camera_optical, - ] - def _publish_tf(self, msg: PoseStamped) -> None: - transforms = self._odom_to_tf(msg) - self.tf.publish(*transforms) + self.tf.publish(Transform.from_pose(self.frame_id, msg)) if self.odom.transport: self.odom.publish(msg) - def publish_camera_info(self) -> None: + def _static_publish(self) -> None: + frame_remap = { + "base_link": self.frame_id, + "camera_link": self.config.camera_link_frame, + "camera_optical": self.config.camera_optical_frame, + } + stamped_statics = [ + Transform( + translation=t.translation, + rotation=t.rotation, + frame_id=frame_remap.get(t.frame_id, t.frame_id), + child_frame_id=frame_remap.get(t.child_frame_id, t.child_frame_id), + ) + for t in self.static_transforms.values() + ] + period = 1.0 / self.config.static_transform_publish_rate while True: + now = time.time() + for st in stamped_statics: + st.ts = now + self.tf.publish(*stamped_statics) self.camera_info.publish(self.camera_info_static) - time.sleep(1.0) + time.sleep(period) @rpc def move(self, twist: Twist, duration: float = 0.0) -> bool: diff --git a/docs/capabilities/memory/plot.md b/docs/capabilities/memory/plot.md index 87328dadf8..5eb3155901 100644 --- a/docs/capabilities/memory/plot.md +++ b/docs/capabilities/memory/plot.md @@ -360,10 +360,7 @@ m.data.save("assets/plants_peak_detections.png") from dimos.perception.detection.type.detection3d.imageDetections3DPC import ( ImageDetections3DPC, ) -from dimos.robot.unitree.go2.connection import ( - _camera_info_static as go2_camerainfo, - BASE_TO_OPTICAL, -) +from dimos.robot.unitree.go2.connection import GO2Connection from dimos.memory2.vis.space.elements import Box3D from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.Transform import Transform @@ -372,7 +369,7 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 # TODO We need a nicer way to get optical transform for image streams # depending on the source def world_to_optical(base_pose): - return -(Transform.from_pose("base_link", base_pose) + BASE_TO_OPTICAL) + return -(Transform.from_pose("base_link", base_pose) + GO2Connection.static_transforms["camera_link"] + GO2Connection.static_transforms["camera_optical"]) drawing = Space() @@ -380,7 +377,7 @@ drawing.add(global_map) drawing.add(detections) -camera_info = go2_camerainfo() +camera_info = GO2Connection.camera_info_static detections3d = (detections .map_data(lambda obs: ImageDetections3DPC.from_2d( From 99623634a889d6d0294af243c90c3b0778f68877 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 23 May 2026 14:33:30 -0700 Subject: [PATCH 02/71] move consts, handle rate=0 --- dimos/constants.py | 4 +++ dimos/robot/unitree/go2/connection.py | 49 +++++++++++++++------------ 2 files changed, 31 insertions(+), 22 deletions(-) diff --git a/dimos/constants.py b/dimos/constants.py index d849f4aaf3..fb88c04426 100644 --- a/dimos/constants.py +++ b/dimos/constants.py @@ -53,3 +53,7 @@ DEFAULT_THREAD_JOIN_TIMEOUT = 2.0 DEFAULT_BUILD_NATIVE = False + +DEFAULT_ROBOT_FRAME = "base_link" +DEFAULT_CAMERA_FRAME = "camera_link" +DEFAULT_CAMERA_OPTICAL_FRAME = "camera_optical" diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index dff3a37fdc..f9bca029db 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -25,7 +25,13 @@ import rerun.blueprint as rrb from dimos.agents.annotation import skill -from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT +from dimos.constants import ( + DEFAULT_CAMERA_FRAME, + DEFAULT_CAMERA_OPTICAL_FRAME, + DEFAULT_CAPACITY_COLOR_IMAGE, + DEFAULT_ROBOT_FRAME, + DEFAULT_THREAD_JOIN_TIMEOUT, +) from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.core import rpc from dimos.core.global_config import GlobalConfig @@ -67,10 +73,10 @@ class Go2Mode(str, Enum): class ConnectionConfig(ModuleConfig): ip: str = Field(default_factory=lambda m: m["g"].robot_ip) mode: Go2Mode = Go2Mode.DEFAULT - frame_id: str | None = "base_link" - camera_link_frame: str = "camera_link" - camera_optical_frame: str = "camera_optical" - static_transform_publish_rate: float = 1.0 + frame_id: str | None = DEFAULT_ROBOT_FRAME + camera_link_frame: str = DEFAULT_CAMERA_FRAME + camera_optical_frame: str = DEFAULT_CAMERA_OPTICAL_FRAME + static_publish_rate: float = 1.0 class Go2ConnectionProtocol(Protocol): @@ -204,14 +210,14 @@ class GO2Connection(Module, Camera, Pointcloud): camera_link=Transform( translation=Vector3(0.3, 0.0, 0.0), rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", - child_frame_id="camera_link", + frame_id=DEFAULT_ROBOT_FRAME, + child_frame_id=DEFAULT_CAMERA_FRAME, ), camera_optical=Transform( translation=Vector3(0.0, 0.0, 0.0), rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id="camera_link", - child_frame_id="camera_optical", + frame_id=DEFAULT_CAMERA_FRAME, + child_frame_id=DEFAULT_CAMERA_OPTICAL_FRAME, ), ) _static_publish_thread: Thread | None = None @@ -284,9 +290,9 @@ def _publish_tf(self, msg: PoseStamped) -> None: def _static_publish(self) -> None: frame_remap = { - "base_link": self.frame_id, - "camera_link": self.config.camera_link_frame, - "camera_optical": self.config.camera_optical_frame, + DEFAULT_ROBOT_FRAME: self.frame_id, + DEFAULT_CAMERA_FRAME: self.config.camera_link_frame, + DEFAULT_CAMERA_OPTICAL_FRAME: self.config.camera_optical_frame, } stamped_statics = [ Transform( @@ -297,14 +303,15 @@ def _static_publish(self) -> None: ) for t in self.static_transforms.values() ] - period = 1.0 / self.config.static_transform_publish_rate - while True: - now = time.time() - for st in stamped_statics: - st.ts = now - self.tf.publish(*stamped_statics) - self.camera_info.publish(self.camera_info_static) - time.sleep(period) + if self.config.static_publish_rate > 0: + period = 1.0 / self.config.static_publish_rate + while True: + now = time.time() + for st in stamped_statics: + st.ts = now + self.tf.publish(*stamped_statics) + self.camera_info.publish(self.camera_info_static) + time.sleep(period) @rpc def move(self, twist: Twist, duration: float = 0.0) -> bool: @@ -359,8 +366,6 @@ def observe(self) -> Image | None: def deploy(dimos: ModuleCoordinator, ip: str, prefix: str = "") -> "ModuleProxy": - from dimos.constants import DEFAULT_CAPACITY_COLOR_IMAGE - connection = dimos.deploy(GO2Connection, ip=ip) connection.pointcloud.transport = pSHMTransport( From 4e6485e06e83ea782778730b9119c4c0eded8e4e Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 23 May 2026 14:37:12 -0700 Subject: [PATCH 03/71] dont mutate --- dimos/robot/unitree/go2/connection.py | 21 +++++++++------------ 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index f9bca029db..e887a1b0bb 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -294,21 +294,18 @@ def _static_publish(self) -> None: DEFAULT_CAMERA_FRAME: self.config.camera_link_frame, DEFAULT_CAMERA_OPTICAL_FRAME: self.config.camera_optical_frame, } - stamped_statics = [ - Transform( - translation=t.translation, - rotation=t.rotation, - frame_id=frame_remap.get(t.frame_id, t.frame_id), - child_frame_id=frame_remap.get(t.child_frame_id, t.child_frame_id), - ) - for t in self.static_transforms.values() - ] if self.config.static_publish_rate > 0: period = 1.0 / self.config.static_publish_rate while True: - now = time.time() - for st in stamped_statics: - st.ts = now + stamped_statics = [ + Transform( + translation=t.translation, + rotation=t.rotation, + frame_id=frame_remap.get(t.frame_id, t.frame_id), + child_frame_id=frame_remap.get(t.child_frame_id, t.child_frame_id), + ) + for t in self.static_transforms.values() + ] self.tf.publish(*stamped_statics) self.camera_info.publish(self.camera_info_static) time.sleep(period) From 99b2d1d406b844f6b666fed72e071399729c7387 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 23 May 2026 15:54:12 -0700 Subject: [PATCH 04/71] fully switch to urdf --- dimos/constants.py | 2 - dimos/memory2/vis/space/rerun.py | 24 ++------ dimos/robot/config.py | 84 ++++++++++++++++++++++++++- dimos/robot/model_parser.py | 19 ++++++ dimos/robot/unitree/go2/config.py | 22 +++++++ dimos/robot/unitree/go2/connection.py | 62 +++++++++----------- dimos/robot/unitree/go2/go2.urdf | 16 +++++ docs/capabilities/memory/plot.md | 3 +- 8 files changed, 177 insertions(+), 55 deletions(-) create mode 100644 dimos/robot/unitree/go2/config.py diff --git a/dimos/constants.py b/dimos/constants.py index fb88c04426..ba5e5a7c91 100644 --- a/dimos/constants.py +++ b/dimos/constants.py @@ -55,5 +55,3 @@ DEFAULT_BUILD_NATIVE = False DEFAULT_ROBOT_FRAME = "base_link" -DEFAULT_CAMERA_FRAME = "camera_link" -DEFAULT_CAMERA_OPTICAL_FRAME = "camera_optical" diff --git a/dimos/memory2/vis/space/rerun.py b/dimos/memory2/vis/space/rerun.py index 435097e77c..d193897a11 100644 --- a/dimos/memory2/vis/space/rerun.py +++ b/dimos/memory2/vis/space/rerun.py @@ -22,11 +22,10 @@ from dimos.memory2.type.observation import Observation from dimos.memory2.vis.color import Color from dimos.memory2.vis.space.elements import Arrow, Box3D, Camera, Point, Polyline, Pose, Text -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.OccupancyGrid import OccupancyGrid from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.robot.unitree.go2.config import Go2Config if TYPE_CHECKING: from dimos.memory2.vis.space.space import Space @@ -39,20 +38,6 @@ def _rgba(el: Any) -> tuple[int, int, int, int]: return c.with_alpha(c.a * opacity).rgba_u8() -# base_link → camera_optical extrinsics (applied at render time for image observations) -_BASE_TO_OPTICAL = Transform( - translation=Vector3(0.3, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", - child_frame_id="camera_link", -) + Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id="camera_link", - child_frame_id="camera_optical", -) - - def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: """Render a Space to a Rerun viewer.""" import rerun as rr @@ -240,8 +225,11 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: data = obs.data img = _as_image(data) if img is not None: - # Apply base→optical extrinsics for camera frustum rendering - world_T_optical = Transform.from_pose("world", obs.pose_stamped) + _BASE_TO_OPTICAL + # TODO: fix there should be a standard way to set camera frames + base_to_optical = Go2Config.static_absolute_transforms[ + "base_link/camera_link/camera_optical" + ] + world_T_optical = Transform.from_pose("world", obs.pose_stamped) + base_to_optical rr.log(path, world_T_optical.to_pose().to_rerun(), static=True) h, w = img.shape[:2] focal = max(w, h) diff --git a/dimos/robot/config.py b/dimos/robot/config.py index d1131b0301..f5a6db1d7a 100644 --- a/dimos/robot/config.py +++ b/dimos/robot/config.py @@ -21,6 +21,7 @@ from __future__ import annotations +from functools import cached_property from pathlib import Path from typing import Any @@ -31,8 +32,9 @@ from dimos.manipulation.planning.spec.config import RobotModelConfig 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.robot.model_parser import ModelDescription, parse_model +from dimos.robot.model_parser import JointDescription, ModelDescription, parse_model class GripperConfig(BaseModel): @@ -189,6 +191,86 @@ def joint_name_mapping(self) -> dict[str, str]: def coordinator_task_name(self) -> str: return f"traj_{self.name}" + @cached_property + def body_frame(self) -> str: + """Returns the robot's root link from the urdf (usually base_link), skips past ``type="floating"`` and returns structural link ("world" can be the true root frame, but that is detached from the robot)""" + model = self._ensure_parsed() + outgoing: dict[str, list[JointDescription]] = {} + for joint in model.joints: + outgoing.setdefault(joint.parent_link, []).append(joint) + current = model.root_link + while True: + floating_joint = next( + (joint for joint in outgoing.get(current, []) if joint.type == "floating"), + None, + ) + if floating_joint is None: + return current + current = floating_joint.child_link + + @cached_property + def static_relative_transforms(self) -> dict[str, Transform]: + """The key is the absolute path to the joint, with the value being the relative tranform (one link, not all links) + + Example:: + print(RobotConfig(model_path="go2.urdf").static_relative_transforms) + # output: + { + "base_link/camera_link": Transform(translation=(0.3, 0, 0), rotation=identity, frame_id="base_link", child_frame_id="camera_link"), + # NOTE: this is the transform from camera_link to camera_optical NOT baselink to camera_optical + "base_link/camera_link/camera_optical": Transform(translation=(0, 0, 0), rotation=(-0.5, 0.5, -0.5, 0.5), frame_id="camera_link", child_frame_id="camera_optical"), + } + """ + model = self._ensure_parsed() + by_child = {joint.child_link: joint for joint in model.joints if joint.child_link} + + def path_to(link: str) -> str: + parts = [link] + current = link + while current in by_child: + current = by_child[current].parent_link + if not current: + break + parts.append(current) + return "/".join(reversed(parts)) + + result: dict[str, Transform] = {} + for joint in model.joints: + if joint.type != "fixed" or not joint.child_link: + continue + roll, pitch, yaw = joint.origin_rpy + tx, ty, tz = joint.origin_xyz + result[path_to(joint.child_link)] = Transform( + translation=Vector3(tx, ty, tz), + rotation=Quaternion.from_euler(Vector3(roll, pitch, yaw)), + frame_id=joint.parent_link, + child_frame_id=joint.child_link, + ) + return result + + @cached_property + def static_absolute_transforms(self) -> dict[str, Transform]: + """Same keys as ``static_relative_transforms`` but each value is the cumulative + transform from the root link down to the leaf link (i.e. all hops composed). + + Example:: + print(RobotConfig(model_path="go2.urdf").static_absolute_transforms) + # output: + { + "base_link/camera_link": Transform(..., frame_id="base_link", child_frame_id="camera_link"), + "base_link/camera_link/camera_optical": Transform(..., frame_id="base_link", child_frame_id="camera_optical"), + } + """ + relative_transforms = self.static_relative_transforms + # Process parents before children so each entry can build on the prior absolute. + result: dict[str, Transform] = {} + for path in sorted(relative_transforms.keys(), key=lambda path: path.count("/")): + parent_path = path.rsplit("/", 1)[0] + parent_absolute = result.get(parent_path) + relative = relative_transforms[path] + result[path] = parent_absolute + relative if parent_absolute is not None else relative + return result + # -- Converter methods ---------------------------------------------------- def to_robot_model_config(self) -> RobotModelConfig: diff --git a/dimos/robot/model_parser.py b/dimos/robot/model_parser.py index 1b760d065d..febca8d8cb 100644 --- a/dimos/robot/model_parser.py +++ b/dimos/robot/model_parser.py @@ -37,6 +37,9 @@ class JointDescription: effort_limit: float | None = None parent_link: str = "" child_link: str = "" + # Joint origin in the parent link's frame. Defaults to identity. + origin_xyz: tuple[float, float, float] = (0.0, 0.0, 0.0) + origin_rpy: tuple[float, float, float] = (0.0, 0.0, 0.0) @dataclass @@ -139,6 +142,13 @@ def _parse_urdf_string(xml_string: str) -> ModelDescription: velocity = _float_or_none(limit_elem.get("velocity")) effort = _float_or_none(limit_elem.get("effort")) + origin_xyz = (0.0, 0.0, 0.0) + origin_rpy = (0.0, 0.0, 0.0) + origin_elem = joint_elem.find("origin") + if origin_elem is not None: + origin_xyz = _parse_triple(origin_elem.get("xyz", "0 0 0")) + origin_rpy = _parse_triple(origin_elem.get("rpy", "0 0 0")) + joints.append( JointDescription( name=name, @@ -149,6 +159,8 @@ def _parse_urdf_string(xml_string: str) -> ModelDescription: effort_limit=effort, parent_link=parent_link, child_link=child_link, + origin_xyz=origin_xyz, + origin_rpy=origin_rpy, ) ) @@ -239,4 +251,11 @@ def _float_or_none(value: str | None) -> float | None: return None +def _parse_triple(value: str) -> tuple[float, float, float]: + parts = value.split() + if len(parts) != 3: + raise ValueError(f"Expected 3 floats, got {value!r}") + return (float(parts[0]), float(parts[1]), float(parts[2])) + + __all__ = ["JointDescription", "ModelDescription", "parse_model"] diff --git a/dimos/robot/unitree/go2/config.py b/dimos/robot/unitree/go2/config.py new file mode 100644 index 0000000000..81d095a40c --- /dev/null +++ b/dimos/robot/unitree/go2/config.py @@ -0,0 +1,22 @@ +# 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 pathlib import Path + +from dimos.robot.config import RobotConfig + +Go2Config = RobotConfig( + name="unitree_go2", + model_path=Path(__file__).parent / "go2.urdf", +) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index e887a1b0bb..a7316be21c 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -26,8 +26,6 @@ from dimos.agents.annotation import skill from dimos.constants import ( - DEFAULT_CAMERA_FRAME, - DEFAULT_CAMERA_OPTICAL_FRAME, DEFAULT_CAPACITY_COLOR_IMAGE, DEFAULT_ROBOT_FRAME, DEFAULT_THREAD_JOIN_TIMEOUT, @@ -47,14 +45,13 @@ from dimos.memory2.replay import Replay, resolve_db_path from dimos.memory2.store.sqlite import SqliteStore from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped -from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 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.robot.unitree.connection import UnitreeWebRTCConnection +from dimos.robot.unitree.go2.config import Go2Config from dimos.utils.decorators.decorators import cached_property, simple_mcache if sys.version_info < (3, 13): @@ -74,9 +71,10 @@ class ConnectionConfig(ModuleConfig): ip: str = Field(default_factory=lambda m: m["g"].robot_ip) mode: Go2Mode = Go2Mode.DEFAULT frame_id: str | None = DEFAULT_ROBOT_FRAME - camera_link_frame: str = DEFAULT_CAMERA_FRAME - camera_optical_frame: str = DEFAULT_CAMERA_OPTICAL_FRAME static_publish_rate: float = 1.0 + static_transforms: dict[str, Transform] = Field( + default_factory=lambda: dict(Go2Config.static_relative_transforms) + ) class Go2ConnectionProtocol(Protocol): @@ -205,21 +203,6 @@ class GO2Connection(Module, Camera, Pointcloud): connection: Go2ConnectionProtocol camera_info_static: CameraInfo = _camera_info_static() - static_transforms = dict( - # key=default child transform name - camera_link=Transform( - translation=Vector3(0.3, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id=DEFAULT_ROBOT_FRAME, - child_frame_id=DEFAULT_CAMERA_FRAME, - ), - camera_optical=Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id=DEFAULT_CAMERA_FRAME, - child_frame_id=DEFAULT_CAMERA_OPTICAL_FRAME, - ), - ) _static_publish_thread: Thread | None = None _latest_video_frame: Image | None = None @@ -289,24 +272,37 @@ def _publish_tf(self, msg: PoseStamped) -> None: self.odom.publish(msg) def _static_publish(self) -> None: - frame_remap = { - DEFAULT_ROBOT_FRAME: self.frame_id, - DEFAULT_CAMERA_FRAME: self.config.camera_link_frame, - DEFAULT_CAMERA_OPTICAL_FRAME: self.config.camera_optical_frame, + config_frame_remap = { + Go2Config.body_frame: self.frame_id, } + # remape the frame_id + stamped_statics = [ + Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=config_frame_remap.get(transform.frame_id, transform.frame_id), + child_frame_id=config_frame_remap.get( + transform.child_frame_id, transform.child_frame_id + ), + ) + for transform in self.config.static_transforms.values() + ] + if self.config.static_publish_rate > 0: period = 1.0 / self.config.static_publish_rate while True: - stamped_statics = [ + now = time.time() + self.tf.publish( + # reconstructed here to refresh the timestamp Transform( - translation=t.translation, - rotation=t.rotation, - frame_id=frame_remap.get(t.frame_id, t.frame_id), - child_frame_id=frame_remap.get(t.child_frame_id, t.child_frame_id), + translation=transform.translation, + rotation=transform.rotation, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ts=now, ) - for t in self.static_transforms.values() - ] - self.tf.publish(*stamped_statics) + for transform in stamped_statics + ) self.camera_info.publish(self.camera_info_static) time.sleep(period) diff --git a/dimos/robot/unitree/go2/go2.urdf b/dimos/robot/unitree/go2/go2.urdf index 4e67e9ca8e..fc83858f7f 100644 --- a/dimos/robot/unitree/go2/go2.urdf +++ b/dimos/robot/unitree/go2/go2.urdf @@ -19,4 +19,20 @@ + + + + + + + + + + + + + + + + diff --git a/docs/capabilities/memory/plot.md b/docs/capabilities/memory/plot.md index 5eb3155901..ac2189aea8 100644 --- a/docs/capabilities/memory/plot.md +++ b/docs/capabilities/memory/plot.md @@ -361,6 +361,7 @@ from dimos.perception.detection.type.detection3d.imageDetections3DPC import ( ImageDetections3DPC, ) from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.robot.unitree.go2.config import Go2Config from dimos.memory2.vis.space.elements import Box3D from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.Transform import Transform @@ -369,7 +370,7 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 # TODO We need a nicer way to get optical transform for image streams # depending on the source def world_to_optical(base_pose): - return -(Transform.from_pose("base_link", base_pose) + GO2Connection.static_transforms["camera_link"] + GO2Connection.static_transforms["camera_optical"]) + return -(Transform.from_pose("base_link", base_pose) + Go2Config.static_absolute_transforms["base_link/camera_link/camera_optical"]) drawing = Space() From 6423f7fa292ae15068e73e46655b929178580fb6 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 23 May 2026 16:37:29 -0700 Subject: [PATCH 05/71] improve --- dimos/memory2/vis/space/rerun.py | 7 ++-- dimos/robot/config.py | 55 ++++++--------------------- dimos/robot/unitree/go2/config.py | 12 ++++++ dimos/robot/unitree/go2/connection.py | 49 +++++++++--------------- dimos/utils/testing/test_moment.py | 13 ++++++- docs/capabilities/memory/plot.md | 4 +- 6 files changed, 61 insertions(+), 79 deletions(-) diff --git a/dimos/memory2/vis/space/rerun.py b/dimos/memory2/vis/space/rerun.py index d193897a11..f00097ea32 100644 --- a/dimos/memory2/vis/space/rerun.py +++ b/dimos/memory2/vis/space/rerun.py @@ -226,9 +226,10 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: img = _as_image(data) if img is not None: # TODO: fix there should be a standard way to set camera frames - base_to_optical = Go2Config.static_absolute_transforms[ - "base_link/camera_link/camera_optical" - ] + base_to_optical = ( + Go2Config.static_transforms["camera_link"] + + Go2Config.static_transforms["camera_optical"] + ) world_T_optical = Transform.from_pose("world", obs.pose_stamped) + base_to_optical rr.log(path, world_T_optical.to_pose().to_rerun(), static=True) h, w = img.shape[:2] diff --git a/dimos/robot/config.py b/dimos/robot/config.py index f5a6db1d7a..4213c37147 100644 --- a/dimos/robot/config.py +++ b/dimos/robot/config.py @@ -191,6 +191,10 @@ def joint_name_mapping(self) -> dict[str, str]: def coordinator_task_name(self) -> str: return f"traj_{self.name}" + @cached_property + def all_frame_ids(self) -> list[str]: + return list(self._ensure_parsed().links) + @cached_property def body_frame(self) -> str: """Returns the robot's root link from the urdf (usually base_link), skips past ``type="floating"`` and returns structural link ("world" can be the true root frame, but that is detached from the robot)""" @@ -209,38 +213,24 @@ def body_frame(self) -> str: current = floating_joint.child_link @cached_property - def static_relative_transforms(self) -> dict[str, Transform]: - """The key is the absolute path to the joint, with the value being the relative tranform (one link, not all links) + def static_transforms(self) -> dict[str, Transform]: + """Key is the child frame, transform is parent to child Example:: - print(RobotConfig(model_path="go2.urdf").static_relative_transforms) + print(RobotConfig(model_path="go2.urdf").static_transforms) # output: { - "base_link/camera_link": Transform(translation=(0.3, 0, 0), rotation=identity, frame_id="base_link", child_frame_id="camera_link"), - # NOTE: this is the transform from camera_link to camera_optical NOT baselink to camera_optical - "base_link/camera_link/camera_optical": Transform(translation=(0, 0, 0), rotation=(-0.5, 0.5, -0.5, 0.5), frame_id="camera_link", child_frame_id="camera_optical"), + "camera_link": Transform(translation=(0.3, 0, 0), rotation=identity, frame_id="base_link", child_frame_id="camera_link"), + "camera_optical": Transform(translation=(0, 0, 0), rotation=(-0.5, 0.5, -0.5, 0.5), frame_id="camera_link", child_frame_id="camera_optical"), } """ - model = self._ensure_parsed() - by_child = {joint.child_link: joint for joint in model.joints if joint.child_link} - - def path_to(link: str) -> str: - parts = [link] - current = link - while current in by_child: - current = by_child[current].parent_link - if not current: - break - parts.append(current) - return "/".join(reversed(parts)) - result: dict[str, Transform] = {} - for joint in model.joints: + for joint in self._ensure_parsed().joints: if joint.type != "fixed" or not joint.child_link: continue roll, pitch, yaw = joint.origin_rpy tx, ty, tz = joint.origin_xyz - result[path_to(joint.child_link)] = Transform( + result[joint.child_link] = Transform( translation=Vector3(tx, ty, tz), rotation=Quaternion.from_euler(Vector3(roll, pitch, yaw)), frame_id=joint.parent_link, @@ -248,29 +238,6 @@ def path_to(link: str) -> str: ) return result - @cached_property - def static_absolute_transforms(self) -> dict[str, Transform]: - """Same keys as ``static_relative_transforms`` but each value is the cumulative - transform from the root link down to the leaf link (i.e. all hops composed). - - Example:: - print(RobotConfig(model_path="go2.urdf").static_absolute_transforms) - # output: - { - "base_link/camera_link": Transform(..., frame_id="base_link", child_frame_id="camera_link"), - "base_link/camera_link/camera_optical": Transform(..., frame_id="base_link", child_frame_id="camera_optical"), - } - """ - relative_transforms = self.static_relative_transforms - # Process parents before children so each entry can build on the prior absolute. - result: dict[str, Transform] = {} - for path in sorted(relative_transforms.keys(), key=lambda path: path.count("/")): - parent_path = path.rsplit("/", 1)[0] - parent_absolute = result.get(parent_path) - relative = relative_transforms[path] - result[path] = parent_absolute + relative if parent_absolute is not None else relative - return result - # -- Converter methods ---------------------------------------------------- def to_robot_model_config(self) -> RobotModelConfig: diff --git a/dimos/robot/unitree/go2/config.py b/dimos/robot/unitree/go2/config.py index 81d095a40c..59e2e9ea3d 100644 --- a/dimos/robot/unitree/go2/config.py +++ b/dimos/robot/unitree/go2/config.py @@ -12,10 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. +from importlib import resources from pathlib import Path +from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.robot.config import RobotConfig +_FRONT_CAMERA_720_YAML = resources.files("dimos.robot.unitree.go2").joinpath( + "front_camera_720.yaml" +) + + +def camera_info_static() -> CameraInfo: + with resources.as_file(_FRONT_CAMERA_720_YAML) as yaml_path: + return CameraInfo.from_yaml(str(yaml_path)) + + Go2Config = RobotConfig( name="unitree_go2", model_path=Path(__file__).parent / "go2.urdf", diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index a7316be21c..528f6563ce 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -13,7 +13,6 @@ # limitations under the License. from enum import Enum -from importlib import resources import sys from threading import Thread import time @@ -51,7 +50,7 @@ from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.robot.unitree.connection import UnitreeWebRTCConnection -from dimos.robot.unitree.go2.config import Go2Config +from dimos.robot.unitree.go2.config import Go2Config, camera_info_static from dimos.utils.decorators.decorators import cached_property, simple_mcache if sys.version_info < (3, 13): @@ -73,7 +72,10 @@ class ConnectionConfig(ModuleConfig): frame_id: str | None = DEFAULT_ROBOT_FRAME static_publish_rate: float = 1.0 static_transforms: dict[str, Transform] = Field( - default_factory=lambda: dict(Go2Config.static_relative_transforms) + default_factory=lambda: dict(Go2Config.static_transforms) + ) + frame_mapping: dict[str, str] = Field( + default_factory=lambda: {each: each for each in Go2Config.all_frame_ids} ) @@ -94,16 +96,6 @@ def enable_rage_mode(self) -> bool: ... def publish_request(self, topic: str, data: dict) -> dict: ... # type: ignore[type-arg] -_FRONT_CAMERA_720_YAML = resources.files("dimos.robot.unitree.go2").joinpath( - "front_camera_720.yaml" -) - - -def _camera_info_static() -> CameraInfo: - with resources.as_file(_FRONT_CAMERA_720_YAML) as yaml_path: - return CameraInfo.from_yaml(str(yaml_path)) - - def make_connection(ip: str | None, cfg: GlobalConfig) -> Go2ConnectionProtocol: connection_type = cfg.unitree_connection_type @@ -202,7 +194,7 @@ class GO2Connection(Module, Camera, Pointcloud): camera_info: Out[CameraInfo] connection: Go2ConnectionProtocol - camera_info_static: CameraInfo = _camera_info_static() + camera_info_static: CameraInfo = camera_info_static() _static_publish_thread: Thread | None = None _latest_video_frame: Image | None = None @@ -272,18 +264,13 @@ def _publish_tf(self, msg: PoseStamped) -> None: self.odom.publish(msg) def _static_publish(self) -> None: - config_frame_remap = { - Go2Config.body_frame: self.frame_id, - } - # remape the frame_id + remap = {Go2Config.body_frame: self.frame_id, **self.config.frame_mapping} stamped_statics = [ Transform( translation=transform.translation, rotation=transform.rotation, - frame_id=config_frame_remap.get(transform.frame_id, transform.frame_id), - child_frame_id=config_frame_remap.get( - transform.child_frame_id, transform.child_frame_id - ), + frame_id=remap.get(transform.frame_id, transform.frame_id), + child_frame_id=remap.get(transform.child_frame_id, transform.child_frame_id), ) for transform in self.config.static_transforms.values() ] @@ -292,16 +279,18 @@ def _static_publish(self) -> None: period = 1.0 / self.config.static_publish_rate while True: now = time.time() + # reconstructed each iteration to refresh the timestamp self.tf.publish( - # reconstructed here to refresh the timestamp - Transform( - translation=transform.translation, - rotation=transform.rotation, - frame_id=transform.frame_id, - child_frame_id=transform.child_frame_id, - ts=now, + *( + Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ts=now, + ) + for transform in stamped_statics ) - for transform in stamped_statics ) self.camera_info.publish(self.camera_info_static) time.sleep(period) diff --git a/dimos/utils/testing/test_moment.py b/dimos/utils/testing/test_moment.py index 7518a1af02..a60d34c085 100644 --- a/dimos/utils/testing/test_moment.py +++ b/dimos/utils/testing/test_moment.py @@ -23,6 +23,7 @@ from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.protocol.tf.tf import TF from dimos.robot.unitree.go2 import connection +from dimos.robot.unitree.go2.config import Go2Config from dimos.utils.data import get_data from dimos.utils.testing.moment import Moment, SensorMoment @@ -51,7 +52,17 @@ def transforms(self) -> list[Transform]: # back and forth through time and the viewer doesn't get confused odom = self.odom.value odom.ts = time.time() - return connection.GO2Connection._odom_to_tf(odom) + statics = [ + Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ts=odom.ts, + ) + for transform in Go2Config.static_transforms.values() + ] + return [Transform.from_pose(Go2Config.body_frame, odom), *statics] def publish(self) -> None: t = TF() diff --git a/docs/capabilities/memory/plot.md b/docs/capabilities/memory/plot.md index ac2189aea8..0c22d19774 100644 --- a/docs/capabilities/memory/plot.md +++ b/docs/capabilities/memory/plot.md @@ -370,7 +370,9 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 # TODO We need a nicer way to get optical transform for image streams # depending on the source def world_to_optical(base_pose): - return -(Transform.from_pose("base_link", base_pose) + Go2Config.static_absolute_transforms["base_link/camera_link/camera_optical"]) + statics = Go2Config.static_transforms + base_to_optical = statics["camera_link"] + statics["camera_optical"] + return -(Transform.from_pose("base_link", base_pose) + base_to_optical) drawing = Space() From 4cfdc837148cdd92643f57e5a280749ce178bc63 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 23 May 2026 17:06:11 -0700 Subject: [PATCH 06/71] camera_info_static --- dimos/utils/testing/test_moment.py | 5 ++--- docs/capabilities/memory/plot.md | 5 ++--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/dimos/utils/testing/test_moment.py b/dimos/utils/testing/test_moment.py index a60d34c085..2028639e28 100644 --- a/dimos/utils/testing/test_moment.py +++ b/dimos/utils/testing/test_moment.py @@ -22,8 +22,7 @@ from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.protocol.tf.tf import TF -from dimos.robot.unitree.go2 import connection -from dimos.robot.unitree.go2.config import Go2Config +from dimos.robot.unitree.go2.config import Go2Config, camera_info_static from dimos.utils.data import get_data from dimos.utils.testing.moment import Moment, SensorMoment @@ -69,7 +68,7 @@ def publish(self) -> None: t.publish(*self.transforms) t.stop() - camera_info = connection._camera_info_static() + camera_info = camera_info_static() camera_info.ts = time.time() camera_info_transport: LCMTransport[CameraInfo] = LCMTransport("/camera_info", CameraInfo) camera_info_transport.publish(camera_info) diff --git a/docs/capabilities/memory/plot.md b/docs/capabilities/memory/plot.md index 0c22d19774..44a6627829 100644 --- a/docs/capabilities/memory/plot.md +++ b/docs/capabilities/memory/plot.md @@ -360,8 +360,7 @@ m.data.save("assets/plants_peak_detections.png") from dimos.perception.detection.type.detection3d.imageDetections3DPC import ( ImageDetections3DPC, ) -from dimos.robot.unitree.go2.connection import GO2Connection -from dimos.robot.unitree.go2.config import Go2Config +from dimos.robot.unitree.go2.config import Go2Config, camera_info_static from dimos.memory2.vis.space.elements import Box3D from dimos.msgs.geometry_msgs.Pose import Pose from dimos.msgs.geometry_msgs.Transform import Transform @@ -380,7 +379,7 @@ drawing.add(global_map) drawing.add(detections) -camera_info = GO2Connection.camera_info_static +camera_info = camera_info_static() detections3d = (detections .map_data(lambda obs: ImageDetections3DPC.from_2d( From ec78a8697505dc7d24209454dfabbf688fcf403a Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sat, 23 May 2026 17:23:17 -0700 Subject: [PATCH 07/71] - --- docs/capabilities/memory/plot.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/docs/capabilities/memory/plot.md b/docs/capabilities/memory/plot.md index 44a6627829..b87ce4a238 100644 --- a/docs/capabilities/memory/plot.md +++ b/docs/capabilities/memory/plot.md @@ -369,9 +369,7 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 # TODO We need a nicer way to get optical transform for image streams # depending on the source def world_to_optical(base_pose): - statics = Go2Config.static_transforms - base_to_optical = statics["camera_link"] + statics["camera_optical"] - return -(Transform.from_pose("base_link", base_pose) + base_to_optical) + return -(Transform.from_pose("base_link", base_pose) + Go2Config.static_transforms["camera_link"] + Go2Config.static_transforms["camera_optical"]) drawing = Space() From 96fa258f78ec09d49eae8b670c2caf03ab1540f0 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sun, 24 May 2026 06:20:25 -0700 Subject: [PATCH 08/71] go2: set odom tf parent frame from config (default "world") --- dimos/robot/unitree/go2/connection.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 528f6563ce..13a2160b35 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -70,6 +70,7 @@ class ConnectionConfig(ModuleConfig): ip: str = Field(default_factory=lambda m: m["g"].robot_ip) mode: Go2Mode = Go2Mode.DEFAULT frame_id: str | None = DEFAULT_ROBOT_FRAME + parent_frame_id: str = "world" static_publish_rate: float = 1.0 static_transforms: dict[str, Transform] = Field( default_factory=lambda: dict(Go2Config.static_transforms) @@ -259,7 +260,15 @@ def stop(self) -> None: super().stop() def _publish_tf(self, msg: PoseStamped) -> None: - self.tf.publish(Transform.from_pose(self.frame_id, msg)) + self.tf.publish( + Transform( + translation=msg.position, + rotation=msg.orientation, + frame_id=self.config.parent_frame_id, + child_frame_id=self.frame_id, + ts=msg.ts, + ) + ) if self.odom.transport: self.odom.publish(msg) From adeec3c544b211964aa1f47b0c08044f563127d3 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sun, 24 May 2026 10:09:55 -0700 Subject: [PATCH 09/71] perception: fix conftest after go2 connection refactor Use Go2Config.static_transforms + local _odom_to_tf helper, and the public camera_info_static() (was the removed _camera_info_static). --- dimos/perception/detection/conftest.py | 38 ++++++++++++++++++++++---- 1 file changed, 33 insertions(+), 5 deletions(-) diff --git a/dimos/perception/detection/conftest.py b/dimos/perception/detection/conftest.py index 5b991a1806..9a60ac81db 100644 --- a/dimos/perception/detection/conftest.py +++ b/dimos/perception/detection/conftest.py @@ -35,11 +35,39 @@ from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC from dimos.protocol.tf.tf import TF -from dimos.robot.unitree.go2 import connection +from dimos.robot.unitree.go2.config import Go2Config, camera_info_static from dimos.robot.unitree.type.odometry import Odometry from dimos.utils.data import get_data +def _odom_to_tf(odom: Odometry) -> list[Transform]: + """Build [world→base_link, base_link→camera_link, camera_link→camera_optical] for tests. + + Replaces the removed GO2Connection._odom_to_tf using Go2Config.static_transforms. + """ + base_link = Transform( + translation=odom.position, + rotation=odom.orientation, + frame_id=odom.frame_id or "world", + child_frame_id="base_link", + ts=odom.ts, + ) + statics = [ + Transform( + translation=t.translation, + rotation=t.rotation, + frame_id=t.frame_id, + child_frame_id=t.child_frame_id, + ts=odom.ts, + ) + for t in ( + Go2Config.static_transforms["camera_link"], + Go2Config.static_transforms["camera_optical"], + ) + ] + return [base_link, *statics] + + class Moment(TypedDict, total=False): odom_frame: Odometry lidar_frame: PointCloud2 @@ -97,7 +125,7 @@ def moment_provider(**kwargs) -> Moment: if odom_frame is None: raise ValueError("No odom frame found") - transforms = connection.GO2Connection._odom_to_tf(odom_frame) + transforms = _odom_to_tf(odom_frame) tf.receive_transform(*transforms) @@ -105,7 +133,7 @@ def moment_provider(**kwargs) -> Moment: "odom_frame": odom_frame, "lidar_frame": lidar_frame, "image_frame": image_frame, - "camera_info": connection._camera_info_static(), + "camera_info": camera_info_static(), "transforms": transforms, "tf": tf, } @@ -246,8 +274,8 @@ def object_db_module(get_moment): c = mock.create_autospec(CameraInfo, spec_set=True, instance=True) module2d = Detection2DModule(detector=lambda: Yolo2DDetector(device="cpu"), camera_info=c) - module3d = Detection3DModule(camera_info=connection._camera_info_static()) - moduleDB = ObjectDBModule(camera_info=connection._camera_info_static()) + module3d = Detection3DModule(camera_info=camera_info_static()) + moduleDB = ObjectDBModule(camera_info=camera_info_static()) # Process 5 frames to build up object history for i in range(5): From a7872417b4de218855c878a708e4a2c0265b3ca9 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 25 May 2026 09:42:45 -0700 Subject: [PATCH 10/71] restore --- dimos/constants.py | 1 + dimos/robot/unitree/go2/connection.py | 28 ++++++++++++++++++++++++++- dimos/utils/testing/test_moment.py | 15 +++----------- 3 files changed, 31 insertions(+), 13 deletions(-) diff --git a/dimos/constants.py b/dimos/constants.py index ba5e5a7c91..2f71d10c14 100644 --- a/dimos/constants.py +++ b/dimos/constants.py @@ -54,4 +54,5 @@ DEFAULT_BUILD_NATIVE = False +DEFAULT_WORLD_FRAME = "world" DEFAULT_ROBOT_FRAME = "base_link" diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 7c0bfe79ee..3c7390fd46 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -28,6 +28,7 @@ DEFAULT_CAPACITY_COLOR_IMAGE, DEFAULT_ROBOT_FRAME, DEFAULT_THREAD_JOIN_TIMEOUT, + DEFAULT_WORLD_FRAME, ) from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.core import rpc @@ -70,7 +71,7 @@ class ConnectionConfig(ModuleConfig): ip: str = Field(default_factory=lambda m: m["g"].robot_ip) mode: Go2Mode = Go2Mode.DEFAULT frame_id: str | None = DEFAULT_ROBOT_FRAME - parent_frame_id: str = "world" + parent_frame_id: str = DEFAULT_WORLD_FRAME static_publish_rate: float = 1.0 static_transforms: dict[str, Transform] = Field( default_factory=lambda: dict(Go2Config.static_transforms) @@ -259,6 +260,31 @@ def stop(self) -> None: super().stop() + @classmethod + def _odom_to_tf(cls: Odometry) -> list[Transform]: + """[world→base_link, base_link→camera_link, camera_link→camera_optical] for tests.""" + base_link = Transform( + translation=cls.position, + rotation=cls.orientation, + frame_id=cls.frame_id or DEFAULT_WORLD_FRAME, + child_frame_id=DEFAULT_ROBOT_FRAME, + ts=cls.ts, + ) + statics = [ + Transform( + translation=t.translation, + rotation=t.rotation, + frame_id=t.frame_id, + child_frame_id=t.child_frame_id, + ts=cls.ts, + ) + for t in ( + Go2Config.static_transforms["camera_link"], + Go2Config.static_transforms["camera_optical"], + ) + ] + return [base_link, *statics] + def _publish_tf(self, msg: PoseStamped) -> None: self.tf.publish( Transform( diff --git a/dimos/utils/testing/test_moment.py b/dimos/utils/testing/test_moment.py index 2028639e28..6b83941610 100644 --- a/dimos/utils/testing/test_moment.py +++ b/dimos/utils/testing/test_moment.py @@ -22,7 +22,8 @@ from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.protocol.tf.tf import TF -from dimos.robot.unitree.go2.config import Go2Config, camera_info_static +from dimos.robot.unitree.go2.config import camera_info_static +from dimos.robot.unitree.go2.connection import connection from dimos.utils.data import get_data from dimos.utils.testing.moment import Moment, SensorMoment @@ -51,17 +52,7 @@ def transforms(self) -> list[Transform]: # back and forth through time and the viewer doesn't get confused odom = self.odom.value odom.ts = time.time() - statics = [ - Transform( - translation=transform.translation, - rotation=transform.rotation, - frame_id=transform.frame_id, - child_frame_id=transform.child_frame_id, - ts=odom.ts, - ) - for transform in Go2Config.static_transforms.values() - ] - return [Transform.from_pose(Go2Config.body_frame, odom), *statics] + return connection.GO2Connection._odom_to_tf(odom) def publish(self) -> None: t = TF() From f273a859fad587eebef2110a2794a9409ed91678 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 25 May 2026 09:43:20 -0700 Subject: [PATCH 11/71] - --- dimos/robot/unitree/go2/connection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 3c7390fd46..0d3501fcd5 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -70,7 +70,7 @@ class Go2Mode(str, Enum): class ConnectionConfig(ModuleConfig): ip: str = Field(default_factory=lambda m: m["g"].robot_ip) mode: Go2Mode = Go2Mode.DEFAULT - frame_id: str | None = DEFAULT_ROBOT_FRAME + frame_id: str = DEFAULT_ROBOT_FRAME parent_frame_id: str = DEFAULT_WORLD_FRAME static_publish_rate: float = 1.0 static_transforms: dict[str, Transform] = Field( From 39509c1d464dd68f165547b6f25d92ea51942337 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 25 May 2026 10:00:12 -0700 Subject: [PATCH 12/71] clean up _odom_to_tf --- dimos/perception/detection/conftest.py | 32 ++--------------- dimos/robot/unitree/dimsim_connection.py | 43 ++-------------------- dimos/robot/unitree/go2/config.py | 46 ++++++++++++++++++++++++ dimos/robot/unitree/go2/connection.py | 25 ------------- dimos/robot/unitree/go2/go2.urdf | 6 ++++ 5 files changed, 56 insertions(+), 96 deletions(-) diff --git a/dimos/perception/detection/conftest.py b/dimos/perception/detection/conftest.py index 9a60ac81db..3cfee54fdf 100644 --- a/dimos/perception/detection/conftest.py +++ b/dimos/perception/detection/conftest.py @@ -35,39 +35,11 @@ from dimos.perception.detection.type.detection3d.imageDetections3DPC import ImageDetections3DPC from dimos.perception.detection.type.detection3d.pointcloud import Detection3DPC from dimos.protocol.tf.tf import TF -from dimos.robot.unitree.go2.config import Go2Config, camera_info_static +from dimos.robot.unitree.go2.config import camera_info_static, odom_to_tf from dimos.robot.unitree.type.odometry import Odometry from dimos.utils.data import get_data -def _odom_to_tf(odom: Odometry) -> list[Transform]: - """Build [world→base_link, base_link→camera_link, camera_link→camera_optical] for tests. - - Replaces the removed GO2Connection._odom_to_tf using Go2Config.static_transforms. - """ - base_link = Transform( - translation=odom.position, - rotation=odom.orientation, - frame_id=odom.frame_id or "world", - child_frame_id="base_link", - ts=odom.ts, - ) - statics = [ - Transform( - translation=t.translation, - rotation=t.rotation, - frame_id=t.frame_id, - child_frame_id=t.child_frame_id, - ts=odom.ts, - ) - for t in ( - Go2Config.static_transforms["camera_link"], - Go2Config.static_transforms["camera_optical"], - ) - ] - return [base_link, *statics] - - class Moment(TypedDict, total=False): odom_frame: Odometry lidar_frame: PointCloud2 @@ -125,7 +97,7 @@ def moment_provider(**kwargs) -> Moment: if odom_frame is None: raise ValueError("No odom frame found") - transforms = _odom_to_tf(odom_frame) + transforms = odom_to_tf(odom_frame) tf.receive_transform(*transforms) diff --git a/dimos/robot/unitree/dimsim_connection.py b/dimos/robot/unitree/dimsim_connection.py index 5afcd1fda7..1ea921a7c0 100644 --- a/dimos/robot/unitree/dimsim_connection.py +++ b/dimos/robot/unitree/dimsim_connection.py @@ -21,14 +21,12 @@ from dimos.core.global_config import GlobalConfig from dimos.core.transport import LCMTransport 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.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 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.protocol.tf.tf import LCMTF +from dimos.robot.unitree.go2.config import odom_to_tf from dimos.simulation.dimsim.dimsim_process import DimSimProcess from dimos.utils.logging_config import setup_logger @@ -101,41 +99,4 @@ def publish_request(self, topic: str, data: dict[str, Any]) -> dict[Any, Any]: return {} def _handle_odom(self, msg: PoseStamped) -> None: - self._tf.publish(*_odom_to_tf(msg)) - - -def _odom_to_tf(odom: PoseStamped) -> list[Transform]: - """Build transform chain from odometry pose. - - Transform tree: world -> base_link -> {camera_link -> camera_optical, lidar_link} - """ - camera_link = Transform( - translation=Vector3(0.3, 0.0, 0.0), # camera 30cm forward - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", - child_frame_id="camera_link", - ts=odom.ts, - ) - - camera_optical = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(-0.5, 0.5, -0.5, 0.5), - frame_id="camera_link", - child_frame_id="camera_optical", - ts=odom.ts, - ) - - lidar_link = Transform( - translation=Vector3(0.0, 0.0, 0.0), - rotation=Quaternion(0.0, 0.0, 0.0, 1.0), - frame_id="base_link", - child_frame_id="lidar_link", - ts=odom.ts, - ) - - return [ - Transform.from_pose("base_link", odom), - camera_link, - camera_optical, - lidar_link, - ] + self._tf.publish(*odom_to_tf(msg)) diff --git a/dimos/robot/unitree/go2/config.py b/dimos/robot/unitree/go2/config.py index 59e2e9ea3d..fd516eb9f1 100644 --- a/dimos/robot/unitree/go2/config.py +++ b/dimos/robot/unitree/go2/config.py @@ -12,11 +12,31 @@ # See the License for the specific language governing permissions and # limitations under the License. +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. from importlib import resources from pathlib import Path +from dimos.constants import ( + DEFAULT_ROBOT_FRAME, + DEFAULT_WORLD_FRAME, +) +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.robot.config import RobotConfig +from dimos.robot.unitree.go2.config import Go2Config, camera_info_static _FRONT_CAMERA_720_YAML = resources.files("dimos.robot.unitree.go2").joinpath( "front_camera_720.yaml" @@ -28,6 +48,32 @@ def camera_info_static() -> CameraInfo: return CameraInfo.from_yaml(str(yaml_path)) +def odom_to_tf(cls: Odometry) -> list[Transform]: + """[world→base_link, base_link→camera_link, camera_link→camera_optical] for tests.""" + base_link = Transform( + translation=cls.position, + rotation=cls.orientation, + frame_id=cls.frame_id or DEFAULT_WORLD_FRAME, + child_frame_id=DEFAULT_ROBOT_FRAME, + ts=cls.ts, + ) + statics = [ + Transform( + translation=t.translation, + rotation=t.rotation, + frame_id=t.frame_id, + child_frame_id=t.child_frame_id, + ts=cls.ts, + ) + for t in ( + Go2Config.static_transforms["camera_link"], + Go2Config.static_transforms["camera_optical"], + Go2Config.static_transforms["lidar_link"], + ) + ] + return [base_link, *statics] + + Go2Config = RobotConfig( name="unitree_go2", model_path=Path(__file__).parent / "go2.urdf", diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 0d3501fcd5..c73cac69aa 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -260,31 +260,6 @@ def stop(self) -> None: super().stop() - @classmethod - def _odom_to_tf(cls: Odometry) -> list[Transform]: - """[world→base_link, base_link→camera_link, camera_link→camera_optical] for tests.""" - base_link = Transform( - translation=cls.position, - rotation=cls.orientation, - frame_id=cls.frame_id or DEFAULT_WORLD_FRAME, - child_frame_id=DEFAULT_ROBOT_FRAME, - ts=cls.ts, - ) - statics = [ - Transform( - translation=t.translation, - rotation=t.rotation, - frame_id=t.frame_id, - child_frame_id=t.child_frame_id, - ts=cls.ts, - ) - for t in ( - Go2Config.static_transforms["camera_link"], - Go2Config.static_transforms["camera_optical"], - ) - ] - return [base_link, *statics] - def _publish_tf(self, msg: PoseStamped) -> None: self.tf.publish( Transform( diff --git a/dimos/robot/unitree/go2/go2.urdf b/dimos/robot/unitree/go2/go2.urdf index fc83858f7f..91c914389b 100644 --- a/dimos/robot/unitree/go2/go2.urdf +++ b/dimos/robot/unitree/go2/go2.urdf @@ -35,4 +35,10 @@ + + + + + + From cc360d2e727ac733628d37c4f923a0f12edb25d1 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 25 May 2026 10:06:10 -0700 Subject: [PATCH 13/71] fix imports --- dimos/robot/unitree/go2/config.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/dimos/robot/unitree/go2/config.py b/dimos/robot/unitree/go2/config.py index fd516eb9f1..2fd77d5edc 100644 --- a/dimos/robot/unitree/go2/config.py +++ b/dimos/robot/unitree/go2/config.py @@ -36,7 +36,6 @@ from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.robot.config import RobotConfig -from dimos.robot.unitree.go2.config import Go2Config, camera_info_static _FRONT_CAMERA_720_YAML = resources.files("dimos.robot.unitree.go2").joinpath( "front_camera_720.yaml" @@ -48,6 +47,12 @@ def camera_info_static() -> CameraInfo: return CameraInfo.from_yaml(str(yaml_path)) +Go2Config = RobotConfig( + name="unitree_go2", + model_path=Path(__file__).parent / "go2.urdf", +) + + def odom_to_tf(cls: Odometry) -> list[Transform]: """[world→base_link, base_link→camera_link, camera_link→camera_optical] for tests.""" base_link = Transform( @@ -72,9 +77,3 @@ def odom_to_tf(cls: Odometry) -> list[Transform]: ) ] return [base_link, *statics] - - -Go2Config = RobotConfig( - name="unitree_go2", - model_path=Path(__file__).parent / "go2.urdf", -) From e07fbb8ddbf418897d0c04d3b47cc6cf15cd5ee9 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 25 May 2026 10:12:48 -0700 Subject: [PATCH 14/71] - --- dimos/utils/testing/test_moment.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/utils/testing/test_moment.py b/dimos/utils/testing/test_moment.py index 6b83941610..f01cfc1f6c 100644 --- a/dimos/utils/testing/test_moment.py +++ b/dimos/utils/testing/test_moment.py @@ -22,8 +22,8 @@ from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.protocol.tf.tf import TF +from dimos.robot.unitree.go2 import connection from dimos.robot.unitree.go2.config import camera_info_static -from dimos.robot.unitree.go2.connection import connection from dimos.utils.data import get_data from dimos.utils.testing.moment import Moment, SensorMoment From 9218f18daeac4ea1dfe4ccd5d82635bf305cb517 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Mon, 25 May 2026 10:46:37 -0700 Subject: [PATCH 15/71] - --- dimos/utils/testing/test_moment.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/dimos/utils/testing/test_moment.py b/dimos/utils/testing/test_moment.py index f01cfc1f6c..ac6b69ced8 100644 --- a/dimos/utils/testing/test_moment.py +++ b/dimos/utils/testing/test_moment.py @@ -22,8 +22,7 @@ from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.protocol.tf.tf import TF -from dimos.robot.unitree.go2 import connection -from dimos.robot.unitree.go2.config import camera_info_static +from dimos.robot.unitree.go2.config import camera_info_static, odom_to_tf from dimos.utils.data import get_data from dimos.utils.testing.moment import Moment, SensorMoment @@ -52,7 +51,7 @@ def transforms(self) -> list[Transform]: # back and forth through time and the viewer doesn't get confused odom = self.odom.value odom.ts = time.time() - return connection.GO2Connection._odom_to_tf(odom) + return odom_to_tf(odom) def publish(self) -> None: t = TF() From b8dcdab6648f9ad0807cd1a7cb998219413b5c98 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 16:39:35 -0700 Subject: [PATCH 16/71] add static publish system --- dimos/core/module.py | 124 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 124 insertions(+) diff --git a/dimos/core/module.py b/dimos/core/module.py index f2aed9d185..d6da815ee3 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -19,6 +19,7 @@ import json import sys import threading +import time from typing import ( TYPE_CHECKING, Any, @@ -40,6 +41,7 @@ from dimos.core.resource import CompositeResource from dimos.core.rpc_client import RpcCall from dimos.core.stream import In, Out, RemoteOut, Transport +from dimos.msgs.geometry_msgs.Transform import Transform from dimos.protocol.rpc.pubsubrpc import LCMRPC from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, DEFAULT_RPC_TIMEOUTS, RPCSpec from dimos.protocol.service.spec import BaseConfig, Configurable @@ -107,6 +109,11 @@ class ModuleConfig(BaseConfig): tf_transport: type[TFSpec] = LCMTF # type: ignore[type-arg] frame_id_prefix: str | None = None frame_id: str | None = None + static_transforms: dict[str, Transform] = Field(default_factory=dict) + # TODO: in the future we should make self.tf.publish error if it tried to publish a transform that references a frame that is not mentioned in this dict (same with self.tf.get) + # TODO: later expose frame remappings, somehow, at the blueprint level + frame_mapping: dict[str, str] = Field(default_factory=dict) + static_publish_rate: float = 1.0 g: GlobalConfig = global_config @@ -146,7 +153,10 @@ def __init__(self, config_args: dict[str, Any]) -> None: self._module_closed_lock = threading.Lock() self._tools = {} self._tools_lock = threading.Lock() + self._static_publish_thread: threading.Thread | None = None + self._static_publish_stop = threading.Event() self._loop, self._loop_thread = get_loop() + self.frame_mapping, self.static_transforms = self._setup_frames() try: self.rpc = self.config.rpc_transport( # type: ignore[call-arg] rpc_timeouts=self.config.rpc_timeouts, @@ -182,6 +192,7 @@ def build(self) -> None: def start(self) -> None: self._start_main() self._auto_bind_handlers() + self._start_static_publish() @rpc def stop(self) -> None: @@ -189,12 +200,121 @@ def stop(self) -> None: super().stop() self._close_module() + def _setup_frames(self) -> dict[str, str]: + frame_mapping_field = type(self.config).model_fields["frame_mapping"] + if not hasattr(frame_mapping_field, "default_factory") or not callable( + frame_mapping_field.default_factory + ): + raise Exception( + f"""In the {self.name!r} module config definition, the frame_remapping needs to be a pydantic field, not a dict""" + ) + existing_frames = frame_mapping_field.default_factory() + + # given something like: + # class MyConfig: + # frame_mapping: dict[str, str] = Field(default_factory=lambda: dict( + # body="base_link", + # parent="world", + # )) + # the "body" is what I call a common_name with "base_link" being the remapped name (the REAL frame id that other modules can query/use) + + # step1 (for static_transforms only) translate urdf_name=>common_name + reverse_mapping = {value: key for key, value in existing_frames.items()} + static_transforms_common_names = { + reverse_mapping.get(urdf_frame_id, urdf_frame_id): Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=reverse_mapping.get(transform.frame_id, transform.frame_id), + child_frame_id=reverse_mapping.get( + transform.child_frame_id, transform.child_frame_id + ), + ) + for urdf_frame_id, transform in self.config.static_transforms.items() + } + # step2 map common_name=>real_frame_id + valid_frame_ids = set(existing_frames.keys()) | set(static_transforms_common_names.keys()) + final_frame_mapping = { + **existing_frames, + **self.config.frame_mapping, + } + for existing_frame, remapped_frame in final_frame_mapping.items(): + if existing_frame not in valid_frame_ids: + raise Exception( + f"""On module {self.name}, tried to map {existing_frame!r} to {remapped_frame!r} but that first frame doesn't exist. The existing ones are: {list(existing_frames.keys())!r} """ + ) + static_transforms_final = { + final_frame_mapping.get(common_frame_id, common_frame_id): Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=final_frame_mapping.get(transform.frame_id, transform.frame_id), + child_frame_id=final_frame_mapping.get( + transform.child_frame_id, transform.child_frame_id + ), + ) + for common_frame_id, transform in static_transforms_common_names.items() + } + return final_frame_mapping, static_transforms_final + + def _start_static_publish(self) -> None: + if not self.config.static_transforms or self.config.static_publish_rate <= 0: + return + self._static_publish_stop.clear() + self._static_publish_thread = threading.Thread( + target=self._static_publish, + daemon=True, + ) + self._static_publish_thread.start() + + # TODO: we're only using this until self.tf.publish_static is implemented + def _static_publish(self) -> None: + period = 1.0 / self.config.static_publish_rate + while not self._static_publish_stop.wait(period): + now = time.time() + self.tf.publish( + *( + Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ts=now, + ) + for transform in self.static_transforms.values() + ) + ) + tfs = list( + Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ts=now, + ) + for transform in self.static_transforms.values() + ) + for each in tfs: + print(f"""each.frame_id = {each.frame_id}""") + print(f"""each.child_frame_id = {each.child_frame_id}""") + print(f"""tfs = {tfs}""") + self._on_static_publish() + + def _on_static_publish(self) -> None: + """ + This is a callback for modules to publish other data (ex: camera info) in the static loop + This should rarely used, but exists for the few cases where it is needed + """ + def _close_module(self) -> None: with self._module_closed_lock: if self._module_closed: return self._module_closed = True + self._static_publish_stop.set() + if self._static_publish_thread and self._static_publish_thread.is_alive(): + self._static_publish_thread.join(timeout=self._loop_thread_timeout) + self._static_publish_thread = None + self._close_all_tools() self._close_rpc() @@ -249,6 +369,8 @@ def __getstate__(self): # type: ignore[no-untyped-def] state.pop("_main_gen", None) state.pop("_tools", None) state.pop("_tools_lock", None) + state.pop("_static_publish_thread", None) + state.pop("_static_publish_stop", None) return state def __setstate__(self, state) -> None: # type: ignore[no-untyped-def] @@ -263,6 +385,8 @@ def __setstate__(self, state) -> None: # type: ignore[no-untyped-def] self._main_gen = None self._tools = {} self._tools_lock = threading.Lock() + self._static_publish_thread = None + self._static_publish_stop = threading.Event() @property def tf(self): # type: ignore[no-untyped-def] From 1f4b99551ccc8011099cbaf06778dc7bccffd65a Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 17:02:35 -0700 Subject: [PATCH 17/71] fixup go2 frame publishing --- dimos/robot/unitree/go2/connection.py | 79 ++++++++------------------- 1 file changed, 23 insertions(+), 56 deletions(-) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index c73cac69aa..119aa373c5 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -14,7 +14,6 @@ from enum import Enum import sys -from threading import Thread import time from typing import TYPE_CHECKING, Any, Protocol @@ -26,8 +25,6 @@ from dimos.agents.annotation import skill from dimos.constants import ( DEFAULT_CAPACITY_COLOR_IMAGE, - DEFAULT_ROBOT_FRAME, - DEFAULT_THREAD_JOIN_TIMEOUT, DEFAULT_WORLD_FRAME, ) from dimos.core.coordination.module_coordinator import ModuleCoordinator @@ -70,15 +67,17 @@ class Go2Mode(str, Enum): class ConnectionConfig(ModuleConfig): ip: str = Field(default_factory=lambda m: m["g"].robot_ip) mode: Go2Mode = Go2Mode.DEFAULT - frame_id: str = DEFAULT_ROBOT_FRAME - parent_frame_id: str = DEFAULT_WORLD_FRAME - static_publish_rate: float = 1.0 + frame_mapping: dict[str, str] = Field( + default_factory=lambda: dict( + body=Go2Config.body_frame, + parent=DEFAULT_WORLD_FRAME, + camera_link="camera_link", + camera_optical="camera_optical", + ) + ) static_transforms: dict[str, Transform] = Field( default_factory=lambda: dict(Go2Config.static_transforms) ) - frame_mapping: dict[str, str] = Field( - default_factory=lambda: {each: each for each in Go2Config.all_frame_ids} - ) class Go2ConnectionProtocol(Protocol): @@ -197,7 +196,6 @@ class GO2Connection(Module, Camera, Pointcloud): connection: Go2ConnectionProtocol camera_info_static: CameraInfo = camera_info_static() - _static_publish_thread: Thread | None = None _latest_video_frame: Image | None = None @classmethod @@ -224,21 +222,20 @@ def start(self) -> None: return self.connection.start() - def onimage(image: Image) -> None: + def on_image(image: Image) -> None: + image.frame_id = self.frame_mapping["camera_optical"] self.color_image.publish(image) self._latest_video_frame = image - self.register_disposable(self.connection.lidar_stream().subscribe(self.lidar.publish)) + def on_lidar(pointcloud: PointCloud2) -> None: + pointcloud.frame_id = self.frame_mapping["body"] + self.lidar.publish(pointcloud) + + self.register_disposable(self.connection.lidar_stream().subscribe(on_lidar)) self.register_disposable(self.connection.odom_stream().subscribe(self._publish_tf)) - self.register_disposable(self.connection.video_stream().subscribe(onimage)) + self.register_disposable(self.connection.video_stream().subscribe(on_image)) self.register_disposable(Disposable(self.cmd_vel.subscribe(self.move))) - self._static_publish_thread = Thread( - target=self._static_publish, - daemon=True, - ) - self._static_publish_thread.start() - self.standup() time.sleep(3) self.connection.balance_stand() @@ -255,56 +252,26 @@ def stop(self) -> None: if self.connection: self.connection.stop() - if self._static_publish_thread and self._static_publish_thread.is_alive(): - self._static_publish_thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) - super().stop() + def _on_static_publish(self) -> None: + self.camera_info_static.frame_id = self.frame_mapping["camera_optical"] + self.camera_info.publish(self.camera_info_static) + def _publish_tf(self, msg: PoseStamped) -> None: self.tf.publish( Transform( translation=msg.position, rotation=msg.orientation, - frame_id=self.config.parent_frame_id, - child_frame_id=self.frame_id, + frame_id=self.frame_mapping["parent"], + child_frame_id=self.frame_mapping["body"], ts=msg.ts, ) ) if self.odom.transport: + msg.frame_id = self.frame_mapping["parent"] self.odom.publish(msg) - def _static_publish(self) -> None: - remap = {Go2Config.body_frame: self.frame_id, **self.config.frame_mapping} - stamped_statics = [ - Transform( - translation=transform.translation, - rotation=transform.rotation, - frame_id=remap.get(transform.frame_id, transform.frame_id), - child_frame_id=remap.get(transform.child_frame_id, transform.child_frame_id), - ) - for transform in self.config.static_transforms.values() - ] - - if self.config.static_publish_rate > 0: - period = 1.0 / self.config.static_publish_rate - while True: - now = time.time() - # reconstructed each iteration to refresh the timestamp - self.tf.publish( - *( - Transform( - translation=transform.translation, - rotation=transform.rotation, - frame_id=transform.frame_id, - child_frame_id=transform.child_frame_id, - ts=now, - ) - for transform in stamped_statics - ) - ) - self.camera_info.publish(self.camera_info_static) - time.sleep(period) - @rpc def move(self, twist: Twist, duration: float = 0.0) -> bool: """Send movement command to robot.""" From 75ed8a72bc85f09dc453db2dd6a5ad0634f43dde Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 22:13:16 -0700 Subject: [PATCH 18/71] fixup tf.get time tolerance for static transforms --- dimos/core/module.py | 47 +++++++----------- dimos/memory2/module.py | 6 +-- dimos/msgs/geometry_msgs/Transform.py | 2 + dimos/protocol/tf/tf.py | 71 ++++++++++++++++++++++++--- 4 files changed, 87 insertions(+), 39 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index d6da815ee3..6de0fd83c4 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -113,7 +113,7 @@ class ModuleConfig(BaseConfig): # TODO: in the future we should make self.tf.publish error if it tried to publish a transform that references a frame that is not mentioned in this dict (same with self.tf.get) # TODO: later expose frame remappings, somehow, at the blueprint level frame_mapping: dict[str, str] = Field(default_factory=dict) - static_publish_rate: float = 1.0 + static_publish_interval: float = 1.0 g: GlobalConfig = global_config @@ -190,9 +190,11 @@ def build(self) -> None: @rpc def start(self) -> None: + # NOTE: there's basically always going to be some inital race around static transform frames and tf.get's + # putting the statics at the top helps mitigate/reduce that + self._start_static_publish() self._start_main() self._auto_bind_handlers() - self._start_static_publish() @rpc def stop(self) -> None: @@ -200,7 +202,7 @@ def stop(self) -> None: super().stop() self._close_module() - def _setup_frames(self) -> dict[str, str]: + def _setup_frames(self) -> tuple[dict[str, str], dict[str, Transform]]: frame_mapping_field = type(self.config).model_fields["frame_mapping"] if not hasattr(frame_mapping_field, "default_factory") or not callable( frame_mapping_field.default_factory @@ -208,7 +210,7 @@ def _setup_frames(self) -> dict[str, str]: raise Exception( f"""In the {self.name!r} module config definition, the frame_remapping needs to be a pydantic field, not a dict""" ) - existing_frames = frame_mapping_field.default_factory() + existing_frames: dict[str, str] = frame_mapping_field.default_factory() # type: ignore[call-arg] # given something like: # class MyConfig: @@ -256,33 +258,25 @@ def _setup_frames(self) -> dict[str, str]: return final_frame_mapping, static_transforms_final def _start_static_publish(self) -> None: - if not self.config.static_transforms or self.config.static_publish_rate <= 0: + self._static_publish() + if not self.config.static_transforms or self.config.static_publish_interval <= 0: return self._static_publish_stop.clear() self._static_publish_thread = threading.Thread( - target=self._static_publish, + target=self._static_publisher, daemon=True, ) self._static_publish_thread.start() - # TODO: we're only using this until self.tf.publish_static is implemented + # TODO: later this should be replaced with latching streams + def _static_publisher(self) -> None: + while not self._static_publish_stop.wait(self.config.static_publish_interval): + self._static_publish() + def _static_publish(self) -> None: - period = 1.0 / self.config.static_publish_rate - while not self._static_publish_stop.wait(period): - now = time.time() - self.tf.publish( - *( - Transform( - translation=transform.translation, - rotation=transform.rotation, - frame_id=transform.frame_id, - child_frame_id=transform.child_frame_id, - ts=now, - ) - for transform in self.static_transforms.values() - ) - ) - tfs = list( + now = time.time() + self.tf.publish_static( + *( Transform( translation=transform.translation, rotation=transform.rotation, @@ -292,11 +286,8 @@ def _static_publish(self) -> None: ) for transform in self.static_transforms.values() ) - for each in tfs: - print(f"""each.frame_id = {each.frame_id}""") - print(f"""each.child_frame_id = {each.child_frame_id}""") - print(f"""tfs = {tfs}""") - self._on_static_publish() + ) + self._on_static_publish() def _on_static_publish(self) -> None: """ diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index b584553bae..800a9aa41d 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -311,11 +311,7 @@ def on_msg(msg: Any) -> None: if not pose: logger.warning( - "[%s] No tf available for frame '%s' at time %s (msg ts: %s), storing without pose", - name, - frame_id, - ts, - getattr(msg, "ts", None), + f"""[{name}] No tf available for frame {name!r} at time {ts} (msg ts: {getattr(msg, "ts", None)}), storing without pose\n{self.tf.tree_str}""" ) stream.append(msg, ts=ts, pose=pose) diff --git a/dimos/msgs/geometry_msgs/Transform.py b/dimos/msgs/geometry_msgs/Transform.py index 9b08c8dadd..972302333d 100644 --- a/dimos/msgs/geometry_msgs/Transform.py +++ b/dimos/msgs/geometry_msgs/Transform.py @@ -48,6 +48,7 @@ def __init__( # type: ignore[no-untyped-def] frame_id: str = "world", child_frame_id: str = "unset", ts: float = 0.0, + static: bool = False, **kwargs, ) -> None: self.frame_id = frame_id @@ -55,6 +56,7 @@ def __init__( # type: ignore[no-untyped-def] self.ts = ts if ts != 0.0 else time.time() self.translation = translation if translation is not None else Vector3() self.rotation = rotation if rotation is not None else Quaternion() + self.static = static def now(self) -> Transform: """Return a copy of this Transform with the current timestamp.""" diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 05c600d733..4d38fe8e73 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -115,6 +115,7 @@ class MultiTBuffer: def __init__(self, buffer_size: float = 10.0) -> None: self.buffers: dict[tuple[str, str], TBuffer] = {} self.buffer_size = buffer_size + self._statics: dict[tuple[str, str], Transform] = {} self._cv = threading.Condition() def receive_transform(self, *args: Transform) -> None: @@ -124,6 +125,8 @@ def receive_transform(self, *args: Transform) -> None: if key not in self.buffers: self.buffers[key] = TBuffer(self.buffer_size) self.buffers[key].add(transform) + if transform.static: + self._statics[key] = transform self._cv.notify_all() def get_frames(self) -> set[str]: @@ -160,13 +163,18 @@ def get_transform( ) with self._cv: - # Check forward direction key = (parent_frame, child_frame) + reverse_key = (child_frame, parent_frame) + + # Check statics first (no tolerance needed) + if key in self._statics: + return self._statics[key] + if reverse_key in self._statics: + return self._statics[reverse_key].inverse() + + # Check dynamic buffers if key in self.buffers: return self.buffers[key].get(time_point, time_tolerance) # type: ignore[arg-type] - - # Check reverse direction and return inverse - reverse_key = (child_frame, parent_frame) if reverse_key in self.buffers: transform = self.buffers[reverse_key].get(time_point, time_tolerance) # type: ignore[arg-type] return transform.inverse() if transform else None @@ -274,6 +282,43 @@ def get_transform_search( return None + @property + def tree_str(self) -> str: + if not self.buffers: + return "(empty)" + + children: dict[str, list[str]] = {} + all_children: set[str] = set() + for parent, child in self.buffers: + children.setdefault(parent, []).append(child) + all_children.add(child) + + roots = [frame for frame in children if frame not in all_children] + if not roots: + roots = sorted(children.keys())[:1] + + for root in roots: + children.setdefault(root, []) + for frame_list in children.values(): + frame_list.sort() + + lines: list[str] = [] + + def walk(frame: str, prefix: str, is_last: bool, is_root: bool) -> None: + connector = "" if is_root else ("└── " if is_last else "├── ") + lines.append(f"{prefix}{connector}{frame}") + child_prefix = prefix if is_root else prefix + (" " if is_last else "│ ") + kids = children.get(frame, []) + for index, kid in enumerate(kids): + walk(kid, child_prefix, index == len(kids) - 1, False) + + for index, root in enumerate(sorted(roots)): + if index > 0: + lines.append("") + walk(root, "", index == len(roots) - 1, True) + + return "\n".join(lines) + def graph(self) -> str: import subprocess @@ -356,14 +401,28 @@ def publish(self, *args: Transform) -> None: self.pubsub.publish(topic, TFMessage(*args)) def publish_static(self, *args: Transform) -> None: - raise NotImplementedError("Static transforms not implemented in PubSubTF.") + # TODO: note this is a stop-gap, we'd rather this be a latched publish but thats not supported at the moment + static_transforms = tuple( + Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ts=transform.ts, + static=True, + ) + for transform in args + ) + self.receive_transform(*static_transforms) + topic = getattr(self.config, "topic", None) + if topic: + self.pubsub.publish(topic, TFMessage(*static_transforms)) def publish_all(self) -> None: """Publish all transforms currently stored in all buffers.""" all_transforms = [] with self._cv: for buffer in self.buffers.values(): - # Get the latest transform from each buffer latest = buffer.get() # get() with no args returns latest if latest: all_transforms.append(latest) From 5ffcc8c1ce49b150e58cc1bb691d4606b890b76a Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 22:19:59 -0700 Subject: [PATCH 19/71] fixup test --- dimos/robot/cli/test_dimos.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/dimos/robot/cli/test_dimos.py b/dimos/robot/cli/test_dimos.py index b925d4ebd9..a12727ac34 100644 --- a/dimos/robot/cli/test_dimos.py +++ b/dimos/robot/cli/test_dimos.py @@ -69,6 +69,7 @@ class TestModuleB(Module): " * testmodulea.default_rpc_timeout: float (default: 120.0)", " * testmodulea.frame_id_prefix: str | None (default: None)", " * testmodulea.frame_id: str | None (default: None)", + " * testmodulea.static_publish_interval: float (default: 1.0)", " * testmodulea.min_interval_sec: float (default: 0.1)", " * testmodulea.entity_prefix: str (default: world)", " * testmodulea.viewer_mode: typing.Literal['native', 'web', 'connect', 'none'] (default: native)", @@ -76,6 +77,7 @@ class TestModuleB(Module): " * testmoduleb.default_rpc_timeout: float (default: 120.0)", " * testmoduleb.frame_id_prefix: str | None (default: None)", " * testmoduleb.frame_id: str | None (default: None)", + " * testmoduleb.static_publish_interval: float (default: 1.0)", " * testmoduleb.memory_limit: str (default: 25%)", " * testmoduleb.ip: str (default: 127.0.0.1)", "", @@ -110,6 +112,7 @@ class TestModuleB(Module): " * testmodulea.default_rpc_timeout: float (default: 120.0)", " * testmodulea.frame_id_prefix: str | None (default: foo)", " * testmodulea.frame_id: str | None (default: None)", + " * testmodulea.static_publish_interval: float (default: 1.0)", " * testmodulea.min_interval_sec: float (default: 0.1)", " * testmodulea.entity_prefix: str (default: world)", " * testmodulea.viewer_mode: typing.Literal['native', 'web', 'connect', 'none'] (default: web)", @@ -117,6 +120,7 @@ class TestModuleB(Module): " * testmoduleb.default_rpc_timeout: float (default: 120.0)", " * testmoduleb.frame_id_prefix: str | None (default: None)", " * testmoduleb.frame_id: str | None (default: None)", + " * testmoduleb.static_publish_interval: float (default: 1.0)", " * testmoduleb.memory_limit: str (default: 25%)", " * testmoduleb.ip: str (default: 1.1.1.1)", "", @@ -140,6 +144,7 @@ class TestModule(Module): " * testmodule.default_rpc_timeout: float (default: 120.0)", " * testmodule.frame_id_prefix: str | None (default: None)", " * testmodule.frame_id: str | None (default: None)", + " * testmodule.static_publish_interval: float (default: 1.0)", " * [Required] testmodule.foo: int", " * testmodule.spam: str (default: eggs)", "", From 0614ea2a1248c4fb310aefd7c628f1474c921506 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 22:26:38 -0700 Subject: [PATCH 20/71] fix leak --- dimos/core/module.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/dimos/core/module.py b/dimos/core/module.py index 6de0fd83c4..9839829186 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -274,6 +274,8 @@ def _static_publisher(self) -> None: self._static_publish() def _static_publish(self) -> None: + if not self.static_transforms: + return now = time.time() self.tf.publish_static( *( From 9b0ae35a19474f78163d0409d6c156c0df2527c0 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 22:41:00 -0700 Subject: [PATCH 21/71] =?UTF-8?q?fix=20review=20nits:=20rename=20cls?= =?UTF-8?q?=E2=86=92odom=20in=20odom=5Fto=5Ftf,=20make=20camera=5Finfo=5Fs?= =?UTF-8?q?tatic=20instance-level?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- dimos/robot/unitree/go2/config.py | 12 ++++++------ dimos/robot/unitree/go2/connection.py | 7 ++++--- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/dimos/robot/unitree/go2/config.py b/dimos/robot/unitree/go2/config.py index 2fd77d5edc..5823f862e9 100644 --- a/dimos/robot/unitree/go2/config.py +++ b/dimos/robot/unitree/go2/config.py @@ -53,14 +53,14 @@ def camera_info_static() -> CameraInfo: ) -def odom_to_tf(cls: Odometry) -> list[Transform]: +def odom_to_tf(odom: Odometry) -> list[Transform]: """[world→base_link, base_link→camera_link, camera_link→camera_optical] for tests.""" base_link = Transform( - translation=cls.position, - rotation=cls.orientation, - frame_id=cls.frame_id or DEFAULT_WORLD_FRAME, + translation=odom.position, + rotation=odom.orientation, + frame_id=odom.frame_id or DEFAULT_WORLD_FRAME, child_frame_id=DEFAULT_ROBOT_FRAME, - ts=cls.ts, + ts=odom.ts, ) statics = [ Transform( @@ -68,7 +68,7 @@ def odom_to_tf(cls: Odometry) -> list[Transform]: rotation=t.rotation, frame_id=t.frame_id, child_frame_id=t.child_frame_id, - ts=cls.ts, + ts=odom.ts, ) for t in ( Go2Config.static_transforms["camera_link"], diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 119aa373c5..ef96776645 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -211,9 +211,10 @@ def rerun_views(cls): # type: ignore[no-untyped-def] def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self.connection = make_connection(self.config.ip, self.config.g) + self._camera_info_static = camera_info_static() if hasattr(self.connection, "camera_info_static"): - self.camera_info_static = self.connection.camera_info_static + self._camera_info_static = self.connection.camera_info_static @rpc def start(self) -> None: @@ -255,8 +256,8 @@ def stop(self) -> None: super().stop() def _on_static_publish(self) -> None: - self.camera_info_static.frame_id = self.frame_mapping["camera_optical"] - self.camera_info.publish(self.camera_info_static) + self._camera_info_static.frame_id = self.frame_mapping["camera_optical"] + self.camera_info.publish(self._camera_info_static) def _publish_tf(self, msg: PoseStamped) -> None: self.tf.publish( From d28982a4c0e4e3ace1323e200bf2cc243aed256a Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 22:44:04 -0700 Subject: [PATCH 22/71] =?UTF-8?q?fix=20bare=20Exception=E2=86=92ValueError?= =?UTF-8?q?=20in=20=5Fsetup=5Fframes,=20propagate=20static=20through=20Tra?= =?UTF-8?q?nsform=20ops?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- dimos/core/module.py | 4 ++-- dimos/msgs/geometry_msgs/Transform.py | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index 9839829186..bb0b3025dc 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -207,7 +207,7 @@ def _setup_frames(self) -> tuple[dict[str, str], dict[str, Transform]]: if not hasattr(frame_mapping_field, "default_factory") or not callable( frame_mapping_field.default_factory ): - raise Exception( + raise ValueError( f"""In the {self.name!r} module config definition, the frame_remapping needs to be a pydantic field, not a dict""" ) existing_frames: dict[str, str] = frame_mapping_field.default_factory() # type: ignore[call-arg] @@ -241,7 +241,7 @@ def _setup_frames(self) -> tuple[dict[str, str], dict[str, Transform]]: } for existing_frame, remapped_frame in final_frame_mapping.items(): if existing_frame not in valid_frame_ids: - raise Exception( + raise ValueError( f"""On module {self.name}, tried to map {existing_frame!r} to {remapped_frame!r} but that first frame doesn't exist. The existing ones are: {list(existing_frames.keys())!r} """ ) static_transforms_final = { diff --git a/dimos/msgs/geometry_msgs/Transform.py b/dimos/msgs/geometry_msgs/Transform.py index 972302333d..3607891c08 100644 --- a/dimos/msgs/geometry_msgs/Transform.py +++ b/dimos/msgs/geometry_msgs/Transform.py @@ -66,6 +66,7 @@ def now(self) -> Transform: frame_id=self.frame_id, child_frame_id=self.child_frame_id, ts=time.time(), + static=self.static, ) def __repr__(self) -> str: @@ -133,6 +134,7 @@ def __add__(self, other: Transform) -> Transform: frame_id=self.frame_id, child_frame_id=other.child_frame_id, ts=self.ts, + static=self.static or other.static, ) def inverse(self) -> Transform: @@ -157,6 +159,7 @@ def inverse(self) -> Transform: frame_id=self.child_frame_id, # Swap frame references child_frame_id=self.frame_id, ts=self.ts, + static=self.static, ) def __neg__(self) -> Transform: From e9cda440cf4afab6bdbf46cdb6fe6cf6f27e25ba Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 22:51:23 -0700 Subject: [PATCH 23/71] misc --- dimos/core/module.py | 3 ++- dimos/memory2/module.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index bb0b3025dc..c053c44117 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -218,7 +218,8 @@ def _setup_frames(self) -> tuple[dict[str, str], dict[str, Transform]]: # body="base_link", # parent="world", # )) - # the "body" is what I call a common_name with "base_link" being the remapped name (the REAL frame id that other modules can query/use) + # the "body" is what I call a common_name + # "base_link" is the REAL frame id that other modules can query/use # step1 (for static_transforms only) translate urdf_name=>common_name reverse_mapping = {value: key for key, value in existing_frames.items()} diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index 800a9aa41d..a5d30ed622 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -311,7 +311,7 @@ def on_msg(msg: Any) -> None: if not pose: logger.warning( - f"""[{name}] No tf available for frame {name!r} at time {ts} (msg ts: {getattr(msg, "ts", None)}), storing without pose\n{self.tf.tree_str}""" + f"""[{name}] No tf available for frame {frame_id!r} at time {ts} (msg ts: {getattr(msg, "ts", None)}), storing without pose\n{self.tf.tree_str}""" ) stream.append(msg, ts=ts, pose=pose) From 76bb8588b52331b4d75c1a722d1a5479c38192b6 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 23:07:57 -0700 Subject: [PATCH 24/71] switch to /tf_static to avoid the LCM field modification --- dimos/msgs/geometry_msgs/Transform.py | 5 -- dimos/protocol/tf/tf.py | 70 ++++++++++++++++++--------- 2 files changed, 47 insertions(+), 28 deletions(-) diff --git a/dimos/msgs/geometry_msgs/Transform.py b/dimos/msgs/geometry_msgs/Transform.py index 3607891c08..9b08c8dadd 100644 --- a/dimos/msgs/geometry_msgs/Transform.py +++ b/dimos/msgs/geometry_msgs/Transform.py @@ -48,7 +48,6 @@ def __init__( # type: ignore[no-untyped-def] frame_id: str = "world", child_frame_id: str = "unset", ts: float = 0.0, - static: bool = False, **kwargs, ) -> None: self.frame_id = frame_id @@ -56,7 +55,6 @@ def __init__( # type: ignore[no-untyped-def] self.ts = ts if ts != 0.0 else time.time() self.translation = translation if translation is not None else Vector3() self.rotation = rotation if rotation is not None else Quaternion() - self.static = static def now(self) -> Transform: """Return a copy of this Transform with the current timestamp.""" @@ -66,7 +64,6 @@ def now(self) -> Transform: frame_id=self.frame_id, child_frame_id=self.child_frame_id, ts=time.time(), - static=self.static, ) def __repr__(self) -> str: @@ -134,7 +131,6 @@ def __add__(self, other: Transform) -> Transform: frame_id=self.frame_id, child_frame_id=other.child_frame_id, ts=self.ts, - static=self.static or other.static, ) def inverse(self) -> Transform: @@ -159,7 +155,6 @@ def inverse(self) -> Transform: frame_id=self.child_frame_id, # Swap frame references child_frame_id=self.frame_id, ts=self.ts, - static=self.static, ) def __neg__(self) -> Transform: diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index 4d38fe8e73..b7224932eb 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -125,10 +125,32 @@ def receive_transform(self, *args: Transform) -> None: if key not in self.buffers: self.buffers[key] = TBuffer(self.buffer_size) self.buffers[key].add(transform) - if transform.static: - self._statics[key] = transform + # pruning happens in self.buffers[key].add(), but we don't want static ones to get pruned + # would be better to fix pruning somehow but the that class does the pruning + # is more generic than TF-speficic pruning, so I'm going to backfill instead for now + # plz improve when the latching-topic change is made + self._backfill_statics() self._cv.notify_all() + def receive_static_transform(self, *args: Transform) -> None: + with self._cv: + for transform in args: + key = (transform.frame_id, transform.child_frame_id) + if key not in self.buffers: + self.buffers[key] = TBuffer(self.buffer_size) + self.buffers[key].save(transform) + self._statics[key] = transform + self._cv.notify_all() + + def _backfill_statics(self) -> None: + """Re-add static transforms to buffers if pruning removed them.""" + for key, transform in self._statics.items(): + buffer = self.buffers.get(key) + if buffer is None or len(buffer) == 0: + if buffer is None: + self.buffers[key] = TBuffer(self.buffer_size) + self.buffers[key].save(transform) + def get_frames(self) -> set[str]: frames = set() with self._cv: @@ -357,6 +379,7 @@ def __str__(self) -> str: class PubSubTFConfig(TFConfig): topic: Topic | None = None # Required field but needs default for dataclass inheritance + static_topic: Topic | None = None pubsub: type[PubSub] | PubSub | None = None # type: ignore[type-arg] autostart: bool = True @@ -386,6 +409,9 @@ def start(self, sub: bool = True) -> None: topic = getattr(self.config, "topic", None) if topic: self.pubsub.subscribe(topic, self.receive_msg) + static_topic = getattr(self.config, "static_topic", None) + if static_topic: + self.pubsub.subscribe(static_topic, self._receive_static_msg) def stop(self) -> None: self.pubsub.stop() @@ -401,34 +427,28 @@ def publish(self, *args: Transform) -> None: self.pubsub.publish(topic, TFMessage(*args)) def publish_static(self, *args: Transform) -> None: - # TODO: note this is a stop-gap, we'd rather this be a latched publish but thats not supported at the moment - static_transforms = tuple( - Transform( - translation=transform.translation, - rotation=transform.rotation, - frame_id=transform.frame_id, - child_frame_id=transform.child_frame_id, - ts=transform.ts, - static=True, - ) - for transform in args - ) - self.receive_transform(*static_transforms) - topic = getattr(self.config, "topic", None) - if topic: - self.pubsub.publish(topic, TFMessage(*static_transforms)) + self.receive_static_transform(*args) + static_topic = getattr(self.config, "static_topic", None) + if static_topic: + self.pubsub.publish(static_topic, TFMessage(*args)) def publish_all(self) -> None: """Publish all transforms currently stored in all buffers.""" - all_transforms = [] + dynamic = [] + static = [] with self._cv: - for buffer in self.buffers.values(): + for key, buffer in self.buffers.items(): latest = buffer.get() # get() with no args returns latest if latest: - all_transforms.append(latest) + if key in self._statics: + static.append(latest) + else: + dynamic.append(latest) - if all_transforms: - self.publish(*all_transforms) + if dynamic: + self.publish(*dynamic) + if static: + self.publish_static(*static) def get( self, @@ -470,9 +490,13 @@ def get_pose( def receive_msg(self, msg: TFMessage, topic: Topic) -> None: self.receive_tfmessage(msg) + def _receive_static_msg(self, msg: TFMessage, topic: Topic) -> None: + self.receive_static_transform(*msg.transforms) + class LCMPubsubConfig(PubSubTFConfig): topic: Topic = field(default_factory=lambda: Topic("/tf", TFMessage)) + static_topic: Topic = field(default_factory=lambda: Topic("/tf_static", TFMessage)) pubsub: type[PubSub] | PubSub | None = LCM # type: ignore[type-arg] autostart: bool = True From e91840d7aab99790467b7c2e293f9e22d94ce49d Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 23:25:58 -0700 Subject: [PATCH 25/71] misc --- dimos/core/module.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/dimos/core/module.py b/dimos/core/module.py index c053c44117..4ff8f48389 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -112,6 +112,7 @@ class ModuleConfig(BaseConfig): static_transforms: dict[str, Transform] = Field(default_factory=dict) # TODO: in the future we should make self.tf.publish error if it tried to publish a transform that references a frame that is not mentioned in this dict (same with self.tf.get) # TODO: later expose frame remappings, somehow, at the blueprint level + # frame mapping is how others can remap frame names, in the future we should have a way to do this remapping at the blueprint level too frame_mapping: dict[str, str] = Field(default_factory=dict) static_publish_interval: float = 1.0 g: GlobalConfig = global_config @@ -208,7 +209,7 @@ def _setup_frames(self) -> tuple[dict[str, str], dict[str, Transform]]: frame_mapping_field.default_factory ): raise ValueError( - f"""In the {self.name!r} module config definition, the frame_remapping needs to be a pydantic field, not a dict""" + f"""In the {self.name!r} module config definition, the frame_mapping needs to be a pydantic field, not a dict""" ) existing_frames: dict[str, str] = frame_mapping_field.default_factory() # type: ignore[call-arg] @@ -295,7 +296,7 @@ def _static_publish(self) -> None: def _on_static_publish(self) -> None: """ This is a callback for modules to publish other data (ex: camera info) in the static loop - This should rarely used, but exists for the few cases where it is needed + This should be rarely used, but exists for the few cases where it is needed """ def _close_module(self) -> None: From 58bf2ff32e9decc7e6fc7317c703fbff5fbd559c Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 23:26:48 -0700 Subject: [PATCH 26/71] add message throttling --- dimos/memory2/module.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index a5d30ed622..43bd2a084f 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -241,6 +241,7 @@ class RecorderConfig(MemoryModuleConfig): overwrite: bool = True default_frame_id: str = "base_link" tf_tolerance: float = 0.5 + tf_warning_interval: float = 5.0 db_path: str | Path = "recording.db" @@ -302,6 +303,8 @@ def _port_to_stream(self, name: str, input_topic: In[Any], stream: Stream[Any]) default_frame_id = self.config.default_frame_id tf_tolerance = self.config.tf_tolerance + tf_warning_interval = self.config.tf_warning_interval + last_warn_times: dict[str, float] = {} def on_msg(msg: Any) -> None: ts = getattr(msg, "ts", None) or time.time() @@ -310,9 +313,13 @@ def on_msg(msg: Any) -> None: pose = transform.to_pose() if transform is not None else None if not pose: - logger.warning( - f"""[{name}] No tf available for frame {frame_id!r} at time {ts} (msg ts: {getattr(msg, "ts", None)}), storing without pose\n{self.tf.tree_str}""" - ) + now = time.monotonic() + # throtle is per-frame so we don't miss anything important + if now - last_warn_times.get(frame_id, 0.0) > tf_warning_interval: + last_warn_times[frame_id] = now + logger.warning( + f"""[{name}] No tf available for frame {frame_id!r} at time {ts} (msg ts: {getattr(msg, "ts", None)}), storing without pose\n{self.tf.tree_str}""" + ) stream.append(msg, ts=ts, pose=pose) self.register_disposable(Disposable(input_topic.subscribe(on_msg))) From 19c9f939ef4c4a1e0ab5dfde2fe8a598009958d2 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 23:33:40 -0700 Subject: [PATCH 27/71] add tests --- dimos/protocol/tf/test_tf.py | 155 +++++++++++++++++++++++++++++++++++ 1 file changed, 155 insertions(+) diff --git a/dimos/protocol/tf/test_tf.py b/dimos/protocol/tf/test_tf.py index 281e99c0d7..524632c028 100644 --- a/dimos/protocol/tf/test_tf.py +++ b/dimos/protocol/tf/test_tf.py @@ -18,8 +18,10 @@ import threading import time +from pydantic import Field import pytest +from dimos.core.module import Module, 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 @@ -805,3 +807,156 @@ def test_get_with_longer_transform_chain(self) -> None: assert result.frame_id == "world" assert result.child_frame_id == "hand" + + +class TestStaticTransforms: + LCM_PROPAGATION_DELAY = 2 + MODULE_START_DELAY = 2 + STRICT_TOLERANCE = 0.001 + STALE_AGE = 1000 + FAR_FUTURE_OFFSET = 9999 + + def test_static_survives_time_tolerance(self) -> None: + """Static transforms should be retrievable regardless of time tolerance.""" + buffer = MultiTBuffer(buffer_size=10.0) + old_time = time.time() - self.STALE_AGE + + camera_link = Transform( + translation=Vector3(0.3, 0.0, 0.1), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="base_link", + child_frame_id="camera_link", + ts=old_time, + ) + buffer.receive_static_transform(camera_link) + + result = buffer.get_transform( + "base_link", + "camera_link", + time_point=time.time(), + time_tolerance=self.STRICT_TOLERANCE, + ) + assert result is not None + assert abs(result.translation.x - 0.3) < 1e-6 + + def test_static_not_pruned_by_dynamic(self) -> None: + """Adding dynamic transforms to a different pair shouldn't prune statics from buffers.""" + buffer = MultiTBuffer(buffer_size=2.0) + + camera_link = Transform( + translation=Vector3(0.3, 0.0, 0.1), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="base_link", + child_frame_id="camera_link", + ts=time.time() - self.STALE_AGE, + ) + buffer.receive_static_transform(camera_link) + + for i in range(20): + buffer.receive_transform( + Transform( + translation=Vector3(float(i), 0.0, 0.0), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="world", + child_frame_id="base_link", + ts=time.time(), + ) + ) + + assert ("base_link", "camera_link") in buffer.buffers + assert len(buffer.buffers[("base_link", "camera_link")]) > 0 + + def test_static_in_bfs_chain(self) -> None: + """BFS should find paths through static frames.""" + buffer = MultiTBuffer(buffer_size=10.0) + now = time.time() + + buffer.receive_transform( + Transform( + translation=Vector3(1.0, 0.0, 0.0), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="world", + child_frame_id="base_link", + ts=now, + ) + ) + buffer.receive_static_transform( + Transform( + translation=Vector3(0.3, 0.0, 0.1), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="base_link", + child_frame_id="camera_link", + ts=now, + ) + ) + + result = buffer.get("world", "camera_link") + assert result is not None + assert abs(result.translation.x - 1.3) < 1e-6 + assert abs(result.translation.z - 0.1) < 1e-6 + + def test_static_publish_over_lcm(self) -> None: + """Static transforms published on /tf_static should be received as statics.""" + broadcaster = TF() + receiver = TF() + + camera_link = Transform( + translation=Vector3(0.3, 0.0, 0.1), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="base_link", + child_frame_id="camera_link", + ts=time.time(), + ) + broadcaster.publish_static(camera_link) + time.sleep(self.LCM_PROPAGATION_DELAY) + + assert ("base_link", "camera_link") in receiver._statics + + result = receiver.get( + "base_link", + "camera_link", + time_point=time.time() + self.FAR_FUTURE_OFFSET, + time_tolerance=self.STRICT_TOLERANCE, + ) + assert result is not None + assert abs(result.translation.x - 0.3) < 1e-6 + + broadcaster.stop() + receiver.stop() + + def test_module_publishes_static_transforms(self) -> None: + """A Module with static_transforms config should publish them via tf.""" + + class TestConfig(ModuleConfig): + frame_mapping: dict[str, str] = Field( + default_factory=lambda: dict( + body="base_link", + parent="world", + camera_link="camera_link", + ) + ) + static_transforms: dict[str, Transform] = Field( + default_factory=lambda: { + "camera_link": Transform( + translation=Vector3(0.3, 0.0, 0.1), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="base_link", + child_frame_id="camera_link", + ), + } + ) + + class TestModule(Module): + config: TestConfig + + module = TestModule() + module.start() + time.sleep(self.MODULE_START_DELAY) + + result = module.tf.get("base_link", "camera_link", time_tolerance=self.STRICT_TOLERANCE) + assert result is not None + assert abs(result.translation.x - 0.3) < 1e-6 + assert abs(result.translation.z - 0.1) < 1e-6 + + module.stop() + module._close_module() From 6a97b9802504e949149c5c2676914c883f2b41fb Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 26 May 2026 23:42:13 -0700 Subject: [PATCH 28/71] misc --- dimos/memory2/module.py | 6 ++++-- dimos/protocol/tf/tf.py | 2 +- dimos/robot/unitree/go2/config.py | 13 ------------- 3 files changed, 5 insertions(+), 16 deletions(-) diff --git a/dimos/memory2/module.py b/dimos/memory2/module.py index 43bd2a084f..a734e43e08 100644 --- a/dimos/memory2/module.py +++ b/dimos/memory2/module.py @@ -24,7 +24,7 @@ from reactivex.disposable import Disposable from dimos.agents.annotation import skill -from dimos.constants import DIMOS_PROJECT_ROOT +from dimos.constants import DEFAULT_WORLD_FRAME, DIMOS_PROJECT_ROOT from dimos.core.core import rpc from dimos.core.module import Module, ModuleConfig from dimos.memory2.embed import EmbedImages @@ -309,7 +309,9 @@ def _port_to_stream(self, name: str, input_topic: In[Any], stream: Stream[Any]) def on_msg(msg: Any) -> None: ts = getattr(msg, "ts", None) or time.time() frame_id = getattr(msg, "frame_id", None) or default_frame_id - transform = self.tf.get("world", frame_id, time_point=ts, time_tolerance=tf_tolerance) + transform = self.tf.get( + DEFAULT_WORLD_FRAME, frame_id, time_point=ts, time_tolerance=tf_tolerance + ) pose = transform.to_pose() if transform is not None else None if not pose: diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index b7224932eb..ed753655de 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -126,7 +126,7 @@ def receive_transform(self, *args: Transform) -> None: self.buffers[key] = TBuffer(self.buffer_size) self.buffers[key].add(transform) # pruning happens in self.buffers[key].add(), but we don't want static ones to get pruned - # would be better to fix pruning somehow but the that class does the pruning + # would be better to fix pruning somehow but the class that does the pruning # is more generic than TF-speficic pruning, so I'm going to backfill instead for now # plz improve when the latching-topic change is made self._backfill_statics() diff --git a/dimos/robot/unitree/go2/config.py b/dimos/robot/unitree/go2/config.py index 5823f862e9..bd60670d8c 100644 --- a/dimos/robot/unitree/go2/config.py +++ b/dimos/robot/unitree/go2/config.py @@ -12,19 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Copyright 2025-2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. from importlib import resources from pathlib import Path From 18525ed4183e997bf5a1e97288a029c16fd3180c Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 27 May 2026 02:39:59 -0700 Subject: [PATCH 29/71] fix tree_str reading buffers without lock (greptile review) --- dimos/protocol/tf/tf.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/dimos/protocol/tf/tf.py b/dimos/protocol/tf/tf.py index ed753655de..52ad45b869 100644 --- a/dimos/protocol/tf/tf.py +++ b/dimos/protocol/tf/tf.py @@ -306,12 +306,15 @@ def get_transform_search( @property def tree_str(self) -> str: - if not self.buffers: + with self._cv: + keys = list(self.buffers.keys()) + + if not keys: return "(empty)" children: dict[str, list[str]] = {} all_children: set[str] = set() - for parent, child in self.buffers: + for parent, child in keys: children.setdefault(parent, []).append(child) all_children.add(child) From c4834694b2441b6b0378f65dbed4fd091c13f5d6 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sun, 31 May 2026 01:40:17 -0700 Subject: [PATCH 30/71] review: address greptile on dimos/memory2/vis/space/rerun.py:237 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix syntax error introduced by manual merge resolution: `rr.log(path, world_T_optic)static=True)` → `rr.log(path, world_T_optical.to_pose().to_rerun(), static=True)` Both sides of the merge had identical content at this line; restoring. --- dimos/memory2/vis/space/rerun.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/memory2/vis/space/rerun.py b/dimos/memory2/vis/space/rerun.py index e0b7ce3025..59e69ce0ed 100644 --- a/dimos/memory2/vis/space/rerun.py +++ b/dimos/memory2/vis/space/rerun.py @@ -234,7 +234,7 @@ def render(space: Space, app_id: str = "space", spawn: bool = True) -> None: + Go2Config.static_transforms["camera_optical"] ) world_T_optical = Transform.from_pose("world", ps) + base_to_optical - rr.log(path, world_T_optic)static=True) + rr.log(path, world_T_optical.to_pose().to_rerun(), static=True) h, w = img.shape[:2] focal = max(w, h) rr.log( From 10b4fbd9d91662fc2eb7110c57b58c49cb613389 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sun, 31 May 2026 10:23:55 -0700 Subject: [PATCH 31/71] fix(mypy): import camera_info_static from go2.config (not connection) --- dimos/mapping/loop_closure/eval.py | 4 ++-- dimos/mapping/loop_closure/utils/markers_rrd.py | 4 ++-- dimos/utils/cli/map.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/dimos/mapping/loop_closure/eval.py b/dimos/mapping/loop_closure/eval.py index ddd9d4cb2a..04ea228359 100644 --- a/dimos/mapping/loop_closure/eval.py +++ b/dimos/mapping/loop_closure/eval.py @@ -41,7 +41,7 @@ from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.sensor_msgs.Image import Image from dimos.perception.fiducial.marker_transformer import DetectMarkers -from dimos.robot.unitree.go2.connection import _camera_info_static +from dimos.robot.unitree.go2.config import camera_info_static from dimos.utils.data import get_data DEFAULT_DATASETS = [f"hk_village{i}" for i in range(1, 7)] @@ -68,7 +68,7 @@ def _eval_recording( ) -> tuple[float, float]: """Returns (pgo_time_s, spread_m) for one recording.""" db_path = get_data(f"{name}.db") - cam_info = _camera_info_static() + cam_info = camera_info_static() store = SqliteStore(path=str(db_path)) with store: diff --git a/dimos/mapping/loop_closure/utils/markers_rrd.py b/dimos/mapping/loop_closure/utils/markers_rrd.py index a98e79ddc3..09d124089f 100644 --- a/dimos/mapping/loop_closure/utils/markers_rrd.py +++ b/dimos/mapping/loop_closure/utils/markers_rrd.py @@ -40,7 +40,7 @@ from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.perception.fiducial.marker_transformer import DetectMarkers -from dimos.robot.unitree.go2.connection import _camera_info_static +from dimos.robot.unitree.go2.config import camera_info_static from dimos.utils.data import resolve_named_path TIMELINE = "ts" @@ -64,7 +64,7 @@ def main( ), ) -> None: db_path = resolve_named_path(dataset, ".db") - cam_info = _camera_info_static() + cam_info = camera_info_static() rr.init("dimos markers", recording_id=db_path.stem) rr.save(str(out)) diff --git a/dimos/utils/cli/map.py b/dimos/utils/cli/map.py index c162a015ef..eaa996740f 100644 --- a/dimos/utils/cli/map.py +++ b/dimos/utils/cli/map.py @@ -233,7 +233,7 @@ def main( from dimos.msgs.sensor_msgs.CameraInfo import CameraInfo from dimos.msgs.sensor_msgs.Image import Image from dimos.perception.fiducial.marker_transformer import DetectMarkers - from dimos.robot.unitree.go2.connection import _camera_info_static + from dimos.robot.unitree.go2.config import camera_info_static from dimos.utils.data import resolve_named_path from dimos.visualization.rerun.init import rerun_init @@ -326,7 +326,7 @@ def main( # (verified: matches lidar_base_pose + BASE_TO_OPTICAL to ~1mm). # No mount composition needed. color_image = store.stream("color_image", Image) - cam_info = CameraInfo.from_yaml(str(camera_info)) if camera_info else _camera_info_static() + cam_info = CameraInfo.from_yaml(str(camera_info)) if camera_info else camera_info_static() xf = DetectMarkers( camera_info=cam_info, marker_length_m=marker_size, From 9c148f3f3c282dffa8eac44ed3a1195e3fa363b1 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 3 Jun 2026 21:46:53 -0700 Subject: [PATCH 32/71] fix: correct camera_info_static import in replay CLI tools --- dimos/mapping/utils/cli/replay.py | 4 ++-- dimos/mapping/utils/cli/replay_marker.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/dimos/mapping/utils/cli/replay.py b/dimos/mapping/utils/cli/replay.py index c78704f725..22af2800a4 100644 --- a/dimos/mapping/utils/cli/replay.py +++ b/dimos/mapping/utils/cli/replay.py @@ -214,13 +214,13 @@ def main( 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.robot.unitree.go2.config import camera_info_static from dimos.utils.data import resolve_named_path db_path = resolve_named_path(dataset, ".db") if out is None: out = Path.cwd() / f"{db_path.stem}.rrd" - cam_info = _camera_info_static() + cam_info = camera_info_static() # Resolve which streams to voxelize: all PointCloud2 streams, or the # explicit --map-source subset. Validate up front so typos fail fast. diff --git a/dimos/mapping/utils/cli/replay_marker.py b/dimos/mapping/utils/cli/replay_marker.py index 4d04815051..07f97cb5a2 100644 --- a/dimos/mapping/utils/cli/replay_marker.py +++ b/dimos/mapping/utils/cli/replay_marker.py @@ -72,13 +72,13 @@ def main( from dimos.msgs.sensor_msgs.Image import Image from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.perception.fiducial.marker_transformer import DetectMarkers - from dimos.robot.unitree.go2.connection import _camera_info_static + from dimos.robot.unitree.go2.config import camera_info_static from dimos.utils.data import resolve_named_path db_path = resolve_named_path(dataset, ".db") if out is None: out = Path.cwd() / f"{db_path.stem}.rrd" - cam_info = _camera_info_static() + cam_info = camera_info_static() rr.init("dimos markers", recording_id=db_path.stem) rr.save(str(out)) From a2200511a9dde025fdfa79fbfab1e511a40c68a7 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Sun, 21 Jun 2026 21:26:41 +0800 Subject: [PATCH 33/71] create TfModule --- dimos/core/module.py | 107 ++----------------- dimos/core/tf_module.py | 142 ++++++++++++++++++++++++++ dimos/protocol/tf/test_tf.py | 10 +- dimos/robot/cli/test_dimos.py | 5 - dimos/robot/unitree/go2/connection.py | 6 +- 5 files changed, 161 insertions(+), 109 deletions(-) create mode 100644 dimos/core/tf_module.py diff --git a/dimos/core/module.py b/dimos/core/module.py index b2c3679d89..847d335de1 100644 --- a/dimos/core/module.py +++ b/dimos/core/module.py @@ -19,7 +19,6 @@ import json import sys import threading -import time from typing import ( TYPE_CHECKING, Any, @@ -41,7 +40,6 @@ from dimos.core.resource import CompositeResource from dimos.core.rpc_client import RpcCall from dimos.core.stream import In, Out, RemoteOut, Transport -from dimos.msgs.geometry_msgs.Transform import Transform from dimos.protocol.rpc.pubsubrpc import LCMRPC from dimos.protocol.rpc.spec import DEFAULT_RPC_TIMEOUT, DEFAULT_RPC_TIMEOUTS, RPCSpec from dimos.protocol.service.spec import BaseConfig, Configurable @@ -111,12 +109,11 @@ class ModuleConfig(BaseConfig): tf_transport: type[TFSpec] = LCMTF # type: ignore[type-arg] frame_id_prefix: str | None = None frame_id: str | None = None - static_transforms: dict[str, Transform] = Field(default_factory=dict) - # TODO: in the future we should make self.tf.publish error if it tried to publish a transform that references a frame that is not mentioned in this dict (same with self.tf.get) # TODO: later expose frame remappings, somehow, at the blueprint level # frame mapping is how others can remap frame names, in the future we should have a way to do this remapping at the blueprint level too + # (static-transform publishing lives on the TfModule subclass, but frame_mapping + # stays here because blueprint-level frame namespacing applies to every module) frame_mapping: dict[str, str] = Field(default_factory=dict) - static_publish_interval: float = 1.0 g: GlobalConfig = global_config @@ -156,10 +153,8 @@ def __init__(self, config_args: dict[str, Any]) -> None: self._module_closed_lock = threading.Lock() self._tools = {} self._tools_lock = threading.Lock() - self._static_publish_thread: threading.Thread | None = None - self._static_publish_stop = threading.Event() self._loop, self._loop_thread = get_loop() - self.frame_mapping, self.static_transforms = self._setup_frames() + self.frame_mapping = self._setup_frame_mapping() try: self.rpc = self.config.rpc_transport( # type: ignore[call-arg] rpc_timeouts=self.config.rpc_timeouts, @@ -193,9 +188,6 @@ def build(self) -> None: @rpc def start(self) -> None: - # NOTE: there's basically always going to be some inital race around static transform frames and tf.get's - # putting the statics at the top helps mitigate/reduce that - self._start_static_publish() self._start_main() self._auto_bind_handlers() @@ -205,7 +197,10 @@ def stop(self) -> None: super().stop() self._close_module() - def _setup_frames(self) -> tuple[dict[str, str], dict[str, Transform]]: + def _setup_frame_mapping(self) -> dict[str, str]: + # frame_mapping resolution stays on the base Module (every module exposes a + # resolved frame_mapping for blueprint-level namespacing). Static-transform + # publishing lives on the TfModule subclass. frame_mapping_field = type(self.config).model_fields["frame_mapping"] if not hasattr(frame_mapping_field, "default_factory") or not callable( frame_mapping_field.default_factory @@ -213,8 +208,6 @@ def _setup_frames(self) -> tuple[dict[str, str], dict[str, Transform]]: raise ValueError( f"""In the {self.name!r} module config definition, the frame_mapping needs to be a pydantic field, not a dict""" ) - existing_frames: dict[str, str] = frame_mapping_field.default_factory() # type: ignore[call-arg] - # given something like: # class MyConfig: # frame_mapping: dict[str, str] = Field(default_factory=lambda: dict( @@ -223,83 +216,14 @@ def _setup_frames(self) -> tuple[dict[str, str], dict[str, Transform]]: # )) # the "body" is what I call a common_name # "base_link" is the REAL frame id that other modules can query/use - - # step1 (for static_transforms only) translate urdf_name=>common_name - reverse_mapping = {value: key for key, value in existing_frames.items()} - static_transforms_common_names = { - reverse_mapping.get(urdf_frame_id, urdf_frame_id): Transform( - translation=transform.translation, - rotation=transform.rotation, - frame_id=reverse_mapping.get(transform.frame_id, transform.frame_id), - child_frame_id=reverse_mapping.get( - transform.child_frame_id, transform.child_frame_id - ), - ) - for urdf_frame_id, transform in self.config.static_transforms.items() - } - # step2 map common_name=>real_frame_id - valid_frame_ids = set(existing_frames.keys()) | set(static_transforms_common_names.keys()) - final_frame_mapping = { - **existing_frames, - **self.config.frame_mapping, - } + existing_frames: dict[str, str] = frame_mapping_field.default_factory() # type: ignore[call-arg] + final_frame_mapping = {**existing_frames, **self.config.frame_mapping} for existing_frame, remapped_frame in final_frame_mapping.items(): - if existing_frame not in valid_frame_ids: + if existing_frame not in existing_frames: raise ValueError( f"""On module {self.name}, tried to map {existing_frame!r} to {remapped_frame!r} but that first frame doesn't exist. The existing ones are: {list(existing_frames.keys())!r} """ ) - static_transforms_final = { - final_frame_mapping.get(common_frame_id, common_frame_id): Transform( - translation=transform.translation, - rotation=transform.rotation, - frame_id=final_frame_mapping.get(transform.frame_id, transform.frame_id), - child_frame_id=final_frame_mapping.get( - transform.child_frame_id, transform.child_frame_id - ), - ) - for common_frame_id, transform in static_transforms_common_names.items() - } - return final_frame_mapping, static_transforms_final - - def _start_static_publish(self) -> None: - self._static_publish() - if not self.config.static_transforms or self.config.static_publish_interval <= 0: - return - self._static_publish_stop.clear() - self._static_publish_thread = threading.Thread( - target=self._static_publisher, - daemon=True, - ) - self._static_publish_thread.start() - - # TODO: later this should be replaced with latching streams - def _static_publisher(self) -> None: - while not self._static_publish_stop.wait(self.config.static_publish_interval): - self._static_publish() - - def _static_publish(self) -> None: - if not self.static_transforms: - return - now = time.time() - self.tf.publish_static( - *( - Transform( - translation=transform.translation, - rotation=transform.rotation, - frame_id=transform.frame_id, - child_frame_id=transform.child_frame_id, - ts=now, - ) - for transform in self.static_transforms.values() - ) - ) - self._on_static_publish() - - def _on_static_publish(self) -> None: - """ - This is a callback for modules to publish other data (ex: camera info) in the static loop - This should be rarely used, but exists for the few cases where it is needed - """ + return final_frame_mapping def _close_module(self) -> None: with self._module_closed_lock: @@ -307,11 +231,6 @@ def _close_module(self) -> None: return self._module_closed = True - self._static_publish_stop.set() - if self._static_publish_thread and self._static_publish_thread.is_alive(): - self._static_publish_thread.join(timeout=self._loop_thread_timeout) - self._static_publish_thread = None - self._close_all_tools() self._close_rpc() @@ -366,8 +285,6 @@ def __getstate__(self): # type: ignore[no-untyped-def] state.pop("_main_gen", None) state.pop("_tools", None) state.pop("_tools_lock", None) - state.pop("_static_publish_thread", None) - state.pop("_static_publish_stop", None) return state def __setstate__(self, state) -> None: # type: ignore[no-untyped-def] @@ -382,8 +299,6 @@ def __setstate__(self, state) -> None: # type: ignore[no-untyped-def] self._main_gen = None self._tools = {} self._tools_lock = threading.Lock() - self._static_publish_thread = None - self._static_publish_stop = threading.Event() @property def tf(self): # type: ignore[no-untyped-def] diff --git a/dimos/core/tf_module.py b/dimos/core/tf_module.py new file mode 100644 index 0000000000..f0955464a1 --- /dev/null +++ b/dimos/core/tf_module.py @@ -0,0 +1,142 @@ +# 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. +import threading +import time +from typing import Any + +from pydantic import Field + +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.msgs.geometry_msgs.Transform import Transform + + +class TfModuleConfig(ModuleConfig): + static_transforms: dict[str, Transform] = Field(default_factory=dict) + # TODO: in the future we should make self.tf.publish error if it tried to publish a transform that references a frame that is not mentioned in this dict (same with self.tf.get) + static_publish_interval: float = 1.0 + + +class TfModule(Module): + """A Module that republishes its config's static (non-moving) transforms on an interval. + + Modules that need to publish fixed frames inherit from this instead of Module + directly, so the base Module stays free of transform-publishing machinery that + most modules don't need. `frame_mapping` resolution still lives on Module (every + module needs it for blueprint-level frame namespacing); only the static-transform + publishing is added here. + """ + + config: TfModuleConfig + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._static_publish_thread: threading.Thread | None = None + # created lazily in start() so the module stays picklable for worker deployment + self._static_publish_stop: threading.Event | None = None + self.static_transforms = self._resolve_static_transforms() + + @rpc + def start(self) -> None: + # NOTE: there's basically always going to be some inital race around static transform frames and tf.get's + # publishing the statics before main starts helps mitigate/reduce that + self._start_static_publish() + super().start() + + @rpc + def stop(self) -> None: + self._stop_static_publish() + super().stop() + + def _resolve_static_transforms(self) -> dict[str, Transform]: + frame_mapping_field = type(self.config).model_fields["frame_mapping"] + existing_frames: dict[str, str] = frame_mapping_field.default_factory() # type: ignore[misc,union-attr,call-arg] + + # step1 translate urdf_name=>common_name (see Module._setup_frame_mapping for what + # "common name" vs "real frame id" means) + reverse_mapping = {value: key for key, value in existing_frames.items()} + static_transforms_common_names = { + reverse_mapping.get(urdf_frame_id, urdf_frame_id): Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=reverse_mapping.get(transform.frame_id, transform.frame_id), + child_frame_id=reverse_mapping.get( + transform.child_frame_id, transform.child_frame_id + ), + ) + for urdf_frame_id, transform in self.config.static_transforms.items() + } + # step2 map common_name=>real_frame_id using the module's resolved frame_mapping + final_frame_mapping = self.frame_mapping + return { + final_frame_mapping.get(common_frame_id, common_frame_id): Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=final_frame_mapping.get(transform.frame_id, transform.frame_id), + child_frame_id=final_frame_mapping.get( + transform.child_frame_id, transform.child_frame_id + ), + ) + for common_frame_id, transform in static_transforms_common_names.items() + } + + def _start_static_publish(self) -> None: + self._static_publish() + if not self.static_transforms or self.config.static_publish_interval <= 0: + return + self._static_publish_stop = threading.Event() + self._static_publish_thread = threading.Thread( + target=self._static_publisher, + daemon=True, + ) + self._static_publish_thread.start() + + # TODO: later this should be replaced with latching streams + def _static_publisher(self) -> None: + stop = self._static_publish_stop + assert stop is not None + while not stop.wait(self.config.static_publish_interval): + self._static_publish() + + def _static_publish(self) -> None: + if not self.static_transforms: + return + now = time.time() + self.tf.publish_static( + *( + Transform( + translation=transform.translation, + rotation=transform.rotation, + frame_id=transform.frame_id, + child_frame_id=transform.child_frame_id, + ts=now, + ) + for transform in self.static_transforms.values() + ) + ) + self._on_static_publish() + + def _on_static_publish(self) -> None: + """ + This is a callback for modules to publish other data (ex: camera info) in the static loop + This should be rarely used, but exists for the few cases where it is needed + """ + + def _stop_static_publish(self) -> None: + if self._static_publish_stop is not None: + self._static_publish_stop.set() + if self._static_publish_thread and self._static_publish_thread.is_alive(): + self._static_publish_thread.join(timeout=self._loop_thread_timeout) + self._static_publish_thread = None + self._static_publish_stop = None diff --git a/dimos/protocol/tf/test_tf.py b/dimos/protocol/tf/test_tf.py index 524632c028..189d290335 100644 --- a/dimos/protocol/tf/test_tf.py +++ b/dimos/protocol/tf/test_tf.py @@ -21,7 +21,7 @@ from pydantic import Field import pytest -from dimos.core.module import Module, ModuleConfig +from dimos.core.tf_module import TfModule, TfModuleConfig from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform @@ -811,7 +811,7 @@ def test_get_with_longer_transform_chain(self) -> None: class TestStaticTransforms: LCM_PROPAGATION_DELAY = 2 - MODULE_START_DELAY = 2 + MODULE_START_DELAY = 0.05 STRICT_TOLERANCE = 0.001 STALE_AGE = 1000 FAR_FUTURE_OFFSET = 9999 @@ -925,9 +925,9 @@ def test_static_publish_over_lcm(self) -> None: receiver.stop() def test_module_publishes_static_transforms(self) -> None: - """A Module with static_transforms config should publish them via tf.""" + """A TfModule with static_transforms config should publish them via tf.""" - class TestConfig(ModuleConfig): + class TestConfig(TfModuleConfig): frame_mapping: dict[str, str] = Field( default_factory=lambda: dict( body="base_link", @@ -946,7 +946,7 @@ class TestConfig(ModuleConfig): } ) - class TestModule(Module): + class TestModule(TfModule): config: TestConfig module = TestModule() diff --git a/dimos/robot/cli/test_dimos.py b/dimos/robot/cli/test_dimos.py index a12727ac34..b925d4ebd9 100644 --- a/dimos/robot/cli/test_dimos.py +++ b/dimos/robot/cli/test_dimos.py @@ -69,7 +69,6 @@ class TestModuleB(Module): " * testmodulea.default_rpc_timeout: float (default: 120.0)", " * testmodulea.frame_id_prefix: str | None (default: None)", " * testmodulea.frame_id: str | None (default: None)", - " * testmodulea.static_publish_interval: float (default: 1.0)", " * testmodulea.min_interval_sec: float (default: 0.1)", " * testmodulea.entity_prefix: str (default: world)", " * testmodulea.viewer_mode: typing.Literal['native', 'web', 'connect', 'none'] (default: native)", @@ -77,7 +76,6 @@ class TestModuleB(Module): " * testmoduleb.default_rpc_timeout: float (default: 120.0)", " * testmoduleb.frame_id_prefix: str | None (default: None)", " * testmoduleb.frame_id: str | None (default: None)", - " * testmoduleb.static_publish_interval: float (default: 1.0)", " * testmoduleb.memory_limit: str (default: 25%)", " * testmoduleb.ip: str (default: 127.0.0.1)", "", @@ -112,7 +110,6 @@ class TestModuleB(Module): " * testmodulea.default_rpc_timeout: float (default: 120.0)", " * testmodulea.frame_id_prefix: str | None (default: foo)", " * testmodulea.frame_id: str | None (default: None)", - " * testmodulea.static_publish_interval: float (default: 1.0)", " * testmodulea.min_interval_sec: float (default: 0.1)", " * testmodulea.entity_prefix: str (default: world)", " * testmodulea.viewer_mode: typing.Literal['native', 'web', 'connect', 'none'] (default: web)", @@ -120,7 +117,6 @@ class TestModuleB(Module): " * testmoduleb.default_rpc_timeout: float (default: 120.0)", " * testmoduleb.frame_id_prefix: str | None (default: None)", " * testmoduleb.frame_id: str | None (default: None)", - " * testmoduleb.static_publish_interval: float (default: 1.0)", " * testmoduleb.memory_limit: str (default: 25%)", " * testmoduleb.ip: str (default: 1.1.1.1)", "", @@ -144,7 +140,6 @@ class TestModule(Module): " * testmodule.default_rpc_timeout: float (default: 120.0)", " * testmodule.frame_id_prefix: str | None (default: None)", " * testmodule.frame_id: str | None (default: None)", - " * testmodule.static_publish_interval: float (default: 1.0)", " * [Required] testmodule.foo: int", " * testmodule.spam: str (default: eggs)", "", diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index ef96776645..a84ec9bf16 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -30,9 +30,9 @@ from dimos.core.coordination.module_coordinator import ModuleCoordinator from dimos.core.core import rpc from dimos.core.global_config import GlobalConfig -from dimos.core.module import Module, ModuleConfig from dimos.core.resource import CompositeResource from dimos.core.stream import In, Out +from dimos.core.tf_module import TfModule, TfModuleConfig from dimos.core.transport import LCMTransport, pSHMTransport from dimos.spec.perception import Camera, Pointcloud from dimos.utils.logging_config import setup_logger @@ -64,7 +64,7 @@ class Go2Mode(str, Enum): RAGE = "rage" -class ConnectionConfig(ModuleConfig): +class ConnectionConfig(TfModuleConfig): ip: str = Field(default_factory=lambda m: m["g"].robot_ip) mode: Go2Mode = Go2Mode.DEFAULT frame_mapping: dict[str, str] = Field( @@ -183,7 +183,7 @@ def publish_request(self, topic: str, data: dict): # type: ignore[no-untyped-de _Config = TypeVar("_Config", bound=ConnectionConfig, default=ConnectionConfig) -class GO2Connection(Module, Camera, Pointcloud): +class GO2Connection(TfModule, Camera, Pointcloud): dedicated_worker = True config: ConnectionConfig From 4e63c00ebab14c6358e9dad2b3a9982c3ab73e31 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 12:41:01 -0700 Subject: [PATCH 34/71] Ray trace local region --- dimos/mapping/ray_tracing/module.py | 17 +- dimos/mapping/ray_tracing/rust/Cargo.lock | 59 ++++ dimos/mapping/ray_tracing/rust/Cargo.toml | 2 + dimos/mapping/ray_tracing/rust/flake.lock | 10 +- dimos/mapping/ray_tracing/rust/flake.nix | 2 +- dimos/mapping/ray_tracing/rust/src/main.rs | 309 +++++++++++------- dimos/mapping/ray_tracing/rust/src/python.rs | 64 ++-- .../ray_tracing/rust/src/voxel_ray_tracer.rs | 233 ++++++++----- dimos/mapping/ray_tracing/test_transformer.py | 49 ++- dimos/mapping/ray_tracing/transformer.py | 69 +++- dimos/mapping/ray_tracing/voxel_map.py | 5 +- dimos/mapping/ray_tracing/voxel_map.pyi | 14 +- pyproject.toml | 1 + 13 files changed, 610 insertions(+), 224 deletions(-) diff --git a/dimos/mapping/ray_tracing/module.py b/dimos/mapping/ray_tracing/module.py index 5b215f9ff6..a7fcab1d2f 100644 --- a/dimos/mapping/ray_tracing/module.py +++ b/dimos/mapping/ray_tracing/module.py @@ -19,6 +19,7 @@ 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.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.spec import mapping @@ -47,10 +48,23 @@ class RayTracingVoxelMapConfig(NativeModuleConfig): graze_cos: float = 0.7 # Only spare a voxel whose neighborhood was hit within this many frames. recency_window: int = 15 + # Integrate every frame, publish the local map and region bounds every + # Nth frame. Zero disables them. + emit_every: int = 1 + # Publish the global map every Nth frame. Zero disables it. + global_emit_every: int = 1 + # Size the local region to this percentile of batch point distances, + # so a stray far hit cannot inflate the region the planner recomputes. + region_percentile: float = 95.0 class RayTracingVoxelMap(NativeModule, mapping.GlobalPointcloud): - """Rust voxel-map module with raycast clearing of dynamic objects.""" + """Rust voxel-map module with raycast clearing of dynamic objects. + + region_bounds describes the cylinder local_map covers, packed into a + PoseStamped. Position holds the center. Orientation holds radius, z_min, + z_max, and zero. It shares the local_map stamp. + """ config: RayTracingVoxelMapConfig @@ -58,6 +72,7 @@ class RayTracingVoxelMap(NativeModule, mapping.GlobalPointcloud): odometry: In[Odometry] global_map: Out[PointCloud2] local_map: Out[PointCloud2] + region_bounds: Out[PoseStamped] # Verify protocol port compliance (mypy will flag missing ports) diff --git a/dimos/mapping/ray_tracing/rust/Cargo.lock b/dimos/mapping/ray_tracing/rust/Cargo.lock index 293fccf43f..122775cdb5 100644 --- a/dimos/mapping/ray_tracing/rust/Cargo.lock +++ b/dimos/mapping/ray_tracing/rust/Cargo.lock @@ -33,6 +33,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "arrayvec" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" + [[package]] name = "autocfg" version = "1.5.1" @@ -63,6 +69,31 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + [[package]] name = "darling" version = "0.20.11" @@ -136,11 +167,13 @@ name = "dimos-voxel-ray-tracing" version = "0.1.0" dependencies = [ "ahash", + "arrayvec", "dimos-module", "lcm-msgs", "nalgebra", "numpy", "pyo3", + "rayon", "serde", "tokio", "tracing", @@ -158,6 +191,12 @@ dependencies = [ "syn", ] +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + [[package]] name = "errno" version = "0.3.14" @@ -701,6 +740,26 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "regex" version = "1.12.3" diff --git a/dimos/mapping/ray_tracing/rust/Cargo.toml b/dimos/mapping/ray_tracing/rust/Cargo.toml index f1a63d34d6..42c60f96bd 100644 --- a/dimos/mapping/ray_tracing/rust/Cargo.toml +++ b/dimos/mapping/ray_tracing/rust/Cargo.toml @@ -21,6 +21,8 @@ lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "r tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] } serde = { version = "1", features = ["derive"] } ahash = "0.8" +arrayvec = "0.7" +rayon = "1" tracing = "0.1" pyo3 = { version = "0.25", features = ["extension-module", "abi3-py310"] } numpy = "0.25" diff --git a/dimos/mapping/ray_tracing/rust/flake.lock b/dimos/mapping/ray_tracing/rust/flake.lock index a548660557..31aab7d531 100644 --- a/dimos/mapping/ray_tracing/rust/flake.lock +++ b/dimos/mapping/ray_tracing/rust/flake.lock @@ -3,11 +3,11 @@ "dimos-repo": { "flake": false, "locked": { - "lastModified": 1779865691, - "narHash": "sha256-2CVWcov7DiC1qX/B/zFKDJiSYsnbrZ3FNT/viprFWTQ=", - "ref": "refs/heads/jeff/feat/g1_raycast", - "rev": "51666bcd298c1d08bdee179f176f45c0a7dd417d", - "revCount": 744, + "lastModified": 1781823023, + "narHash": "sha256-YkqAYJlqFI52p1ZfM0G29+85eSykkIgXU7xs7oqgopo=", + "ref": "refs/heads/andrew/feat/kronk-nav-3", + "rev": "8ef1d416ddf141e354902ab86cedcd748220e0e9", + "revCount": 882, "type": "git", "url": "file:../../../.." }, diff --git a/dimos/mapping/ray_tracing/rust/flake.nix b/dimos/mapping/ray_tracing/rust/flake.nix index bb2fbfd52d..54721c779d 100644 --- a/dimos/mapping/ray_tracing/rust/flake.nix +++ b/dimos/mapping/ray_tracing/rust/flake.nix @@ -34,7 +34,7 @@ cargoRoot = "dimos/mapping/ray_tracing/rust"; buildAndTestSubdir = "dimos/mapping/ray_tracing/rust"; - cargoHash = "sha256-g30NaoLdtWT5YBsEnE4Xv+EMnI5HHFtZAUtdEL/VbKQ="; + cargoHash = "sha256-0d0dlNDvDplA7oWTyUWOCOlS74Zie8uMQ+ps6lXntOI="; meta.mainProgram = "voxel_ray_tracing"; }; diff --git a/dimos/mapping/ray_tracing/rust/src/main.rs b/dimos/mapping/ray_tracing/rust/src/main.rs index d8de5f42ab..be9db8aba8 100644 --- a/dimos/mapping/ray_tracing/rust/src/main.rs +++ b/dimos/mapping/ray_tracing/rust/src/main.rs @@ -6,11 +6,13 @@ use std::time::Duration; use ahash::AHashSet; use dimos_module::{error_throttled, run, warn_throttled, Input, LcmTransport, Module, Output}; use dimos_voxel_ray_tracing::voxel_ray_tracer::{ - iter_global_points, update_map, Config, LocalBounds, VoxelKey, VoxelMap, + batch_local_bounds, iter_global_points, update_map, Config, LocalBounds, VoxelKey, VoxelMap, }; +use lcm_msgs::geometry_msgs::{Point, Pose, PoseStamped, Quaternion}; use lcm_msgs::nav_msgs::Odometry; use lcm_msgs::sensor_msgs::{PointCloud2, PointField}; use lcm_msgs::std_msgs::{Header, Time}; +use nalgebra::{UnitQuaternion, Vector3}; #[derive(Module)] struct RayTracingVoxelMap { @@ -26,27 +28,40 @@ struct RayTracingVoxelMap { #[output(encode = PointCloud2::encode)] local_map: Output, + // Region the local_map covers, packed into a PoseStamped. Position holds + // the cylinder center and orientation holds radius, z_min, z_max, zero. + // Stamped identically to local_map so consumers can pair them. + #[output(encode = PoseStamped::encode)] + region_bounds: Output, + #[config] config: Config, map: VoxelMap, - last_origin: Option<(f32, f32, f32)>, + last_pose: Option<(Vector3, UnitQuaternion)>, + frame_count: u32, + batch_points: Vec<(f32, f32, f32)>, + batch_origins: Vec<(f32, f32, f32)>, } impl RayTracingVoxelMap { async fn on_odometry(&mut self, msg: Odometry) { - self.last_origin = Some(( - msg.pose.pose.position.x as f32, - msg.pose.pose.position.y as f32, - msg.pose.pose.position.z as f32, + let p = &msg.pose.pose.position; + let q = &msg.pose.pose.orientation; + self.last_pose = Some(( + Vector3::new(p.x as f32, p.y as f32, p.z as f32), + UnitQuaternion::from_quaternion(nalgebra::Quaternion::new( + q.w as f32, q.x as f32, q.y as f32, q.z as f32, + )), )); } async fn on_lidar(&mut self, msg: PointCloud2) { - let Some(origin) = self.last_origin else { + let Some((translation, rotation)) = self.last_pose else { // Need at least one odometry sample before we can raycast. return; }; + let origin = (translation.x, translation.y, translation.z); let voxel_size = self.config.voxel_size; @@ -65,45 +80,106 @@ impl RayTracingVoxelMap { return; } + // Register sensor-frame clouds into the world by the odom pose. + let rot = rotation.to_rotation_matrix(); + let points: Vec<(f32, f32, f32)> = points + .iter() + .map(|&(x, y, z)| { + let p = rot * Vector3::new(x, y, z) + translation; + (p.x, p.y, p.z) + }) + .collect(); + + // The integrated points are world-frame either way. + let out_frame_id = "world"; + let live = update_map(&mut self.map, origin, &points, &self.config); - let half = voxel_size * 0.5; - let mut z_min = f32::INFINITY; - let mut z_max = f32::NEG_INFINITY; - let mut r_xy_max_sq = 0.0_f32; - for &(kx, ky, kz) in &live { - let cx = kx as f32 * voxel_size + half; - let cy = ky as f32 * voxel_size + half; - let cz = kz as f32 * voxel_size + half; - z_min = z_min.min(cz); - z_max = z_max.max(cz); - let dx = cx - origin.0; - let dy = cy - origin.1; - r_xy_max_sq = r_xy_max_sq.max(dx * dx + dy * dy); + // The batch only feeds the local region bounds, so don't let it grow + // when the local map is disabled. + if self.config.emit_every > 0 { + self.batch_points.extend_from_slice(&points); + self.batch_origins.push(origin); + } + + self.frame_count += 1; + if self + .frame_count + .is_multiple_of(self.config.global_emit_every) + { + let cloud = build_global_cloud( + &self.map, + &live, + voxel_size, + out_frame_id, + msg.header.stamp.clone(), + ); + if let Err(e) = self.global_map.publish(&cloud).await { + error_throttled!( + Duration::from_secs(1), + error = %e, + "Updated global voxel map failed to publish", + ); + } } + if !self.frame_count.is_multiple_of(self.config.emit_every) { + return; + } + + // Percentile-bounded cylinder over the batch plus the clearing margin. + let margin = self.config.shadow_depth + voxel_size; + let (cx, cy, radius, z_min, z_max) = batch_local_bounds( + &self.batch_points, + &self.batch_origins, + self.config.region_percentile, + margin, + ); + self.batch_points.clear(); + self.batch_origins.clear(); let cylinder = LocalBounds { - origin_x: origin.0, - origin_y: origin.1, - r_xy_max_sq, + origin_x: cx, + origin_y: cy, + r_xy_max_sq: radius * radius, z_min, z_max, }; - let (global_cloud, local_cloud) = build_pointclouds( + let bounds_msg = PoseStamped { + header: Header { + seq: 0, + stamp: msg.header.stamp.clone(), + frame_id: out_frame_id.to_string(), + }, + pose: Pose { + position: Point { + x: cx as f64, + y: cy as f64, + z: 0.0, + }, + orientation: Quaternion { + x: radius as f64, + y: z_min as f64, + z: z_max as f64, + w: 0.0, + }, + }, + }; + if let Err(e) = self.region_bounds.publish(&bounds_msg).await { + error_throttled!( + Duration::from_secs(1), + error = %e, + "Region bounds failed to publish", + ); + } + + let local_cloud = build_local_cloud( &self.map, &live, voxel_size, &cylinder, - &msg.header.frame_id, + out_frame_id, msg.header.stamp, ); - if let Err(e) = self.global_map.publish(&global_cloud).await { - error_throttled!( - Duration::from_secs(1), - error = %e, - "Updated global voxel map failed to publish", - ); - } if let Err(e) = self.local_map.publish(&local_cloud).await { error_throttled!( Duration::from_secs(1), @@ -180,93 +256,103 @@ fn read_f32_le(buf: &[u8], off: usize) -> f32 { f32::from_le_bytes(bytes) } -fn build_pointclouds( - map: &VoxelMap, - live: &AHashSet, +fn write_point(data: &mut Vec, n: &mut i32, 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()); + data.extend_from_slice(&0.0_f32.to_le_bytes()); + *n += 1; +} + +/// Centers of live voxels not yet healthy enough to appear in the map. +fn unhealthy_live_centers<'a>( + map: &'a VoxelMap, + live: &'a AHashSet, voxel_size: f32, - cylinder: &LocalBounds, - frame_id: &str, - stamp: Time, -) -> (PointCloud2, PointCloud2) { +) -> impl Iterator + 'a { let half = voxel_size * 0.5; - let mut global_data = Vec::with_capacity((map.voxels.len() + live.len()) * 16); - let mut local_data = Vec::with_capacity(live.len() * 2 * 16); - let mut global_n: i32 = 0; - let mut local_n: i32 = 0; - - let write_point = |data: &mut Vec, n: &mut i32, 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()); - data.extend_from_slice(&0.0_f32.to_le_bytes()); - *n += 1; - }; - - // add healthy voxels to global, and local if necessary - for (x, y, z) in iter_global_points(map, voxel_size) { - write_point(&mut global_data, &mut global_n, x, y, z); - if cylinder.contains(x, y, z) { - write_point(&mut local_data, &mut local_n, x, y, z); - } - } - - // add live voxels to both if they aren't already there - for &(kx, ky, kz) in live { + live.iter().filter_map(move |&(kx, ky, kz)| { if matches!(map.voxels.get(&(kx, ky, kz)), Some(c) if c.health > 0) { - continue; + return None; } - let x = kx as f32 * voxel_size + half; - let y = ky as f32 * voxel_size + half; - let z = kz as f32 * voxel_size + half; - write_point(&mut global_data, &mut global_n, x, y, z); - write_point(&mut local_data, &mut local_n, x, y, z); - } + Some(( + kx as f32 * voxel_size + half, + ky as f32 * voxel_size + half, + kz as f32 * voxel_size + half, + )) + }) +} +fn make_cloud(data: Vec, n: i32, frame_id: &str, stamp: Time) -> PointCloud2 { let make_field = |name: &str, off: i32| PointField { name: name.into(), offset: off, datatype: PointField::FLOAT32 as u8, count: 1, }; - let fields = vec![ - make_field("x", 0), - make_field("y", 4), - make_field("z", 8), - make_field("intensity", 12), - ]; - - let global_cloud = PointCloud2 { - header: Header { - seq: 0, - stamp: stamp.clone(), - frame_id: frame_id.into(), - }, - height: 1, - width: global_n, - fields: fields.clone(), - is_bigendian: false, - point_step: 16, - row_step: 16 * global_n, - data: global_data, - is_dense: true, - }; - let local_cloud = PointCloud2 { + PointCloud2 { header: Header { seq: 0, stamp, frame_id: frame_id.into(), }, height: 1, - width: local_n, - fields, + width: n, + fields: vec![ + make_field("x", 0), + make_field("y", 4), + make_field("z", 8), + make_field("intensity", 12), + ], is_bigendian: false, point_step: 16, - row_step: 16 * local_n, - data: local_data, + row_step: 16 * n, + data, is_dense: true, - }; + } +} + +/// All healthy voxels plus this frame's live voxels. +fn build_global_cloud( + map: &VoxelMap, + live: &AHashSet, + voxel_size: f32, + frame_id: &str, + stamp: Time, +) -> PointCloud2 { + let mut data = Vec::with_capacity((map.voxels.len() + live.len()) * 16); + let mut n: i32 = 0; + for (x, y, z) in iter_global_points(map, voxel_size) { + write_point(&mut data, &mut n, x, y, z); + } + for (x, y, z) in unhealthy_live_centers(map, live, voxel_size) { + write_point(&mut data, &mut n, x, y, z); + } + make_cloud(data, n, frame_id, stamp) +} - (global_cloud, local_cloud) +/// Healthy voxels and this frame's live voxels, all inside the cylinder. +fn build_local_cloud( + map: &VoxelMap, + live: &AHashSet, + voxel_size: f32, + cylinder: &LocalBounds, + frame_id: &str, + stamp: Time, +) -> PointCloud2 { + let mut data = Vec::with_capacity(live.len() * 2 * 16); + let mut n: i32 = 0; + for (x, y, z) in iter_global_points(map, voxel_size) { + if cylinder.contains(x, y, z) { + write_point(&mut data, &mut n, x, y, z); + } + } + for (x, y, z) in unhealthy_live_centers(map, live, voxel_size) { + if cylinder.contains(x, y, z) { + write_point(&mut data, &mut n, x, y, z); + } + } + make_cloud(data, n, frame_id, stamp) } #[tokio::main] @@ -315,8 +401,8 @@ mod tests { z_min: 0.0, z_max: 1.0, }; - let (global, local) = - build_pointclouds(&map, &live, 1.0, &cylinder, "world", Time::default()); + let global = build_global_cloud(&map, &live, 1.0, "world", Time::default()); + let local = build_local_cloud(&map, &live, 1.0, &cylinder, "world", Time::default()); assert!(cloud_points(&global).contains(&voxel_center(0, 0, 0))); assert!(cloud_points(&local).contains(&voxel_center(0, 0, 0))); } @@ -333,8 +419,8 @@ mod tests { z_min: -10.0, z_max: 10.0, }; - let (global, local) = - build_pointclouds(&map, &live, 1.0, &cylinder, "world", Time::default()); + let global = build_global_cloud(&map, &live, 1.0, "world", Time::default()); + let local = build_local_cloud(&map, &live, 1.0, &cylinder, "world", Time::default()); assert!(cloud_points(&global).contains(&voxel_center(5, 0, 0))); assert!(!cloud_points(&local).contains(&voxel_center(5, 0, 0))); assert_eq!(local.width, 0); @@ -352,28 +438,31 @@ mod tests { z_min: 0.0, z_max: 1.0, }; - let (global, local) = - build_pointclouds(&map, &live, 1.0, &cylinder, "world", Time::default()); + let global = build_global_cloud(&map, &live, 1.0, "world", Time::default()); + let local = build_local_cloud(&map, &live, 1.0, &cylinder, "world", Time::default()); assert!(cloud_points(&global).contains(&voxel_center(0, 0, 5))); assert!(!cloud_points(&local).contains(&voxel_center(0, 0, 5))); assert_eq!(local.width, 0); } #[test] - fn local_map_always_includes_live_voxels() { + fn live_voxels_follow_the_cylinder_in_local_map() { let map = VoxelMap::default(); let mut live: AHashSet = AHashSet::new(); + live.insert((1, 0, 0)); live.insert((10, 10, 10)); let cylinder = LocalBounds { origin_x: 0.0, origin_y: 0.0, - r_xy_max_sq: 0.0, + r_xy_max_sq: 4.0, z_min: 0.0, - z_max: 0.0, + z_max: 1.0, }; - let (global, local) = - build_pointclouds(&map, &live, 1.0, &cylinder, "world", Time::default()); + let global = build_global_cloud(&map, &live, 1.0, "world", Time::default()); + let local = build_local_cloud(&map, &live, 1.0, &cylinder, "world", Time::default()); + assert!(cloud_points(&global).contains(&voxel_center(1, 0, 0))); assert!(cloud_points(&global).contains(&voxel_center(10, 10, 10))); - assert!(cloud_points(&local).contains(&voxel_center(10, 10, 10))); + assert!(cloud_points(&local).contains(&voxel_center(1, 0, 0))); + assert!(!cloud_points(&local).contains(&voxel_center(10, 10, 10))); } } diff --git a/dimos/mapping/ray_tracing/rust/src/python.rs b/dimos/mapping/ray_tracing/rust/src/python.rs index ff49f5ed04..99f81fc24b 100644 --- a/dimos/mapping/ray_tracing/rust/src/python.rs +++ b/dimos/mapping/ray_tracing/rust/src/python.rs @@ -8,9 +8,46 @@ use pyo3::prelude::*; use validator::Validate; use crate::voxel_ray_tracer::{ - iter_global_normals, iter_global_points, update_map, Config, LocalBounds, VoxelMap, + batch_local_bounds, iter_global_normals, iter_global_points, update_map, Config, LocalBounds, + VoxelMap, }; +fn extract_tuples(arr: &Bound<'_, PyAny>, name: &str) -> PyResult> { + let arr: PyReadonlyArray2<'_, f32> = arr.extract().map_err(|_| { + PyValueError::new_err(format!("{name} must be a (N, 3) float32 numpy array")) + })?; + let shape = arr.shape(); + if shape[1] != 3 { + return Err(PyValueError::new_err(format!( + "{name} must be (N, 3) float32, got shape {:?}", + shape + ))); + } + let view = arr.as_array(); + Ok((0..shape[0]) + .filter_map(|i| { + let x = view[[i, 0]]; + let y = view[[i, 1]]; + let z = view[[i, 2]]; + (x.is_finite() && y.is_finite() && z.is_finite()).then_some((x, y, z)) + }) + .collect()) +} + +/// Local region a batch of frames observed, as (cx, cy, radius, z_min, z_max). +/// Non-finite points are ignored. +#[pyfunction] +fn local_bounds( + points: &Bound<'_, PyAny>, + origins: &Bound<'_, PyAny>, + percentile: f32, + margin: f32, +) -> PyResult<(f32, f32, f32, f32, f32)> { + let pts = extract_tuples(points, "points")?; + let origs = extract_tuples(origins, "origins")?; + Ok(batch_local_bounds(&pts, &origs, percentile, margin)) +} + #[pyclass] pub struct VoxelRayMapper { config: Config, @@ -54,6 +91,9 @@ impl VoxelRayMapper { max_health, graze_cos, recency_window, + emit_every: 1, + global_emit_every: 1, + region_percentile: 95.0, }; config .validate() @@ -70,26 +110,7 @@ impl VoxelRayMapper { points: &Bound<'_, PyAny>, origin: (f32, f32, f32), ) -> PyResult<()> { - let points: PyReadonlyArray2<'_, f32> = points - .extract() - .map_err(|_| PyValueError::new_err("points must be a (N, 3) float32 numpy array"))?; - let shape = points.shape(); - if shape[1] != 3 { - return Err(PyValueError::new_err(format!( - "points must be (N, 3) float32, got shape {:?}", - shape - ))); - } - let arr = points.as_array(); - let n = shape[0]; - let pts: Vec<(f32, f32, f32)> = (0..n) - .filter_map(|i| { - let x = arr[[i, 0]]; - let y = arr[[i, 1]]; - let z = arr[[i, 2]]; - (x.is_finite() && y.is_finite() && z.is_finite()).then_some((x, y, z)) - }) - .collect(); + let pts = extract_tuples(points, "points")?; let cfg = &self.config; let map = &mut self.map; @@ -205,5 +226,6 @@ impl VoxelRayMapper { #[pymodule] fn dimos_voxel_ray_tracing(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + m.add_function(wrap_pyfunction!(local_bounds, m)?)?; Ok(()) } diff --git a/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs b/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs index 1549e5a1ab..770412eb58 100644 --- a/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs +++ b/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs @@ -2,8 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 use ahash::{AHashMap, AHashSet}; +use arrayvec::ArrayVec; use dimos_module::native_config; use nalgebra::{Matrix3, Vector3}; +use rayon::prelude::*; use validator::ValidationError; pub type VoxelKey = (i32, i32, i32); @@ -32,6 +34,17 @@ pub struct Config { /// Only spare a voxel whose neighborhood was hit within this many frames. /// A stale voxel can be cleared, even if it's a grazing hit. Large disables it. pub recency_window: u32, + /// Integrate every frame, publish the local map and region bounds every + /// Nth frame. Zero disables them. + #[validate(range(min = 0))] + pub emit_every: u32, + /// Publish the global map every Nth frame. Zero disables it. + #[validate(range(min = 0))] + pub global_emit_every: u32, + /// Size the local region to this percentile of batch point distances, + /// so a stray far hit cannot inflate the region the planner recomputes. + #[validate(range(min = 0.0, max = 100.0))] + pub region_percentile: f32, } fn validate_health_range(cfg: &Config) -> Result<(), ValidationError> { @@ -83,7 +96,7 @@ impl VoxelMap { .voxels .keys() .copied() - .map(|k| (k, pooled_normal(&self.voxels, k, voxel_size))) + .map(|k| (k, pooled_normal_and_recency(&self.voxels, k, voxel_size).0)) .collect(); for (k, n) in updates { self.voxels.get_mut(&k).unwrap().normal = n; @@ -93,6 +106,7 @@ impl VoxelMap { const NORMAL_MIN_POINTS: u32 = 3; const NORMAL_NEIGHBOR_RADIUS: i32 = 1; +const NEIGHBORHOOD_CAP: usize = (2 * NORMAL_NEIGHBOR_RADIUS as usize + 1).pow(3); const NORMAL_REWEIGHT_ITERS: u32 = 3; /// Neighbor weight falloff with plane distance, as a fraction of voxel size. const NORMAL_PLANE_SIGMA_FRAC: f32 = 0.5; @@ -188,15 +202,17 @@ struct Neighbor { centroid: Vector3, } -/// Find voxel's normal from its neighborhood. -fn pooled_normal( +/// Find voxel's normal and the most recent frame any voxel in its +/// neighborhood was hit, from one scan of the neighborhood. +fn pooled_normal_and_recency( voxels: &AHashMap, key: VoxelKey, voxel_size: f32, -) -> Option> { +) -> (Option>, u32) { let r = NORMAL_NEIGHBOR_RADIUS; - let mut nbs: Vec = Vec::new(); + let mut nbs: ArrayVec = ArrayVec::new(); let mut n_raw: u32 = 0; + let mut recency = 0; for dx in -r..=r { for dy in -r..=r { for dz in -r..=r { @@ -204,6 +220,7 @@ fn pooled_normal( let Some(v) = voxels.get(&nk) else { continue; }; + recency = recency.max(v.last_hit); if v.num_pts == 0 { continue; } @@ -224,12 +241,12 @@ fn pooled_normal( } } if n_raw < NORMAL_MIN_POINTS { - return None; + return (None, recency); } let sigma = NORMAL_PLANE_SIGMA_FRAC * voxel_size; let two_sig2 = 2.0 * sigma * sigma; - let mut weights = vec![1.0_f32; nbs.len()]; + let mut weights = [1.0_f32; NEIGHBORHOOD_CAP]; let mut cov = Matrix3::zeros(); for _ in 0..NORMAL_REWEIGHT_ITERS { let (mut wn, mut s, mut t) = (0.0_f32, Vector3::zeros(), Matrix3::zeros()); @@ -260,25 +277,9 @@ fn pooled_normal( // get rid of planes if we had to discard too many points to get a plane let kept: f32 = nbs.iter().zip(&weights).map(|(nb, &w)| w * nb.n).sum(); if kept < NORMAL_MIN_SUPPORT * n_raw as f32 { - return None; - } - fit_normal(cov) -} - -/// Most recent frame any voxel in the neighborhood was hit. -fn neighborhood_recency(voxels: &AHashMap, key: VoxelKey) -> u32 { - let r = NORMAL_NEIGHBOR_RADIUS; - let mut best = 0; - for dx in -r..=r { - for dy in -r..=r { - for dz in -r..=r { - if let Some(v) = voxels.get(&(key.0 + dx, key.1 + dy, key.2 + dz)) { - best = best.max(v.last_hit); - } - } - } + return (None, recency); } - best + (fit_normal(cov), recency) } /// Refit the cached normal and neighborhood recency of every voxel whose @@ -301,14 +302,11 @@ fn refresh_voxels( } } let updates: Vec<(VoxelKey, Option>, u32)> = dirty - .iter() + .par_iter() .filter(|k| map.voxels.contains_key(k)) .map(|&k| { - ( - k, - pooled_normal(&map.voxels, k, voxel_size), - neighborhood_recency(&map.voxels, k), - ) + let (normal, recency) = pooled_normal_and_recency(&map.voxels, k, voxel_size); + (k, normal, recency) }) .collect(); for (k, n, rec) in updates { @@ -322,16 +320,12 @@ fn refresh_voxels( /// Spare a clearing miss only when a grazing ray skims a recently hit planar /// surface. Stale or voxels with no normal are left to the health checks. fn should_spare( - voxels: &AHashMap, - key: VoxelKey, + c: &Voxel, ray_unit: Vector3, graze_cos: f32, frame: u32, recency_window: u32, ) -> bool { - let Some(c) = voxels.get(&key) else { - return false; - }; match c.normal { Some(n) => { frame.saturating_sub(c.recency) <= recency_window && ray_unit.dot(&n).abs() < graze_cos @@ -359,6 +353,48 @@ impl LocalBounds { } } +/// The local region a batch of frames observed, as (cx, cy, radius, z_min, +/// z_max). A cylinder centered on the mean origin, sized to a percentile of +/// the point distances so a stray far hit cannot inflate it. Points must be +/// finite. An empty batch yields a zero-radius region at the mean origin. +pub fn batch_local_bounds( + points: &[(f32, f32, f32)], + origins: &[(f32, f32, f32)], + percentile_pct: f32, + margin: f32, +) -> (f32, f32, f32, f32, f32) { + let n = origins.len().max(1) as f64; + let cx = (origins.iter().map(|o| o.0 as f64).sum::() / n) as f32; + let cy = (origins.iter().map(|o| o.1 as f64).sum::() / n) as f32; + if points.is_empty() { + let cz = (origins.iter().map(|o| o.2 as f64).sum::() / n) as f32; + return (cx, cy, 0.0, cz, cz); + } + + let mut dist: Vec = points.iter().map(|p| (p.0 - cx).hypot(p.1 - cy)).collect(); + let mut zs: Vec = points.iter().map(|p| p.2).collect(); + let radius = percentile(&mut dist, percentile_pct) + margin; + let z_min = percentile(&mut zs, 100.0 - percentile_pct) - margin; + let z_max = percentile(&mut zs, percentile_pct) + margin; + (cx, cy, radius, z_min, z_max) +} + +fn percentile(values: &mut [f32], p: f32) -> f32 { + let n = values.len(); + if n == 1 { + return values[0]; + } + let rank = (p as f64 / 100.0).clamp(0.0, 1.0) * (n - 1) as f64; + let lo = rank.floor() as usize; + let frac = (rank - lo as f64) as f32; + let (_, &mut v_lo, rest) = values.select_nth_unstable_by(lo, |a, b| a.total_cmp(b)); + if frac == 0.0 || rest.is_empty() { + return v_lo; + } + let v_hi = rest.iter().copied().fold(f32::INFINITY, f32::min); + v_lo + frac * (v_hi - v_lo) +} + pub fn iter_global_points( map: &VoxelMap, voxel_size: f32, @@ -423,35 +459,46 @@ pub fn update_map( let frame = map.frame; let hits = live_voxels(points, cfg.voxel_size); - let mut misses: AHashSet = AHashSet::new(); let origin_voxel = world_to_voxel(origin.0, origin.1, origin.2, inv); let step = cfg.ray_subsample as usize; - for (i, &p) in points.iter().enumerate() { - if i % step != 0 { - continue; - } - let dx = p.0 - origin.0; - let dy = p.1 - origin.1; - let dz = p.2 - origin.2; - if dx * dx + dy * dy + dz * dz > max_range_sq { - continue; - } - let endpoint = world_to_voxel(p.0, p.1, p.2, inv); - find_misses_along_ray( - &mut misses, - &map.voxels, - origin, - p, - cfg.voxel_size, - cfg.shadow_depth, - cfg.grace_depth, - cfg.graze_cos, - frame, - cfg.recency_window, - origin_voxel, - endpoint, - ); - } + let voxels = &map.voxels; + let misses: AHashSet = points + .par_iter() + .enumerate() + .fold(AHashSet::new, |mut misses, (i, &p)| { + if i % step != 0 { + return misses; + } + let dx = p.0 - origin.0; + let dy = p.1 - origin.1; + let dz = p.2 - origin.2; + if dx * dx + dy * dy + dz * dz > max_range_sq { + return misses; + } + let endpoint = world_to_voxel(p.0, p.1, p.2, inv); + find_misses_along_ray( + &mut misses, + voxels, + origin, + p, + cfg.voxel_size, + cfg.shadow_depth, + cfg.grace_depth, + cfg.graze_cos, + frame, + cfg.recency_window, + origin_voxel, + endpoint, + ); + misses + }) + .reduce(AHashSet::new, |mut a, mut b| { + if a.len() < b.len() { + std::mem::swap(&mut a, &mut b); + } + a.extend(b); + a + }); // add new hits for v in &hits { @@ -619,17 +666,10 @@ fn find_misses_along_ray( continue; } - if map_voxels.contains_key(&(x, y, z)) - && !should_spare( - map_voxels, - (x, y, z), - ray_unit, - graze_cos, - frame, - recency_window, - ) - { - misses.insert((x, y, z)); + if let Some(c) = map_voxels.get(&(x, y, z)) { + if !should_spare(c, ray_unit, graze_cos, frame, recency_window) { + misses.insert((x, y, z)); + } } } } @@ -649,6 +689,9 @@ mod tests { max_health: 1, graze_cos: 0.5, recency_window: 60, + emit_every: 1, + global_emit_every: 1, + region_percentile: 95.0, } } @@ -696,6 +739,33 @@ mod tests { assert_eq!(misses, expected); } + #[test] + fn batch_bounds_ignore_far_outlier() { + let origins = [(1.0, 1.0, 0.5), (3.0, 1.0, 0.5)]; + let mut points: Vec<(f32, f32, f32)> = (0..99) + .map(|i| { + let a = i as f32 / 99.0 * std::f32::consts::TAU; + (2.0 + a.cos(), 1.0 + a.sin(), (i % 10) as f32 * 0.1) + }) + .collect(); + points.push((60.0, 1.0, 30.0)); + let (cx, cy, radius, z_min, z_max) = batch_local_bounds(&points, &origins, 95.0, 0.3); + assert_eq!(cx, 2.0); + assert_eq!(cy, 1.0); + assert!(radius < 2.0, "outlier inflated radius to {radius}"); + assert!(z_max < 2.0, "outlier inflated z_max to {z_max}"); + assert!((-0.5..=0.0).contains(&z_min), "z_min out of range: {z_min}"); + } + + #[test] + fn batch_bounds_empty_points_zero_radius() { + let origins = [(1.0, 2.0, 3.0)]; + let (cx, cy, radius, z_min, z_max) = batch_local_bounds(&[], &origins, 95.0, 0.3); + assert_eq!((cx, cy, radius), (1.0, 2.0, 0.0)); + assert_eq!(z_min, 3.0); + assert_eq!(z_max, 3.0); + } + #[test] fn hits_insert_voxels() { let cfg = basic_config(); @@ -808,6 +878,9 @@ mod tests { max_health: 1, graze_cos: 0.5, recency_window: 60, + emit_every: 1, + global_emit_every: 1, + region_percentile: 95.0, }; // Build the floor over a y band so it is a 2d plane, not a wire. let max_x = 25.0_f32; @@ -960,6 +1033,9 @@ mod tests { max_health: 1, graze_cos: 0.5, recency_window: 60, + emit_every: 1, + global_emit_every: 1, + region_percentile: 95.0, }; // Staircase @@ -1031,6 +1107,9 @@ mod tests { max_health: 1, graze_cos: 0.5, recency_window: 60, + emit_every: 1, + global_emit_every: 1, + region_percentile: 95.0, }; // Flat floor from the sensor out to a vertical wall. @@ -1090,6 +1169,9 @@ mod tests { max_health: 1, graze_cos, recency_window: 60, + emit_every: 1, + global_emit_every: 1, + region_percentile: 95.0, }; // Staircase topped by a flat landing and a back wall. @@ -1218,6 +1300,9 @@ mod tests { max_health: 1, graze_cos: 0.5, recency_window, + emit_every: 1, + global_emit_every: 1, + region_percentile: 95.0, }; let (mut map, _) = build_surface(&floor, voxel_size, cfg.max_health); let row: Vec = map diff --git a/dimos/mapping/ray_tracing/test_transformer.py b/dimos/mapping/ray_tracing/test_transformer.py index ce622f738e..b58fcca5f2 100644 --- a/dimos/mapping/ray_tracing/test_transformer.py +++ b/dimos/mapping/ray_tracing/test_transformer.py @@ -50,6 +50,21 @@ def test_emit_every_n_yields_on_cadence_and_flushes_remainder() -> None: assert [r.tags["frame_count"] for r in results] == [3, 6, 7] +def test_negative_emit_every_is_rejected() -> None: + with pytest.raises(ValueError, match="emit_every"): + RayTraceMap(emit_every=-1) + + +def test_pose_propagates_to_emitted_obs() -> None: + pose = (1.5, -2.0, 0.5) + obs = _obs(_cube(), ts=1.0, pose=pose) + + [emitted] = list(RayTraceMap()(iter([obs]))) + + assert emitted.pose_tuple is not None + assert emitted.pose_tuple[:3] == pose + + def test_poseless_obs_are_skipped() -> None: points = _cube() poseless = Observation(id=1, ts=0.0, pose=None, _data=PointCloud2.from_numpy(points)) @@ -60,6 +75,34 @@ def test_poseless_obs_are_skipped() -> None: assert [r.tags["frame_count"] for r in results] == [1] -def test_negative_emit_every_is_rejected() -> None: - with pytest.raises(ValueError): - RayTraceMap(emit_every=-1) +def _ring( + center: tuple[float, float], radius: float, z: float, n: int = 100 +) -> NDArray[np.float32]: + angles = np.linspace(0.0, 2.0 * np.pi, n, endpoint=False) + xs = center[0] + radius * np.cos(angles) + ys = center[1] + radius * np.sin(angles) + zs = np.full_like(xs, z) + return np.stack([xs, ys, zs], axis=1).astype(np.float32) + + +def test_emit_local_tags_region_bounds_around_registered_origin() -> None: + margin = 0.2 + 0.1 + # Sensor-frame ring centered on the sensor; the pose registers it to (2, 3, 0.5). + obs = _obs(_ring((0.0, 0.0), radius=1.0, z=0.0), ts=1.0, pose=(2.0, 3.0, 0.5)) + + [emitted] = list(RayTraceMap(emit_local=True)(iter([obs]))) + + cx, cy, radius, z_min, z_max = emitted.tags["region_bounds"] + assert (cx, cy) == pytest.approx((2.0, 3.0)) + assert radius == pytest.approx(1.0 + margin) + assert z_min == pytest.approx(0.5 - margin) + assert z_max == pytest.approx(0.5 + margin) + + +def test_emit_local_empty_frame_yields_zero_radius_region_at_robot() -> None: + empty = np.empty((0, 3), dtype=np.float32) + obs = _obs(empty, ts=1.0, pose=(1.0, 2.0, 3.0)) + + [emitted] = list(RayTraceMap(emit_local=True)(iter([obs]))) + + assert emitted.tags["region_bounds"] == pytest.approx((1.0, 2.0, 0.0, 3.0, 3.0)) diff --git a/dimos/mapping/ray_tracing/transformer.py b/dimos/mapping/ray_tracing/transformer.py index 847d83a5d3..56302ed687 100644 --- a/dimos/mapping/ray_tracing/transformer.py +++ b/dimos/mapping/ray_tracing/transformer.py @@ -16,21 +16,33 @@ from typing import TYPE_CHECKING, Any +import numpy as np import open3d as o3d # type: ignore[import-untyped] import open3d.core as o3c # type: ignore[import-untyped] -from dimos.mapping.ray_tracing.voxel_map import VoxelRayMapper +from dimos.mapping.ray_tracing.voxel_map import VoxelRayMapper, local_bounds from dimos.memory2.transform import Transformer +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.utils.logging_config import setup_logger if TYPE_CHECKING: from collections.abc import Iterator + from numpy.typing import NDArray + from dimos.memory2.type.observation import Observation +logger = setup_logger() + class RayTraceMap(Transformer[PointCloud2, PointCloud2]): - """Accumulate world-frame lidar into a voxel map with raycast clearing.""" + """Accumulate lidar into a voxel map with raycast clearing. + + Each cloud is sensor-frame and registered into the world by its odometry pose. + """ def __init__( self, @@ -38,6 +50,8 @@ def __init__( voxel_size: float = 0.1, max_range: float = 30.0, emit_every: int = 1, + emit_local: bool = False, + region_percentile: float = 95.0, **mapper_kwargs: Any, ) -> None: if emit_every < 0: @@ -45,16 +59,46 @@ def __init__( self.voxel_size = voxel_size self.max_range = max_range self.emit_every = emit_every + self.emit_local = emit_local + self.region_percentile = region_percentile self._mapper_kwargs = mapper_kwargs + def _local_bounds( + self, + batch_points: list[NDArray[np.float32]], + batch_origins: list[tuple[float, float, float]], + last_obs: Observation[PointCloud2], + ) -> tuple[float, float, float, float, float]: + """Robot-centered cylinder sized to a percentile of the observed points. + + An empty batch yields a zero-radius region at the robot. + """ + if not batch_origins: + pose = last_obs.pose_tuple + assert pose is not None, "poseless obs are skipped upstream" + rx, ry, rz = pose[:3] + return rx, ry, 0.0, rz, rz + + points = np.concatenate(batch_points, axis=0) + origins = np.asarray(batch_origins, dtype=np.float32) + margin = self._mapper_kwargs.get("shadow_depth", 0.2) + self.voxel_size + return local_bounds(points, origins, self.region_percentile, margin) + def _make_obs( self, mapper: VoxelRayMapper, last_obs: Observation[PointCloud2], count: int, + batch_points: list[NDArray[np.float32]], + batch_origins: list[tuple[float, float, float]], ) -> Observation[PointCloud2]: tags = {**last_obs.tags, "frame_count": count} - positions = mapper.global_map() + if self.emit_local: + cx, cy, radius, z_min, z_max = self._local_bounds(batch_points, batch_origins, last_obs) + positions = mapper.local_map((cx, cy, 0.0), radius, z_min, z_max) + tags["region_bounds"] = (cx, cy, radius, z_min, z_max) + else: + positions = mapper.global_map() pcd = o3d.t.geometry.PointCloud() pcd.point["positions"] = o3c.Tensor.from_numpy(positions) cloud = PointCloud2(pointcloud=pcd, frame_id="world", ts=last_obs.ts) @@ -69,17 +113,28 @@ def __call__( ) last_obs: Observation[PointCloud2] | None = None count = 0 + batch_points: list[NDArray[np.float32]] = [] + batch_origins: list[tuple[float, float, float]] = [] for obs in upstream: if obs.pose_tuple is None: + logger.debug("RayTraceMap: obs %s has no pose; skipping", obs.id) continue - x, y, z, *_ = obs.pose_tuple - mapper.add_frame(obs.data.points_f32(), (x, y, z)) + x, y, z, qx, qy, qz, qw = obs.pose_tuple + # Sensor-frame cloud: register into the world by the odom pose. + tf = Transform(translation=Vector3(x, y, z), rotation=Quaternion(qx, qy, qz, qw)) + pts = obs.data.transform(tf).points_f32() + mapper.add_frame(pts, (x, y, z)) + if self.emit_local and pts.size: + batch_points.append(pts) + batch_origins.append((x, y, z)) last_obs = obs count += 1 if self.emit_every > 0 and count % self.emit_every == 0: - yield self._make_obs(mapper, last_obs, count) + yield self._make_obs(mapper, last_obs, count, batch_points, batch_origins) + batch_points = [] + batch_origins = [] if last_obs is not None and (self.emit_every == 0 or count % self.emit_every != 0): - yield self._make_obs(mapper, last_obs, count) + yield self._make_obs(mapper, last_obs, count, batch_points, batch_origins) diff --git a/dimos/mapping/ray_tracing/voxel_map.py b/dimos/mapping/ray_tracing/voxel_map.py index fae8ee61a4..9b052bd3e5 100644 --- a/dimos/mapping/ray_tracing/voxel_map.py +++ b/dimos/mapping/ray_tracing/voxel_map.py @@ -17,7 +17,10 @@ from __future__ import annotations try: - from dimos_voxel_ray_tracing import VoxelRayMapper # noqa: F401 (re-exported) + from dimos_voxel_ray_tracing import ( # noqa: F401 (re-exported) + VoxelRayMapper, + local_bounds, + ) except ImportError as e: raise ImportError( "dimos_voxel_ray_tracing is not built. Run: " diff --git a/dimos/mapping/ray_tracing/voxel_map.pyi b/dimos/mapping/ray_tracing/voxel_map.pyi index 32915ad3fa..278bbde231 100644 --- a/dimos/mapping/ray_tracing/voxel_map.pyi +++ b/dimos/mapping/ray_tracing/voxel_map.pyi @@ -71,4 +71,16 @@ class VoxelRayMapper: def __len__(self) -> int: ... def __repr__(self) -> str: ... -__all__ = ["VoxelRayMapper"] +def local_bounds( + points: NDArray[np.float32], + origins: NDArray[np.float32], + percentile: float, + margin: float, +) -> tuple[float, float, float, float, float]: + """Local region a batch of frames observed, as (cx, cy, radius, z_min, z_max). + + Non-finite points are ignored. + """ + ... + +__all__ = ["VoxelRayMapper", "local_bounds"] diff --git a/pyproject.toml b/pyproject.toml index e4f03e5406..52e79530dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -505,6 +505,7 @@ module = [ "cyclonedds", "cyclonedds.*", "dimos_lcm.*", + "dimos_voxel_ray_tracing", "etils", "faster_whisper", "geometry_msgs.*", From 49c5ab606d5573a113ea949f52dfb88b22bcf572 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 12:53:29 -0700 Subject: [PATCH 35/71] Clean up --- dimos/mapping/ray_tracing/rust/src/main.rs | 4 +++- dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs | 10 ++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/dimos/mapping/ray_tracing/rust/src/main.rs b/dimos/mapping/ray_tracing/rust/src/main.rs index be9db8aba8..2933ff26ab 100644 --- a/dimos/mapping/ray_tracing/rust/src/main.rs +++ b/dimos/mapping/ray_tracing/rust/src/main.rs @@ -28,7 +28,9 @@ struct RayTracingVoxelMap { #[output(encode = PointCloud2::encode)] local_map: Output, - // Region the local_map covers, packed into a PoseStamped. Position holds + // Full region of the local_map, represented as a PoseStamped. We need + // this to update the planner artifacts because local_map only includes + // points that exist (not points that have been removed) Position holds // the cylinder center and orientation holds radius, z_min, z_max, zero. // Stamped identically to local_map so consumers can pair them. #[output(encode = PoseStamped::encode)] diff --git a/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs b/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs index 770412eb58..da28006f82 100644 --- a/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs +++ b/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs @@ -34,8 +34,7 @@ pub struct Config { /// Only spare a voxel whose neighborhood was hit within this many frames. /// A stale voxel can be cleared, even if it's a grazing hit. Large disables it. pub recency_window: u32, - /// Integrate every frame, publish the local map and region bounds every - /// Nth frame. Zero disables them. + /// Publish the accumulated local map and region bounds every Nth frame. Zero disables them. #[validate(range(min = 0))] pub emit_every: u32, /// Publish the global map every Nth frame. Zero disables it. @@ -353,10 +352,9 @@ impl LocalBounds { } } -/// The local region a batch of frames observed, as (cx, cy, radius, z_min, -/// z_max). A cylinder centered on the mean origin, sized to a percentile of -/// the point distances so a stray far hit cannot inflate it. Points must be -/// finite. An empty batch yields a zero-radius region at the mean origin. +/// A cylinder (cx, cy, radius, z_min, z_max) on the mean origin, sized to a +/// percentile of the point distances so a stray far hit cannot inflate it. +/// Points must be finite. An empty batch yields a zero-radius region. pub fn batch_local_bounds( points: &[(f32, f32, f32)], origins: &[(f32, f32, f32)], From 331687ff93b55071c9b152b7d945f2cad7949c4f Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 12:59:26 -0700 Subject: [PATCH 36/71] Speed improvement --- dimos/mapping/ray_tracing/test_transformer.py | 22 +++++++++++++++++++ dimos/mapping/ray_tracing/transformer.py | 9 ++++++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/dimos/mapping/ray_tracing/test_transformer.py b/dimos/mapping/ray_tracing/test_transformer.py index b58fcca5f2..20a5c6b622 100644 --- a/dimos/mapping/ray_tracing/test_transformer.py +++ b/dimos/mapping/ray_tracing/test_transformer.py @@ -106,3 +106,25 @@ def test_emit_local_empty_frame_yields_zero_radius_region_at_robot() -> None: [emitted] = list(RayTraceMap(emit_local=True)(iter([obs]))) assert emitted.tags["region_bounds"] == pytest.approx((1.0, 2.0, 0.0, 3.0, 3.0)) + + +def test_registers_sensor_frame_cloud_by_pose() -> None: + margin = 0.2 + 0.1 + s = 2.0**-0.5 + # 90-degree pitch maps sensor +x to world -z, then translate by (5, 0, 2), + # landing the point at world (5, 0, 1). + point = np.array([[1.0, 0.0, 0.0]], dtype=np.float32) + obs = Observation( + id=0, + ts=1.0, + pose=(5.0, 0.0, 2.0, 0.0, s, 0.0, s), + _data=PointCloud2.from_numpy(point), + ) + + [emitted] = list(RayTraceMap(emit_local=True)(iter([obs]))) + + cx, cy, radius, z_min, z_max = emitted.tags["region_bounds"] + assert (cx, cy) == pytest.approx((5.0, 0.0)) + assert radius == pytest.approx(0.0 + margin) + assert z_min == pytest.approx(1.0 - margin) + assert z_max == pytest.approx(1.0 + margin) diff --git a/dimos/mapping/ray_tracing/transformer.py b/dimos/mapping/ray_tracing/transformer.py index 56302ed687..35cc40dd5f 100644 --- a/dimos/mapping/ray_tracing/transformer.py +++ b/dimos/mapping/ray_tracing/transformer.py @@ -122,8 +122,13 @@ def __call__( continue x, y, z, qx, qy, qz, qw = obs.pose_tuple # Sensor-frame cloud: register into the world by the odom pose. - tf = Transform(translation=Vector3(x, y, z), rotation=Quaternion(qx, qy, qz, qw)) - pts = obs.data.transform(tf).points_f32() + # Apply it to the f32 array directly to skip an Open3D float64 round-trip. + mat = Transform( + translation=Vector3(x, y, z), rotation=Quaternion(qx, qy, qz, qw) + ).to_matrix() + rot = mat[:3, :3].astype(np.float32) + trans = mat[:3, 3].astype(np.float32) + pts = obs.data.points_f32() @ rot.T + trans mapper.add_frame(pts, (x, y, z)) if self.emit_local and pts.size: batch_points.append(pts) From 583e39ae017030132bb35e3ff1799d26c88b657a Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 13:45:36 -0700 Subject: [PATCH 37/71] Incremental planner builds --- dimos/mapping/ray_tracing/module.py | 15 +- .../nav_3d/mls_planner/goal_relay.py | 52 + .../nav_3d/mls_planner/mls_planner.pyi | 40 +- .../nav_3d/mls_planner/mls_planner_native.py | 20 +- .../nav_3d/mls_planner/rust/src/adjacency.rs | 90 +- .../nav_3d/mls_planner/rust/src/dijkstra.rs | 240 ++++- .../nav_3d/mls_planner/rust/src/edges.rs | 207 +++- .../nav_3d/mls_planner/rust/src/main.rs | 324 ++++-- .../mls_planner/rust/src/mls_planner.rs | 988 +++++++++++++++++- .../nav_3d/mls_planner/rust/src/nodes.rs | 494 ++++++++- .../nav_3d/mls_planner/rust/src/planner.rs | 724 +++++++++++-- .../nav_3d/mls_planner/rust/src/python.rs | 135 ++- .../nav_3d/mls_planner/rust/src/surfaces.rs | 138 ++- .../nav_3d/mls_planner/test_transformer.py | 88 +- .../nav_3d/mls_planner/transformer.py | 29 +- .../nav_3d/mls_planner/utils/plan_rrd.py | 293 +++++- pyproject.toml | 1 + 17 files changed, 3414 insertions(+), 464 deletions(-) create mode 100644 dimos/navigation/nav_3d/mls_planner/goal_relay.py diff --git a/dimos/mapping/ray_tracing/module.py b/dimos/mapping/ray_tracing/module.py index a7fcab1d2f..851ff8069d 100644 --- a/dimos/mapping/ray_tracing/module.py +++ b/dimos/mapping/ray_tracing/module.py @@ -44,12 +44,13 @@ class RayTracingVoxelMapConfig(NativeModuleConfig): # Bounds for the health of voxels. Positive health means voxel is occupied. min_health: int = -2 max_health: int = 1 - # Spare a clearing miss when |ray dot surface normal| is below this. + # Don't clear a miss when abs of ray dot normal is below this, clear it when above. + # Higher clears only on direct hits, lower clears on slight grazes too. graze_cos: float = 0.7 # Only spare a voxel whose neighborhood was hit within this many frames. + # A stale voxel can be cleared, even if it's a grazing hit. Large disables it. recency_window: int = 15 - # Integrate every frame, publish the local map and region bounds every - # Nth frame. Zero disables them. + # Publish the accumulated local map and region bounds every Nth frame. Zero disables them. emit_every: int = 1 # Publish the global map every Nth frame. Zero disables it. global_emit_every: int = 1 @@ -59,12 +60,7 @@ class RayTracingVoxelMapConfig(NativeModuleConfig): class RayTracingVoxelMap(NativeModule, mapping.GlobalPointcloud): - """Rust voxel-map module with raycast clearing of dynamic objects. - - region_bounds describes the cylinder local_map covers, packed into a - PoseStamped. Position holds the center. Orientation holds radius, z_min, - z_max, and zero. It shares the local_map stamp. - """ + """Rust voxel-map module with raycast clearing of dynamic objects.""" config: RayTracingVoxelMapConfig @@ -75,6 +71,5 @@ class RayTracingVoxelMap(NativeModule, mapping.GlobalPointcloud): region_bounds: Out[PoseStamped] -# Verify protocol port compliance (mypy will flag missing ports) if TYPE_CHECKING: RayTracingVoxelMap() diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py new file mode 100644 index 0000000000..0fab1cff5c --- /dev/null +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -0,0 +1,52 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from 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.nav_msgs.Odometry import Odometry + + +class GoalRelayConfig(ModuleConfig): + pass + + +class GoalRelay(Module): + """Convert Odometry to PoseStamped""" + + config: GoalRelayConfig + + odometry: In[Odometry] + goal: In[PointStamped] + + start_pose: Out[PoseStamped] + goal_pose: Out[PoseStamped] + + @rpc + def start(self) -> None: + super().start() + self.register_disposable(Disposable(self.odometry.subscribe(self._on_odometry))) + self.register_disposable(Disposable(self.goal.subscribe(self._on_goal))) + + def _on_odometry(self, msg: Odometry) -> None: + self.start_pose.publish(msg.to_pose_stamped()) + + def _on_goal(self, point: PointStamped) -> None: + self.goal_pose.publish(point.to_pose_stamped()) diff --git a/dimos/navigation/nav_3d/mls_planner/mls_planner.pyi b/dimos/navigation/nav_3d/mls_planner/mls_planner.pyi index dfff32d2c8..07e0d2ea3b 100644 --- a/dimos/navigation/nav_3d/mls_planner/mls_planner.pyi +++ b/dimos/navigation/nav_3d/mls_planner/mls_planner.pyi @@ -23,20 +23,44 @@ class MLSPlanner: *, voxel_size: float, robot_height: float, - surface_dilation_passes: int = 3, - surface_erosion_passes: int = 3, + surface_closing_radius: float = 0.15, node_spacing_m: float = 1.0, - node_wall_buffer_m: float = 0.3, - node_step_threshold_m: float = 0.25, + wall_clearance_m: float = 0.3, + wall_buffer_m: float = 0.75, + wall_buffer_weight: float = 100.0, + step_threshold_m: float = 0.25, + step_penalty_weight: float = 4.0, ) -> None: ... def update_global_map(self, points: NDArray[np.float32]) -> None: """Voxelize the map and rebuild surfaces, nodes, and edges. Shape (N, 3) float32.""" ... + def update_region( + self, + points: NDArray[np.float32], + origin: tuple[float, float], + radius: float, + z_min: float, + z_max: float, + ) -> None: + """Replace the cylindrical region with a local map slice and rebuild. + + Points are (N, 3) float32. + """ + ... + def surface_map(self) -> NDArray[np.float32]: """Standable surface cells as (M, 3) float32 centers.""" ... + def surface_clearance_map(self) -> NDArray[np.float32]: + """Surface cells as (M, 4) float32 rows of [x, y, z, clearance]. + + Clearance is the horizontal distance to the nearest untraversable edge. + Unreached cells report +inf. + """ + ... + def nodes(self) -> NDArray[np.float32]: """Graph node positions as (K, 3) float32.""" ... @@ -53,6 +77,14 @@ class MLSPlanner: """Plan a path between start and goal. Returns (W, 3) float32, or None if unreachable.""" ... + def voxel_count(self) -> int: + """Number of occupied voxels in the current map.""" + ... + + def voxel_map(self) -> NDArray[np.float32]: + """Accumulated occupied voxel centers as (N, 3) float32, for visualization.""" + ... + def clear(self) -> None: """Drop the graph and buffered state.""" ... diff --git a/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py b/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py index 1a435fa0cb..edc8bc4c1b 100644 --- a/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py +++ b/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py @@ -34,19 +34,29 @@ class MLSPlannerNativeConfig(NativeModuleConfig): voxel_size: float = 0.1 robot_height: float = 1.5 - surface_dilation_passes: int = 3 - surface_erosion_passes: int = 3 + surface_closing_radius: float = 0.3 node_spacing_m: float = 1.0 - node_wall_buffer_m: float = 0.3 - node_step_threshold_m: float = 0.25 + wall_clearance_m: float = 0.3 + wall_buffer_m: float = 0.75 + wall_buffer_weight: float = 100.0 + step_threshold_m: float = 0.25 + step_penalty_weight: float = 4.0 + goal_tolerance: float = 0.3 + viz_publish_hz: float = 2.0 class MLSPlannerNative(NativeModule): - """Rust-backed MLS planner.""" + """Rust-backed MLS planner. + + Feed either global_map, which rebuilds fully per message, or the local_map + plus region_bounds pair from RayTracingVoxelMap for incremental updates. + """ config: MLSPlannerNativeConfig global_map: In[PointCloud2] + local_map: In[PointCloud2] + region_bounds: In[PoseStamped] start_pose: In[PoseStamped] goal_pose: In[PoseStamped] diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/adjacency.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/adjacency.rs index a4408acae7..baf81839cc 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/adjacency.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/adjacency.rs @@ -6,7 +6,7 @@ //! Uses a "slot map" to store cells. When inserting, either expand the map //! or reuse a freed location marked with a tombstone. -use ahash::AHashMap; +use ahash::{AHashMap, AHashSet}; use rayon::prelude::*; use crate::voxel::VoxelKey; @@ -21,9 +21,20 @@ pub const NO_CELL: CellId = u32::MAX; const TOMBSTONE: VoxelKey = (i32::MIN, i32::MIN, i32::MIN); const NEIGHBORS_4: [(i32, i32); 4] = [(-1, 0), (1, 0), (0, -1), (0, 1)]; +/// Vertical extent of a `dz`-cell change in meters, the step-penalty input. +#[inline] +pub fn rise(dz: i32, voxel_size: f32) -> f32 { + dz.unsigned_abs() as f32 * voxel_size +} + #[derive(Clone, Copy, Debug)] pub struct Edge { pub dest: CellId, + /// Geometric cost, set at build time and never mutated. + pub base_cost: f32, + /// Vertical change of the edge in meters, for the step penalty. + pub rise: f32, + /// base_cost scaled by the wall penalty plus the step penalty. pub cost: f32, } @@ -114,9 +125,19 @@ impl SurfaceCells { &self.edges[id as usize] } + #[inline] + pub fn edges_mut(&mut self, id: CellId) -> &mut [Edge] { + &mut self.edges[id as usize] + } + #[cfg(test)] pub fn add_edge(&mut self, src: CellId, dest: CellId, cost: f32) { - self.edges[src as usize].push(Edge { dest, cost }); + self.edges[src as usize].push(Edge { + dest, + base_cost: cost, + rise: 0.0, + cost, + }); } /// Iterate live cells: (id, outgoing edges). @@ -209,12 +230,75 @@ pub fn build_surface_cells( .get(&(ix + dx, iy + dy, nz)) .expect("neighbor cell exists in lookup"); let cost = ((dx * dx + dy * dy + dz * dz) as f32).sqrt() * voxel_size; - local.push(Edge { dest, cost }); + local.push(Edge { + dest, + base_cost: cost, + rise: rise(dz, voxel_size), + cost, + }); } } }); } +/// Recompute outgoing edges for the seed cells and their surface neighbors, +/// matching build_surface_cells. Call after an incremental insert or remove so +/// the affected region matches a full rebuild. +pub fn rebuild_edges_around( + cells: &mut SurfaceCells, + surface_lookup: &SurfaceLookup, + seeds: &[VoxelKey], + voxel_size: f32, + step_threshold_cells: i32, +) { + let mut affected: AHashSet = AHashSet::new(); + for &(ix, iy, iz) in seeds { + if let Some(id) = cells.id((ix, iy, iz)) { + affected.insert(id); + } + for (dx, dy) in NEIGHBORS_4 { + let Some(nzs) = surface_lookup.get(&(ix + dx, iy + dy)) else { + continue; + }; + for &nz in nzs { + if (nz - iz).abs() > step_threshold_cells { + continue; + } + if let Some(id) = cells.id((ix + dx, iy + dy, nz)) { + affected.insert(id); + } + } + } + } + + for id in affected { + let (ix, iy, iz) = cells.coord(id); + let mut edges: Vec = Vec::new(); + for (dx, dy) in NEIGHBORS_4 { + let Some(nzs) = surface_lookup.get(&(ix + dx, iy + dy)) else { + continue; + }; + for &nz in nzs { + let dz = nz - iz; + if dz.abs() > step_threshold_cells { + continue; + } + let dest = cells + .id((ix + dx, iy + dy, nz)) + .expect("neighbor cell exists in lookup"); + let cost = ((dx * dx + dy * dy + dz * dz) as f32).sqrt() * voxel_size; + edges.push(Edge { + dest, + base_cost: cost, + rise: rise(dz, voxel_size), + cost, + }); + } + } + cells.edges[id as usize] = edges; + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/dijkstra.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/dijkstra.rs index 59a2890c89..76dfc745be 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/dijkstra.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/dijkstra.rs @@ -7,7 +7,10 @@ use std::cmp::Ordering; use std::collections::BinaryHeap; +use ahash::AHashSet; + use crate::adjacency::{CellId, SurfaceCells, NO_CELL}; +use crate::voxel::VoxelKey; #[derive(Default)] pub struct DijkstraState { @@ -28,37 +31,136 @@ impl DijkstraState { self.source.resize(n, 0); self.heap.clear(); } + + /// Grow the vecs to hold `n` slots without disturbing existing labels. + /// New slots default to unreached. + fn ensure_capacity(&mut self, n: usize) { + if self.dist.len() < n { + self.dist.resize(n, f32::INFINITY); + self.pred.resize(n, NO_CELL); + self.source.resize(n, 0); + } + } } -/// Multi-source dijkstra. -/// -/// Labels each node with distance to nearest source, the source id, and the path. -pub fn dijkstra(cells: &SurfaceCells, sources: &[CellId], state: &mut DijkstraState) { +/// Which edge weight a search uses. +#[derive(Clone, Copy)] +pub enum Weight { + /// Geometric distance, for the wall-distance field. + Base, + /// Wall-safe penalized cost, for the node Voronoi. + Penalized, +} + +impl Weight { + #[inline] + fn of(self, edge: &crate::adjacency::Edge) -> f32 { + match self { + Weight::Base => edge.base_cost, + Weight::Penalized => edge.cost, + } + } +} + +/// Multi-source dijkstra labeling each cell with its nearest source and path. +pub fn dijkstra( + cells: &SurfaceCells, + sources: &[CellId], + state: &mut DijkstraState, + weight: Weight, +) { state.reset(cells.slot_capacity()); - for (label, &s) in sources.iter().enumerate() { + for &s in sources { if !cells.is_live(s) { continue; } state.dist[s as usize] = 0.0; - state.source[s as usize] = label as u32; - state.heap.push(Scored(0.0, s)); + state.source[s as usize] = s; + state.heap.push(Scored(0.0, cells.coord(s), s)); } - while let Some(Scored(d, u)) = state.heap.pop() { + while let Some(Scored(d, _, u)) = state.heap.pop() { let cur = state.dist[u as usize]; if d > cur { continue; } let su = state.source[u as usize]; for edge in cells.neighbors(u) { - let nd = d + edge.cost; + let nd = d + weight.of(edge); let v = edge.dest as usize; if nd < state.dist[v] { state.dist[v] = nd; state.pred[v] = u; state.source[v] = su; - state.heap.push(Scored(nd, edge.dest)); + state + .heap + .push(Scored(nd, cells.coord(edge.dest), edge.dest)); + } + } + } +} + +/// Multi-source dijkstra that re-labels only cells in `window`, seeded from +/// in-window sources and the cached frontier just outside it. Correct while the +/// window margin exceeds the reach of the change. +pub fn dijkstra_region( + cells: &SurfaceCells, + sources: &[CellId], + window: &AHashSet, + state: &mut DijkstraState, + weight: Weight, +) { + state.ensure_capacity(cells.slot_capacity()); + state.heap.clear(); + + for &w in window { + let i = w as usize; + state.dist[i] = f32::INFINITY; + state.pred[i] = NO_CELL; + state.source[i] = 0; + } + + for &s in sources { + if !cells.is_live(s) || !window.contains(&s) { + continue; + } + state.dist[s as usize] = 0.0; + state.source[s as usize] = s; + state.heap.push(Scored(0.0, cells.coord(s), s)); + } + + let mut frontier: AHashSet = AHashSet::new(); + for &w in window { + for edge in cells.neighbors(w) { + let n = edge.dest; + if !window.contains(&n) && state.dist[n as usize].is_finite() { + frontier.insert(n); + } + } + } + for &n in &frontier { + state + .heap + .push(Scored(state.dist[n as usize], cells.coord(n), n)); + } + + while let Some(Scored(d, _, u)) = state.heap.pop() { + if d > state.dist[u as usize] { + continue; + } + let su = state.source[u as usize]; + for edge in cells.neighbors(u) { + let v = edge.dest; + if !window.contains(&v) { + continue; + } + let nd = d + weight.of(edge); + if nd < state.dist[v as usize] { + state.dist[v as usize] = nd; + state.pred[v as usize] = u; + state.source[v as usize] = su; + state.heap.push(Scored(nd, cells.coord(v), v)); } } } @@ -70,9 +172,11 @@ pub fn dijkstra(cells: &SurfaceCells, sources: &[CellId], state: &mut DijkstraSt pub fn walk_preds(state: &DijkstraState, start: CellId) -> Vec { let mut cells = vec![start]; let mut cur = start; + let mut seen: AHashSet = AHashSet::new(); + seen.insert(start); loop { let p = state.pred[cur as usize]; - if p == NO_CELL { + if p == NO_CELL || !seen.insert(p) { break; } cur = p; @@ -81,7 +185,7 @@ pub fn walk_preds(state: &DijkstraState, start: CellId) -> Vec { cells } -struct Scored(f32, CellId); +struct Scored(f32, VoxelKey, CellId); impl PartialEq for Scored { fn eq(&self, other: &Self) -> bool { @@ -104,7 +208,95 @@ impl Ord for Scored { #[cfg(test)] mod tests { use super::*; - use crate::adjacency::SurfaceCells; + use crate::adjacency::{ + build_surface_cells, build_surface_lookup, SurfaceCells, SurfaceLookup, + }; + use crate::voxel::VoxelKey; + + fn grid(n: i32) -> SurfaceCells { + let cells: Vec = (0..n) + .flat_map(|x| (0..n).map(move |y| (x, y, 0))) + .collect(); + let mut lookup = SurfaceLookup::new(); + build_surface_lookup(&cells, &mut lookup); + let mut sc = SurfaceCells::default(); + build_surface_cells(&mut sc, &lookup, 0.1, 2); + sc + } + + fn root_of(state: &DijkstraState, start: CellId) -> CellId { + let mut cur = start; + while state.pred[cur as usize] != NO_CELL { + cur = state.pred[cur as usize]; + } + cur + } + + #[test] + fn region_window_all_equals_full() { + let sc = grid(10); + let sources = [sc.id((0, 0, 0)).unwrap(), sc.id((9, 9, 0)).unwrap()]; + + let mut full = DijkstraState::default(); + dijkstra(&sc, &sources, &mut full, Weight::Penalized); + + let window: AHashSet = sc.ids().collect(); + let mut region = DijkstraState::default(); + dijkstra_region(&sc, &sources, &window, &mut region, Weight::Penalized); + + for id in sc.ids() { + assert_eq!( + region.dist[id as usize], + full.dist[id as usize], + "dist mismatch at {:?}", + sc.coord(id) + ); + } + } + + #[test] + fn region_partial_window_reproduces_cached_distances() { + let sc = grid(12); + let sources = [sc.id((0, 0, 0)).unwrap(), sc.id((11, 11, 0)).unwrap()]; + + let mut full = DijkstraState::default(); + dijkstra(&sc, &sources, &mut full, Weight::Penalized); + + // Seed the regional state with the full result as the cache, then + // recompute an interior block. Nothing changed, so the block must come + // back identical and every cell must still trace to a real source. + let mut region = DijkstraState { + dist: full.dist.clone(), + pred: full.pred.clone(), + source: full.source.clone(), + ..Default::default() + }; + + let window: AHashSet = sc + .ids() + .filter(|&id| { + let (x, y, _) = sc.coord(id); + (3..=8).contains(&x) && (3..=8).contains(&y) + }) + .collect(); + dijkstra_region(&sc, &sources, &window, &mut region, Weight::Penalized); + + for &id in &window { + assert_eq!( + region.dist[id as usize], + full.dist[id as usize], + "dist mismatch at {:?}", + sc.coord(id) + ); + let root = root_of(®ion, id); + assert!( + sources.contains(&root), + "cell {:?} traces to non-source {:?}", + sc.coord(id), + sc.coord(root) + ); + } + } fn chain(n: i32) -> (SurfaceCells, Vec) { let mut sc = SurfaceCells::default(); @@ -120,7 +312,7 @@ mod tests { fn single_source_dist_and_pred() { let (sc, ids) = chain(5); let mut st = DijkstraState::default(); - dijkstra(&sc, &[ids[0]], &mut st); + dijkstra(&sc, &[ids[0]], &mut st, Weight::Penalized); for (i, &id) in ids.iter().enumerate().take(5) { assert_eq!(st.dist[id as usize], i as f32); assert_eq!(st.source[id as usize], 0); @@ -140,13 +332,13 @@ mod tests { fn multi_source_labels_by_nearest() { let (sc, ids) = chain(5); let mut st = DijkstraState::default(); - dijkstra(&sc, &[ids[0], ids[4]], &mut st); - assert_eq!(st.source[ids[0] as usize], 0); - assert_eq!(st.source[ids[1] as usize], 0); - assert_eq!(st.source[ids[3] as usize], 1); - assert_eq!(st.source[ids[4] as usize], 1); + dijkstra(&sc, &[ids[0], ids[4]], &mut st, Weight::Penalized); + assert_eq!(st.source[ids[0] as usize], ids[0]); + assert_eq!(st.source[ids[1] as usize], ids[0]); + assert_eq!(st.source[ids[3] as usize], ids[4]); + assert_eq!(st.source[ids[4] as usize], ids[4]); let s2 = st.source[ids[2] as usize]; - assert!(s2 == 0 || s2 == 1); + assert!(s2 == ids[0] || s2 == ids[4]); assert_eq!(st.dist[ids[0] as usize], 0.0); assert_eq!(st.dist[ids[1] as usize], 1.0); assert_eq!(st.dist[ids[2] as usize], 2.0); @@ -166,7 +358,7 @@ mod tests { sc.add_edge(c, d, 1.0); sc.add_edge(d, c, 1.0); let mut st = DijkstraState::default(); - dijkstra(&sc, &[a], &mut st); + dijkstra(&sc, &[a], &mut st, Weight::Penalized); assert_eq!(st.dist[a as usize], 0.0); assert_eq!(st.dist[b as usize], 1.0); assert!(!st.dist[c as usize].is_finite()); @@ -186,7 +378,7 @@ mod tests { sc.add_edge(c, b, 1.0); sc.add_edge(b, c, 1.0); let mut st = DijkstraState::default(); - dijkstra(&sc, &[a], &mut st); + dijkstra(&sc, &[a], &mut st, Weight::Penalized); assert_eq!(st.dist[b as usize], 2.0); assert_eq!(st.pred[b as usize], c); } @@ -195,9 +387,9 @@ mod tests { fn buffer_reuse_does_not_leak_prior_state() { let (sc1, ids1) = chain(5); let mut st = DijkstraState::default(); - dijkstra(&sc1, &[ids1[0]], &mut st); + dijkstra(&sc1, &[ids1[0]], &mut st, Weight::Penalized); let (sc2, ids2) = chain(3); - dijkstra(&sc2, &[ids2[0]], &mut st); + dijkstra(&sc2, &[ids2[0]], &mut st, Weight::Penalized); for (i, &id) in ids2.iter().enumerate().take(3) { assert_eq!(st.dist[id as usize], i as f32); } diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs index 8884b96af0..22ccb8dbd2 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs @@ -8,17 +8,18 @@ //! the Voronoi region. We use the boundaries of these regions to build the //! edges between start nodes. -use ahash::AHashMap; +use ahash::{AHashMap, AHashSet}; use rayon::prelude::*; -use crate::adjacency::{CellId, Edge, SurfaceCells, SurfaceLookup, NO_CELL}; -use crate::dijkstra::{dijkstra, walk_preds, DijkstraState}; +use crate::adjacency::{CellId, SurfaceCells, SurfaceLookup, NO_CELL}; +use crate::dijkstra::{dijkstra, dijkstra_region, walk_preds, DijkstraState, Weight}; use crate::nodes::NodeData; use crate::voxel::VoxelKey; -/// Index into planner graph nodes -pub type NodeId = u32; -pub const NO_NODE: NodeId = u32::MAX; +/// A node is identified by the CellId it sits on, stable across incremental +/// updates so cached edges and the Voronoi forest survive a regional rebuild. +pub type NodeId = CellId; +pub const NO_NODE: NodeId = NO_CELL; /// Index into planner graph node edges pub type NodeEdgeIdx = u32; @@ -40,8 +41,13 @@ pub struct PlannerGraph { pub surface_lookup: SurfaceLookup, pub nodes: Vec, pub node_edges: Vec, - pub node_adj: Vec>, + pub node_adj: AHashMap>, + /// Each cell's nearest node and the predecessor back toward it. The planner + /// walks these to expand a node-to-node edge into its cell path. pub cell_state: DijkstraState, + /// Each cell's distance to the nearest wall, recomputed only in the changed + /// window and reused elsewhere. + pub wall_state: DijkstraState, } impl PlannerGraph { @@ -59,7 +65,7 @@ pub fn build_node_edges( nodes: &[NodeData], state: &mut DijkstraState, out_edges: &mut Vec, - out_adj: &mut Vec>, + out_adj: &mut AHashMap>, ) { out_edges.clear(); out_adj.clear(); @@ -70,34 +76,94 @@ pub fn build_node_edges( } let source_cells: Vec = nodes.iter().map(|n| n.cell_id).collect(); - dijkstra(cells, &source_cells, state); + dijkstra(cells, &source_cells, state, Weight::Penalized); best_boundary_edges(cells, state, out_edges); - out_adj.resize_with(nodes.len(), Vec::new); - for v in out_adj.iter_mut() { - v.clear(); + rebuild_node_adj(out_edges, out_adj); +} + +/// Rebuild the per-node edge index from the edge list. +fn rebuild_node_adj(edges: &[NodeEdge], out_adj: &mut AHashMap>) { + out_adj.clear(); + for (edge_idx, edge) in edges.iter().enumerate() { + out_adj + .entry(edge.a) + .or_default() + .push(edge_idx as NodeEdgeIdx); + out_adj + .entry(edge.b) + .or_default() + .push(edge_idx as NodeEdgeIdx); + } +} + +/// Incremental build_node_edges. Only the cells in the window changed, so redo +/// the Voronoi there, keep cached edges whose boundary lies outside the window, +/// and rescan the window for new node-to-node crossings. +pub fn build_node_edges_region( + cells: &SurfaceCells, + nodes: &[NodeData], + window: &AHashSet, + state: &mut DijkstraState, + out_edges: &mut Vec, + out_adj: &mut AHashMap>, +) { + let source_cells: Vec = nodes.iter().map(|n| n.cell_id).collect(); + if source_cells.is_empty() { + state.reset(cells.slot_capacity()); + out_edges.clear(); + out_adj.clear(); + return; } - for (edge_idx, edge) in out_edges.iter().enumerate() { - out_adj[edge.a as usize].push(edge_idx as NodeEdgeIdx); - out_adj[edge.b as usize].push(edge_idx as NodeEdgeIdx); + dijkstra_region(cells, &source_cells, window, state, Weight::Penalized); + + let live_node: AHashSet = source_cells.iter().copied().collect(); + + let mut merged: AHashMap<(NodeId, NodeId), NodeEdge> = AHashMap::new(); + for e in out_edges.iter() { + if window.contains(&e.boundary_u) || window.contains(&e.boundary_v) { + continue; + } + if !live_node.contains(&e.a) || !live_node.contains(&e.b) { + continue; + } + merged.insert((e.a, e.b), *e); } + + let win_cells: Vec = window.iter().copied().collect(); + merge_min(&mut merged, boundary_edge_map(cells, state, &win_cells)); + + out_edges.clear(); + out_edges.extend(merged.into_values()); + out_edges.par_sort_unstable_by_key(|e| (e.a, e.b)); + rebuild_node_adj(out_edges, out_adj); } fn best_boundary_edges(cells: &SurfaceCells, state: &DijkstraState, out: &mut Vec) { - let cell_entries: Vec<(CellId, &[Edge])> = cells.iter().collect(); + let scan: Vec = cells.ids().collect(); + let merged = boundary_edge_map(cells, state, &scan); + out.clear(); + out.extend(merged.into_values()); + out.par_sort_unstable_by_key(|e| (e.a, e.b)); +} - let merged: AHashMap<(NodeId, NodeId), NodeEdge> = cell_entries - .par_iter() +/// Cheapest Voronoi-boundary crossing per adjacent node pair over the scanned cells. +fn boundary_edge_map( + cells: &SurfaceCells, + state: &DijkstraState, + scan: &[CellId], +) -> AHashMap<(NodeId, NodeId), NodeEdge> { + scan.par_iter() .fold( AHashMap::<(NodeId, NodeId), NodeEdge>::new, - |mut local, (u, edges)| { - let du = state.dist[*u as usize]; + |mut local, &u| { + let du = state.dist[u as usize]; if !du.is_finite() { return local; } - let sa = state.source[*u as usize]; - for edge in *edges { + let sa = state.source[u as usize]; + for edge in cells.neighbors(u) { let v = edge.dest; let dv = state.dist[v as usize]; if !dv.is_finite() { @@ -108,13 +174,15 @@ fn best_boundary_edges(cells: &SurfaceCells, state: &DijkstraState, out: &mut Ve continue; } let cost = du + edge.cost + dv; - + // Skip impassable crossings + if !cost.is_finite() { + continue; + } let (key_a, key_b, bu, bv) = if sa < sb { - (sa, sb, *u, v) + (sa, sb, u, v) } else { - (sb, sa, v, *u) + (sb, sa, v, u) }; - let entry = local.entry((key_a, key_b)).or_insert(NodeEdge { a: key_a, b: key_b, @@ -131,19 +199,23 @@ fn best_boundary_edges(cells: &SurfaceCells, state: &DijkstraState, out: &mut Ve local }, ) - .reduce(AHashMap::<(NodeId, NodeId), NodeEdge>::new, |mut a, b| { - for (k, v_edge) in b { - let entry = a.entry(k).or_insert(v_edge); - if v_edge.cost < entry.cost { - *entry = v_edge; - } - } + .reduce(AHashMap::new, |mut a, b| { + merge_min(&mut a, b); a - }); + }) +} - out.clear(); - out.extend(merged.into_values()); - out.par_sort_unstable_by_key(|e| (e.a, e.b)); +/// Keep the lower-cost edge for each node pair when merging two maps. +fn merge_min( + into: &mut AHashMap<(NodeId, NodeId), NodeEdge>, + from: AHashMap<(NodeId, NodeId), NodeEdge>, +) { + for (k, edge) in from { + let entry = into.entry(k).or_insert(edge); + if edge.cost < entry.cost { + *entry = edge; + } + } } /// Walk every node-graph edge and emit one segment per consecutive cell @@ -211,16 +283,69 @@ mod tests { let pg = setup(&strip_cells(), &[(3, 0, 0), (15, 0, 0)]); assert_eq!(pg.node_edges.len(), 1); let e = &pg.node_edges[0]; - assert_eq!((e.a, e.b), (0, 1)); - assert_eq!(pg.node_adj[0], vec![0]); - assert_eq!(pg.node_adj[1], vec![0]); + let a = pg.cells.id((3, 0, 0)).unwrap(); + let b = pg.cells.id((15, 0, 0)).unwrap(); + assert_eq!((e.a.min(e.b), e.a.max(e.b)), (a.min(b), a.max(b))); + assert_eq!(pg.node_adj[&a], vec![0]); + assert_eq!(pg.node_adj[&b], vec![0]); } #[test] fn three_nodes_in_line_form_a_chain() { let pg = setup(&strip_cells(), &[(3, 0, 0), (10, 0, 0), (17, 0, 0)]); + let c = |k| pg.cells.id(k).unwrap(); let pairs: Vec<(NodeId, NodeId)> = pg.node_edges.iter().map(|e| (e.a, e.b)).collect(); - assert_eq!(pairs, vec![(0, 1), (1, 2)]); + assert_eq!( + pairs, + vec![ + (c((3, 0, 0)), c((10, 0, 0))), + (c((10, 0, 0)), c((17, 0, 0))) + ] + ); + } + + #[test] + fn infinite_crossing_is_not_an_edge() { + // The only crossing between the two nodes is impassable, so no edge. + let surface: Vec = (0..6).map(|x| (x, 0, 0)).collect(); + let mut plg = PlannerGraph::new(); + build_surface_lookup(&surface, &mut plg.surface_lookup); + build_surface_cells(&mut plg.cells, &plg.surface_lookup, VOXEL, 2); + + let c2 = plg.cells.id((2, 0, 0)).unwrap(); + let c3 = plg.cells.id((3, 0, 0)).unwrap(); + for e in plg.cells.edges_mut(c2) { + if e.dest == c3 { + e.cost = f32::INFINITY; + } + } + for e in plg.cells.edges_mut(c3) { + if e.dest == c2 { + e.cost = f32::INFINITY; + } + } + + plg.nodes = [(0, 0, 0), (5, 0, 0)] + .iter() + .map(|&c| NodeData { + cell_id: plg.cells.id(c).unwrap(), + pos: surface_point_xyz(c.0, c.1, c.2, VOXEL), + }) + .collect(); + build_node_edges( + &plg.cells, + &plg.nodes, + &mut plg.cell_state, + &mut plg.node_edges, + &mut plg.node_adj, + ); + + assert!( + plg.node_edges.is_empty(), + "an infinite crossing is not an edge" + ); + // Walking boundaries must not panic on an unset boundary cell. + edges_to_segments(&plg.cells, &plg.cell_state, &plg.node_edges); } #[test] diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs index c26f62e954..91452ad1f0 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs @@ -1,23 +1,51 @@ // Copyright 2026 Dimensional Inc. // SPDX-License-Identifier: Apache-2.0 +use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use dimos_mls_planner::edges::{edges_to_segments, PlannerGraph}; -use dimos_mls_planner::mls_planner::{Config, Planner}; +use dimos_mls_planner::mls_planner::{Config, Planner, RegionBounds}; use dimos_mls_planner::voxel::surface_point_xyz; use dimos_module::{error_throttled, run, warn_throttled, Input, LcmTransport, Module, Output}; use lcm_msgs::geometry_msgs::{Point, Pose, PoseStamped, Quaternion}; use lcm_msgs::nav_msgs::Path; use lcm_msgs::sensor_msgs::{PointCloud2, PointField}; use lcm_msgs::std_msgs::{Header, Time}; -use tracing::info; +use tokio::sync::Notify; +use tracing::debug; + +/// A point in the planner's world frame. +type Xyz = (f32, f32, f32); + +/// State shared between the handle loop and the worker. The handle loop writes +/// the newest value, the worker reads it. +type Shared = Arc>>; + +/// A map input handed from the handle loop to the worker. Only the newest is +/// kept, so a dropped intermediate frame is harmless. +enum MapUpdate { + Region { + cloud: PointCloud2, + bounds: PoseStamped, + }, + Global { + cloud: PointCloud2, + }, +} #[derive(Module)] +#[module(setup = spawn_worker)] struct MlsPlanner { #[input(decode = PointCloud2::decode, handler = on_global_map)] global_map: Input, + #[input(decode = PointCloud2::decode, handler = on_local_map)] + local_map: Input, + + #[input(decode = PoseStamped::decode, handler = on_region_bounds)] + region_bounds: Input, + #[input(decode = PoseStamped::decode, handler = on_start_pose)] start_pose: Input, @@ -39,79 +67,236 @@ struct MlsPlanner { #[config] config: Config, - planner: Planner, - latest_start: Option<(f32, f32, f32)>, + // These live only on the handle loop. We hold a local map and its bounds + // here until their stamps match, then hand the paired snapshot to the worker. + pending_local: Option, + pending_bounds: Option, + + // Shared with the worker task. The handle loop only writes these and wakes + // the worker, so it never blocks on the heavy map processing. + pending: Shared, + latest_start: Shared, + active_goal: Shared, + wake: Arc, } impl MlsPlanner { + /// Spawn the background task that handles map updates and planning. + async fn spawn_worker(&mut self) { + let worker = Worker { + pending: Arc::clone(&self.pending), + latest_start: Arc::clone(&self.latest_start), + active_goal: Arc::clone(&self.active_goal), + wake: Arc::clone(&self.wake), + config: self.config.clone(), + surface_map: self.surface_map.clone(), + nodes: self.nodes.clone(), + node_edges: self.node_edges.clone(), + path: self.path.clone(), + }; + tokio::spawn(worker.run()); + } + async fn on_global_map(&mut self, msg: PointCloud2) { - let points = match extract_xyz(&msg) { - Ok(p) => p, - Err(e) => { - warn_throttled!( - Duration::from_secs(1), - error = %e, - "Failed to extract lidar points, dropped a cloud.", - ); - return; - } + self.hand_off(MapUpdate::Global { cloud: msg }); + } + + async fn on_local_map(&mut self, msg: PointCloud2) { + self.pending_local = Some(msg); + self.try_pair(); + } + + async fn on_region_bounds(&mut self, msg: PoseStamped) { + self.pending_bounds = Some(msg); + self.try_pair(); + } + + /// Hand a local map and its stamp-matching bounds to the worker once both + /// are in hand. + fn try_pair(&mut self) { + let (Some(bounds_msg), Some(cloud)) = (&self.pending_bounds, &self.pending_local) else { + return; }; - if points.is_empty() { + if !same_stamp(&bounds_msg.header.stamp, &cloud.header.stamp) { return; } + let bounds = self.pending_bounds.take().expect("checked above"); + let cloud = self.pending_local.take().expect("checked above"); + self.hand_off(MapUpdate::Region { cloud, bounds }); + } - let t = Instant::now(); - self.planner.update_global_map(&points, &self.config); - let rebuild_ms = ms(t.elapsed()); + /// Store the newest map input and wake the worker. + fn hand_off(&self, update: MapUpdate) { + *self.pending.lock().expect("pending mutex") = Some(update); + self.wake.notify_one(); + } + + /// Record the latest start pose. The worker reads it when it replans. No + /// wake here, so odometry never drives replanning. + async fn on_start_pose(&mut self, msg: PoseStamped) { + let p = &msg.pose.position; + *self.latest_start.lock().expect("start mutex") = + Some((p.x as f32, p.y as f32, p.z as f32)); + } + + /// Set the active goal, or clear it when the goal is non-finite, which is + /// the cancel signal. Wake the worker so a new click replans right away. + async fn on_goal_pose(&mut self, msg: PoseStamped) { + let p = &msg.pose.position; + let goal = (p.x as f32, p.y as f32, p.z as f32); + *self.active_goal.lock().expect("goal mutex") = + (goal.0.is_finite() && goal.1.is_finite() && goal.2.is_finite()).then_some(goal); + self.wake.notify_one(); + } +} +/// Owns the planner graph and does every map mutation, graph publish, and +/// replan off the handle loop. Woken by the handlers via `wake`. +struct Worker { + pending: Shared, + latest_start: Shared, + active_goal: Shared, + wake: Arc, + config: Config, + surface_map: Output, + nodes: Output, + node_edges: Output, + path: Output, +} + +impl Worker { + async fn run(self) { + let mut planner = Planner::default(); + let mut last_path_at: Option = None; + let mut last_viz_at: Option = None; + loop { + self.wake.notified().await; + // Take the newest map input, dropping any intermediates. + let update = self.pending.lock().expect("pending mutex").take(); + if let Some(update) = update { + self.apply_update(&mut planner, update, &mut last_viz_at) + .await; + } + self.maybe_replan(&mut planner, &mut last_path_at).await; + } + } + + /// Update the graph every cycle so planning sees fresh surfaces, then publish + /// the viz clouds rate-capped to viz_publish_hz, since building them is costly + /// and nothing on the planning path reads them. + async fn apply_update( + &self, + planner: &mut Planner, + update: MapUpdate, + last_viz_at: &mut Option, + ) { + let now = Instant::now(); + let viz_due = self.config.viz_publish_hz > 0.0 && { + let viz_interval = Duration::from_secs_f32(1.0 / self.config.viz_publish_hz); + last_viz_at.is_none_or(|t| now.duration_since(t) >= viz_interval) + }; + + let messages = tokio::task::block_in_place(|| { + let updated = self.ingest(planner, update); + (updated && viz_due).then(|| self.build_graph_messages(planner)) + }); + + if let Some((surface, node_cloud, edges)) = messages { + publish_cloud(&self.surface_map, &surface).await; + publish_cloud(&self.nodes, &node_cloud).await; + publish_path(&self.node_edges, &edges).await; + *last_viz_at = Some(now); + } + } + + /// Mutate the graph from a map update. Returns false if the cloud was + /// unusable. + fn ingest(&self, planner: &mut Planner, update: MapUpdate) -> bool { + match update { + MapUpdate::Region { cloud, bounds } => { + let points = match extract_xyz(&cloud) { + Ok(p) => p, + Err(e) => { + warn_throttled!( + Duration::from_secs(1), + error = %e, + "Failed to extract local map points, dropped a region update.", + ); + return false; + } + }; + let bounds = RegionBounds { + origin_x: bounds.pose.position.x as f32, + origin_y: bounds.pose.position.y as f32, + radius: bounds.pose.orientation.x as f32, + z_min: bounds.pose.orientation.y as f32, + z_max: bounds.pose.orientation.z as f32, + }; + + let update_start = Instant::now(); + planner.update_region(&points, &bounds, &self.config); + debug!( + update_ms = update_start.elapsed().as_secs_f64() * 1e3, + local_points = points.len(), + "local region processed" + ); + } + MapUpdate::Global { cloud } => { + let points = match extract_xyz(&cloud) { + Ok(p) => p, + Err(e) => { + warn_throttled!( + Duration::from_secs(1), + error = %e, + "Failed to extract lidar points, dropped a cloud.", + ); + return false; + } + }; + if points.is_empty() { + return false; + } + planner.update_global_map(&points, &self.config); + debug!(global_map_points = points.len(), "global_map processed"); + } + } + true + } + + fn build_graph_messages(&self, planner: &Planner) -> (PointCloud2, PointCloud2, Path) { let voxel_size = self.config.voxel_size; let frame = &self.config.world_frame; - let graph = self.planner.graph(); + let graph = planner.graph(); - let surface_points: Vec<(f32, f32, f32)> = self - .planner + let surface_points: Vec = planner .surface() - .iter() - .map(|&(ix, iy, iz)| surface_point_xyz(ix, iy, iz, voxel_size)) + .map(|(ix, iy, iz)| surface_point_xyz(ix, iy, iz, voxel_size)) .collect(); - publish_cloud(&self.surface_map, &surface_points, frame, now()).await; - - let node_points: Vec<(f32, f32, f32)> = graph.nodes.iter().map(|n| n.pos).collect(); - publish_cloud(&self.nodes, &node_points, frame, now()).await; - - let edges_path = build_segments_path(graph, voxel_size, frame, now()); - publish_path(&self.node_edges, &edges_path).await; - - info!( - global_map_points = points.len(), - voxels = self.planner.voxel_count(), - surface_cells = self.planner.surface().len(), - nodes = graph.nodes.len(), - edges = graph.node_edges.len(), - rebuild_ms, - "global_map processed", - ); - } + let surface = build_pc2_xyz(&surface_points, frame, now()); - async fn on_start_pose(&mut self, msg: PoseStamped) { - let p = &msg.pose.position; - self.latest_start = Some((p.x as f32, p.y as f32, p.z as f32)); - // Drop any previous plan so the visualizer doesn't show a stale path - // rooted at the old start. - publish_path(&self.path, &empty_path(&self.config.world_frame, now())).await; + let node_points: Vec = graph.nodes.iter().map(|n| n.pos).collect(); + let node_cloud = build_pc2_xyz(&node_points, frame, now()); + + let edges = build_segments_path(graph, voxel_size, frame, now()); + (surface, node_cloud, edges) } - async fn on_goal_pose(&mut self, msg: PoseStamped) { - let Some(start) = self.latest_start else { - tracing::warn!("MLSPlanner received goal before start; skipping"); + /// Replan from the latest start to the active goal. This gates and does the + /// IO. The planning itself lives in Planner::plan. + async fn maybe_replan(&self, planner: &mut Planner, last_path_at: &mut Option) { + let start = *self.latest_start.lock().expect("start mutex"); + let goal = *self.active_goal.lock().expect("goal mutex"); + let (Some(start), Some(goal)) = (start, goal) else { return; }; + if is_at_goal(start, goal, self.config.goal_tolerance) { + *self.active_goal.lock().expect("goal mutex") = None; + return; + } - let p = &msg.pose.position; - let goal = (p.x as f32, p.y as f32, p.z as f32); - - let t_plan = Instant::now(); - let waypoints = match self.planner.plan(start, goal, &self.config) { + let plan_start = Instant::now(); + let waypoints = tokio::task::block_in_place(|| planner.plan(start, goal, &self.config)); + let waypoints = match waypoints { Some(wp) => wp, None => { tracing::warn!(?start, ?goal, "no path between start and goal"); @@ -119,27 +304,32 @@ impl MlsPlanner { return; } }; - let plan_ms = ms(t_plan.elapsed()); + let plan_ms = plan_start.elapsed().as_secs_f64() * 1e3; + let produced = Instant::now(); + let since_last_ms = last_path_at.map_or(-1.0, |t| (produced - t).as_secs_f64() * 1e3); + *last_path_at = Some(produced); let stamp = now(); let path_msg = build_path_from_waypoints(&waypoints, &self.config.world_frame, stamp); - info!(waypoints = waypoints.len(), plan_ms, "path planned"); + debug!( + waypoints = waypoints.len(), + plan_ms, since_last_ms, "path planned" + ); publish_path(&self.path, &path_msg).await; } } -fn ms(d: Duration) -> f64 { - d.as_secs_f64() * 1000.0 +/// True when start is within `tol` of goal in the ground plane. +fn is_at_goal(start: Xyz, goal: Xyz, tol: f32) -> bool { + (start.0 - goal.0).hypot(start.1 - goal.1) < tol +} + +fn same_stamp(a: &Time, b: &Time) -> bool { + a.sec == b.sec && a.nsec == b.nsec } -async fn publish_cloud( - out: &Output, - points: &[(f32, f32, f32)], - frame_id: &str, - stamp: Time, -) { - let cloud = build_pc2_xyz(points, frame_id, stamp); - if let Err(e) = out.publish(&cloud).await { +async fn publish_cloud(out: &Output, cloud: &PointCloud2) { + if let Err(e) = out.publish(cloud).await { error_throttled!( Duration::from_secs(1), error = %e, diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs index 8e72f1517c..0d9c41e739 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs @@ -5,31 +5,106 @@ use ahash::AHashSet; use dimos_module::native_config; +use rayon::prelude::*; +use validator::ValidationError; -use crate::adjacency::{build_surface_cells, build_surface_lookup}; -use crate::edges::{build_node_edges, PlannerGraph}; -use crate::nodes::place_nodes; +use crate::adjacency::{build_surface_cells, build_surface_lookup, rebuild_edges_around, CellId}; +use crate::edges::{build_node_edges, build_node_edges_region, PlannerGraph}; +use crate::nodes::{place_nodes, place_nodes_region}; use crate::planner; -use crate::surfaces::{extract_surfaces, ColumnIz}; +use crate::surfaces::{ + add_to_by_col, extract_surfaces, extract_surfaces_region, remove_from_by_col, ColumnIz, +}; use crate::voxel::{voxelize, VoxelKey}; #[native_config] +#[derive(Clone)] +#[validate(schema(function = "validate_wall_buffer"))] pub struct Config { pub world_frame: String, #[validate(range(exclusive_min = 0.0))] pub voxel_size: f32, #[validate(range(exclusive_min = 0.0))] pub robot_height: f32, - #[validate(range(min = 0))] - pub surface_dilation_passes: u32, - #[validate(range(min = 0))] - pub surface_erosion_passes: u32, + /// Radius in meters of the morphological closing that fills small holes in + /// the extracted surface. Fills holes up to twice this wide. + #[validate(range(min = 0.0))] + pub surface_closing_radius: f32, #[validate(range(exclusive_min = 0.0))] pub node_spacing_m: f32, + /// Hard clearance. Cells closer than this to a wall or edge are impassable. + #[validate(range(min = 0.0))] + pub wall_clearance_m: f32, + /// Width of the soft standoff zone beyond the clearance. Paths prefer to stay + /// clearance + buffer from walls. + #[validate(range(min = 0.0))] + pub wall_buffer_m: f32, + /// Peak soft wall penalty at the clearance edge: the cost multiplier there is + /// 1 + this, decaying to 1 at the outer edge of the buffer zone. + #[validate(range(min = 0.0))] + pub wall_buffer_weight: f32, + /// Max traversable vertical step. Taller steps are impassable. #[validate(range(min = 0.0))] - pub node_wall_buffer_m: f32, + pub step_threshold_m: f32, + /// Soft cost added per meter of vertical climb. #[validate(range(min = 0.0))] - pub node_step_threshold_m: f32, + pub step_penalty_weight: f32, + /// Ground-plane distance from goal at which the planner stops replanning. + #[validate(range(exclusive_min = 0.0))] + pub goal_tolerance: f32, + /// Rate cap for republishing the surface_map / nodes / node_edges viz + /// artifacts. 0 disables them entirely. The path output is unthrottled. + #[validate(range(min = 0.0))] + pub viz_publish_hz: f32, +} + +/// The soft wall penalty needs a non-zero zone to act in. +fn validate_wall_buffer(config: &Config) -> Result<(), ValidationError> { + if config.wall_buffer_weight > 0.0 && config.wall_buffer_m == 0.0 { + return Err(ValidationError::new( + "wall_buffer_weight requires wall_buffer_m > 0", + )); + } + Ok(()) +} + +impl Config { + /// Compute number of dilations and erosions to do based on closing radius + pub fn closing_passes(&self) -> u32 { + (self.surface_closing_radius / self.voxel_size).ceil() as u32 + } +} + +/// Cylindrical region the planner re-derives from a local map slice. +pub struct RegionBounds { + pub origin_x: f32, + pub origin_y: f32, + pub radius: f32, + pub z_min: f32, + pub z_max: f32, +} + +impl RegionBounds { + fn contains_voxel(&self, (kx, ky, kz): VoxelKey, voxel_size: f32) -> bool { + let half = voxel_size * 0.5; + let z = kz as f32 * voxel_size + half; + if z < self.z_min || z > self.z_max { + return false; + } + let dx = kx as f32 * voxel_size + half - self.origin_x; + let dy = ky as f32 * voxel_size + half - self.origin_y; + dx * dx + dy * dy <= self.radius * self.radius + } + + /// Inclusive voxel-column bounding box of the cylinder in the xy plane. + fn column_bbox(&self, voxel_size: f32) -> (i32, i32, i32, i32) { + let inv = 1.0 / voxel_size; + let x0 = ((self.origin_x - self.radius) * inv).floor() as i32; + let x1 = ((self.origin_x + self.radius) * inv).floor() as i32; + let y0 = ((self.origin_y - self.radius) * inv).floor() as i32; + let y1 = ((self.origin_y + self.radius) * inv).floor() as i32; + (x0, x1, y0, y1) + } } #[derive(Default)] @@ -37,43 +112,279 @@ pub struct Planner { graph: PlannerGraph, voxel_map: AHashSet, by_col: ColumnIz, - surface: Vec, } impl Planner { pub fn update_global_map(&mut self, points: &[(f32, f32, f32)], config: &Config) { let voxel_size = config.voxel_size; let clearance = (config.robot_height / voxel_size).ceil() as i32; - let step = (config.node_step_threshold_m / voxel_size).floor() as i32; self.voxel_map.clear(); for &p in points { self.voxel_map.insert(voxelize(p, voxel_size)); } + let mut surface: Vec = Vec::new(); extract_surfaces( &self.voxel_map, clearance, - config.surface_dilation_passes, - config.surface_erosion_passes, + config.closing_passes(), &mut self.by_col, - &mut self.surface, + &mut surface, + ); + build_surface_lookup(&surface, &mut self.graph.surface_lookup); + + self.rebuild_graph(config); + } + + /// Update planner artifacts within a local region instead of recomputing + /// the entire planner on the entire map. + pub fn update_region( + &mut self, + local_points: &[(f32, f32, f32)], + bounds: &RegionBounds, + config: &Config, + ) { + let voxel_size = config.voxel_size; + let clearance = (config.robot_height / voxel_size).ceil() as i32; + let pad = (2 * config.closing_passes()) as i32; + + let changed = self.replace_region_voxels(local_points, bounds, voxel_size); + + // No voxel changed, so surfaces and the graph are untouched. + let Some((bx0, bx1, by0, by1)) = changed else { + return; + }; + + // A changed voxel column shifts surfaces only within pad of it, so the + // write-back box is the changed-column bbox grown by pad. + let write = (bx0 - pad, bx1 + pad, by0 - pad, by1 + pad); + let new_cells = + extract_surfaces_region(&self.by_col, clearance, config.closing_passes(), write); + let (added, removed) = self.replace_surface_region(write, &new_cells); + + self.rebuild_region_graph(added, removed, config); + } + + /// Patch cells for the changed surface, then re-place nodes and edges over + /// the change window. A no-op when no surface cell changed. + fn rebuild_region_graph( + &mut self, + added: Vec, + removed: Vec, + config: &Config, + ) { + let step = (config.step_threshold_m / config.voxel_size).floor() as i32; + for &c in &removed { + self.graph.cells.remove(c); + } + for &c in &added { + self.graph.cells.insert(c); + } + let mut seeds = added; + seeds.extend_from_slice(&removed); + if seeds.is_empty() { + return; + } + + rebuild_edges_around( + &mut self.graph.cells, + &self.graph.surface_lookup, + &seeds, + config.voxel_size, + step, + ); + let window = self.node_window(&seeds, config); + place_nodes_region( + &mut self.graph.cells, + &window, + config.voxel_size, + config.node_spacing_m, + config.wall_clearance_m, + config.wall_buffer_m, + config.wall_buffer_weight, + config.step_penalty_weight, + &mut self.graph.wall_state, + &mut self.graph.nodes, + ); + build_node_edges_region( + &self.graph.cells, + &self.graph.nodes, + &window, + &mut self.graph.cell_state, + &mut self.graph.node_edges, + &mut self.graph.node_adj, ); + } + + /// Replace the cylinder's voxels with the local map points, ignoring + /// points outside it. Returns the column bbox of changed voxels, or None + /// if nothing changed. + fn replace_region_voxels( + &mut self, + local_points: &[(f32, f32, f32)], + bounds: &RegionBounds, + voxel_size: f32, + ) -> Option<(i32, i32, i32, i32)> { + let new_set: AHashSet = local_points + .iter() + .map(|&p| voxelize(p, voxel_size)) + .collect(); + + let (x0, x1, y0, y1) = bounds.column_bbox(voxel_size); + let by_col = &self.by_col; + let stale: Vec = (x0..(x1 + 1)) + .into_par_iter() + .flat_map_iter(|ix| { + let mut local: Vec = Vec::new(); + for iy in y0..=y1 { + let Some(zs) = by_col.get(&(ix, iy)) else { + continue; + }; + for &iz in zs { + let k = (ix, iy, iz); + if bounds.contains_voxel(k, voxel_size) && !new_set.contains(&k) { + local.push(k); + } + } + } + local + }) + .collect(); + + let mut bb = ChangeBounds::new(); + for &k in &stale { + bb.add(k.0, k.1); + self.voxel_map.remove(&k); + remove_from_by_col(&mut self.by_col, k); + } + for &k in &new_set { + if !bounds.contains_voxel(k, voxel_size) { + continue; + } + if self.voxel_map.insert(k) { + bb.add(k.0, k.1); + add_to_by_col(&mut self.by_col, k); + } + } + bb.bounds() + } + + /// Replace the surface_lookup entries for columns in the write box with + /// the freshly extracted cells. Returns the added and removed cells so + /// only the affected parts of the graph get patched. + fn replace_surface_region( + &mut self, + write: (i32, i32, i32, i32), + new_cells: &[VoxelKey], + ) -> (Vec, Vec) { + let (x0, x1, y0, y1) = write; + let mut old: AHashSet = AHashSet::new(); + for ix in x0..=x1 { + for iy in y0..=y1 { + if let Some(zs) = self.graph.surface_lookup.remove(&(ix, iy)) { + for iz in zs { + old.insert((ix, iy, iz)); + } + } + } + } + let new: AHashSet = new_cells.iter().copied().collect(); + + let mut touched: AHashSet<(i32, i32)> = AHashSet::new(); + for &(ix, iy, iz) in new_cells { + self.graph + .surface_lookup + .entry((ix, iy)) + .or_default() + .push(iz); + touched.insert((ix, iy)); + } + for col in touched { + if let Some(zs) = self.graph.surface_lookup.get_mut(&col) { + zs.sort_unstable(); + zs.dedup(); + } + } + + let added: Vec = new.iter().filter(|c| !old.contains(c)).copied().collect(); + let removed: Vec = old.iter().filter(|c| !new.contains(c)).copied().collect(); + (added, removed) + } + + /// Rebuild all cells from surface_lookup, then nodes and edges. + fn rebuild_graph(&mut self, config: &Config) { + let voxel_size = config.voxel_size; + let step = (config.step_threshold_m / voxel_size).floor() as i32; - build_surface_lookup(&self.surface, &mut self.graph.surface_lookup); build_surface_cells( &mut self.graph.cells, &self.graph.surface_lookup, voxel_size, step, ); + self.rebuild_nodes(config); + } + + /// Live cells within the changed-cell bbox grown by the node-graph margin, + /// which covers the reach of any node, edge, or Voronoi change. + fn node_window(&self, changed: &[VoxelKey], config: &Config) -> AHashSet { + // A few extra cells beyond the morphology, wall-buffer, and spacing reach. + const SLACK_CELLS: i32 = 2; + let voxel_size = config.voxel_size; + let pad = (2 * config.closing_passes()) as i32; + let buffer_cells = + ((config.wall_clearance_m + config.wall_buffer_m) / voxel_size).ceil() as i32; + let spacing_cells = (config.node_spacing_m / voxel_size).ceil() as i32; + let margin = pad + buffer_cells + spacing_cells + SLACK_CELLS; + + let mut bb = ChangeBounds::new(); + for &(ix, iy, _) in changed { + bb.add(ix, iy); + } + let Some((min_x, max_x, min_y, max_y)) = bb.bounds() else { + return AHashSet::new(); + }; + let (x0, x1, y0, y1) = ( + min_x - margin, + max_x + margin, + min_y - margin, + max_y + margin, + ); + let lookup = &self.graph.surface_lookup; + let cells = &self.graph.cells; + let ids: Vec = (x0..(x1 + 1)) + .into_par_iter() + .flat_map_iter(|ix| { + let mut local: Vec = Vec::new(); + for iy in y0..=y1 { + let Some(zs) = lookup.get(&(ix, iy)) else { + continue; + }; + for &iz in zs { + if let Some(id) = cells.id((ix, iy, iz)) { + local.push(id); + } + } + } + local + }) + .collect(); + ids.into_iter().collect() + } + + /// Full rebuild of nodes and node edges from the current cells. + fn rebuild_nodes(&mut self, config: &Config) { place_nodes( &mut self.graph.cells, - voxel_size, + config.voxel_size, config.node_spacing_m, - config.node_wall_buffer_m, - &mut self.graph.cell_state, + config.wall_clearance_m, + config.wall_buffer_m, + config.wall_buffer_weight, + config.step_penalty_weight, + &mut self.graph.wall_state, &mut self.graph.nodes, ); @@ -95,24 +406,645 @@ impl Planner { if self.graph.nodes.is_empty() { return None; } - planner::plan( - &self.graph, - start, - goal, - config.voxel_size, - config.robot_height, - ) + planner::plan(&self.graph, start, goal, config) } pub fn graph(&self) -> &PlannerGraph { &self.graph } - pub fn surface(&self) -> &[VoxelKey] { - &self.surface + pub fn surface(&self) -> impl Iterator + '_ { + self.graph + .surface_lookup + .iter() + .flat_map(|(&(ix, iy), zs)| zs.iter().map(move |&iz| (ix, iy, iz))) + } + + /// Surface cells paired with their wall clearance, the distance to the + /// nearest untraversable edge. Unreached cells report +inf. + pub fn surface_clearance(&self) -> Vec<(VoxelKey, f32)> { + let dist = &self.graph.wall_state.dist; + self.graph + .cells + .ids() + .map(|id| { + let d = dist.get(id as usize).copied().unwrap_or(f32::INFINITY); + (self.graph.cells.coord(id), d) + }) + .collect() } pub fn voxel_count(&self) -> usize { self.voxel_map.len() } + + pub fn voxel_keys(&self) -> impl Iterator + '_ { + self.voxel_map.iter().copied() + } +} + +/// Running inclusive xy bounding box of changed columns. +struct ChangeBounds { + min_x: i32, + max_x: i32, + min_y: i32, + max_y: i32, + any: bool, +} + +impl ChangeBounds { + fn new() -> Self { + Self { + min_x: i32::MAX, + max_x: i32::MIN, + min_y: i32::MAX, + max_y: i32::MIN, + any: false, + } + } + + fn add(&mut self, ix: i32, iy: i32) { + self.any = true; + self.min_x = self.min_x.min(ix); + self.max_x = self.max_x.max(ix); + self.min_y = self.min_y.min(iy); + self.max_y = self.max_y.max(iy); + } + + fn bounds(&self) -> Option<(i32, i32, i32, i32)> { + self.any + .then_some((self.min_x, self.max_x, self.min_y, self.max_y)) + } +} + +#[cfg(test)] +mod region_tests { + use super::*; + use std::collections::{BTreeMap, BTreeSet}; + + fn test_config() -> Config { + Config { + world_frame: String::new(), + voxel_size: 0.1, + robot_height: 0.5, + surface_closing_radius: 0.3, + node_spacing_m: 1.0, + wall_clearance_m: 0.0, + wall_buffer_m: 0.3, + wall_buffer_weight: 1.0, + step_threshold_m: 0.25, + step_penalty_weight: 0.0, + goal_tolerance: 0.3, + viz_publish_hz: 2.0, + } + } + + /// Floor slab with a wall down the middle, as world-frame point centers. + fn world_points() -> Vec<(f32, f32, f32)> { + let vs = 0.1_f32; + let half = vs * 0.5; + let mut pts = Vec::new(); + for ix in 0..40 { + for iy in 0..40 { + pts.push((ix as f32 * vs + half, iy as f32 * vs + half, half)); + } + } + // a wall column from z=0 up, to create wall-adjacency for nodes + for iy in 0..40 { + for iz in 0..15 { + pts.push(( + 20.0 * vs + half, + iy as f32 * vs + half, + iz as f32 * vs + half, + )); + } + } + pts + } + + fn surface_set(p: &Planner) -> BTreeSet { + p.surface().collect() + } + + fn voxel_set(p: &Planner) -> BTreeSet { + p.voxel_map.iter().copied().collect() + } + + /// Cell adjacency keyed by coordinate, independent of CellId. + fn cell_edges(p: &Planner) -> BTreeMap> { + let cells = &p.graph.cells; + let mut out: BTreeMap> = BTreeMap::new(); + for (id, edges) in cells.iter() { + let src = cells.coord(id); + let set = out.entry(src).or_default(); + for e in edges { + set.insert((cells.coord(e.dest), e.cost.to_bits())); + } + } + out + } + + fn node_coords(p: &Planner) -> BTreeSet { + p.graph + .nodes + .iter() + .map(|n| p.graph.cells.coord(n.cell_id)) + .collect() + } + + fn node_edge_pairs(p: &Planner) -> BTreeSet<(VoxelKey, VoxelKey, u32)> { + let cells = &p.graph.cells; + p.graph + .node_edges + .iter() + .map(|e| { + let a = cells.coord(e.a); + let b = cells.coord(e.b); + let (lo, hi) = if a <= b { (a, b) } else { (b, a) }; + (lo, hi, e.cost.to_bits()) + }) + .collect() + } + + #[test] + fn region_update_removes_stale_voxels() { + let cfg = test_config(); + let bounds = RegionBounds { + origin_x: 2.0, + origin_y: 2.0, + radius: 1.0, + z_min: -1.0, + z_max: 2.0, + }; + let all = world_points(); + + let mut full = Planner::default(); + full.update_global_map(&all, &cfg); + + let inside: Vec<_> = all + .iter() + .copied() + .filter(|&p| bounds.contains_voxel(voxelize(p, cfg.voxel_size), cfg.voxel_size)) + .collect(); + let outside: Vec<_> = all + .iter() + .copied() + .filter(|&p| !bounds.contains_voxel(voxelize(p, cfg.voxel_size), cfg.voxel_size)) + .collect(); + + // Seed the cylinder with a stack of junk voxels not present in the + // world, so update_region must clear them and the surface they induce. + let mut seeded = outside.clone(); + for iz in 3..8 { + seeded.push((2.05, 2.05, iz as f32 * cfg.voxel_size + 0.05)); + } + let mut region = Planner::default(); + region.update_global_map(&seeded, &cfg); + region.update_region(&inside, &bounds, &cfg); + + assert_eq!(voxel_set(®ion), voxel_set(&full), "voxel mismatch"); + assert_eq!(surface_set(®ion), surface_set(&full), "surface mismatch"); + assert_eq!( + cell_edges(®ion), + cell_edges(&full), + "cell edges mismatch" + ); + assert_eq!(node_coords(®ion), node_coords(&full), "node mismatch"); + assert_eq!( + node_edge_pairs(®ion), + node_edge_pairs(&full), + "node edge mismatch" + ); + } + + /// A point outside the region bounds must not enter the planner's voxel + /// map, where it could never be cleared and would inflate the rebuild box. + #[test] + fn region_update_ignores_points_outside_bounds() { + let cfg = test_config(); + let bounds = RegionBounds { + origin_x: 2.0, + origin_y: 2.0, + radius: 1.0, + z_min: -1.0, + z_max: 2.0, + }; + let inside = (2.05, 2.05, 0.05); + let outside = (10.05, 10.05, 0.05); + + let mut p = Planner::default(); + p.update_region(&[inside, outside], &bounds, &cfg); + + assert!(p.voxel_map.contains(&voxelize(inside, cfg.voxel_size))); + assert!(!p.voxel_map.contains(&voxelize(outside, cfg.voxel_size))); + } + + /// Floor 8m x 8m with a wall at x=4m that only a gap at y in [3.5, 4.5] + /// passes through, so crossing the wall is a non-trivial route. + fn big_world() -> Vec<(f32, f32, f32)> { + let vs = 0.1_f32; + let half = vs * 0.5; + let mut pts = Vec::new(); + for ix in 0..80 { + for iy in 0..80 { + pts.push((ix as f32 * vs + half, iy as f32 * vs + half, half)); + } + } + for iy in 0..80 { + if (35..45).contains(&iy) { + continue; + } + for iz in 0..15 { + pts.push(( + 40.0 * vs + half, + iy as f32 * vs + half, + iz as f32 * vs + half, + )); + } + } + pts + } + + fn slice(all: &[(f32, f32, f32)], b: &RegionBounds, vs: f32) -> Vec<(f32, f32, f32)> { + all.iter() + .copied() + .filter(|&p| b.contains_voxel(voxelize(p, vs), vs)) + .collect() + } + + fn path_len(w: &[(f32, f32, f32)]) -> f32 { + w.windows(2) + .map(|p| { + let dx = p[1].0 - p[0].0; + let dy = p[1].1 - p[0].1; + let dz = p[1].2 - p[0].2; + (dx * dx + dy * dy + dz * dz).sqrt() + }) + .sum() + } + + type Pose = (f32, f32, f32); + const PLAN_PAIRS: [(Pose, Pose); 4] = [ + ((0.5, 0.5, 0.05), (7.5, 7.5, 0.05)), + ((0.5, 7.5, 0.05), (7.5, 0.5, 0.05)), + ((0.5, 0.5, 0.05), (0.5, 7.5, 0.05)), + ((7.5, 0.5, 0.05), (7.5, 7.5, 0.05)), + ]; + + fn assert_plans_equivalent(full: &Planner, region: &Planner, cfg: &Config) { + for (s, g) in PLAN_PAIRS { + let pf = full.plan(s, g, cfg); + let pr = region.plan(s, g, cfg); + assert_eq!( + pf.is_some(), + pr.is_some(), + "path existence differs for {s:?} -> {g:?}" + ); + if let (Some(pf), Some(pr)) = (pf, pr) { + let (lf, lr) = (path_len(&pf), path_len(&pr)); + assert!(lr <= lf * 1.6 + 0.5, "region path too long: {lr} vs {lf}"); + assert!(lf <= lr * 1.6 + 0.5, "full path too long: {lf} vs {lr}"); + } + } + } + + /// Re-observing the same geometry must change nothing: no voxel, surface, + /// cell, node, or edge moves. This is the anti-jitter guarantee, far nodes + /// stay put when their region is re-seen, matching a full rebuild. + #[test] + fn region_reobserve_leaves_graph_bit_identical() { + let cfg = test_config(); + let all = big_world(); + let vs = cfg.voxel_size; + + let mut p = Planner::default(); + p.update_global_map(&all, &cfg); + let before_cells = cell_edges(&p); + let before_nodes = node_coords(&p); + let before_edges = node_edge_pairs(&p); + + for &(cx, cy) in &[(2.0, 2.0), (4.0, 4.0), (6.0, 3.0), (1.5, 7.0), (7.0, 7.0)] { + let b = RegionBounds { + origin_x: cx, + origin_y: cy, + radius: 1.2, + z_min: -1.0, + z_max: 2.0, + }; + p.update_region(&slice(&all, &b, vs), &b, &cfg); + } + + assert_eq!( + cell_edges(&p), + before_cells, + "cells changed on re-observation" + ); + assert_eq!( + node_coords(&p), + before_nodes, + "nodes moved on re-observation" + ); + assert_eq!( + node_edge_pairs(&p), + before_edges, + "edges changed on re-observation" + ); + } + + /// Build the planner purely from streamed local cylinders, as the live + /// pipeline does, and require equivalent planning to a one-shot full build. + #[test] + fn region_stream_only_plans_like_full() { + let cfg = test_config(); + let all = big_world(); + let vs = cfg.voxel_size; + + let mut full = Planner::default(); + full.update_global_map(&all, &cfg); + + let mut region = Planner::default(); + let mut cx = 0.5; + while cx <= 7.5 { + let mut cy = 0.5; + while cy <= 7.5 { + let b = RegionBounds { + origin_x: cx, + origin_y: cy, + radius: 1.5, + z_min: -1.0, + z_max: 2.0, + }; + let s = slice(&all, &b, vs); + if !s.is_empty() { + region.update_region(&s, &b, &cfg); + } + cy += 1.0; + } + cx += 1.0; + } + + assert_eq!( + voxel_set(®ion), + voxel_set(&full), + "stream did not reconstruct the map" + ); + assert_plans_equivalent(&full, ®ion, &cfg); + } + + /// Floor split by a wall with a narrow 1-cell gap near x=1.0 and a wide gap + /// near x=4.5. Start and goal straddle the narrow gap. + fn two_gap_world() -> Vec<(f32, f32, f32)> { + let vs = 0.1_f32; + let half = vs * 0.5; + let mut pts = Vec::new(); + for ix in 0..60 { + for iy in 0..40 { + pts.push((ix as f32 * vs + half, iy as f32 * vs + half, half)); + } + } + for ix in 0..60 { + if ix == 10 || (40..50).contains(&ix) { + continue; + } + for iz in 0..7 { + pts.push(( + ix as f32 * vs + half, + 20.0 * vs + half, + iz as f32 * vs + half, + )); + } + } + pts + } + + /// The hard clearance floor must make the narrow gap impassable, forcing + /// the longer detour through the wide gap. + #[test] + fn hard_clearance_floor_avoids_narrow_gap() { + let mut cfg = test_config(); + cfg.node_spacing_m = 0.8; + let pts = two_gap_world(); + let start = (1.0, 1.0, 0.05); + let goal = (1.0, 3.5, 0.05); + let max_x = |w: &[(f32, f32, f32)]| w.iter().map(|p| p.0).fold(f32::MIN, f32::max); + + // No clearance: the shortest route slips straight through the narrow gap. + cfg.wall_clearance_m = 0.0; + let mut open = Planner::default(); + open.update_global_map(&pts, &cfg); + let wp_open = open.plan(start, goal, &cfg).expect("open plan exists"); + + // Clearance wider than the narrow gap: it is impassable, so detour wide. + cfg.wall_clearance_m = 0.2; + let mut safe = Planner::default(); + safe.update_global_map(&pts, &cfg); + let wp_safe = safe.plan(start, goal, &cfg).expect("safe plan exists"); + + assert!(max_x(&wp_open) < 2.0, "open path should use the near gap"); + assert!( + max_x(&wp_safe) > 3.5, + "safe path should detour to the wide gap: max_x={}", + max_x(&wp_safe) + ); + assert!( + path_len(&wp_safe) > path_len(&wp_open) * 1.5, + "safe route should be substantially longer: {} vs {}", + path_len(&wp_safe), + path_len(&wp_open) + ); + } + + /// Every cell the smoothed path crosses, between waypoints included, must + /// clear the hard wall distance. + #[test] + fn final_path_clears_wall_distance() { + let mut cfg = test_config(); + cfg.wall_clearance_m = 0.2; + cfg.wall_buffer_m = 0.5; + let all = big_world(); + let mut p = Planner::default(); + p.update_global_map(&all, &cfg); + + let wp = p + .plan((0.7, 4.0, 0.05), (7.3, 4.0, 0.05), &cfg) + .expect("plan exists"); + let clearance: std::collections::HashMap = + p.surface_clearance().into_iter().collect(); + let vs = cfg.voxel_size; + let key = |x: f32, y: f32, z: f32| { + ( + (x / vs).floor() as i32, + (y / vs).floor() as i32, + (z / vs).round() as i32 - 1, + ) + }; + + // Interior waypoints are exact cell centers; sample between them too. + let interior = &wp[1..wp.len() - 1]; + assert!(interior.len() >= 2, "expected a multi-cell path"); + for pair in interior.windows(2) { + let (a, b) = (pair[0], pair[1]); + for k in 0..=24 { + let t = k as f32 / 24.0; + let x = a.0 + t * (b.0 - a.0); + let y = a.1 + t * (b.1 - a.1); + let z = a.2 + t * (b.2 - a.2); + if let Some(&c) = clearance.get(&key(x, y, z)) { + assert!( + c >= cfg.wall_clearance_m - 1e-4, + "path point ({x:.2},{y:.2}) sits {c:.3} from a wall, under the {} clearance", + cfg.wall_clearance_m + ); + } + } + } + } + + /// Solid 0.3 m block, taller than the step threshold; the path must route + /// around it and never climb on. + fn block_world() -> Vec<(f32, f32, f32)> { + let vs = 0.1_f32; + let half = vs * 0.5; + let mut pts = Vec::new(); + for ix in 0..40 { + for iy in 0..12 { + pts.push((ix as f32 * vs + half, iy as f32 * vs + half, half)); + } + } + // A solid block, 0.3 m tall, blocking the iy 0..6 lane around ix 18..22. + for ix in 18..22 { + for iy in 0..6 { + for iz in 0..4 { + pts.push(( + ix as f32 * vs + half, + iy as f32 * vs + half, + iz as f32 * vs + half, + )); + } + } + } + pts + } + + #[test] + fn final_path_never_climbs_over_threshold_step() { + let mut cfg = test_config(); + cfg.surface_closing_radius = 0.0; + cfg.wall_clearance_m = 0.0; + cfg.wall_buffer_m = 0.0; + cfg.node_spacing_m = 0.5; + let pts = block_world(); + let mut p = Planner::default(); + p.update_global_map(&pts, &cfg); + + let wp = p + .plan((1.0, 0.5, 0.05), (3.9, 0.5, 0.05), &cfg) + .expect("plan exists"); + + // The block top is at z = 0.4; the floor surface point is z = 0.1. No + // interior waypoint may land on the block. + for w in &wp[1..wp.len() - 1] { + assert!( + w.2 < 0.25, + "path climbed onto the 0.3 m block at {w:?}, exceeding the step threshold" + ); + } + // It had to detour out of the blocked lane (iy < 0.6). + let max_y = wp.iter().map(|p| p.1).fold(f32::MIN, f32::max); + assert!( + max_y > 0.6, + "path did not detour around the block: max_y={max_y}" + ); + } + + /// Flat floor with a crossable 0.2 m ridge blocking ix 15 except a flat gap + /// at iy 10..12. Crossing is short but climbs two steps; the detour is flat. + /// Route choice is read from the xy lane, since smoothing flattens the ridge + /// waypoints away. + fn ridge_world() -> Vec<(f32, f32, f32)> { + let vs = 0.1_f32; + let half = vs * 0.5; + let mut pts = Vec::new(); + for ix in 0..40 { + for iy in 0..12 { + pts.push((ix as f32 * vs + half, iy as f32 * vs + half, half)); + } + } + // A 0.2 m ridge cap at ix 15, iy 0..10: a 2-cell step up and back down. + for iy in 0..10 { + pts.push((15.0 * vs + half, iy as f32 * vs + half, 2.0 * vs + half)); + } + pts + } + + #[test] + fn step_penalty_diverts_path_around_ridge() { + let mut cfg = test_config(); + cfg.surface_closing_radius = 0.0; + cfg.wall_clearance_m = 0.0; + cfg.wall_buffer_m = 0.0; + cfg.node_spacing_m = 0.5; + let pts = ridge_world(); + let start = (1.0, 0.5, 0.05); + let goal = (2.9, 0.5, 0.05); + let max_y = |w: &[(f32, f32, f32)]| w.iter().map(|p| p.1).fold(f32::MIN, f32::max); + + // No step penalty: the short route crosses the ridge low. + cfg.step_penalty_weight = 0.0; + let mut cheap = Planner::default(); + cheap.update_global_map(&pts, &cfg); + let wp_cheap = cheap.plan(start, goal, &cfg).expect("plan exists"); + + // Heavy step penalty: the flat detour to the iy 10 gap wins. + cfg.step_penalty_weight = 30.0; + let mut avoid = Planner::default(); + avoid.update_global_map(&pts, &cfg); + let wp_avoid = avoid.plan(start, goal, &cfg).expect("plan exists"); + + assert!( + max_y(&wp_cheap) < 0.6, + "with no step penalty the path should cross the ridge low: max_y={}", + max_y(&wp_cheap) + ); + assert!( + max_y(&wp_avoid) > 0.9, + "with a heavy step penalty the path should detour to the flat gap: max_y={}", + max_y(&wp_avoid) + ); + } + + #[test] + fn goal_on_subclearance_spur_still_plans() { + let mut cfg = test_config(); + cfg.surface_closing_radius = 0.0; + cfg.wall_clearance_m = 0.3; + cfg.wall_buffer_m = 0.0; + cfg.wall_buffer_weight = 0.0; + cfg.node_spacing_m = 0.5; + + let vs = 0.1_f32; + let half = vs * 0.5; + let mut pts = Vec::new(); + for ix in 0..10 { + for iy in 0..10 { + pts.push((ix as f32 * vs + half, iy as f32 * vs + half, half)); + } + } + // A 1-wide spur off the open area: every spur cell is wall-adjacent so + // none clears the clearance and the penalized Voronoi cannot own them. + for ix in 10..16 { + pts.push((ix as f32 * vs + half, 5.0 * vs + half, half)); + } + + let mut p = Planner::default(); + p.update_global_map(&pts, &cfg); + + let start = (0.45, 0.45, 0.0); + let goal = (15.0 * vs + half, 5.0 * vs + half, 0.0); + let wp = p + .plan(start, goal, &cfg) + .expect("goal on a sub-clearance spur still reaches its component node"); + let last = *wp.last().expect("path has waypoints"); + assert!((last.0 - goal.0).abs() < 1e-3 && (last.1 - goal.1).abs() < 1e-3); + } } diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/nodes.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/nodes.rs index 50c0463763..93ce063bdf 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/nodes.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/nodes.rs @@ -5,11 +5,13 @@ //! nodes at local maxima via NMS, and rescale cell-edge costs to push paths //! toward corridor centers. -use ahash::AHashMap; +use std::cmp::Ordering; + +use ahash::{AHashMap, AHashSet}; use rayon::prelude::*; use crate::adjacency::{CellId, Edge, SurfaceCells}; -use crate::dijkstra::{dijkstra, DijkstraState}; +use crate::dijkstra::{dijkstra, dijkstra_region, DijkstraState, Weight}; use crate::voxel::{surface_point_xyz, VoxelKey}; #[derive(Clone, Copy, Debug)] @@ -22,11 +24,15 @@ pub struct NodeData { /// /// Runs multi source dijkstra using edges as sources, then distribute nodes /// using a grid based NMS. +#[allow(clippy::too_many_arguments)] pub fn place_nodes( cells: &mut SurfaceCells, voxel_size: f32, node_spacing_m: f32, - node_wall_buffer_m: f32, + wall_clearance_m: f32, + wall_buffer_m: f32, + wall_buffer_weight: f32, + step_penalty_weight: f32, state: &mut DijkstraState, out_nodes: &mut Vec, ) { @@ -37,20 +43,54 @@ pub fn place_nodes( let mut wall_seeds: Vec = Vec::new(); collect_wall_adjacent_cells(cells, &mut wall_seeds); - dijkstra(cells, &wall_seeds, state); + dijkstra(cells, &wall_seeds, state, Weight::Base); - let mut candidates: Vec = cells + // Floor is the hard clearance; NMS already prefers the clearest cells. + let node_floor = wall_clearance_m; + let candidates: Vec = cells .ids() - .filter(|&id| state.dist[id as usize] >= node_wall_buffer_m) + .filter(|&id| state.dist[id as usize] >= node_floor) .collect(); - candidates.par_sort_unstable_by(|&a, &b| { - state.dist[b as usize] - .total_cmp(&state.dist[a as usize]) - .then(a.cmp(&b)) - }); + place_from_candidates( + cells, + candidates, + &state.dist, + &[], + voxel_size, + node_spacing_m, + out_nodes, + ); + + let domain: Vec = cells.ids().collect(); + ensure_node_per_component(cells, &state.dist, voxel_size, &domain, out_nodes); - let survivors = nms_grid(cells, &candidates, voxel_size, node_spacing_m); + apply_wall_safe_penalty( + cells, + &state.dist, + wall_clearance_m, + wall_buffer_m, + wall_buffer_weight, + step_penalty_weight, + ); +} +/// Sort candidates by descending wall distance, thin them with NMS against the +/// seed nodes, and append the survivors as nodes. +fn place_from_candidates( + cells: &SurfaceCells, + mut candidates: Vec, + dist: &[f32], + seeds: &[CellId], + voxel_size: f32, + node_spacing_m: f32, + out_nodes: &mut Vec, +) { + candidates.par_sort_unstable_by(|&a, &b| { + dist[b as usize] + .total_cmp(&dist[a as usize]) + .then(cells.coord(a).cmp(&cells.coord(b))) + }); + let survivors = nms_grid(cells, &candidates, seeds, voxel_size, node_spacing_m); out_nodes.reserve(survivors.len()); for &id in &survivors { let (ix, iy, iz) = cells.coord(id); @@ -59,31 +99,135 @@ pub fn place_nodes( pos: surface_point_xyz(ix, iy, iz, voxel_size), }); } +} + +/// Regional counterpart to place_nodes: recompute the wall-distance field and +/// node placement inside the window, keeping cached nodes outside it as NMS +/// seeds so spacing holds across the seam. +#[allow(clippy::too_many_arguments)] +pub fn place_nodes_region( + cells: &mut SurfaceCells, + window: &AHashSet, + voxel_size: f32, + node_spacing_m: f32, + wall_clearance_m: f32, + wall_buffer_m: f32, + wall_buffer_weight: f32, + step_penalty_weight: f32, + wall_state: &mut DijkstraState, + nodes: &mut Vec, +) { + let mut wall_seeds: Vec = Vec::new(); + collect_wall_adjacent_in_window(cells, window, &mut wall_seeds); + dijkstra_region(cells, &wall_seeds, window, wall_state, Weight::Base); - apply_wall_safe_penalty(cells, &state.dist, node_wall_buffer_m); + nodes.retain(|n| cells.is_live(n.cell_id) && !window.contains(&n.cell_id)); + let kept: Vec = nodes.iter().map(|n| n.cell_id).collect(); + + let node_floor = wall_clearance_m; + let candidates: Vec = window + .iter() + .copied() + .filter(|&id| cells.is_live(id) && wall_state.dist[id as usize] >= node_floor) + .collect(); + place_from_candidates( + cells, + candidates, + &wall_state.dist, + &kept, + voxel_size, + node_spacing_m, + nodes, + ); + + let domain: Vec = window + .iter() + .copied() + .filter(|&id| cells.is_live(id)) + .collect(); + ensure_node_per_component(cells, &wall_state.dist, voxel_size, &domain, nodes); + + apply_wall_safe_penalty_region( + cells, + &wall_state.dist, + wall_clearance_m, + wall_buffer_m, + wall_buffer_weight, + step_penalty_weight, + window, + ); } -/// Cells missing any of their 4 xy-direction neighbors are treated as -/// boundaries. Direction membership is tracked with a 4-bit mask so the -/// 349k-cell case avoids per-cell hashset allocation. +/// Wall-adjacency over a cell subset, matching collect_wall_adjacent_cells. +fn collect_wall_adjacent_in_window( + cells: &SurfaceCells, + window: &AHashSet, + out: &mut Vec, +) { + out.clear(); + for &id in window { + if cells.is_live(id) && is_wall_adjacent(cells, id) { + out.push(id); + } + } +} + +/// A cell is wall-adjacent when it is missing at least one of its 4 xy-direction +/// neighbors. Membership is tracked with a 4-bit mask to avoid per-cell +/// allocation on the 349k-cell case. +fn is_wall_adjacent(cells: &SurfaceCells, id: CellId) -> bool { + let (cx, cy, _) = cells.coord(id); + let mut mask: u8 = 0; + for e in cells.neighbors(id) { + let (nx, ny, _) = cells.coord(e.dest); + mask |= match (nx - cx, ny - cy) { + (-1, 0) => 1, + (1, 0) => 2, + (0, -1) => 4, + (0, 1) => 8, + _ => 0, + }; + } + mask != 0b1111 +} + +/// Rescale edge costs for the window and its neighbors, whose wall distance may +/// have changed. Idempotent via base_cost. +fn apply_wall_safe_penalty_region( + cells: &mut SurfaceCells, + dist: &[f32], + clearance_m: f32, + buffer_m: f32, + buffer_weight: f32, + step_weight: f32, + window: &AHashSet, +) { + let mut affected: AHashSet = AHashSet::with_capacity(window.len() * 2); + for &w in window { + affected.insert(w); + for e in cells.neighbors(w) { + affected.insert(e.dest); + } + } + for id in affected { + scale_edges( + cells.edges_mut(id), + id, + dist, + clearance_m, + buffer_m, + buffer_weight, + step_weight, + ); + } +} + +/// Wall-adjacent cells over the whole graph. Falls back to a single cell so a +/// fully-enclosed map still seeds the wall-distance field. fn collect_wall_adjacent_cells(cells: &SurfaceCells, out: &mut Vec) { out.clear(); - for (id, edges) in cells.iter() { - let (cx, cy, _) = cells.coord(id); - - // Check if all 4 neighbors are present - let mut mask: u8 = 0; - for e in edges { - let (nx, ny, _) = cells.coord(e.dest); - mask |= match (nx - cx, ny - cy) { - (-1, 0) => 1, - (1, 0) => 2, - (0, -1) => 4, - (0, 1) => 8, - _ => 0, - }; - } - if mask != 0b1111 { + for id in cells.ids() { + if is_wall_adjacent(cells, id) { out.push(id); } } @@ -95,9 +239,13 @@ fn collect_wall_adjacent_cells(cells: &SurfaceCells, out: &mut Vec) { } /// Space out nodes based on minimum distance. +/// +/// The seed nodes suppress nearby candidates without being emitted, keeping a +/// regional re-placement consistent with cached nodes outside the window. fn nms_grid( cells: &SurfaceCells, candidates_sorted: &[CellId], + seeds: &[CellId], voxel_size: f32, node_spacing_m: f32, ) -> Vec { @@ -113,6 +261,9 @@ fn nms_grid( }; let mut bins: AHashMap<(i32, i32, i32), Vec> = AHashMap::new(); + for &s in seeds { + bins.entry(bin_of(cells.coord(s))).or_default().push(s); + } let mut survivors: Vec = Vec::new(); for &id in candidates_sorted { let coord = cells.coord(id); @@ -144,23 +295,183 @@ fn nms_grid( survivors } -/// Scale every edge cost by the average of its endpoint penalties, which -/// pushes shortest paths away from walls. Unreached cells have -/// dist == +INFINITY which collapses to penalty 1.0. -fn apply_wall_safe_penalty(cells: &mut SurfaceCells, dist: &[f32], buffer_m: f32) { +/// Scale each edge by its endpoints' average wall penalty and add the step +/// penalty. Unreached cells (dist +INFINITY) collapse the wall penalty to 1.0. +fn apply_wall_safe_penalty( + cells: &mut SurfaceCells, + dist: &[f32], + clearance_m: f32, + buffer_m: f32, + buffer_weight: f32, + step_weight: f32, +) { let mut edge_lists: Vec<(CellId, &mut Vec)> = cells.iter_edges_mut().collect(); edge_lists.par_iter_mut().for_each(|(src, edges)| { - let pu = penalty_of(dist[*src as usize], buffer_m); - for edge in edges.iter_mut() { - let pv = penalty_of(dist[edge.dest as usize], buffer_m); - edge.cost *= (pu + pv) / 2.0; - } + scale_edges( + edges, + *src, + dist, + clearance_m, + buffer_m, + buffer_weight, + step_weight, + ); }); } +/// Rescale one cell's outgoing edges from base_cost. Idempotent, so a regional +/// repass cannot compound the penalty. +#[inline] +fn scale_edges( + edges: &mut [Edge], + src: CellId, + dist: &[f32], + clearance_m: f32, + buffer_m: f32, + buffer_weight: f32, + step_weight: f32, +) { + let pu = penalty_of(dist[src as usize], clearance_m, buffer_m, buffer_weight); + for edge in edges.iter_mut() { + let pv = penalty_of( + dist[edge.dest as usize], + clearance_m, + buffer_m, + buffer_weight, + ); + edge.cost = edge.base_cost * (pu + pv) / 2.0 + step_weight * edge.rise; + } +} + +/// Lateral wall multiplier at wall distance d. Infinite inside the clearance, +/// then 1 + weight at the clearance edge decaying convexly to 1 at +/// clearance_m + buffer_m, and 1 beyond. #[inline] -fn penalty_of(d: f32, buffer_m: f32) -> f32 { - (1.0 + (buffer_m - d) / buffer_m).max(1.0) +pub(crate) fn penalty_of(d: f32, clearance_m: f32, buffer_m: f32, weight: f32) -> f32 { + if d < clearance_m { + return f32::INFINITY; + } + let outer = clearance_m + buffer_m; + if d >= outer { + return 1.0; + } + let band = buffer_m.max(1e-3); + let t = (outer - d) / band; // 0 at the outer edge, 1 at the clearance edge + 1.0 + weight * t * t +} + +/// Seed a node in every connected component in `domain` that the clearance +/// floor left empty, so a thin or sparse component is still reachable. `domain` +/// is every live cell for a full rebuild, or the window for an incremental one. +fn ensure_node_per_component( + cells: &SurfaceCells, + dist: &[f32], + voxel_size: f32, + domain: &[CellId], + out_nodes: &mut Vec, +) { + if domain.is_empty() { + return; + } + let node_cells: AHashSet = out_nodes.iter().map(|n| n.cell_id).collect(); + let in_domain: AHashSet = domain.iter().copied().collect(); + + let mut uf = UnionFind::default(); + for &id in domain { + uf.make(id); + } + for &id in domain { + for e in cells.neighbors(id) { + if in_domain.contains(&e.dest) { + uf.union(id, e.dest); + } + } + } + + // Served: the component holds or borders a node, including one outside domain. + let mut served: AHashSet = AHashSet::new(); + for &id in domain { + let touches_node = node_cells.contains(&id) + || cells + .neighbors(id) + .iter() + .any(|e| node_cells.contains(&e.dest)); + if touches_node { + served.insert(uf.find(id)); + } + } + + // Clearest cell per still-unserved component. + let mut best: AHashMap = AHashMap::new(); + for &id in domain { + let root = uf.find(id); + if served.contains(&root) { + continue; + } + match best.get(&root) { + Some(&cur) if !is_clearer(cells, dist, id, cur) => {} + _ => { + best.insert(root, id); + } + } + } + + out_nodes.reserve(best.len()); + for &id in best.values() { + let (ix, iy, iz) = cells.coord(id); + out_nodes.push(NodeData { + cell_id: id, + pos: surface_point_xyz(ix, iy, iz, voxel_size), + }); + } +} + +/// Better fallback seed: farther from a wall, ties broken by coordinate. +fn is_clearer(cells: &SurfaceCells, dist: &[f32], a: CellId, b: CellId) -> bool { + match dist[a as usize].total_cmp(&dist[b as usize]) { + Ordering::Greater => true, + Ordering::Less => false, + Ordering::Equal => cells.coord(a) < cells.coord(b), + } +} + +/// Union-find keyed by CellId so the incremental path pays for the window only. +#[derive(Default)] +struct UnionFind { + parent: AHashMap, +} + +impl UnionFind { + fn make(&mut self, x: CellId) { + self.parent.entry(x).or_insert(x); + } + + fn find(&mut self, x: CellId) -> CellId { + let mut root = x; + while let Some(&p) = self.parent.get(&root) { + if p == root { + break; + } + root = p; + } + let mut cur = x; + while let Some(&p) = self.parent.get(&cur) { + if p == root { + break; + } + self.parent.insert(cur, root); + cur = p; + } + root + } + + fn union(&mut self, a: CellId, b: CellId) { + let ra = self.find(a); + let rb = self.find(b); + if ra != rb { + self.parent.insert(ra, rb); + } + } } #[cfg(test)] @@ -193,7 +504,9 @@ mod tests { let mut sc = build_cells(&open_patch(0, 0, 10), 2); let mut state = DijkstraState::default(); let mut nodes = Vec::new(); - place_nodes(&mut sc, VOXEL, 1.0, 0.3, &mut state, &mut nodes); + place_nodes( + &mut sc, VOXEL, 1.0, 0.0, 0.3, 1.0, 0.0, &mut state, &mut nodes, + ); assert!(!nodes.is_empty()); for n in &nodes { let (ix, iy, _) = sc.coord(n.cell_id); @@ -212,10 +525,34 @@ mod tests { let mut sc = build_cells(&cells_in, 2); let mut state = DijkstraState::default(); let mut nodes = Vec::new(); - place_nodes(&mut sc, VOXEL, 1.0, 0.3, &mut state, &mut nodes); + place_nodes( + &mut sc, VOXEL, 1.0, 0.0, 0.3, 1.0, 0.0, &mut state, &mut nodes, + ); assert!(!nodes.is_empty()); } + #[test] + fn each_disconnected_component_gets_a_node() { + // Two 1-wide strips far apart: every cell is wall-adjacent so none + // clears the 0.5 m clearance floor, yet each disconnected strip must + // still get exactly one node. + let mut cells_in: Vec = (0..8).map(|ix| (ix, 0, 0)).collect(); + cells_in.extend((0..8).map(|ix| (ix, 20, 0))); + let mut sc = build_cells(&cells_in, 2); + let mut state = DijkstraState::default(); + let mut nodes = Vec::new(); + place_nodes( + &mut sc, VOXEL, 1.0, 0.5, 0.3, 1.0, 0.0, &mut state, &mut nodes, + ); + assert_eq!( + nodes.len(), + 2, + "each disconnected component needs its own node" + ); + let ys: Vec = nodes.iter().map(|n| sc.coord(n.cell_id).1).collect(); + assert!(ys.contains(&0) && ys.contains(&20)); + } + #[test] fn nms_enforces_spacing() { let mut cells_in = open_patch(0, 0, 10); @@ -223,7 +560,9 @@ mod tests { let mut sc = build_cells(&cells_in, 2); let mut state = DijkstraState::default(); let mut nodes = Vec::new(); - place_nodes(&mut sc, VOXEL, 1.0, 0.3, &mut state, &mut nodes); + place_nodes( + &mut sc, VOXEL, 1.0, 0.0, 0.3, 1.0, 0.0, &mut state, &mut nodes, + ); assert!(nodes.len() >= 2); for i in 0..nodes.len() { for j in (i + 1)..nodes.len() { @@ -238,13 +577,74 @@ mod tests { } } + #[test] + fn penalty_ramps_across_buffer_zone() { + // clearance 0.1, soft zone 0.4 wide, so the outer edge is at 0.5. + let (clearance, buffer, w) = (0.1, 0.4, 4.0); + assert!(penalty_of(0.05, clearance, buffer, w).is_infinite()); + assert!((penalty_of(0.1, clearance, buffer, w) - 5.0).abs() < 1e-6); + assert!((penalty_of(0.5, clearance, buffer, w) - 1.0).abs() < 1e-6); + assert!((penalty_of(1.0, clearance, buffer, w) - 1.0).abs() < 1e-6); + assert!((penalty_of(0.3, clearance, buffer, w) - 2.0).abs() < 1e-6); + } + + #[test] + fn wall_penalty_doubles_cost_at_the_wall() { + // On a 1-wide strip every cell is wall-adjacent (d = 0), so with zero + // clearance the ramp peaks at 2 and edge cost is twice the geometric. + let cells_in: Vec = (0..10).map(|ix| (ix, 0, 0)).collect(); + let mut sc = build_cells(&cells_in, 2); + let mut state = DijkstraState::default(); + let mut nodes = Vec::new(); + place_nodes( + &mut sc, VOXEL, 1.0, 0.0, 0.3, 1.0, 0.0, &mut state, &mut nodes, + ); + let id = sc.id((5, 0, 0)).unwrap(); + assert!((sc.neighbors(id)[0].cost - 2.0 * VOXEL).abs() < 1e-5); + } + + #[test] + fn step_penalty_adds_to_vertical_edges() { + // A 2-cell rise (0.2 m) between adjacent cells. With weight 10 the edge + // gains 10 * 0.2 = 2.0 on top of its geometric and wall cost. + let cells_in: Vec = vec![(0, 0, 0), (1, 0, 2), (2, 0, 2)]; + let cost_with = |step_weight: f32| { + let mut sc = build_cells(&cells_in, 2); + let mut state = DijkstraState::default(); + let mut nodes = Vec::new(); + place_nodes( + &mut sc, + VOXEL, + 1.0, + 0.0, + 0.3, + 1.0, + step_weight, + &mut state, + &mut nodes, + ); + let id = sc.id((0, 0, 0)).unwrap(); + sc.neighbors(id) + .iter() + .find(|e| sc.coord(e.dest) == (1, 0, 2)) + .unwrap() + .cost + }; + assert!( + (cost_with(10.0) - cost_with(0.0) - 10.0 * 0.2).abs() < 1e-4, + "step penalty must add weight * rise" + ); + } + #[test] fn wall_cells_scale_outbound_cost() { let cells_in: Vec = (0..10).map(|ix| (ix, 0, 0)).collect(); let mut sc = build_cells(&cells_in, 2); let mut state = DijkstraState::default(); let mut nodes = Vec::new(); - place_nodes(&mut sc, VOXEL, 1.0, 0.3, &mut state, &mut nodes); + place_nodes( + &mut sc, VOXEL, 1.0, 0.0, 0.3, 1.0, 0.0, &mut state, &mut nodes, + ); let id0 = sc.id((0, 0, 0)).unwrap(); let outbound = sc.neighbors(id0); assert!(!outbound.is_empty()); diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs index e3bb8e818a..85f12345b3 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs @@ -2,50 +2,63 @@ // SPDX-License-Identifier: Apache-2.0 use std::cmp::Ordering; -use std::collections::BinaryHeap; +use std::collections::{BinaryHeap, VecDeque}; -use ahash::AHashMap; +use ahash::{AHashMap, AHashSet}; -use crate::adjacency::{CellId, SurfaceLookup}; +use crate::adjacency::{rise, CellId, SurfaceCells, SurfaceLookup}; use crate::dijkstra::walk_preds; use crate::edges::{NodeEdgeIdx, NodeId, PlannerGraph, NO_NODE}; +use crate::mls_planner::Config; +use crate::nodes::penalty_of; use crate::voxel::{surface_point_xyz, VoxelKey}; -/// Snap a pose to the best surface cell. -pub fn snap_pose_to_cell( +/// Robot-rooted candidate search radius, in multiples of node spacing. +const CANDIDATE_RADIUS_FACTOR: f32 = 3.0; + +/// Horizontal search radius when snapping a pose to the surface. +const SNAP_SEARCH_RADIUS_M: f32 = 1.5; + +/// Max snap candidates tried when connecting the start. +const MAX_SNAP_ATTEMPTS: usize = 64; + +/// Surface cells near the pose, nearest first in xy. +pub fn snap_candidates( surface_lookup: &SurfaceLookup, pose: (f32, f32, f32), voxel_size: f32, tolerance_m: f32, -) -> Option { +) -> Vec { let ix = (pose.0 / voxel_size).floor() as i32; let iy = (pose.1 / voxel_size).floor() as i32; let target_iz = (pose.2 / voxel_size).floor() as i32 - 1; let tol_cells = (tolerance_m / voxel_size).ceil() as i32; + let search_radius = (SNAP_SEARCH_RADIUS_M / voxel_size).ceil() as i32; - if let Some(cell) = best_iz_in_column(surface_lookup, ix, iy, target_iz, tol_cells) { - return Some(cell); - } - - const SEARCH_RADIUS: i32 = 5; - let mut best: Option<(i32, VoxelKey)> = None; - for dix in -SEARCH_RADIUS..=SEARCH_RADIUS { - for diy in -SEARCH_RADIUS..=SEARCH_RADIUS { - if dix == 0 && diy == 0 { - continue; - } - let Some(cell) = + let mut found: Vec<(i32, VoxelKey)> = Vec::new(); + for dix in -search_radius..=search_radius { + for diy in -search_radius..=search_radius { + if let Some(cell) = best_iz_in_column(surface_lookup, ix + dix, iy + diy, target_iz, tol_cells) - else { - continue; - }; - let d2 = dix * dix + diy * diy; - if best.is_none_or(|(bd, _)| d2 < bd) { - best = Some((d2, cell)); + { + found.push((dix * dix + diy * diy, cell)); } } } - best.map(|(_, c)| c) + found.sort_by_key(|&(d2, _)| d2); + found.into_iter().map(|(_, c)| c).collect() +} + +/// Snap a pose to the nearest surface cell. +pub fn snap_pose_to_cell( + surface_lookup: &SurfaceLookup, + pose: (f32, f32, f32), + voxel_size: f32, + tolerance_m: f32, +) -> Option { + snap_candidates(surface_lookup, pose, voxel_size, tolerance_m) + .into_iter() + .next() } fn best_iz_in_column( @@ -77,93 +90,303 @@ pub fn plan( plg: &PlannerGraph, start_pose: (f32, f32, f32), goal_pose: (f32, f32, f32), - voxel_size: f32, - z_tolerance_m: f32, + config: &Config, ) -> Option> { - let start_coord = - snap_pose_to_cell(&plg.surface_lookup, start_pose, voxel_size, z_tolerance_m)?; - let goal_coord = snap_pose_to_cell(&plg.surface_lookup, goal_pose, voxel_size, z_tolerance_m)?; - let start_cell = plg.cells.id(start_coord)?; - let goal_cell = plg.cells.id(goal_coord)?; - - let node_idx_by_cell: AHashMap = plg - .nodes - .iter() - .enumerate() - .map(|(i, n)| (n.cell_id, i as NodeId)) - .collect(); + let voxel_size = config.voxel_size; + let z_tolerance_m = config.robot_height; + let start_candidates = + snap_candidates(&plg.surface_lookup, start_pose, voxel_size, z_tolerance_m); + if start_candidates.is_empty() { + tracing::warn!( + ?start_pose, + "plan failed: start does not snap to any surface cell" + ); + return None; + } + let Some(goal_coord) = + snap_pose_to_cell(&plg.surface_lookup, goal_pose, voxel_size, z_tolerance_m) + else { + tracing::warn!( + ?goal_pose, + "plan failed: goal does not snap to any surface cell" + ); + return None; + }; + let Some(goal_cell) = plg.cells.id(goal_coord) else { + tracing::warn!(?goal_coord, "plan failed: goal cell is not in the graph"); + return None; + }; + + let node_cells: AHashSet = plg.nodes.iter().map(|n| n.cell_id).collect(); + + // The penalized Voronoi cannot own sub-clearance goals, so fall back to the + // nearest node by hops. A real failure then reports as disconnected below. + let mut goal_segment = walk_preds(&plg.cell_state, goal_cell); + let mut goal_node = *goal_segment + .last() + .expect("walk_preds returns at least the start cell"); + if !node_cells.contains(&goal_node) { + let Some((node, path)) = nearest_node(&plg.cells, goal_cell, &node_cells) else { + tracing::warn!( + ?goal_coord, + "plan failed: goal's connected component has no graph node" + ); + return None; + }; + goal_node = node; + goal_segment = path; + } - let start_segment = walk_preds(&plg.cell_state, start_cell); - let goal_segment = walk_preds(&plg.cell_state, goal_cell); - let start_node = *node_idx_by_cell.get(start_segment.last()?)?; - let goal_node = *node_idx_by_cell.get(goal_segment.last()?)?; - - let node_seq = shortest_path_nodes(plg, start_node, goal_node)?; - Some(assemble_waypoints( - plg, - &node_seq, - start_pose, - &start_segment, - goal_pose, - &goal_segment, + // Rooted at the goal so one pass covers every node's cost-to-go. + let (cost_to_go, pred_to_goal) = node_dijkstra(plg, goal_node); + + let radius = (config.node_spacing_m * CANDIDATE_RADIUS_FACTOR).max(voxel_size); + let mut entry: Option<(Vec, Vec)> = None; + for &candidate in start_candidates.iter().take(MAX_SNAP_ATTEMPTS) { + let Some(start_cell) = plg.cells.id(candidate) else { + continue; + }; + entry = select_entry( + plg, + start_cell, + goal_node, + &cost_to_go, + &pred_to_goal, + &node_cells, + radius, + ); + if entry.is_some() { + break; + } + } + let Some((lead_in, node_seq)) = entry else { + tracing::warn!( + candidates = start_candidates.len().min(MAX_SNAP_ATTEMPTS), + reachable_nodes = cost_to_go.len(), + total_nodes = plg.nodes.len(), + "plan failed: start and goal lie on separate connected surface components", + ); + return None; + }; + + // Max traversable step in cells, floored to match the graph adjacency. + let step_cells = (config.step_threshold_m / voxel_size).floor() as i32; + + let wall_cost = WallCost { + clearance_m: config.wall_clearance_m, + buffer_m: config.wall_buffer_m, + buffer_weight: config.wall_buffer_weight, voxel_size, + }; + let cells = assemble_cells(plg, &node_seq, &lead_in, &goal_segment); + let cells = string_pull(plg, &cells, step_cells, &wall_cost); + Some(cells_to_waypoints( + plg, &cells, start_pose, goal_pose, voxel_size, )) } -pub fn shortest_path_nodes(plg: &PlannerGraph, start: NodeId, goal: NodeId) -> Option> { - if start == goal { - return Some(vec![start]); +/// Pick the entry node by connect cost plus cost-to-go, with its on-surface +/// lead-in and the node sequence to the goal. +fn select_entry( + plg: &PlannerGraph, + start_cell: CellId, + goal_node: NodeId, + cost_to_go: &AHashMap, + pred_to_goal: &AHashMap, + node_cells: &AHashSet, + radius_m: f32, +) -> Option<(Vec, Vec)> { + let (connect_dist, connect_pred) = robot_search(&plg.cells, start_cell, radius_m); + + let mut entry_node = NO_NODE; + let mut best_score = f32::INFINITY; + for node in &plg.nodes { + let Some(&connect) = connect_dist.get(&node.cell_id) else { + continue; + }; + let Some(&ctg) = cost_to_go.get(&node.cell_id) else { + continue; + }; + let score = connect + ctg; + if score < best_score { + best_score = score; + entry_node = node.cell_id; + } + } + + if best_score.is_finite() { + let mut lead = walk_local_preds(&connect_pred, entry_node); + lead.reverse(); + return Some((lead, follow_preds(entry_node, goal_node, pred_to_goal)?)); } - let n = plg.nodes.len(); - let mut dist = vec![f32::INFINITY; n]; - let mut pred = vec![NO_NODE; n]; - dist[start as usize] = 0.0; + + let start_segment = walk_preds(&plg.cell_state, start_cell); + let region_node = *start_segment.last()?; + if !node_cells.contains(®ion_node) + || !cost_to_go.get(®ion_node).is_some_and(|c| c.is_finite()) + { + return None; + } + Some(( + start_segment, + follow_preds(region_node, goal_node, pred_to_goal)?, + )) +} + +/// Bounded Dijkstra from the robot cell, visiting cells within the radius. +/// Returns per-cell distance and predecessor maps. +fn robot_search( + cells: &SurfaceCells, + source: CellId, + radius_m: f32, +) -> (AHashMap, AHashMap) { + let mut dist: AHashMap = AHashMap::new(); + let mut pred: AHashMap = AHashMap::new(); let mut heap: BinaryHeap = BinaryHeap::new(); - heap.push(Scored(0.0, start)); + dist.insert(source, 0.0); + heap.push(Scored(0.0, source)); while let Some(Scored(d, u)) = heap.pop() { - if d > dist[u as usize] { + if d > radius_m { + break; + } + if d > dist.get(&u).copied().unwrap_or(f32::INFINITY) { continue; } - if u == goal { - break; + for edge in cells.neighbors(u) { + let nd = d + edge.cost; + if nd < dist.get(&edge.dest).copied().unwrap_or(f32::INFINITY) { + dist.insert(edge.dest, nd); + pred.insert(edge.dest, u); + heap.push(Scored(nd, edge.dest)); + } } - for &edge_idx in &plg.node_adj[u as usize] { + } + (dist, pred) +} + +/// Nearest node to `from` by hops, ignoring edge cost so it reaches a node +/// across cells the wall penalty makes impassable. Returns the node and the +/// path from `from`. +fn nearest_node( + cells: &SurfaceCells, + from: CellId, + node_cells: &AHashSet, +) -> Option<(NodeId, Vec)> { + if node_cells.contains(&from) { + return Some((from, vec![from])); + } + let mut pred: AHashMap = AHashMap::new(); + let mut seen: AHashSet = AHashSet::new(); + let mut queue: VecDeque = VecDeque::new(); + seen.insert(from); + queue.push_back(from); + + while let Some(u) = queue.pop_front() { + for edge in cells.neighbors(u) { + let v = edge.dest; + if !seen.insert(v) { + continue; + } + pred.insert(v, u); + if node_cells.contains(&v) { + let mut path = vec![v]; + let mut cur = v; + while let Some(&p) = pred.get(&cur) { + cur = p; + path.push(cur); + } + path.reverse(); + return Some((v, path)); + } + queue.push_back(v); + } + } + None +} + +/// Walk predecessors back to the search source. +fn walk_local_preds(pred: &AHashMap, from: CellId) -> Vec { + let mut path = vec![from]; + let mut cur = from; + while let Some(&p) = pred.get(&cur) { + cur = p; + path.push(cur); + } + path +} + +/// Cost-to-go to source for every reachable node, with a predecessor pointing +/// one hop toward it. Nodes are keyed by their CellId. Unreachable nodes are +/// simply absent from the maps. +fn node_dijkstra( + plg: &PlannerGraph, + source: NodeId, +) -> (AHashMap, AHashMap) { + let mut dist: AHashMap = AHashMap::new(); + let mut pred: AHashMap = AHashMap::new(); + dist.insert(source, 0.0); + let mut heap: BinaryHeap = BinaryHeap::new(); + heap.push(Scored(0.0, source)); + + while let Some(Scored(d, u)) = heap.pop() { + if d > dist.get(&u).copied().unwrap_or(f32::INFINITY) { + continue; + } + let Some(adj) = plg.node_adj.get(&u) else { + continue; + }; + for &edge_idx in adj { let edge = &plg.node_edges[edge_idx as usize]; let neighbor = if edge.a == u { edge.b } else { edge.a }; let nd = d + edge.cost; - if nd < dist[neighbor as usize] { - dist[neighbor as usize] = nd; - pred[neighbor as usize] = u; + if nd < dist.get(&neighbor).copied().unwrap_or(f32::INFINITY) { + dist.insert(neighbor, nd); + pred.insert(neighbor, u); heap.push(Scored(nd, neighbor)); } } } + (dist, pred) +} - if !dist[goal as usize].is_finite() { - return None; +/// Build the node sequence by following goal-pointing predecessors. +fn follow_preds( + from: NodeId, + goal: NodeId, + pred: &AHashMap, +) -> Option> { + let mut seq = vec![from]; + let mut cur = from; + while cur != goal { + let &next = pred.get(&cur)?; + cur = next; + seq.push(cur); } - let mut path = vec![goal]; - let mut cur = goal; - while pred[cur as usize] != NO_NODE { - cur = pred[cur as usize]; - path.push(cur); + Some(seq) +} + +/// Append a cell, cancelling an out-and-back spur when the next cell retraces +/// the second-to-last. +fn push_cell(cells: &mut Vec, c: CellId) { + if cells.len() >= 2 && cells[cells.len() - 2] == c { + cells.pop(); + } else if cells.last() != Some(&c) { + cells.push(c); } - path.reverse(); - Some(path) } -fn assemble_waypoints( +/// Build the cell path from the entry lead-in through the node edges to the goal. +fn assemble_cells( plg: &PlannerGraph, node_seq: &[NodeId], - start_pose: (f32, f32, f32), - start_segment: &[CellId], - goal_pose: (f32, f32, f32), + lead_in: &[CellId], goal_segment: &[CellId], - voxel_size: f32, -) -> Vec<(f32, f32, f32)> { +) -> Vec { let mut cells: Vec = Vec::new(); - cells.extend_from_slice(start_segment); + for &c in lead_in { + push_cell(&mut cells, c); + } for pair in node_seq.windows(2) { let (a, b) = (pair[0], pair[1]); @@ -181,21 +404,29 @@ fn assemble_waypoints( let to_b = walk_preds(&plg.cell_state, end_side); for c in from_a.into_iter().chain(to_b) { - if cells.last() != Some(&c) { - cells.push(c); - } + push_cell(&mut cells, c); } } for &c in goal_segment.iter().rev() { - if cells.last() != Some(&c) { - cells.push(c); - } + push_cell(&mut cells, c); } + cells +} + +/// Convert the cell path to world waypoints, with the raw start and goal poses +/// as the endpoints. +fn cells_to_waypoints( + plg: &PlannerGraph, + cells: &[CellId], + start_pose: (f32, f32, f32), + goal_pose: (f32, f32, f32), + voxel_size: f32, +) -> Vec<(f32, f32, f32)> { let mut waypoints: Vec<(f32, f32, f32)> = Vec::with_capacity(cells.len() + 2); waypoints.push(start_pose); - for id in cells { + for &id in cells { let (ix, iy, iz) = plg.cells.coord(id); waypoints.push(surface_point_xyz(ix, iy, iz, voxel_size)); } @@ -203,8 +434,134 @@ fn assemble_waypoints( waypoints } +/// Clearance and step limits the smoother holds the path to. +struct WallCost { + clearance_m: f32, + buffer_m: f32, + buffer_weight: f32, + voxel_size: f32, +} + +/// Replace runs of cells with straight chords that come no closer to a wall and +/// climb no more than the run they replace. +fn string_pull( + plg: &PlannerGraph, + cells: &[CellId], + step_cells: i32, + wc: &WallCost, +) -> Vec { + if cells.len() <= 2 { + return cells.to_vec(); + } + let metrics = |from: CellId, to: CellId| { + segment_metrics( + plg, + plg.cells.coord(from), + plg.cells.coord(to), + step_cells, + wc, + ) + }; + let mut out = vec![cells[0]]; + let mut anchor = 0; + while anchor + 1 < cells.len() { + let mut best = anchor + 1; + let mut rough_pen = 1.0_f32; + let mut rough_rise = 0.0_f32; + let mut j = anchor + 1; + while j < cells.len() { + if let Some((pen, rise)) = metrics(cells[j - 1], cells[j]) { + rough_pen = rough_pen.max(pen); + rough_rise += rise; + } + match metrics(cells[anchor], cells[j]) { + Some((pen, rise)) if pen <= rough_pen + 1e-3 && rise <= rough_rise + 1e-3 => { + best = j + } + _ => break, + } + j += 1; + } + out.push(cells[best]); + anchor = best; + } + out +} + +/// Worst wall penalty and total climb along the straight segment a -> b. None if +/// it leaves the surface, exceeds step_cells, or enters the hard clearance. +fn segment_metrics( + plg: &PlannerGraph, + a: VoxelKey, + b: VoxelKey, + step_cells: i32, + wc: &WallCost, +) -> Option<(f32, f32)> { + let (dx, dy, dz) = (b.0 - a.0, b.1 - a.1, b.2 - a.2); + let samples = dx.abs().max(dy.abs()) * 2; + if samples == 0 { + // A same-column vertical chord is not traversable. + return (dz == 0).then_some((1.0, 0.0)); + } + let (mut last_ix, mut last_iy) = (i32::MIN, i32::MIN); + let mut prev_iz: Option = None; + let mut max_pen = 1.0_f32; + let mut rise_cells = 0i32; + for k in 0..=samples { + let t = k as f32 / samples as f32; + let ix = (a.0 as f32 + t * dx as f32).round() as i32; + let iy = (a.1 as f32 + t * dy as f32).round() as i32; + if ix == last_ix && iy == last_iy { + continue; + } + last_ix = ix; + last_iy = iy; + let iz_line = a.2 as f32 + t * dz as f32; + let zs = plg.surface_lookup.get(&(ix, iy))?; + // Surface cell in this column nearest the interpolated segment height. + let mut nearest: Option<(f32, i32)> = None; + for &iz in zs { + let d = (iz as f32 - iz_line).abs(); + if nearest.is_none_or(|(bd, _)| d < bd) { + nearest = Some((d, iz)); + } + } + let (d, iz) = nearest?; + if d > step_cells as f32 { + return None; + } + // Tally climb and reject an untraversable step between columns. + if let Some(p) = prev_iz { + let step = (iz - p).abs(); + if step > step_cells { + return None; + } + rise_cells += step; + } + prev_iz = Some(iz); + // Columns on the surface but not in the graph carry no wall penalty. + let p = match plg.cells.id((ix, iy, iz)) { + Some(id) => { + let wall_dist = plg + .wall_state + .dist + .get(id as usize) + .copied() + .unwrap_or(f32::INFINITY); + penalty_of(wall_dist, wc.clearance_m, wc.buffer_m, wc.buffer_weight) + } + None => 1.0, + }; + if !p.is_finite() { + return None; + } + max_pen = max_pen.max(p); + } + Some((max_pen, rise(rise_cells, wc.voxel_size))) +} + fn edge_between(plg: &PlannerGraph, a: NodeId, b: NodeId) -> Option { - for &edge_idx in &plg.node_adj[a as usize] { + for &edge_idx in plg.node_adj.get(&a)? { let edge = &plg.node_edges[edge_idx as usize]; let other = if edge.a == a { edge.b } else { edge.a }; if other == b { @@ -271,6 +628,28 @@ mod tests { (0..n).map(|x| (x, 0, 0)).collect() } + fn plan_simple( + plg: &PlannerGraph, + start: (f32, f32, f32), + goal: (f32, f32, f32), + ) -> Option> { + let config = Config { + world_frame: "world".into(), + voxel_size: VOXEL, + robot_height: Z_TOL, + surface_closing_radius: 0.0, + node_spacing_m: 1.0, + wall_clearance_m: 0.2, + wall_buffer_m: 0.5, + wall_buffer_weight: 4.0, + step_threshold_m: 0.25, + step_penalty_weight: 4.0, + goal_tolerance: 0.3, + viz_publish_hz: 2.0, + }; + plan(plg, start, goal, &config) + } + #[test] fn snap_picks_in_column_cell() { let mut lookup = SurfaceLookup::new(); @@ -299,23 +678,25 @@ mod tests { #[test] fn plan_returns_none_if_start_cant_snap() { let plg = graph_with_nodes(&strip(20), &[(10, 0, 0)]); - let result = plan(&plg, (0.5, 0.0, 10.0), (1.0, 0.0, 0.1), VOXEL, Z_TOL); + let result = plan_simple(&plg, (0.5, 0.0, 10.0), (1.0, 0.0, 0.1)); assert!(result.is_none()); } #[test] fn plan_returns_none_if_disconnected() { + // The gap must exceed SNAP_SEARCH_RADIUS_M so no start candidate + // can relocate onto the goal island. let mut cells: Vec = (0..5).map(|x| (x, 0, 0)).collect(); - cells.extend((10..15).map(|x| (x, 0, 0))); - let plg = graph_with_nodes(&cells, &[(2, 0, 0), (12, 0, 0)]); - let result = plan(&plg, (0.25, 0.0, 0.1), (1.25, 0.0, 0.1), VOXEL, Z_TOL); + cells.extend((30..35).map(|x| (x, 0, 0))); + let plg = graph_with_nodes(&cells, &[(2, 0, 0), (32, 0, 0)]); + let result = plan_simple(&plg, (0.25, 0.0, 0.1), (3.25, 0.0, 0.1)); assert!(result.is_none()); } #[test] fn plan_same_start_and_goal_passes_through_snap_cell() { let plg = graph_with_nodes(&strip(20), &[(10, 0, 0)]); - let wp = plan(&plg, (1.0, 0.0, 0.05), (1.0, 0.0, 0.05), VOXEL, Z_TOL).unwrap(); + let wp = plan_simple(&plg, (1.0, 0.0, 0.05), (1.0, 0.0, 0.05)).unwrap(); assert_eq!(wp.first(), Some(&(1.0, 0.0, 0.05))); assert_eq!(wp.last(), Some(&(1.0, 0.0, 0.05))); let snap = surface_point_xyz(10, 0, 0, VOXEL); @@ -325,7 +706,8 @@ mod tests { #[test] fn plan_traces_surface_from_pose_to_first_node() { let plg = graph_with_nodes(&strip(20), &[(3, 0, 0), (15, 0, 0)]); - let wp = plan(&plg, (0.2, 0.0, 0.05), (1.7, 0.0, 0.05), VOXEL, Z_TOL).unwrap(); + let wp = plan_simple(&plg, (0.2, 0.0, 0.05), (1.7, 0.0, 0.05)).unwrap(); + // First waypoint is the robot's own snapped cell, not a jump ahead. let start_cell_pos = surface_point_xyz(2, 0, 0, VOXEL); let goal_cell_pos = surface_point_xyz(17, 0, 0, VOXEL); assert_eq!(wp[1], start_cell_pos); @@ -333,16 +715,150 @@ mod tests { } #[test] - fn plan_three_nodes_visits_them_all() { + fn plan_lead_in_does_not_backtrack_to_region_node() { + // Robot at cell 5 is in node 3's region but sits between it and node 15. + let plg = graph_with_nodes(&strip(20), &[(3, 0, 0), (15, 0, 0)]); + let wp = plan_simple(&plg, (0.55, 0.0, 0.05), (1.7, 0.0, 0.05)).unwrap(); + let xs: Vec = wp[1..wp.len() - 1] + .iter() + .map(|w| (w.0 / VOXEL).floor() as i32) + .collect(); + assert_eq!(xs.first(), Some(&5)); + assert!( + xs.windows(2).all(|p| p[1] >= p[0]), + "lead-in walked backward: {xs:?}" + ); + } + + fn waypoint_key(w: &(f32, f32, f32)) -> VoxelKey { + ( + (w.0 / VOXEL).floor() as i32, + (w.1 / VOXEL).floor() as i32, + (w.2 / VOXEL).round() as i32 - 1, + ) + } + + #[test] + fn plan_path_segments_stay_on_the_surface() { let plg = graph_with_nodes(&strip(20), &[(3, 0, 0), (10, 0, 0), (17, 0, 0)]); - let wp = plan(&plg, (0.2, 0.0, 0.05), (1.9, 0.0, 0.05), VOXEL, Z_TOL).unwrap(); - let node_xy: Vec<(f32, f32)> = plg.nodes.iter().map(|n| (n.pos.0, n.pos.1)).collect(); - for &(nx, ny) in &node_xy { + let wp = plan_simple(&plg, (0.2, 0.0, 0.05), (1.9, 0.0, 0.05)).unwrap(); + // Smoothed waypoints are no longer cell-adjacent, but each segment + // between them must still stay on the surface. + let step_cells = (0.25f32 / VOXEL).floor() as i32; + for w in &wp[1..wp.len() - 1] { + assert!( + plg.cells.id(waypoint_key(w)).is_some(), + "waypoint {w:?} is off the surface" + ); + } + let wc = WallCost { + clearance_m: 0.2, + buffer_m: 0.5, + buffer_weight: 4.0, + voxel_size: VOXEL, + }; + for pair in wp[1..wp.len() - 1].windows(2) { assert!( - wp.iter() - .any(|w| (w.0 - nx).abs() < 1e-5 && (w.1 - ny).abs() < 1e-5), - "node ({nx}, {ny}) should appear among waypoints" + segment_metrics( + &plg, + waypoint_key(&pair[0]), + waypoint_key(&pair[1]), + step_cells, + &wc + ) + .is_some(), + "segment {:?} -> {:?} leaves the surface", + pair[0], + pair[1] ); } } + + #[test] + fn string_pull_straightens_open_area() { + // Filled rectangle: every straight segment is on-surface, so the diagonal + // path collapses instead of staircasing through the nodes. + let mut cells: Vec = Vec::new(); + for x in 0..10 { + for y in 0..6 { + cells.push((x, y, 0)); + } + } + let plg = graph_with_nodes(&cells, &[(2, 2, 0), (7, 3, 0)]); + let wp = plan_simple(&plg, (0.05, 0.05, 0.05), (0.85, 0.55, 0.05)).unwrap(); + let interior = wp.len() - 2; + assert!( + interior <= 4, + "path not straightened: {interior} interior points" + ); + } + + #[test] + fn segment_metrics_rejects_vertical_chord() { + let plg = PlannerGraph::new(); + let wc = WallCost { + clearance_m: 0.2, + buffer_m: 0.3, + buffer_weight: 4.0, + voxel_size: VOXEL, + }; + assert!(segment_metrics(&plg, (5, 5, 0), (5, 5, 4), 2, &wc).is_none()); + assert_eq!( + segment_metrics(&plg, (5, 5, 0), (5, 5, 0), 2, &wc), + Some((1.0, 0.0)) + ); + } + + #[test] + fn string_pull_refuses_shortcut_through_sub_clearance_cell() { + // Straight strip: with open clearance the run collapses to its + // endpoints. Drop one mid cell below the hard clearance and the shortcut + // spanning it is refused, so the smoothed path retains that cell. + let mut plg = PlannerGraph::new(); + build_surface_lookup(&strip(10), &mut plg.surface_lookup); + build_surface_cells(&mut plg.cells, &plg.surface_lookup, VOXEL, 2); + let path: Vec = (0..10).map(|x| plg.cells.id((x, 0, 0)).unwrap()).collect(); + + let wc = WallCost { + clearance_m: 0.2, + buffer_m: 0.5, + buffer_weight: 4.0, + voxel_size: VOXEL, + }; + plg.wall_state.dist = vec![f32::INFINITY; plg.cells.slot_capacity()]; + let open = string_pull(&plg, &path, 1, &wc); + assert_eq!(open.len(), 2, "open strip should collapse to its endpoints"); + + let mid = plg.cells.id((5, 0, 0)).unwrap(); + plg.wall_state.dist[mid as usize] = 0.1; // below the 0.2 clearance + let guarded = string_pull(&plg, &path, 1, &wc); + assert!( + guarded.len() > 2, + "shortcut across a sub-clearance cell must be refused: {guarded:?}" + ); + assert!( + guarded.contains(&mid), + "smoothed path must still traverse the low-clearance cell" + ); + } + + #[test] + fn plan_enters_on_goalward_node_not_nearest() { + // Robot sits past node 2 toward the goal. Entry must skip it for node 10. + let plg = graph_with_nodes(&strip(20), &[(2, 0, 0), (10, 0, 0)]); + let wp = plan_simple(&plg, (0.45, 0.0, 0.05), (1.25, 0.0, 0.05)).unwrap(); + let nearest = surface_point_xyz(2, 0, 0, VOXEL); + assert!( + !wp.iter().any(|w| (w.0 - nearest.0).abs() < 1e-5), + "path doubled back to the nearest node: {wp:?}" + ); + let xs: Vec = wp[1..wp.len() - 1] + .iter() + .map(|w| (w.0 / VOXEL).floor() as i32) + .collect(); + assert!( + xs.windows(2).all(|p| p[1] >= p[0]), + "path stepped backward: {xs:?}" + ); + } } diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs index 58021d5d7f..b92b9ef46e 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs @@ -8,8 +8,9 @@ use pyo3::prelude::*; use validator::Validate; use crate::edges::edges_to_segments; -use crate::mls_planner::{Config, Planner}; +use crate::mls_planner::{Config, Planner, RegionBounds}; use crate::voxel::surface_point_xyz; +use crate::voxel::VoxelKey; #[pyclass] pub struct MLSPlanner { @@ -20,34 +21,46 @@ pub struct MLSPlanner { #[pymethods] impl MLSPlanner { #[new] + #[allow(clippy::too_many_arguments)] #[pyo3(signature = ( *, voxel_size, robot_height, - surface_dilation_passes = 3, - surface_erosion_passes = 3, + surface_closing_radius = 0.3, node_spacing_m = 1.0, - node_wall_buffer_m = 0.3, - node_step_threshold_m = 0.25, + wall_clearance_m = 0.3, + wall_buffer_m = 0.75, + wall_buffer_weight = 100.0, + step_threshold_m = 0.25, + step_penalty_weight = 4.0, ))] fn new( voxel_size: f32, robot_height: f32, - surface_dilation_passes: u32, - surface_erosion_passes: u32, + surface_closing_radius: f32, node_spacing_m: f32, - node_wall_buffer_m: f32, - node_step_threshold_m: f32, + wall_clearance_m: f32, + wall_buffer_m: f32, + wall_buffer_weight: f32, + step_threshold_m: f32, + step_penalty_weight: f32, ) -> PyResult { let config = Config { world_frame: String::new(), voxel_size, robot_height, - surface_dilation_passes, - surface_erosion_passes, + surface_closing_radius, node_spacing_m, - node_wall_buffer_m, - node_step_threshold_m, + wall_clearance_m, + wall_buffer_m, + wall_buffer_weight, + step_threshold_m, + step_penalty_weight, + // Only the binary's replan loop reads goal_tolerance. This + // in-process binding plans on demand and never consults it. + goal_tolerance: 1.0, + // Only the binary's worker publishes viz artifacts. Unused here. + viz_publish_hz: 1.0, }; config .validate() @@ -86,12 +99,56 @@ impl MLSPlanner { Ok(()) } + #[pyo3(signature = (points, origin, radius, z_min, z_max))] + fn update_region( + &mut self, + py: Python<'_>, + points: &Bound<'_, PyAny>, + origin: (f32, f32), + radius: f32, + z_min: f32, + z_max: f32, + ) -> PyResult<()> { + let points: PyReadonlyArray2<'_, f32> = points + .extract() + .map_err(|_| PyValueError::new_err("points must be a (N, 3) float32 numpy array"))?; + let shape = points.shape(); + if shape[1] != 3 { + return Err(PyValueError::new_err(format!( + "points must be (N, 3) float32, got shape {:?}", + shape + ))); + } + let arr = points.as_array(); + let n = shape[0]; + let pts: Vec<(f32, f32, f32)> = (0..n) + .filter_map(|i| { + let x = arr[[i, 0]]; + let y = arr[[i, 1]]; + let z = arr[[i, 2]]; + (x.is_finite() && y.is_finite() && z.is_finite()).then_some((x, y, z)) + }) + .collect(); + + let bounds = RegionBounds { + origin_x: origin.0, + origin_y: origin.1, + radius, + z_min, + z_max, + }; + let config = &self.config; + let planner = &mut self.planner; + py.allow_threads(move || planner.update_region(&pts, &bounds, config)); + Ok(()) + } + fn surface_map<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2> { let voxel_size = self.config.voxel_size; - let surface = self.planner.surface(); + let surface: Vec = self.planner.surface().collect(); let positions: Vec = py.allow_threads(|| { let mut out: Vec = Vec::with_capacity(surface.len() * 3); - for &(ix, iy, iz) in surface { + for (ix, iy, iz) in surface { let (x, y, z) = surface_point_xyz(ix, iy, iz, voxel_size); out.push(x); out.push(y); @@ -105,6 +162,28 @@ impl MLSPlanner { .into_pyarray(py) } + /// Surface cells as (M, 4) float32 rows of x, y, z, clearance, where + /// clearance is the distance to the nearest untraversable edge. + fn surface_clearance_map<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2> { + let voxel_size = self.config.voxel_size; + let cells = self.planner.surface_clearance(); + let values: Vec = py.allow_threads(|| { + let mut out: Vec = Vec::with_capacity(cells.len() * 4); + for ((ix, iy, iz), clearance) in cells { + let (x, y, z) = surface_point_xyz(ix, iy, iz, voxel_size); + out.push(x); + out.push(y); + out.push(z); + out.push(clearance); + } + out + }); + let n = values.len() / 4; + Array2::from_shape_vec((n, 4), values) + .expect("4 elements pushed per cell") + .into_pyarray(py) + } + fn nodes<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2> { let graph = self.planner.graph(); let positions: Vec = py.allow_threads(|| { @@ -164,6 +243,30 @@ impl MLSPlanner { ) } + fn voxel_count(&self) -> usize { + self.planner.voxel_count() + } + + /// Accumulated occupied voxel centers as (N, 3) float32, for visualization. + fn voxel_map<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2> { + let vs = self.config.voxel_size; + let half = vs * 0.5; + let keys: Vec<(i32, i32, i32)> = self.planner.voxel_keys().collect(); + let positions: Vec = py.allow_threads(|| { + let mut out: Vec = Vec::with_capacity(keys.len() * 3); + for (kx, ky, kz) in keys { + out.push(kx as f32 * vs + half); + out.push(ky as f32 * vs + half); + out.push(kz as f32 * vs + half); + } + out + }); + let n = positions.len() / 3; + Array2::from_shape_vec((n, 3), positions) + .expect("3 elements pushed per voxel") + .into_pyarray(py) + } + fn clear(&mut self) { self.planner = Planner::default(); } @@ -173,7 +276,7 @@ impl MLSPlanner { format!( "MLSPlanner(voxel_size={}, surface_cells={}, nodes={}, edges={})", self.config.voxel_size, - self.planner.surface().len(), + self.planner.surface().count(), graph.nodes.len(), graph.node_edges.len(), ) diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/surfaces.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/surfaces.rs index 7451b959c4..19e72de8c2 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/surfaces.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/surfaces.rs @@ -36,8 +36,7 @@ fn is_standable(ix: i32, iy: i32, iz: i32, by_col: &ColumnIz, clearance_cells: i pub fn extract_surfaces( voxel_map: &AHashSet, clearance_cells: i32, - dilation_passes: u32, - erosion_passes: u32, + closing_passes: u32, by_col: &mut ColumnIz, out: &mut Vec, ) { @@ -61,27 +60,92 @@ pub fn extract_surfaces( .par_iter() .flat_map_iter(|((ix, iy), zs)| { let mut local: Vec = Vec::new(); - for w in zs.windows(2) { - if w[1] - w[0] > clearance_cells { - local.push((*ix, *iy, w[0])); + standable_in_column(*ix, *iy, zs, clearance_cells, &mut local); + local + }) + .collect(); + drop(entries); + + close_surface_holes(standable, by_col, closing_passes, clearance_cells, out); +} + +/// Standable cells in one column: any cell with robot clearance above, plus +/// the topmost cell. +fn standable_in_column( + ix: i32, + iy: i32, + zs: &[i32], + clearance_cells: i32, + out: &mut Vec, +) { + for w in zs.windows(2) { + if w[1] - w[0] > clearance_cells { + out.push((ix, iy, w[0])); + } + } + if let Some(&last_iz) = zs.last() { + out.push((ix, iy, last_iz)); + } +} + +/// Insert a voxel into the per-column index, keeping each column sorted. +pub fn add_to_by_col(by_col: &mut ColumnIz, (ix, iy, iz): VoxelKey) { + let zs = by_col.entry((ix, iy)).or_default(); + if let Err(pos) = zs.binary_search(&iz) { + zs.insert(pos, iz); + } +} + +/// Remove a voxel from the per-column index, dropping emptied columns. +pub fn remove_from_by_col(by_col: &mut ColumnIz, (ix, iy, iz): VoxelKey) { + if let Some(zs) = by_col.get_mut(&(ix, iy)) { + if let Ok(pos) = zs.binary_search(&iz) { + zs.remove(pos); + } + if zs.is_empty() { + by_col.remove(&(ix, iy)); + } + } +} + +/// Re-extract surface cells whose columns fall in the inclusive write box. +/// Reads by_col over the box plus the morphology halo so closing at the +/// boundary matches a full rebuild, then filters back to the box. by_col must +/// already be current. +pub fn extract_surfaces_region( + by_col: &ColumnIz, + clearance_cells: i32, + closing_passes: u32, + write: (i32, i32, i32, i32), +) -> Vec { + let (wx0, wx1, wy0, wy1) = write; + let pad = (2 * closing_passes) as i32; + + let standable: Vec = ((wx0 - pad)..(wx1 + pad + 1)) + .into_par_iter() + .flat_map_iter(|ix| { + let mut local: Vec = Vec::new(); + for iy in (wy0 - pad)..=(wy1 + pad) { + if let Some(zs) = by_col.get(&(ix, iy)) { + standable_in_column(ix, iy, zs, clearance_cells, &mut local); } } - if let Some(&last_iz) = zs.last() { - local.push((*ix, *iy, last_iz)); - } local }) .collect(); - drop(entries); + let mut closed: Vec = Vec::new(); close_surface_holes( standable, by_col, - dilation_passes, - erosion_passes, + closing_passes, clearance_cells, - out, + &mut closed, ); + closed + .into_iter() + .filter(|&(ix, iy, _)| ix >= wx0 && ix <= wx1 && iy >= wy0 && iy <= wy1) + .collect() } /// Dilation and erosion on all xy slices of the extracted surfaces @@ -89,12 +153,11 @@ pub fn extract_surfaces( fn close_surface_holes( standable: Vec, by_col: &ColumnIz, - dilation_passes: u32, - erosion_passes: u32, + closing_passes: u32, clearance_cells: i32, out: &mut Vec, ) { - if standable.is_empty() || (dilation_passes == 0 && erosion_passes == 0) { + if standable.is_empty() || closing_passes == 0 { out.extend(standable); return; } @@ -105,16 +168,11 @@ fn close_surface_holes( } let slices: Vec<(i32, Vec<(i32, i32)>)> = by_z.into_iter().collect(); - out.par_extend(slices.par_iter().flat_map_iter(|(iz, xys)| { - close_at_z( - xys, - *iz, - by_col, - dilation_passes, - erosion_passes, - clearance_cells, - ) - })); + out.par_extend( + slices.par_iter().flat_map_iter(|(iz, xys)| { + close_at_z(xys, *iz, by_col, closing_passes, clearance_cells) + }), + ); } /// Close holes on an xy slice of the surfaces. @@ -122,11 +180,10 @@ fn close_at_z( xys: &[(i32, i32)], iz: i32, by_col: &ColumnIz, - dilation_passes: u32, - erosion_passes: u32, + closing_passes: u32, clearance_cells: i32, ) -> Vec { - let pad = (dilation_passes + erosion_passes) as i32; + let pad = (2 * closing_passes) as i32; let mut min_x = i32::MAX; let mut max_x = i32::MIN; let mut min_y = i32::MAX; @@ -148,12 +205,9 @@ fn close_at_z( img.put_pixel((ix - x0) as u32, (iy - y0) as u32, ON); } - if dilation_passes > 0 { - img = dilate(&img, Norm::L1, dilation_passes.min(u8::MAX as u32) as u8); - } - if erosion_passes > 0 { - img = erode(&img, Norm::L1, erosion_passes.min(u8::MAX as u32) as u8); - } + let k = closing_passes.min(u8::MAX as u32) as u8; + img = dilate(&img, Norm::L1, k); + img = erode(&img, Norm::L1, k); let mut out = Vec::new(); for py in 0..h { @@ -181,35 +235,35 @@ mod tests { cells.iter().copied().collect() } - fn run(cells: &[VoxelKey], clearance: i32, dil: u32, ero: u32) -> Vec { + fn run(cells: &[VoxelKey], clearance: i32, closing: u32) -> Vec { let map = voxel_map(cells); let mut by_col = ColumnIz::new(); let mut out = Vec::new(); - extract_surfaces(&map, clearance, dil, ero, &mut by_col, &mut out); + extract_surfaces(&map, clearance, closing, &mut by_col, &mut out); out } #[test] fn empty_input() { - assert!(run(&[], 5, 0, 0).is_empty()); + assert!(run(&[], 5, 0).is_empty()); } #[test] fn single_cell_is_topmost_surface() { - let s = run(&[(0, 0, 0)], 5, 0, 0); + let s = run(&[(0, 0, 0)], 5, 0); assert_eq!(s, vec![(0, 0, 0)]); } #[test] fn stacked_cells_within_headroom_only_topmost_is_surface() { let cells: Vec = (0..5).map(|z| (0, 0, z)).collect(); - let s = run(&cells, 5, 0, 0); + let s = run(&cells, 5, 0); assert_eq!(s, vec![(0, 0, 4)]); } #[test] fn gap_larger_than_headroom_makes_lower_cell_standable() { - let mut s = run(&[(0, 0, 0), (0, 0, 10)], 5, 0, 0); + let mut s = run(&[(0, 0, 0), (0, 0, 10)], 5, 0); s.sort(); assert_eq!(s, vec![(0, 0, 0), (0, 0, 10)]); } @@ -229,7 +283,7 @@ mod tests { .into_iter() .map(|(dx, dy)| (dx, dy, 0)) .collect(); - let s = run(&cells, 5, 3, 3); + let s = run(&cells, 5, 3); assert!( s.contains(&(0, 0, 0)), "closing should fill the center hole" @@ -252,7 +306,7 @@ mod tests { .map(|(dx, dy)| (dx, dy, 0)) .collect(); cells.push((0, 0, 1)); - let s = run(&cells, 5, 3, 3); + let s = run(&cells, 5, 3); assert!(!s.contains(&(0, 0, 0))); } } diff --git a/dimos/navigation/nav_3d/mls_planner/test_transformer.py b/dimos/navigation/nav_3d/mls_planner/test_transformer.py index a38d63b9d6..3fca1fb081 100644 --- a/dimos/navigation/nav_3d/mls_planner/test_transformer.py +++ b/dimos/navigation/nav_3d/mls_planner/test_transformer.py @@ -25,8 +25,18 @@ from dimos.navigation.nav_3d.mls_planner.transformer import MLSPlan -def _obs(points: NDArray[np.float32], pose: tuple[float, float, float]) -> Observation[PointCloud2]: - return Observation(id=0, ts=0.0, pose=pose, _data=PointCloud2.from_numpy(points)) +def _obs( + points: NDArray[np.float32], + pose: tuple[float, float, float], + region_bounds: tuple[float, float, float, float, float], +) -> Observation[PointCloud2]: + return Observation( + id=0, + ts=0.0, + pose=pose, + tags={"region_bounds": region_bounds}, + _data=PointCloud2.from_numpy(points), + ) def _flat_floor(half_extent: float = 3.0, spacing: float = 0.1) -> NDArray[np.float32]: @@ -37,34 +47,74 @@ def _flat_floor(half_extent: float = 3.0, spacing: float = 0.1) -> NDArray[np.fl def test_flat_floor_yields_populated_path_and_planned_true() -> None: - obs = _obs(_flat_floor(), pose=(-2.0, -2.0, 1.0)) + obs = _obs( + _flat_floor(), + pose=(-2.0, -2.0, 1.0), + region_bounds=(0.0, 0.0, 5.0, -1.0, 2.0), + ) [out] = list(MLSPlan(goal=(2.0, 2.0, 0.0), voxel_size=0.2, robot_height=1.0)(iter([obs]))) assert out.tags["planned"] is True assert len(out.data.poses) >= 2 - assert out.tags["voxel_map"] is obs.data - assert out.tags["nodes"].shape[1] == 3 - assert out.tags["surface_map"].shape[1] == 3 + assert out.tags["voxels"] > 0 + assert out.tags["voxel_map"].shape == (out.tags["voxels"], 3) + assert out.tags["surface_clearance"].shape[1] == 4 + assert set(out.tags["timings"]) == {"update_ms", "plan_ms", "total_ms"} + + +def test_poseless_obs_is_skipped() -> None: + points = _flat_floor() + poseless = Observation( + id=1, + ts=0.0, + pose=None, + tags={"region_bounds": (0.0, 0.0, 5.0, -1.0, 2.0)}, + _data=PointCloud2.from_numpy(points), + ) + posed = _obs(points, pose=(-2.0, -2.0, 1.0), region_bounds=(0.0, 0.0, 5.0, -1.0, 2.0)) + results = list( + MLSPlan(goal=(2.0, 2.0, 0.0), voxel_size=0.2, robot_height=1.0)(iter([poseless, posed])) + ) -def test_no_route_yields_empty_path_with_planned_false() -> None: - rng = np.random.default_rng(0) - obs = _obs(rng.random((50, 3)).astype(np.float32), pose=(0.0, 0.0, 0.0)) + assert len(results) == 1 - [out] = list(MLSPlan(goal=(100.0, 100.0, 100.0))(iter([obs]))) - assert out.tags["planned"] is False - assert out.data.poses == [] +def test_missing_region_bounds_raises() -> None: + obs = Observation( + id=0, + ts=0.0, + pose=(0.0, 0.0, 0.0), + _data=PointCloud2.from_numpy(np.zeros((1, 3), dtype=np.float32)), + ) + with pytest.raises(ValueError, match="local map slices"): + list(MLSPlan(goal=(1.0, 1.0, 0.0))(iter([obs]))) -def test_poseless_obs_is_skipped_and_following_posed_obs_plans() -> None: - poseless = Observation(id=1, ts=0.0, pose=None, _data=PointCloud2.from_numpy(_flat_floor())) - posed = _obs(_flat_floor(), pose=(-2.0, -2.0, 1.0)) - results = list( - MLSPlan(goal=(2.0, 2.0, 0.0), voxel_size=0.2, robot_height=1.0)(iter([poseless, posed])) +def test_start_z_is_dropped_by_robot_height() -> None: + obs = _obs( + np.zeros((1, 3), dtype=np.float32), + pose=(1.0, 2.0, 3.0), + region_bounds=(1.0, 2.0, 5.0, -1.0, 5.0), ) - assert len(results) == 1 - assert results[0].tags["planned"] is True + [out] = list(MLSPlan(goal=(10.0, 10.0, 0.0), robot_height=0.4)(iter([obs]))) + + assert out.tags["start"] == (1.0, 2.0, 3.0 - 0.4) + + +def test_no_route_yields_empty_path_with_planned_false() -> None: + # Random points form no traversable surface, so planning fails. + rng = np.random.default_rng(0) + obs = _obs( + rng.random((50, 3)).astype(np.float32), + pose=(0.0, 0.0, 0.0), + region_bounds=(0.5, 0.5, 2.0, -1.0, 2.0), + ) + + [out] = list(MLSPlan(goal=(100.0, 100.0, 100.0))(iter([obs]))) + + assert out.tags["planned"] is False + assert out.data.poses == [] diff --git a/dimos/navigation/nav_3d/mls_planner/transformer.py b/dimos/navigation/nav_3d/mls_planner/transformer.py index 3448bebf22..3f9bc09b7a 100644 --- a/dimos/navigation/nav_3d/mls_planner/transformer.py +++ b/dimos/navigation/nav_3d/mls_planner/transformer.py @@ -14,6 +14,7 @@ from __future__ import annotations +import time from typing import TYPE_CHECKING, Any from dimos.memory2.transform import Transformer @@ -21,6 +22,7 @@ from dimos.msgs.nav_msgs.Path import Path from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.navigation.nav_3d.mls_planner.mls_planner import MLSPlanner +from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: from collections.abc import Iterator @@ -30,6 +32,8 @@ from dimos.memory2.type.observation import Observation +logger = setup_logger() + class MLSPlan(Transformer[PointCloud2, Path]): """Plan paths from current pose to a fixed goal over an accumulating voxel map.""" @@ -72,24 +76,41 @@ def __call__( ) for obs in upstream: if obs.pose_tuple is None: + logger.debug("MLSPlan: obs %s has no pose; skipping", obs.id) continue x, y, z, *_ = obs.pose_tuple start = (float(x), float(y), float(z) - self.robot_height) - voxel_map = obs.data - planner.update_global_map(voxel_map.points_f32()) + bounds = obs.tags.get("region_bounds") + if bounds is None: + raise ValueError( + "MLSPlan consumes local map slices; construct RayTraceMap(emit_local=True)" + ) + ox, oy, radius, z_min, z_max = bounds + t_update = time.perf_counter() + planner.update_region(obs.data.points_f32(), (ox, oy), radius, z_min, z_max) + t_plan = time.perf_counter() waypoints = planner.plan(start, self.goal) + t_done = time.perf_counter() path = self._path_from_waypoints(waypoints, obs.ts) + timings = { + "update_ms": (t_plan - t_update) * 1000, + "plan_ms": (t_done - t_plan) * 1000, + "total_ms": (t_done - t_update) * 1000, + } + yield obs.derive( data=path, tags={ **obs.tags, - "voxel_map": voxel_map, - "surface_map": planner.surface_map(), + "voxel_map": planner.voxel_map(), + "surface_clearance": planner.surface_clearance_map(), "nodes": planner.nodes(), "node_edges": planner.node_edges(), "start": start, "planned": waypoints is not None, + "timings": timings, + "voxels": planner.voxel_count(), }, ) diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index fea6c89aa6..3bb539eaae 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -12,38 +12,91 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Dev replay util: run a lidar+odometry .db through RayTraceMap and MLSPlan into rerun.""" +"""Replay a lidar+odometry .db through RayTraceMap and the MLS planner into rerun. + +Pass one or more --config clearance,buffer,weight to overlay each as a colored path. +""" from __future__ import annotations from pathlib import Path as FsPath -from typing import TYPE_CHECKING +from time import perf_counter +import numpy as np +from numpy.typing import NDArray import rerun as rr import typer from dimos.mapping.ray_tracing.transformer import RayTraceMap +from dimos.memory2.store.memory import MemoryStore from dimos.memory2.store.sqlite import SqliteStore +from dimos.memory2.stream import Stream from dimos.memory2.transform import FnTransformer from dimos.memory2.type.observation import Observation from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.nav_msgs.Path import Path from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation -from dimos.navigation.nav_3d.mls_planner.transformer import MLSPlan +from dimos.navigation.nav_3d.mls_planner.mls_planner import MLSPlanner from dimos.utils.data import resolve_named_path -if TYPE_CHECKING: - import numpy as np - from numpy.typing import NDArray - TIMELINE = "ts" +TIMING_KEYS = ["update_ms", "plan_ms", "total_ms"] +SIZE_KEYS = ["voxels", "surface_cells", "nodes", "edges"] + +# Distinct path colors for overlaid configurations, config 0 first. +PATH_PALETTE = [ + [0, 255, 0], + [255, 0, 255], + [0, 200, 255], + [255, 180, 0], + [255, 80, 80], + [160, 120, 255], + [120, 255, 200], + [255, 255, 120], +] + + +def _parse_configs( + specs: list[str] | None, + clearance: float, + buffer: float, + weight: float, +) -> list[tuple[float, float, float]]: + """Each spec is 'clearance,buffer,weight'. Falls back to the single flags.""" + if not specs: + return [(clearance, buffer, weight)] + out: list[tuple[float, float, float]] = [] + for spec in specs: + parts = spec.replace(" ", "").split(",") + if len(parts) != 3: + raise typer.BadParameter(f"--config must be 'clearance,buffer,weight'; got {spec!r}") + c, b, w = (float(p) for p in parts) + out.append((c, b, w)) + return out + + +def _print_summary(streams: dict[str, dict[str, Stream[float]]]) -> None: + print("\nper-frame summary (mean / p50 / p95 / max):") + for kind, by_key in streams.items(): + for key, stream in by_key.items(): + values = [obs.data for obs in stream] + if not values: + continue + arr = np.asarray(values, dtype=np.float64) + mean, p50, p95, peak = ( + arr.mean(), + np.percentile(arr, 50), + np.percentile(arr, 95), + arr.max(), + ) + print(f" {kind}/{key:<14} {mean:9.2f} {p50:9.2f} {p95:9.2f} {peak:9.2f}") + + PairObs = Observation[tuple[Observation[PointCloud2], Observation[Odometry]]] def _attach_pose_from_odom(pair_obs: PairObs) -> Observation[PointCloud2]: - lidar_obs = pair_obs.data[0] - odom_obs = pair_obs.data[1] + lidar_obs, odom_obs = pair_obs.data odom = odom_obs.data pose_tuple = ( float(odom.position.x), @@ -68,12 +121,60 @@ def _log_edges(edges: NDArray[np.float32], entity: str) -> None: rr.log(entity, rr.LineStrips3D(segments)) -def _log_path(path: Path, entity: str) -> None: - if not path.poses: +def _log_path_wp(waypoints: NDArray[np.float32] | None, entity: str, color: list[int]) -> None: + if waypoints is None or len(waypoints) == 0: rr.log(entity, rr.LineStrips3D([])) return - points = [(float(p.position.x), float(p.position.y), float(p.position.z)) for p in path.poses] - rr.log(entity, rr.LineStrips3D([points], colors=[[0, 255, 0]], radii=0.05)) + points = [(float(p[0]), float(p[1]), float(p[2])) for p in waypoints] + rr.log(entity, rr.LineStrips3D([points], colors=[color], radii=0.05)) + + +def _clearance_colors(clearance: NDArray[np.float32], clamp_m: float) -> NDArray[np.uint8]: + """Map per-cell wall clearance to a blue ramp, clamped so it resolves near walls.""" + norm = np.clip(np.nan_to_num(clearance / clamp_m, nan=1.0, posinf=1.0), 0.0, 1.0) + blocked = np.array([4.0, 8.0, 48.0], dtype=np.float64) + clear = np.array([150.0, 200.0, 255.0], dtype=np.float64) + rgb: NDArray[np.float64] = blocked + norm[:, None] * (clear - blocked) + return rgb.astype(np.uint8) + + +def _log_shared( + start: tuple[float, float, float], + planner: MLSPlanner, + render_voxel: float, + clearance_clamp: float, +) -> tuple[NDArray[np.float32], NDArray[np.float32], NDArray[np.float32]]: + """Log the map artifacts shared by every config from a reference planner. + + Returns (surface, nodes, edges) for metric sizing. + """ + rr.log("world/start", rr.Points3D([start], colors=[[0, 255, 0]], radii=0.1)) + + voxel_map = planner.voxel_map() + if voxel_map.size: + rr.log( + "world/voxel_map", + rr.Points3D(voxel_map, colors=[[180, 125, 125]], radii=render_voxel / 2), + ) + + surface = planner.surface_clearance_map() + if surface.size: + rr.log( + "world/surface_map", + rr.Points3D( + surface[:, :3], + colors=_clearance_colors(surface[:, 3], clearance_clamp), + radii=render_voxel / 2, + ), + ) + + nodes = planner.nodes() + if nodes.size: + rr.log("world/nodes", rr.Points3D(nodes, colors=[[255, 200, 0]], radii=0.05)) + + edges = planner.node_edges() + _log_edges(edges, "world/node_edges") + return surface, nodes, edges def main( @@ -90,19 +191,62 @@ def main( align_tol: float = typer.Option(0.05, "--align-tol", help="Lidar/odom alignment tolerance (s)"), voxel_size: float = typer.Option(0.1, "--voxel-size", help="Voxel edge length (m)"), max_range: float = typer.Option(30.0, "--max-range", help="Max ray cast distance (m)"), + registered_clouds: bool = typer.Option( + True, + "--registered-clouds/--no-registered-clouds", + help="Clouds are already world-registered (fastlio). Use --no-registered-clouds for " + "sensor-frame clouds (pointlio); the odom pose then registers each frame.", + ), ray_subsample: int = typer.Option(1, "--ray-subsample", help="Keep every Nth ray"), emit_every: int = typer.Option(1, "--emit-every", help="Replan every N lidar frames"), - robot_height: float = typer.Option(0.3, "--robot-height", help="Robot height (m)"), + robot_height: float = typer.Option(1.0, "--robot-height", help="Robot height (m)"), + surface_closing_radius: float = typer.Option( + 0.3, + "--surface-closing-radius", + help="Hole-fill radius (m); morphological closing fills holes up to twice this wide", + ), node_spacing: float = typer.Option(1.0, "--node-spacing", help="Graph node spacing (m)"), + wall_clearance: float = typer.Option( + 0.3, + "--wall-clearance", + help="Hard clearance; cells closer to a wall or edge are impassable (m)", + ), + wall_buffer: float = typer.Option( + 0.75, "--wall-buffer", help="Width of the soft standoff zone beyond clearance (m)" + ), + wall_buffer_weight: float = typer.Option( + 100.0, "--wall-buffer-weight", help="Peak soft wall penalty at the clearance edge" + ), + step_height: float = typer.Option( + 0.25, + "--step-height", + help="Max traversable vertical step (m); taller steps are impassable", + ), + step_penalty_weight: float = typer.Option( + 4.0, "--step-penalty-weight", help="Soft cost per meter of vertical climb" + ), + config: list[str] = typer.Option( + None, + "--config", + help="Repeatable 'clearance,buffer,weight' to overlay as colored paths; " + "overrides the single --wall-* flags", + ), goal: tuple[float, float, float] = typer.Option( - (0.0, 0.0, 0.0), - "--goal", - help="Planner goal xyz. Default is dataset-specific; override per recording.", + (0.0, 0.0, 0.0), "--goal", help="Planner goal xyz; override per recording" ), live: bool = typer.Option( False, "--live", help="Also spawn the rerun viewer when --out is set" ), render_voxel: float = typer.Option(0.05, "--render-voxel", help="Rerun voxel render size (m)"), + clearance_clamp: float = typer.Option( + 1.0, "--clearance-clamp", help="Max clearance (m) for the surface color scale" + ), + from_time: float | None = typer.Option( + None, "--from-time", help="Start timestamp (s); default is the stream start" + ), + to_time: float | None = typer.Option( + None, "--to-time", help="End timestamp (s); default is the stream end" + ), ) -> None: db_path = resolve_named_path(dataset, ".db") @@ -120,60 +264,109 @@ def main( store = SqliteStore(path=str(db_path)) with store: lidar = store.stream(lidar_stream, PointCloud2).order_by("ts") + if from_time is not None: + lidar = lidar.from_time(from_time) + if to_time is not None: + lidar = lidar.to_time(to_time) odom = store.stream(odom_stream, Odometry).order_by("ts") pose_tagged = lidar.align(odom, tolerance=align_tol).transform( FnTransformer(_attach_pose_from_odom) ) - pipeline = pose_tagged.transform( + ray_pipeline = pose_tagged.transform( RayTraceMap( voxel_size=voxel_size, max_range=max_range, ray_subsample=ray_subsample, emit_every=emit_every, + emit_local=True, + registered_clouds=registered_clouds, ) - ).transform( - MLSPlan( - goal=goal, + ) + + configs = _parse_configs(config, wall_clearance, wall_buffer, wall_buffer_weight) + planners: list[tuple[str, list[int], MLSPlanner]] = [] + for i, (clr, buf, wgt) in enumerate(configs): + planner = MLSPlanner( voxel_size=voxel_size, robot_height=robot_height, + surface_closing_radius=surface_closing_radius, node_spacing_m=node_spacing, + wall_clearance_m=clr, + wall_buffer_m=buf, + wall_buffer_weight=wgt, + step_threshold_m=step_height, + step_penalty_weight=step_penalty_weight, ) - ) + color = PATH_PALETTE[i % len(PATH_PALETTE)] + label = f"cfg{i}_c{clr:g}_b{buf:g}_w{wgt:g}" + planners.append((label, color, planner)) + print(f"config {i}: clearance={clr} buffer={buf} weight={wgt} color={color} -> {label}") rr.log("world/goal", rr.Points3D([goal], colors=[[255, 0, 0]], radii=0.1), static=True) - for obs in pipeline: - rr.set_time(TIMELINE, timestamp=obs.ts) + metrics = MemoryStore() + timing_streams = {k: metrics.stream(f"timing_{k}", float) for k in TIMING_KEYS} + size_streams = {k: metrics.stream(f"size_{k}", float) for k in SIZE_KEYS} - start = obs.tags["start"] - rr.log("world/start", rr.Points3D([start], colors=[[0, 255, 0]], radii=0.1)) + try: + frame = 0 + for ray_obs in ray_pipeline: + if ray_obs.pose_tuple is None: + continue + bounds = ray_obs.tags.get("region_bounds") + if bounds is None: + raise ValueError("plan_rrd needs RayTraceMap(emit_local=True)") + px, py, pz, *_ = ray_obs.pose_tuple + start = (float(px), float(py), float(pz) - robot_height) + ox, oy, radius, z_min, z_max = bounds + pts = ray_obs.data.points_f32() + rr.set_time(TIMELINE, timestamp=ray_obs.ts) - voxel_map = obs.tags["voxel_map"] - rr.log("world/voxel_map", voxel_map.to_rerun(voxel_size=render_voxel)) + ref_timing: dict[str, float] = {} + surface = nodes = edges = np.empty((0,), dtype=np.float32) + for j, (label, color, planner) in enumerate(planners): + t0 = perf_counter() + planner.update_region(pts, (ox, oy), radius, z_min, z_max) + t1 = perf_counter() + waypoints = planner.plan(start, goal) + t2 = perf_counter() + _log_path_wp(waypoints, f"world/paths/{label}", color) + if j == 0: + ref_timing = { + "update_ms": (t1 - t0) * 1000, + "plan_ms": (t2 - t1) * 1000, + "total_ms": (t2 - t0) * 1000, + } + surface, nodes, edges = _log_shared( + start, planner, render_voxel, clearance_clamp + ) - surface = obs.tags["surface_map"] - if surface.size: - rr.log( - "world/surface_map", - rr.Points3D(surface, colors=[[120, 120, 200]], radii=render_voxel / 2), - ) - - nodes = obs.tags["nodes"] - if nodes.size: - rr.log("world/nodes", rr.Points3D(nodes, colors=[[255, 200, 0]], radii=0.05)) - - _log_edges(obs.tags["node_edges"], "world/node_edges") - _log_path(obs.data, "world/path") + for key, value in ref_timing.items(): + timing_streams[key].append(float(value), ts=ray_obs.ts) + rr.log(f"metrics/timing/{key}", rr.Scalars(value)) + sizes = { + "voxels": planners[0][2].voxel_count(), + "surface_cells": len(surface), + "nodes": len(nodes), + "edges": len(edges), + } + for key, value in sizes.items(): + size_streams[key].append(float(value), ts=ray_obs.ts) + rr.log(f"metrics/size/{key}", rr.Scalars(value)) - count = obs.tags.get("frame_count", "?") - planned = obs.tags.get("planned", False) - print( - f"frame_count={count} planned={planned} waypoints={len(obs.data.poses)}", - end="\r", - flush=True, - ) - print() + frame += 1 + print( + f"frame={frame} configs={len(planners)} " + f"rebuild(ref)={ref_timing['total_ms'] - ref_timing['plan_ms']:.1f}ms " + f"plan(ref)={ref_timing['plan_ms']:.1f}ms", + end="\r", + flush=True, + ) + except KeyboardInterrupt: + print("\ninterrupted; reporting metrics for completed frames") + finally: + _print_summary({"timing": timing_streams, "size": size_streams}) if out is not None: print(f"wrote {out}") diff --git a/pyproject.toml b/pyproject.toml index 52e79530dd..cac7a453c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -505,6 +505,7 @@ module = [ "cyclonedds", "cyclonedds.*", "dimos_lcm.*", + "dimos_mls_planner", "dimos_voxel_ray_tracing", "etils", "faster_whisper", From c3c7a9a91e908a6a06004fbe934b844dbd649c51 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 15:40:33 -0700 Subject: [PATCH 38/71] Remove registered flag from transform --- dimos/mapping/ray_tracing/rust/flake.lock | 10 +- dimos/mapping/ray_tracing/rust/src/main.rs | 181 ++++++++++-------- dimos/mapping/ray_tracing/test_transformer.py | 2 +- .../nav_3d/mls_planner/goal_relay.py | 52 ----- .../nav_3d/mls_planner/rust/src/dijkstra.rs | 10 + .../nav_3d/mls_planner/rust/src/main.rs | 19 ++ .../mls_planner/rust/src/mls_planner.rs | 19 +- .../nav_3d/mls_planner/rust/src/nodes.rs | 2 +- .../nav_3d/mls_planner/utils/plan_rrd.py | 7 - 9 files changed, 157 insertions(+), 145 deletions(-) delete mode 100644 dimos/navigation/nav_3d/mls_planner/goal_relay.py diff --git a/dimos/mapping/ray_tracing/rust/flake.lock b/dimos/mapping/ray_tracing/rust/flake.lock index 31aab7d531..a548660557 100644 --- a/dimos/mapping/ray_tracing/rust/flake.lock +++ b/dimos/mapping/ray_tracing/rust/flake.lock @@ -3,11 +3,11 @@ "dimos-repo": { "flake": false, "locked": { - "lastModified": 1781823023, - "narHash": "sha256-YkqAYJlqFI52p1ZfM0G29+85eSykkIgXU7xs7oqgopo=", - "ref": "refs/heads/andrew/feat/kronk-nav-3", - "rev": "8ef1d416ddf141e354902ab86cedcd748220e0e9", - "revCount": 882, + "lastModified": 1779865691, + "narHash": "sha256-2CVWcov7DiC1qX/B/zFKDJiSYsnbrZ3FNT/viprFWTQ=", + "ref": "refs/heads/jeff/feat/g1_raycast", + "rev": "51666bcd298c1d08bdee179f176f45c0a7dd417d", + "revCount": 744, "type": "git", "url": "file:../../../.." }, diff --git a/dimos/mapping/ray_tracing/rust/src/main.rs b/dimos/mapping/ray_tracing/rust/src/main.rs index 2933ff26ab..d7dc7ac919 100644 --- a/dimos/mapping/ray_tracing/rust/src/main.rs +++ b/dimos/mapping/ray_tracing/rust/src/main.rs @@ -105,89 +105,77 @@ impl RayTracingVoxelMap { } self.frame_count += 1; - if self - .frame_count - .is_multiple_of(self.config.global_emit_every) - { - let cloud = build_global_cloud( - &self.map, - &live, - voxel_size, - out_frame_id, - msg.header.stamp.clone(), + let local_due = self.frame_count.is_multiple_of(self.config.emit_every); + + let cylinder = if local_due { + let margin = self.config.shadow_depth + voxel_size; + let (cx, cy, radius, z_min, z_max) = batch_local_bounds( + &self.batch_points, + &self.batch_origins, + self.config.region_percentile, + margin, ); - if let Err(e) = self.global_map.publish(&cloud).await { + self.batch_points.clear(); + self.batch_origins.clear(); + + let bounds_msg = PoseStamped { + header: Header { + seq: 0, + stamp: msg.header.stamp.clone(), + frame_id: out_frame_id.to_string(), + }, + pose: Pose { + position: Point { + x: cx as f64, + y: cy as f64, + z: 0.0, + }, + orientation: Quaternion { + x: radius as f64, + y: z_min as f64, + z: z_max as f64, + w: 0.0, + }, + }, + }; + if let Err(e) = self.region_bounds.publish(&bounds_msg).await { error_throttled!( Duration::from_secs(1), error = %e, - "Updated global voxel map failed to publish", + "Region bounds failed to publish", ); } - } - if !self.frame_count.is_multiple_of(self.config.emit_every) { - return; - } - - // Percentile-bounded cylinder over the batch plus the clearing margin. - let margin = self.config.shadow_depth + voxel_size; - let (cx, cy, radius, z_min, z_max) = batch_local_bounds( - &self.batch_points, - &self.batch_origins, - self.config.region_percentile, - margin, - ); - self.batch_points.clear(); - self.batch_origins.clear(); - let cylinder = LocalBounds { - origin_x: cx, - origin_y: cy, - r_xy_max_sq: radius * radius, - z_min, - z_max, - }; - - let bounds_msg = PoseStamped { - header: Header { - seq: 0, - stamp: msg.header.stamp.clone(), - frame_id: out_frame_id.to_string(), - }, - pose: Pose { - position: Point { - x: cx as f64, - y: cy as f64, - z: 0.0, - }, - orientation: Quaternion { - x: radius as f64, - y: z_min as f64, - z: z_max as f64, - w: 0.0, - }, - }, + Some(LocalBounds { + origin_x: cx, + origin_y: cy, + r_xy_max_sq: radius * radius, + z_min, + z_max, + }) + } else { + None }; - if let Err(e) = self.region_bounds.publish(&bounds_msg).await { - error_throttled!( - Duration::from_secs(1), - error = %e, - "Region bounds failed to publish", - ); - } - let local_cloud = build_local_cloud( - &self.map, - &live, - voxel_size, - &cylinder, - out_frame_id, - msg.header.stamp, - ); - if let Err(e) = self.local_map.publish(&local_cloud).await { - error_throttled!( - Duration::from_secs(1), - error = %e, - "Updated local voxel map failed to publish", - ); + let global_due = self + .frame_count + .is_multiple_of(self.config.global_emit_every); + + // build and publish the clouds + let stamp = msg.header.stamp; + if let Some(cyl) = &cylinder { + if global_due { + let (global, local) = + build_global_and_local(&self.map, &live, voxel_size, cyl, out_frame_id, stamp); + publish_cloud(&self.global_map, &global).await; + publish_cloud(&self.local_map, &local).await; + } else { + let local = + build_local_cloud(&self.map, &live, voxel_size, cyl, out_frame_id, stamp); + publish_cloud(&self.local_map, &local).await; + } + } else if global_due { + let global = build_global_cloud(&self.map, &live, voxel_size, out_frame_id, stamp); + publish_cloud(&self.global_map, &global).await; } } } @@ -357,6 +345,49 @@ fn build_local_cloud( make_cloud(data, n, frame_id, stamp) } +/// Build the global and local clouds in one pass over the map, so a frame that +/// emits both does not scan the voxel map twice. +fn build_global_and_local( + map: &VoxelMap, + live: &AHashSet, + voxel_size: f32, + cylinder: &LocalBounds, + frame_id: &str, + stamp: Time, +) -> (PointCloud2, PointCloud2) { + let mut g_data = Vec::with_capacity((map.voxels.len() + live.len()) * 16); + let mut g_n: i32 = 0; + let mut l_data = Vec::with_capacity(live.len() * 2 * 16); + let mut l_n: i32 = 0; + let mut push = |x: f32, y: f32, z: f32| { + write_point(&mut g_data, &mut g_n, x, y, z); + if cylinder.contains(x, y, z) { + write_point(&mut l_data, &mut l_n, x, y, z); + } + }; + for (x, y, z) in iter_global_points(map, voxel_size) { + push(x, y, z); + } + for (x, y, z) in unhealthy_live_centers(map, live, voxel_size) { + push(x, y, z); + } + ( + make_cloud(g_data, g_n, frame_id, stamp.clone()), + make_cloud(l_data, l_n, frame_id, stamp), + ) +} + +async fn publish_cloud(out: &Output, cloud: &PointCloud2) { + if let Err(e) = out.publish(cloud).await { + error_throttled!( + Duration::from_secs(1), + error = %e, + topic = %out.topic, + "Voxel map failed to publish", + ); + } +} + #[tokio::main] async fn main() { let transport = LcmTransport::new() diff --git a/dimos/mapping/ray_tracing/test_transformer.py b/dimos/mapping/ray_tracing/test_transformer.py index 20a5c6b622..ec52ec8f37 100644 --- a/dimos/mapping/ray_tracing/test_transformer.py +++ b/dimos/mapping/ray_tracing/test_transformer.py @@ -87,7 +87,7 @@ def _ring( def test_emit_local_tags_region_bounds_around_registered_origin() -> None: margin = 0.2 + 0.1 - # Sensor-frame ring centered on the sensor; the pose registers it to (2, 3, 0.5). + # Sensor-frame ring centered on the sensor. The pose registers it to (2, 3, 0.5). obs = _obs(_ring((0.0, 0.0), radius=1.0, z=0.0), ts=1.0, pose=(2.0, 3.0, 0.5)) [emitted] = list(RayTraceMap(emit_local=True)(iter([obs]))) diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py deleted file mode 100644 index 0fab1cff5c..0000000000 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ /dev/null @@ -1,52 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from __future__ import annotations - -from 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.nav_msgs.Odometry import Odometry - - -class GoalRelayConfig(ModuleConfig): - pass - - -class GoalRelay(Module): - """Convert Odometry to PoseStamped""" - - config: GoalRelayConfig - - odometry: In[Odometry] - goal: In[PointStamped] - - start_pose: Out[PoseStamped] - goal_pose: Out[PoseStamped] - - @rpc - def start(self) -> None: - super().start() - self.register_disposable(Disposable(self.odometry.subscribe(self._on_odometry))) - self.register_disposable(Disposable(self.goal.subscribe(self._on_goal))) - - def _on_odometry(self, msg: Odometry) -> None: - self.start_pose.publish(msg.to_pose_stamped()) - - def _on_goal(self, point: PointStamped) -> None: - self.goal_pose.publish(point.to_pose_stamped()) diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/dijkstra.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/dijkstra.rs index 76dfc745be..25f7250528 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/dijkstra.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/dijkstra.rs @@ -232,6 +232,16 @@ mod tests { cur } + #[test] + fn walk_preds_breaks_on_pred_cycle() { + let mut state = DijkstraState::default(); + state.reset(2); + state.pred[0] = 1; + state.pred[1] = 0; + let path = walk_preds(&state, 0); + assert_eq!(path, vec![0, 1]); + } + #[test] fn region_window_all_equals_full() { let sc = grid(10); diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs index 91452ad1f0..ccf5f8e2b9 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs @@ -526,3 +526,22 @@ async fn main() { .expect("failed to create LCM transport"); run::(transport).await; } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn is_at_goal_respects_tolerance_and_ignores_z() { + assert!(is_at_goal((0.0, 0.0, 0.0), (0.05, 0.0, 9.0), 0.1)); + assert!(!is_at_goal((0.0, 0.0, 0.0), (0.2, 0.0, 0.0), 0.1)); + } + + #[test] + fn same_stamp_compares_sec_and_nsec() { + let a = Time { sec: 5, nsec: 7 }; + assert!(same_stamp(&a, &Time { sec: 5, nsec: 7 })); + assert!(!same_stamp(&a, &Time { sec: 5, nsec: 8 })); + assert!(!same_stamp(&a, &Time { sec: 6, nsec: 7 })); + } +} diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs index 0d9c41e739..7b72b705a8 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs @@ -879,7 +879,7 @@ mod region_tests { ) }; - // Interior waypoints are exact cell centers; sample between them too. + // Interior waypoints are exact cell centers. Sample between them too. let interior = &wp[1..wp.len() - 1]; assert!(interior.len() >= 2, "expected a multi-cell path"); for pair in interior.windows(2) { @@ -900,7 +900,7 @@ mod region_tests { } } - /// Solid 0.3 m block, taller than the step threshold; the path must route + /// Solid 0.3 m block, taller than the step threshold. The path must route /// around it and never climb on. fn block_world() -> Vec<(f32, f32, f32)> { let vs = 0.1_f32; @@ -941,7 +941,7 @@ mod region_tests { .plan((1.0, 0.5, 0.05), (3.9, 0.5, 0.05), &cfg) .expect("plan exists"); - // The block top is at z = 0.4; the floor surface point is z = 0.1. No + // The block top is at z = 0.4. The floor surface point is z = 0.1. No // interior waypoint may land on the block. for w in &wp[1..wp.len() - 1] { assert!( @@ -958,7 +958,7 @@ mod region_tests { } /// Flat floor with a crossable 0.2 m ridge blocking ix 15 except a flat gap - /// at iy 10..12. Crossing is short but climbs two steps; the detour is flat. + /// at iy 10..12. Crossing is short but climbs two steps. The detour is flat. /// Route choice is read from the xy lane, since smoothing flattens the ridge /// waypoints away. fn ridge_world() -> Vec<(f32, f32, f32)> { @@ -1047,4 +1047,15 @@ mod region_tests { let last = *wp.last().expect("path has waypoints"); assert!((last.0 - goal.0).abs() < 1e-3 && (last.1 - goal.1).abs() < 1e-3); } + + #[test] + fn wall_buffer_weight_requires_buffer() { + let mut cfg = test_config(); + cfg.wall_buffer_weight = 1.0; + cfg.wall_buffer_m = 0.0; + assert!(validate_wall_buffer(&cfg).is_err()); + + cfg.wall_buffer_m = 0.3; + assert!(validate_wall_buffer(&cfg).is_ok()); + } } diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/nodes.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/nodes.rs index 93ce063bdf..1854a7986c 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/nodes.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/nodes.rs @@ -45,7 +45,7 @@ pub fn place_nodes( collect_wall_adjacent_cells(cells, &mut wall_seeds); dijkstra(cells, &wall_seeds, state, Weight::Base); - // Floor is the hard clearance; NMS already prefers the clearest cells. + // Floor is the hard clearance. NMS already prefers the clearest cells. let node_floor = wall_clearance_m; let candidates: Vec = cells .ids() diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index 3bb539eaae..ff3301f466 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -191,12 +191,6 @@ def main( align_tol: float = typer.Option(0.05, "--align-tol", help="Lidar/odom alignment tolerance (s)"), voxel_size: float = typer.Option(0.1, "--voxel-size", help="Voxel edge length (m)"), max_range: float = typer.Option(30.0, "--max-range", help="Max ray cast distance (m)"), - registered_clouds: bool = typer.Option( - True, - "--registered-clouds/--no-registered-clouds", - help="Clouds are already world-registered (fastlio). Use --no-registered-clouds for " - "sensor-frame clouds (pointlio); the odom pose then registers each frame.", - ), ray_subsample: int = typer.Option(1, "--ray-subsample", help="Keep every Nth ray"), emit_every: int = typer.Option(1, "--emit-every", help="Replan every N lidar frames"), robot_height: float = typer.Option(1.0, "--robot-height", help="Robot height (m)"), @@ -280,7 +274,6 @@ def main( ray_subsample=ray_subsample, emit_every=emit_every, emit_local=True, - registered_clouds=registered_clouds, ) ) From 1e5e7eb9abc450902e68a61861ca96d9d01ca49d Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 17:17:55 -0700 Subject: [PATCH 39/71] Teardown and other fixes --- dimos/mapping/ray_tracing/module.py | 3 +- dimos/mapping/ray_tracing/rust/src/main.rs | 32 +++++-- .../nav_3d/mls_planner/rust/src/main.rs | 91 ++++++++++++++++--- .../nav_3d/mls_planner/rust/src/planner.rs | 19 ++-- .../nav_3d/mls_planner/rust/src/python.rs | 69 ++++++-------- 5 files changed, 142 insertions(+), 72 deletions(-) diff --git a/dimos/mapping/ray_tracing/module.py b/dimos/mapping/ray_tracing/module.py index 851ff8069d..5e9bf20e6c 100644 --- a/dimos/mapping/ray_tracing/module.py +++ b/dimos/mapping/ray_tracing/module.py @@ -54,8 +54,7 @@ class RayTracingVoxelMapConfig(NativeModuleConfig): emit_every: int = 1 # Publish the global map every Nth frame. Zero disables it. global_emit_every: int = 1 - # Size the local region to this percentile of batch point distances, - # so a stray far hit cannot inflate the region the planner recomputes. + # Size the local region to this percentile of batch point distances. region_percentile: float = 95.0 diff --git a/dimos/mapping/ray_tracing/rust/src/main.rs b/dimos/mapping/ray_tracing/rust/src/main.rs index d7dc7ac919..60bae5ad73 100644 --- a/dimos/mapping/ray_tracing/rust/src/main.rs +++ b/dimos/mapping/ray_tracing/rust/src/main.rs @@ -28,11 +28,8 @@ struct RayTracingVoxelMap { #[output(encode = PointCloud2::encode)] local_map: Output, - // Full region of the local_map, represented as a PoseStamped. We need - // this to update the planner artifacts because local_map only includes - // points that exist (not points that have been removed) Position holds - // the cylinder center and orientation holds radius, z_min, z_max, zero. - // Stamped identically to local_map so consumers can pair them. + // Cylinder bounds of the local map. Position holds the center, orientation + // holds radius, z_min, z_max. Stamped like local_map so consumers pair them. #[output(encode = PoseStamped::encode)] region_bounds: Output, @@ -105,7 +102,7 @@ impl RayTracingVoxelMap { } self.frame_count += 1; - let local_due = self.frame_count.is_multiple_of(self.config.emit_every); + let local_due = emit_due(self.frame_count, self.config.emit_every); let cylinder = if local_due { let margin = self.config.shadow_depth + voxel_size; @@ -156,11 +153,8 @@ impl RayTracingVoxelMap { None }; - let global_due = self - .frame_count - .is_multiple_of(self.config.global_emit_every); + let global_due = emit_due(self.frame_count, self.config.global_emit_every); - // build and publish the clouds let stamp = msg.header.stamp; if let Some(cyl) = &cylinder { if global_due { @@ -180,6 +174,11 @@ impl RayTracingVoxelMap { } } +/// Whether the Nth-frame output fires this frame. Zero disables it. +fn emit_due(frame_count: u32, every: u32) -> bool { + every != 0 && frame_count.is_multiple_of(every) +} + struct ExtractError(&'static str); impl std::fmt::Display for ExtractError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -422,6 +421,19 @@ mod tests { ) } + #[test] + fn emit_due_fires_every_nth_frame_and_zero_disables() { + assert!(emit_due(1, 1)); + assert!(emit_due(2, 1)); + assert!(!emit_due(1, 2)); + assert!(emit_due(2, 2)); + assert!(!emit_due(5, 3)); + assert!(emit_due(6, 3)); + for n in 1..10 { + assert!(!emit_due(n, 0)); + } + } + #[test] fn local_map_includes_voxel_inside_cylinder() { let mut map = VoxelMap::default(); diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs index ccf5f8e2b9..abd860b51d 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs @@ -35,7 +35,7 @@ enum MapUpdate { } #[derive(Module)] -#[module(setup = spawn_worker)] +#[module(setup = spawn_worker, teardown = stop_worker)] struct MlsPlanner { #[input(decode = PointCloud2::decode, handler = on_global_map)] global_map: Input, @@ -78,6 +78,9 @@ struct MlsPlanner { latest_start: Shared, active_goal: Shared, wake: Arc, + + // Handle to the background worker, aborted on teardown. + worker: Option>, } impl MlsPlanner { @@ -94,7 +97,14 @@ impl MlsPlanner { node_edges: self.node_edges.clone(), path: self.path.clone(), }; - tokio::spawn(worker.run()); + self.worker = Some(tokio::spawn(worker.run())); + } + + /// Abort the background worker on teardown. + async fn stop_worker(&mut self) { + if let Some(handle) = self.worker.take() { + handle.abort(); + } } async fn on_global_map(&mut self, msg: PointCloud2) { @@ -114,10 +124,7 @@ impl MlsPlanner { /// Hand a local map and its stamp-matching bounds to the worker once both /// are in hand. fn try_pair(&mut self) { - let (Some(bounds_msg), Some(cloud)) = (&self.pending_bounds, &self.pending_local) else { - return; - }; - if !same_stamp(&bounds_msg.header.stamp, &cloud.header.stamp) { + if !stamps_paired(self.pending_bounds.as_ref(), self.pending_local.as_ref()) { return; } let bounds = self.pending_bounds.take().expect("checked above"); @@ -139,17 +146,29 @@ impl MlsPlanner { Some((p.x as f32, p.y as f32, p.z as f32)); } - /// Set the active goal, or clear it when the goal is non-finite, which is - /// the cancel signal. Wake the worker so a new click replans right away. + /// Set or cancel the active goal from a click, then wake the worker. async fn on_goal_pose(&mut self, msg: PoseStamped) { - let p = &msg.pose.position; - let goal = (p.x as f32, p.y as f32, p.z as f32); - *self.active_goal.lock().expect("goal mutex") = - (goal.0.is_finite() && goal.1.is_finite() && goal.2.is_finite()).then_some(goal); + *self.active_goal.lock().expect("goal mutex") = goal_position(&msg.pose.position); self.wake.notify_one(); } } +/// True when bounds and a local cloud are both present with matching stamps, +/// so they form one paired region update. +fn stamps_paired(bounds: Option<&PoseStamped>, cloud: Option<&PointCloud2>) -> bool { + match (bounds, cloud) { + (Some(b), Some(c)) => same_stamp(&b.header.stamp, &c.header.stamp), + _ => false, + } +} + +/// The goal position, or None when any coordinate is non-finite, which is the +/// cancel signal. +fn goal_position(p: &Point) -> Option { + let goal = (p.x as f32, p.y as f32, p.z as f32); + (goal.0.is_finite() && goal.1.is_finite() && goal.2.is_finite()).then_some(goal) +} + /// Owns the planner graph and does every map mutation, graph publish, and /// replan off the handle loop. Woken by the handlers via `wake`. struct Worker { @@ -544,4 +563,52 @@ mod tests { assert!(!same_stamp(&a, &Time { sec: 5, nsec: 8 })); assert!(!same_stamp(&a, &Time { sec: 6, nsec: 7 })); } + + fn stamped(stamp: Time) -> Header { + Header { + stamp, + ..Default::default() + } + } + + fn bounds_at(stamp: Time) -> PoseStamped { + PoseStamped { + header: stamped(stamp), + ..Default::default() + } + } + + fn cloud_at(stamp: Time) -> PointCloud2 { + PointCloud2 { + header: stamped(stamp), + ..Default::default() + } + } + + #[test] + fn stamps_paired_only_when_both_present_and_stamps_match() { + let s = Time { sec: 2, nsec: 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 }); + assert!(!stamps_paired(Some(&b), Some(&other))); + + assert!(!stamps_paired(Some(&b), None)); + assert!(!stamps_paired(None, Some(&c))); + assert!(!stamps_paired(None, None)); + } + + fn point(x: f64, y: f64, z: f64) -> Point { + Point { x, y, z } + } + + #[test] + fn goal_position_passes_finite_and_cancels_on_non_finite() { + assert_eq!(goal_position(&point(1.0, 2.0, 3.0)), Some((1.0, 2.0, 3.0))); + assert_eq!(goal_position(&point(f64::NAN, 0.0, 0.0)), None); + assert_eq!(goal_position(&point(0.0, f64::INFINITY, 0.0)), None); + assert_eq!(goal_position(&point(0.0, 0.0, f64::NEG_INFINITY)), None); + } } diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs index 85f12345b3..a43f3157ad 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs @@ -200,17 +200,24 @@ fn select_entry( let mut entry_node = NO_NODE; let mut best_score = f32::INFINITY; - for node in &plg.nodes { - let Some(&connect) = connect_dist.get(&node.cell_id) else { + // Scan the bounded reachable set, not every node. Tie-break by cell id for + // deterministic order. + for (&cell, &connect) in &connect_dist { + if !node_cells.contains(&cell) { continue; - }; - let Some(&ctg) = cost_to_go.get(&node.cell_id) else { + } + let Some(&ctg) = cost_to_go.get(&cell) else { continue; }; let score = connect + ctg; - if score < best_score { + let better = match score.partial_cmp(&best_score) { + Some(std::cmp::Ordering::Less) => true, + Some(std::cmp::Ordering::Equal) => cell < entry_node, + _ => false, + }; + if better { best_score = score; - entry_node = node.cell_id; + entry_node = cell; } } diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs index b92b9ef46e..45d9b46791 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs @@ -18,6 +18,31 @@ pub struct MLSPlanner { planner: Planner, } +/// Extract a (N, 3) float32 numpy array into xyz tuples, dropping any row with +/// a non-finite coordinate. +fn extract_points(points: &Bound<'_, PyAny>) -> PyResult> { + let points: PyReadonlyArray2<'_, f32> = points + .extract() + .map_err(|_| PyValueError::new_err("points must be a (N, 3) float32 numpy array"))?; + let shape = points.shape(); + if shape[1] != 3 { + return Err(PyValueError::new_err(format!( + "points must be (N, 3) float32, got shape {:?}", + shape + ))); + } + let arr = points.as_array(); + let n = shape[0]; + Ok((0..n) + .filter_map(|i| { + let x = arr[[i, 0]]; + let y = arr[[i, 1]]; + let z = arr[[i, 2]]; + (x.is_finite() && y.is_finite() && z.is_finite()).then_some((x, y, z)) + }) + .collect()) +} + #[pymethods] impl MLSPlanner { #[new] @@ -72,27 +97,7 @@ impl MLSPlanner { } fn update_global_map(&mut self, py: Python<'_>, points: &Bound<'_, PyAny>) -> PyResult<()> { - let points: PyReadonlyArray2<'_, f32> = points - .extract() - .map_err(|_| PyValueError::new_err("points must be a (N, 3) float32 numpy array"))?; - let shape = points.shape(); - if shape[1] != 3 { - return Err(PyValueError::new_err(format!( - "points must be (N, 3) float32, got shape {:?}", - shape - ))); - } - let arr = points.as_array(); - let n = shape[0]; - let pts: Vec<(f32, f32, f32)> = (0..n) - .filter_map(|i| { - let x = arr[[i, 0]]; - let y = arr[[i, 1]]; - let z = arr[[i, 2]]; - (x.is_finite() && y.is_finite() && z.is_finite()).then_some((x, y, z)) - }) - .collect(); - + let pts = extract_points(points)?; let config = &self.config; let planner = &mut self.planner; py.allow_threads(move || planner.update_global_map(&pts, config)); @@ -109,27 +114,7 @@ impl MLSPlanner { z_min: f32, z_max: f32, ) -> PyResult<()> { - let points: PyReadonlyArray2<'_, f32> = points - .extract() - .map_err(|_| PyValueError::new_err("points must be a (N, 3) float32 numpy array"))?; - let shape = points.shape(); - if shape[1] != 3 { - return Err(PyValueError::new_err(format!( - "points must be (N, 3) float32, got shape {:?}", - shape - ))); - } - let arr = points.as_array(); - let n = shape[0]; - let pts: Vec<(f32, f32, f32)> = (0..n) - .filter_map(|i| { - let x = arr[[i, 0]]; - let y = arr[[i, 1]]; - let z = arr[[i, 2]]; - (x.is_finite() && y.is_finite() && z.is_finite()).then_some((x, y, z)) - }) - .collect(); - + let pts = extract_points(points)?; let bounds = RegionBounds { origin_x: origin.0, origin_y: origin.1, From a5c7fd4f6cc5d479e64ec41c76eab39e3043939c Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 18:12:39 -0700 Subject: [PATCH 40/71] Remove useless code --- dimos/mapping/ray_tracing/test_transformer.py | 10 +++++----- dimos/mapping/ray_tracing/transformer.py | 13 ++++--------- .../nav_3d/mls_planner/test_transformer.py | 12 ------------ dimos/navigation/nav_3d/mls_planner/transformer.py | 7 +------ .../navigation/nav_3d/mls_planner/utils/plan_rrd.py | 5 +---- 5 files changed, 11 insertions(+), 36 deletions(-) diff --git a/dimos/mapping/ray_tracing/test_transformer.py b/dimos/mapping/ray_tracing/test_transformer.py index ec52ec8f37..2d6c91a83b 100644 --- a/dimos/mapping/ray_tracing/test_transformer.py +++ b/dimos/mapping/ray_tracing/test_transformer.py @@ -85,12 +85,12 @@ def _ring( return np.stack([xs, ys, zs], axis=1).astype(np.float32) -def test_emit_local_tags_region_bounds_around_registered_origin() -> None: +def test_tags_region_bounds_around_registered_origin() -> None: margin = 0.2 + 0.1 # Sensor-frame ring centered on the sensor. The pose registers it to (2, 3, 0.5). obs = _obs(_ring((0.0, 0.0), radius=1.0, z=0.0), ts=1.0, pose=(2.0, 3.0, 0.5)) - [emitted] = list(RayTraceMap(emit_local=True)(iter([obs]))) + [emitted] = list(RayTraceMap()(iter([obs]))) cx, cy, radius, z_min, z_max = emitted.tags["region_bounds"] assert (cx, cy) == pytest.approx((2.0, 3.0)) @@ -99,11 +99,11 @@ def test_emit_local_tags_region_bounds_around_registered_origin() -> None: assert z_max == pytest.approx(0.5 + margin) -def test_emit_local_empty_frame_yields_zero_radius_region_at_robot() -> None: +def test_empty_frame_yields_zero_radius_region_at_robot() -> None: empty = np.empty((0, 3), dtype=np.float32) obs = _obs(empty, ts=1.0, pose=(1.0, 2.0, 3.0)) - [emitted] = list(RayTraceMap(emit_local=True)(iter([obs]))) + [emitted] = list(RayTraceMap()(iter([obs]))) assert emitted.tags["region_bounds"] == pytest.approx((1.0, 2.0, 0.0, 3.0, 3.0)) @@ -121,7 +121,7 @@ def test_registers_sensor_frame_cloud_by_pose() -> None: _data=PointCloud2.from_numpy(point), ) - [emitted] = list(RayTraceMap(emit_local=True)(iter([obs]))) + [emitted] = list(RayTraceMap()(iter([obs]))) cx, cy, radius, z_min, z_max = emitted.tags["region_bounds"] assert (cx, cy) == pytest.approx((5.0, 0.0)) diff --git a/dimos/mapping/ray_tracing/transformer.py b/dimos/mapping/ray_tracing/transformer.py index 35cc40dd5f..f532936b28 100644 --- a/dimos/mapping/ray_tracing/transformer.py +++ b/dimos/mapping/ray_tracing/transformer.py @@ -50,7 +50,6 @@ def __init__( voxel_size: float = 0.1, max_range: float = 30.0, emit_every: int = 1, - emit_local: bool = False, region_percentile: float = 95.0, **mapper_kwargs: Any, ) -> None: @@ -59,7 +58,6 @@ def __init__( self.voxel_size = voxel_size self.max_range = max_range self.emit_every = emit_every - self.emit_local = emit_local self.region_percentile = region_percentile self._mapper_kwargs = mapper_kwargs @@ -93,12 +91,9 @@ def _make_obs( batch_origins: list[tuple[float, float, float]], ) -> Observation[PointCloud2]: tags = {**last_obs.tags, "frame_count": count} - if self.emit_local: - cx, cy, radius, z_min, z_max = self._local_bounds(batch_points, batch_origins, last_obs) - positions = mapper.local_map((cx, cy, 0.0), radius, z_min, z_max) - tags["region_bounds"] = (cx, cy, radius, z_min, z_max) - else: - positions = mapper.global_map() + cx, cy, radius, z_min, z_max = self._local_bounds(batch_points, batch_origins, last_obs) + 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) @@ -130,7 +125,7 @@ def __call__( trans = mat[:3, 3].astype(np.float32) pts = obs.data.points_f32() @ rot.T + trans mapper.add_frame(pts, (x, y, z)) - if self.emit_local and pts.size: + if pts.size: batch_points.append(pts) batch_origins.append((x, y, z)) last_obs = obs diff --git a/dimos/navigation/nav_3d/mls_planner/test_transformer.py b/dimos/navigation/nav_3d/mls_planner/test_transformer.py index 3fca1fb081..2cd84e03bb 100644 --- a/dimos/navigation/nav_3d/mls_planner/test_transformer.py +++ b/dimos/navigation/nav_3d/mls_planner/test_transformer.py @@ -81,18 +81,6 @@ def test_poseless_obs_is_skipped() -> None: assert len(results) == 1 -def test_missing_region_bounds_raises() -> None: - obs = Observation( - id=0, - ts=0.0, - pose=(0.0, 0.0, 0.0), - _data=PointCloud2.from_numpy(np.zeros((1, 3), dtype=np.float32)), - ) - - with pytest.raises(ValueError, match="local map slices"): - list(MLSPlan(goal=(1.0, 1.0, 0.0))(iter([obs]))) - - def test_start_z_is_dropped_by_robot_height() -> None: obs = _obs( np.zeros((1, 3), dtype=np.float32), diff --git a/dimos/navigation/nav_3d/mls_planner/transformer.py b/dimos/navigation/nav_3d/mls_planner/transformer.py index 3f9bc09b7a..64dff8bab1 100644 --- a/dimos/navigation/nav_3d/mls_planner/transformer.py +++ b/dimos/navigation/nav_3d/mls_planner/transformer.py @@ -81,12 +81,7 @@ def __call__( x, y, z, *_ = obs.pose_tuple start = (float(x), float(y), float(z) - self.robot_height) - bounds = obs.tags.get("region_bounds") - if bounds is None: - raise ValueError( - "MLSPlan consumes local map slices; construct RayTraceMap(emit_local=True)" - ) - ox, oy, radius, z_min, z_max = bounds + 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) t_plan = time.perf_counter() diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index ff3301f466..bc955ca998 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -273,7 +273,6 @@ def main( max_range=max_range, ray_subsample=ray_subsample, emit_every=emit_every, - emit_local=True, ) ) @@ -307,9 +306,7 @@ def main( for ray_obs in ray_pipeline: if ray_obs.pose_tuple is None: continue - bounds = ray_obs.tags.get("region_bounds") - if bounds is None: - raise ValueError("plan_rrd needs RayTraceMap(emit_local=True)") + bounds = ray_obs.tags["region_bounds"] px, py, pz, *_ = ray_obs.pose_tuple start = (float(px), float(py), float(pz) - robot_height) ox, oy, radius, z_min, z_max = bounds From 1ceef52c59715a666a8774ac3f73a2fcc566fdd9 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 19:57:47 -0700 Subject: [PATCH 41/71] Bug fix --- .../nav_3d/mls_planner/rust/src/main.rs | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs index abd860b51d..22b0d9d4bc 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs @@ -303,15 +303,20 @@ impl Worker { /// Replan from the latest start to the active goal. This gates and does the /// IO. The planning itself lives in Planner::plan. async fn maybe_replan(&self, planner: &mut Planner, last_path_at: &mut Option) { - let start = *self.latest_start.lock().expect("start mutex"); - let goal = *self.active_goal.lock().expect("goal mutex"); - let (Some(start), Some(goal)) = (start, goal) else { + let Some(start) = *self.latest_start.lock().expect("start mutex") else { return; }; - if is_at_goal(start, goal, self.config.goal_tolerance) { - *self.active_goal.lock().expect("goal mutex") = None; - return; - } + let goal = { + let mut guard = self.active_goal.lock().expect("goal mutex"); + let Some(goal) = *guard else { + return; + }; + if is_at_goal(start, goal, self.config.goal_tolerance) { + *guard = None; + return; + } + goal + }; let plan_start = Instant::now(); let waypoints = tokio::task::block_in_place(|| planner.plan(start, goal, &self.config)); From fa5a6b29cb1836b772d9a0f43df5664f7b61473e Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 20:32:12 -0700 Subject: [PATCH 42/71] Minor fixes --- dimos/mapping/ray_tracing/transformer.py | 4 +++- dimos/navigation/nav_3d/mls_planner/rust/src/main.rs | 2 +- dimos/navigation/nav_3d/mls_planner/rust/src/python.rs | 3 +-- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/dimos/mapping/ray_tracing/transformer.py b/dimos/mapping/ray_tracing/transformer.py index f532936b28..cec1670baf 100644 --- a/dimos/mapping/ray_tracing/transformer.py +++ b/dimos/mapping/ray_tracing/transformer.py @@ -37,6 +37,8 @@ logger = setup_logger() +DEFAULT_SHADOW_DEPTH = 0.2 + class RayTraceMap(Transformer[PointCloud2, PointCloud2]): """Accumulate lidar into a voxel map with raycast clearing. @@ -79,7 +81,7 @@ def _local_bounds( points = np.concatenate(batch_points, axis=0) origins = np.asarray(batch_origins, dtype=np.float32) - margin = self._mapper_kwargs.get("shadow_depth", 0.2) + self.voxel_size + margin = self._mapper_kwargs.get("shadow_depth", DEFAULT_SHADOW_DEPTH) + self.voxel_size return local_bounds(points, origins, self.region_percentile, margin) def _make_obs( diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs index 22b0d9d4bc..70a1f21358 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs @@ -343,7 +343,7 @@ impl Worker { } } -/// True when start is within `tol` of goal in the ground plane. +/// True if within tolerance of the goal on the ground plane. fn is_at_goal(start: Xyz, goal: Xyz, tol: f32) -> bool { (start.0 - goal.0).hypot(start.1 - goal.1) < tol } diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs index 45d9b46791..3b6e0ade0e 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs @@ -9,8 +9,7 @@ use validator::Validate; use crate::edges::edges_to_segments; use crate::mls_planner::{Config, Planner, RegionBounds}; -use crate::voxel::surface_point_xyz; -use crate::voxel::VoxelKey; +use crate::voxel::{surface_point_xyz, VoxelKey}; #[pyclass] pub struct MLSPlanner { From 84f7a127d092512459dec8e83f632675c95697cb Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 21:25:50 -0700 Subject: [PATCH 43/71] Add blueprint --- .../navigation/basic_path_follower/module.py | 196 ++++++++++++++++++ .../basic_path_follower/test_module.py | 31 +++ .../nav_3d/mls_planner/goal_relay.py | 52 +++++ dimos/robot/all_blueprints.py | 3 + dimos/robot/unitree/connection.py | 24 +++ .../navigation/unitree_go2_nav_3d.py | 128 ++++++++++++ dimos/robot/unitree/go2/connection.py | 28 ++- 7 files changed, 453 insertions(+), 9 deletions(-) create mode 100644 dimos/navigation/basic_path_follower/module.py create mode 100644 dimos/navigation/basic_path_follower/test_module.py create mode 100644 dimos/navigation/nav_3d/mls_planner/goal_relay.py create mode 100644 dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py diff --git a/dimos/navigation/basic_path_follower/module.py b/dimos/navigation/basic_path_follower/module.py new file mode 100644 index 0000000000..e4c591e1da --- /dev/null +++ b/dimos/navigation/basic_path_follower/module.py @@ -0,0 +1,196 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import math +from threading import Event, RLock, Thread +import time +from typing import Any + +from dimos_lcm.std_msgs import Bool # type: ignore[import-untyped] +import numpy as np +from reactivex.disposable import Disposable + +from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.nav_msgs.Path import Path +from dimos.utils.logging_config import setup_logger +from dimos.utils.trigonometry import angle_diff + +logger = setup_logger() + + +class BasicPathFollowerConfig(ModuleConfig): + speed: float = 0.5 + control_frequency: float = 10.0 + goal_tolerance: float = 0.3 + # Lookahead grows with speed (time_s * speed), clamped to these bounds. + lookahead_time_s: float = 1.5 + min_lookahead_m: float = 0.4 + max_lookahead_m: float = 1.5 + heading_gain: float = 1.5 + max_angular: float = 1.0 + + +def lookahead_distance(speed: float, time_s: float, lo: float, hi: float) -> float: + return min(hi, max(lo, time_s * speed)) + + +class BasicPathFollower(Module): + """Follow a planned path by chasing a lookahead point with P-controlled heading. + + Publishes nav_cmd_vel until the last waypoint is within goal tolerance, then + goal_reached. stop_movement cancels the current path. + """ + + config: BasicPathFollowerConfig + + path: In[Path] + odometry: In[Odometry] + stop_movement: In[Bool] + + nav_cmd_vel: Out[Twist] + goal_reached: Out[Bool] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._lock = RLock() + self._current_odom: PoseStamped | None = None + self._waypoints: np.ndarray | None = None + self._stop_event = Event() + self._thread: Thread | None = None + + @rpc + def start(self) -> None: + super().start() + self.register_disposable(Disposable(self.odometry.subscribe(self._on_odometry))) + self.register_disposable(Disposable(self.path.subscribe(self._on_path))) + if self.stop_movement.transport is not None: + self.register_disposable(Disposable(self.stop_movement.subscribe(self._on_stop))) + self._thread = Thread(target=self._follow, daemon=True) + self._thread.start() + + @rpc + 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()) + super().stop() + + def _on_odometry(self, msg: Odometry) -> None: + with self._lock: + self._current_odom = msg.to_pose_stamped() + + def _on_path(self, path: Path) -> None: + # An empty path is the planner clearing its plan, not a stop. Keep the + # last path. + if len(path.poses) == 0: + return + waypoints = np.array([[p.position.x, p.position.y] for p in path.poses]) + with self._lock: + self._waypoints = waypoints + + def _on_stop(self, msg: Bool) -> None: + if msg.data: + with self._lock: + self._waypoints = None + self.nav_cmd_vel.publish(Twist()) + + def _follow(self) -> None: + period = 1.0 / self.config.control_frequency + while not self._stop_event.is_set(): + start_time = time.perf_counter() + with self._lock: + odom = self._current_odom + waypoints = self._waypoints + if odom is not None and waypoints is not None: + self._step(odom, waypoints) + elapsed = time.perf_counter() - start_time + self._stop_event.wait(max(0.0, period - elapsed)) + + def _step(self, odom: PoseStamped, waypoints: np.ndarray) -> None: + position = np.array([odom.position.x, odom.position.y]) + if float(np.linalg.norm(waypoints[-1] - position)) < self.config.goal_tolerance: + self.nav_cmd_vel.publish(Twist()) + with self._lock: + if self._waypoints is waypoints: + self._waypoints = None + self.goal_reached.publish(Bool(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]), + odom.orientation.euler[2], + ) + + angular = max( + -self.config.max_angular, + min(self.config.max_angular, self.config.heading_gain * yaw_error), + ) + # Taper forward speed by cos(yaw_error): decelerate into turns and pivot + # in place near 90 degrees, rather than stop and lurch. + linear = self.config.speed * max(0.0, math.cos(yaw_error)) + self.nav_cmd_vel.publish(Twist(Vector3(linear, 0, 0), Vector3(0, 0, angular))) + + def _lookahead_point(self, waypoints: np.ndarray, position: np.ndarray) -> np.ndarray: + if len(waypoints) == 1: + return np.asarray(waypoints[0]) + + # Interpolate along the path rather than snap to a waypoint, which would + # cut corners. + seg_idx, start = self._project_onto_path(waypoints, position) + remaining = lookahead_distance( + self.config.speed, + self.config.lookahead_time_s, + self.config.min_lookahead_m, + self.config.max_lookahead_m, + ) + for i in range(seg_idx, len(waypoints) - 1): + end = waypoints[i + 1] + seg = end - start + seg_len = float(np.linalg.norm(seg)) + if seg_len >= remaining: + return np.asarray(start + (remaining / seg_len) * seg) + remaining -= seg_len + start = end + return np.asarray(waypoints[-1]) + + def _project_onto_path( + self, waypoints: np.ndarray, position: np.ndarray + ) -> tuple[int, np.ndarray]: + best_idx = 0 + best_point = np.asarray(waypoints[0]) + best_dist = math.inf + for i in range(len(waypoints) - 1): + a = waypoints[i] + ab = waypoints[i + 1] - a + denom = float(np.dot(ab, ab)) + t = 0.0 if denom == 0 else float(np.clip(np.dot(position - a, ab) / denom, 0.0, 1.0)) + proj = a + t * ab + dist = float(np.linalg.norm(position - proj)) + if dist < best_dist: + best_dist = dist + best_idx = i + best_point = proj + return best_idx, best_point diff --git a/dimos/navigation/basic_path_follower/test_module.py b/dimos/navigation/basic_path_follower/test_module.py new file mode 100644 index 0000000000..a5ad8d5782 --- /dev/null +++ b/dimos/navigation/basic_path_follower/test_module.py @@ -0,0 +1,31 @@ +# 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.navigation.basic_path_follower.module import lookahead_distance + + +def test_lookahead_floor_at_low_speed(): + assert lookahead_distance(0.1, 1.5, 0.4, 1.5) == 0.4 + + +def test_lookahead_scales_in_linear_region(): + assert lookahead_distance(0.5, 1.5, 0.4, 1.5) == 0.75 + + +def test_lookahead_clamped_at_ceiling(): + assert lookahead_distance(2.0, 1.5, 0.4, 1.5) == 1.5 + + +def test_lookahead_monotonic_in_speed(): + assert lookahead_distance(0.8, 1.5, 0.4, 1.5) > lookahead_distance(0.45, 1.5, 0.4, 1.5) diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py new file mode 100644 index 0000000000..905c4f8973 --- /dev/null +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -0,0 +1,52 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from 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.nav_msgs.Odometry import Odometry + + +class GoalRelayConfig(ModuleConfig): + pass + + +class GoalRelay(Module): + """Adapt odometry and goal points to the planner's PoseStamped inputs.""" + + config: GoalRelayConfig + + odometry: In[Odometry] + goal: In[PointStamped] + + start_pose: Out[PoseStamped] + goal_pose: Out[PoseStamped] + + @rpc + def start(self) -> None: + super().start() + self.register_disposable(Disposable(self.odometry.subscribe(self._on_odometry))) + self.register_disposable(Disposable(self.goal.subscribe(self._on_goal))) + + def _on_odometry(self, msg: Odometry) -> None: + self.start_pose.publish(msg.to_pose_stamped()) + + def _on_goal(self, point: PointStamped) -> None: + self.goal_pose.publish(point.to_pose_stamped()) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 22bf7d27e1..c4af00d4d0 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -115,6 +115,7 @@ "unitree-go2-keyboard-teleop": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_keyboard_teleop:unitree_go2_keyboard_teleop", "unitree-go2-markers": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_markers", "unitree-go2-memory": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_memory", + "unitree-go2-nav-3d": "dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_nav_3d:unitree_go2_nav_3d", "unitree-go2-relocalization": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_relocalization", "unitree-go2-ros": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2_ros:unitree_go2_ros", "unitree-go2-security": "dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_security:unitree_go2_security", @@ -139,6 +140,7 @@ "arm-teleop-module": "dimos.teleop.quest.quest_extensions.ArmTeleopModule", "b-box-navigation-module": "dimos.navigation.bbox_navigation.BBoxNavigationModule", "b1-connection-module": "dimos.robot.unitree.b1.connection.B1ConnectionModule", + "basic-path-follower": "dimos.navigation.basic_path_follower.module.BasicPathFollower", "camera-module": "dimos.hardware.sensors.camera.module.CameraModule", "cartesian-motion-controller": "dimos.manipulation.control.servo_control.cartesian_motion_controller.CartesianMotionController", "click-start-goal-router": "dimos.navigation.cmu_nav.modules.click_start_goal_router.click_start_goal_router.ClickStartGoalRouter", @@ -170,6 +172,7 @@ "go2-fleet-connection": "dimos.robot.unitree.go2.fleet_connection.Go2FleetConnection", "go2-memory": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2.Go2Memory", "go2-teleop-module": "dimos.teleop.quest.quest_extensions.Go2TeleopModule", + "goal-relay": "dimos.navigation.nav_3d.mls_planner.goal_relay.GoalRelay", "google-maps-skill-container": "dimos.agents.skills.google_maps_skill_container.GoogleMapsSkillContainer", "gps-nav-skill-container": "dimos.agents.skills.gps_nav_skill.GpsNavSkillContainer", "grasping-module": "dimos.manipulation.grasping.grasping.GraspingModule", diff --git a/dimos/robot/unitree/connection.py b/dimos/robot/unitree/connection.py index f7040d9823..e439ed3672 100644 --- a/dimos/robot/unitree/connection.py +++ b/dimos/robot/unitree/connection.py @@ -329,6 +329,30 @@ def set_obstacle_avoidance(self, enabled: bool = True) -> None: {"api_id": 1001, "parameter": {"enable": int(enabled)}}, ) + def set_motion_mode(self, name: str) -> None: + """Select the robot's top-level motion mode via the motion switcher. + + 'mcf' is the AI/sport controller (firmware >= 1.1.7) that handles + terrain, including walking up and down stairs. 'normal' is the basic + controller. Stair traversal lives in mcf, so the nav stack selects it. + """ + # api_id 1001 = CheckMode, 1002 = SelectMode, param {"name": }. + current = None + try: + resp = self.publish_request(RTC_TOPIC["MOTION_SWITCHER"], {"api_id": 1001}) + current = json.loads(resp["data"]["data"]).get("name") + except (KeyError, TypeError, ValueError) as e: + print(f"Motion mode check failed: {e}") + print(f"Motion mode: current={current!r} requested={name!r}") + if current == name: + return + self.publish_request( + RTC_TOPIC["MOTION_SWITCHER"], + {"api_id": 1002, "parameter": {"name": name}}, + ) + # Switching controllers makes the robot re-stand; give it a moment. + time.sleep(5) + def free_walk(self) -> bool: """Activate FreeWalk locomotion mode — enables walking and velocity commands.""" return bool(self.publish_request(RTC_TOPIC["SPORT_MOD"], {"api_id": SPORT_CMD["FreeWalk"]})) diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py new file mode 100644 index 0000000000..eb36e6dbec --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +# 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. + +"""Go2 3D nav stack: pointlio odometry, raytraced voxel map, MLS planning, path following. + +Expects a mid360 + pointlio running against a lidar mounted on the robot. +The go2's built-in lidar and odometry are remapped aside and unused for +navigation. The map, paths, and robot pose all live in pointlio's world frame. +""" + +import os +from typing import Any + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.global_config import global_config +from dimos.hardware.sensors.lidar.pointlio.module import PointLio +from dimos.mapping.ray_tracing.module import RayTracingVoxelMap +from dimos.navigation.basic_path_follower.module import BasicPathFollower +from dimos.navigation.movement_manager.movement_manager import MovementManager +from dimos.navigation.nav_3d.mls_planner.goal_relay import GoalRelay +from dimos.navigation.nav_3d.mls_planner.mls_planner_native import MLSPlannerNative +from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import rerun_config +from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.visualization.vis_module import vis_module + +voxel_size = 0.1 +# Height of the head-mounted lidar above the ground while standing. +go2_lidar_height = 0.5 + + +def _render_global_map(msg: Any) -> Any: + return msg.to_rerun() + + +def _render_path(msg: Any) -> Any: + # The planner clears its plan with an empty path on every start-pose change. + # Logging those would blank the line. Drop them so the last path stays shown. + if len(msg.poses) == 0: + return None + return msg + + +def _static_robot_body(rr: Any) -> list[Any]: + """Go2-shaped box on pointlio's body frame, counter-rotated for the lidar pitch.""" + return [ + rr.Boxes3D(half_sizes=[0.35, 0.155, 0.2], colors=[(0, 255, 127)]), + rr.Transform3D( + parent_frame="tf#/body", + rotation=rr.RotationAxisAngle(axis=(0, 1, 0), degrees=-45.0), + ), + ] + + +_nav_rerun_config = { + **rerun_config, + "max_hz": { + **rerun_config["max_hz"], + "world/global_map": 1.0, + "world/local_map": 1.0, + }, + "memory_limit": "256MB", + # base_link tf comes from the go2 internal odometry, which is not the map + # frame. Anchor the robot box to pointlio's body frame instead and hide the + # camera frustum that rides base_link. + "static": {"world/tf/body": _static_robot_body}, + "visual_override": { + **rerun_config["visual_override"], + "world/global_map": _render_global_map, + "world/path": _render_path, + "world/camera_info": None, + "world/color_image": None, + "world/lidar": None, + "world/surface_map": None, + "world/nodes": None, + "world/node_edges": None, + }, +} + +unitree_go2_nav_3d = autoconnect( + vis_module(viewer_backend=global_config.viewer, rerun_config=_nav_rerun_config), + # "mcf" for stair traversal + GO2Connection.blueprint(lidar=False, camera=False, motion_mode="mcf").remappings( + [ + (GO2Connection, "lidar", "lidar_l1"), + (GO2Connection, "odom", "odom_go2"), + ] + ), + PointLio.blueprint( + host_ip=os.getenv("LIDAR_HOST_IP", "192.168.1.5"), + lidar_ip=os.getenv("LIDAR_IP", "192.168.1.155"), + body_frame_id="body", + ), + RayTracingVoxelMap.blueprint( + voxel_size=voxel_size, + emit_every=2, + global_emit_every=50, + max_health=10, + graze_cos=0.85, + ), + # global_map is remapped off so the planner runs purely on the + # incremental local_map + region_bounds pair. + MLSPlannerNative.blueprint( + world_frame="odom", + voxel_size=voxel_size, + robot_height=go2_lidar_height, + wall_clearance_m=0.2, + wall_buffer_m=0.75, + wall_buffer_weight=100.0, + step_threshold_m=0.25, + step_penalty_weight=1.0, + viz_publish_hz=0.0, + ).remappings([(MLSPlannerNative, "global_map", "global_map_unused")]), + GoalRelay.blueprint(), + BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6), + MovementManager.blueprint(), +).global_config(n_workers=10, robot_model="unitree_go2", obstacle_avoidance=False) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 185546217b..c8bab179d7 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -67,6 +67,10 @@ class Go2Mode(str, Enum): class ConnectionConfig(ModuleConfig): ip: str = Field(default_factory=lambda m: m["g"].robot_ip) mode: Go2Mode = Go2Mode.DEFAULT + lidar: bool = True + camera: bool = True + # "mcf" for stair traversal, "normal" for basic, None to leave it as is + motion_mode: str | None = None # Per-device AES-128 key (Go2 fw >=1.1.15); defaults from GlobalConfig. aes_128_key: str | None = Field(default_factory=lambda m: m["g"].unitree_aes_128_key) @@ -177,6 +181,9 @@ def balance_stand(self) -> bool: def set_obstacle_avoidance(self, enabled: bool = True) -> None: pass + def set_motion_mode(self, name: str) -> None: + pass + def enable_rage_mode(self) -> bool: return True @@ -249,16 +256,21 @@ def onimage(image: Image) -> None: self.color_image.publish(image) self._latest_video_frame = image - self.register_disposable(self.connection.lidar_stream().subscribe(self.lidar.publish)) + if self.config.lidar: + self.register_disposable(self.connection.lidar_stream().subscribe(self.lidar.publish)) self.register_disposable(self.connection.odom_stream().subscribe(self._publish_tf)) - self.register_disposable(self.connection.video_stream().subscribe(onimage)) self.register_disposable(Disposable(self.cmd_vel.subscribe(self.move))) - self._camera_info_thread = Thread( - target=self.publish_camera_info, - daemon=True, - ) - self._camera_info_thread.start() + if self.config.camera: + self.register_disposable(self.connection.video_stream().subscribe(onimage)) + self._camera_info_thread = Thread( + target=self.publish_camera_info, + daemon=True, + ) + self._camera_info_thread.start() + + if self.config.motion_mode and isinstance(self.connection, UnitreeWebRTCConnection): + self.connection.set_motion_mode(self.config.motion_mode) self.standup() time.sleep(3) @@ -269,8 +281,6 @@ def onimage(image: Image) -> None: self.connection.set_obstacle_avoidance(self.config.g.obstacle_avoidance) - # self.record("go2_bigoffice") - @rpc def stop(self) -> None: self.liedown() From 601df177703d5077ad31b8211310dfc95332fa48 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 23 Jun 2026 11:06:22 -0700 Subject: [PATCH 44/71] Clean up --- dimos/robot/unitree/connection.py | 8 ++------ .../go2/blueprints/basic/unitree_go2_basic.py | 2 +- .../blueprints/navigation/unitree_go2_nav_3d.py | 14 ++------------ 3 files changed, 5 insertions(+), 19 deletions(-) diff --git a/dimos/robot/unitree/connection.py b/dimos/robot/unitree/connection.py index e439ed3672..d9ffb08c22 100644 --- a/dimos/robot/unitree/connection.py +++ b/dimos/robot/unitree/connection.py @@ -330,11 +330,9 @@ def set_obstacle_avoidance(self, enabled: bool = True) -> None: ) def set_motion_mode(self, name: str) -> None: - """Select the robot's top-level motion mode via the motion switcher. + """Select the top-level motion controller via the motion switcher. - 'mcf' is the AI/sport controller (firmware >= 1.1.7) that handles - terrain, including walking up and down stairs. 'normal' is the basic - controller. Stair traversal lives in mcf, so the nav stack selects it. + mcf is the AI/sport controller that traverses stairs. normal is basic. """ # api_id 1001 = CheckMode, 1002 = SelectMode, param {"name": }. current = None @@ -343,14 +341,12 @@ def set_motion_mode(self, name: str) -> None: current = json.loads(resp["data"]["data"]).get("name") except (KeyError, TypeError, ValueError) as e: print(f"Motion mode check failed: {e}") - print(f"Motion mode: current={current!r} requested={name!r}") if current == name: return self.publish_request( RTC_TOPIC["MOTION_SWITCHER"], {"api_id": 1002, "parameter": {"name": name}}, ) - # Switching controllers makes the robot re-stand; give it a moment. time.sleep(5) def free_walk(self) -> bool: diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index 8515f6b973..3bf935e33a 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -95,7 +95,7 @@ def _go2_rerun_blueprint() -> Any: ) -rerun_config = { +rerun_config: dict[str, Any] = { "blueprint": _go2_rerun_blueprint, # Custom converters for specific rerun entity paths # Normally all these would be specified in their respectative modules diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index eb36e6dbec..514d334222 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -13,14 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Go2 3D nav stack: pointlio odometry, raytraced voxel map, MLS planning, path following. +"""3d navigation on Go2 with ray tracing and MLS planning""" -Expects a mid360 + pointlio running against a lidar mounted on the robot. -The go2's built-in lidar and odometry are remapped aside and unused for -navigation. The map, paths, and robot pose all live in pointlio's world frame. -""" - -import os from typing import Any from dimos.core.coordination.blueprints import autoconnect @@ -97,11 +91,7 @@ def _static_robot_body(rr: Any) -> list[Any]: (GO2Connection, "odom", "odom_go2"), ] ), - PointLio.blueprint( - host_ip=os.getenv("LIDAR_HOST_IP", "192.168.1.5"), - lidar_ip=os.getenv("LIDAR_IP", "192.168.1.155"), - body_frame_id="body", - ), + PointLio.blueprint(body_frame_id="body"), RayTracingVoxelMap.blueprint( voxel_size=voxel_size, emit_every=2, From 166e160e8ec525c291c6d292b3370d6d039639e5 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 23 Jun 2026 11:30:56 -0700 Subject: [PATCH 45/71] Fix flake --- dimos/mapping/ray_tracing/rust/flake.lock | 11 ++++++----- dimos/mapping/ray_tracing/rust/flake.nix | 4 ++-- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/dimos/mapping/ray_tracing/rust/flake.lock b/dimos/mapping/ray_tracing/rust/flake.lock index a548660557..071efa165c 100644 --- a/dimos/mapping/ray_tracing/rust/flake.lock +++ b/dimos/mapping/ray_tracing/rust/flake.lock @@ -3,15 +3,16 @@ "dimos-repo": { "flake": false, "locked": { - "lastModified": 1779865691, - "narHash": "sha256-2CVWcov7DiC1qX/B/zFKDJiSYsnbrZ3FNT/viprFWTQ=", - "ref": "refs/heads/jeff/feat/g1_raycast", - "rev": "51666bcd298c1d08bdee179f176f45c0a7dd417d", - "revCount": 744, + "lastModified": 1782159599, + "narHash": "sha256-U2FwqtFzYen8Or50oTYGrqSi8vRePlS3m581ti2T56Q=", + "ref": "main", + "rev": "0bfe9cbbc79c52afa1ccea783147438f89bd7cf2", + "revCount": 836, "type": "git", "url": "file:../../../.." }, "original": { + "ref": "main", "type": "git", "url": "file:../../../.." } diff --git a/dimos/mapping/ray_tracing/rust/flake.nix b/dimos/mapping/ray_tracing/rust/flake.nix index 54721c779d..897d8818a2 100644 --- a/dimos/mapping/ray_tracing/rust/flake.nix +++ b/dimos/mapping/ray_tracing/rust/flake.nix @@ -6,8 +6,8 @@ flake-utils.url = "github:numtide/flake-utils"; # Relative git+file: will be deprecated (nix#12281) but there's no # viable alternative for reaching local path deps outside the flake dir currently - # presumably an alternative will be added before this is removed - dimos-repo = { url = "git+file:../../../.."; flake = false; }; + # presumably an alternative will be added before this is removed. + dimos-repo = { url = "git+file:../../../..?ref=main"; flake = false; }; }; outputs = { self, nixpkgs, flake-utils, dimos-repo }: From c794ffa24c0987c5d487d47db57604d2de5cf2fd Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 23 Jun 2026 11:52:17 -0700 Subject: [PATCH 46/71] More planning --- .../unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 514d334222..90c3fe1ec5 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -94,7 +94,7 @@ def _static_robot_body(rr: Any) -> list[Any]: PointLio.blueprint(body_frame_id="body"), RayTracingVoxelMap.blueprint( voxel_size=voxel_size, - emit_every=2, + emit_every=1, global_emit_every=50, max_health=10, graze_cos=0.85, From 8ec993c9f525ee608696b679f3ae99d85cdcbe95 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 23 Jun 2026 12:17:56 -0700 Subject: [PATCH 47/71] Step defaults and voxel size --- .../unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 90c3fe1ec5..18aacbd17d 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -29,7 +29,7 @@ from dimos.robot.unitree.go2.connection import GO2Connection from dimos.visualization.vis_module import vis_module -voxel_size = 0.1 +voxel_size = 0.08 # Height of the head-mounted lidar above the ground while standing. go2_lidar_height = 0.5 @@ -108,7 +108,7 @@ def _static_robot_body(rr: Any) -> list[Any]: wall_clearance_m=0.2, wall_buffer_m=0.75, wall_buffer_weight=100.0, - step_threshold_m=0.25, + step_threshold_m=0.16, step_penalty_weight=1.0, viz_publish_hz=0.0, ).remappings([(MLSPlannerNative, "global_map", "global_map_unused")]), From cc10edb15bd16b9e1a488c249fb7207d581a165c Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 23 Jun 2026 14:42:06 -0700 Subject: [PATCH 48/71] Validate cached paths on plan fail --- .../sensors/lidar/pointlio/cpp/main.cpp | 9 -- .../navigation/basic_path_follower/module.py | 7 +- .../nav_3d/mls_planner/rust/src/main.rs | 16 ++- .../mls_planner/rust/src/mls_planner.rs | 33 ++++- .../nav_3d/mls_planner/rust/src/planner.rs | 116 ++++++++++++++++-- .../nav_3d/mls_planner/rust/src/python.rs | 2 +- .../navigation/unitree_go2_nav_3d.py | 4 +- 7 files changed, 156 insertions(+), 31 deletions(-) diff --git a/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp b/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp index 72e7a25094..0a8603fd97 100644 --- a/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp +++ b/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp @@ -224,8 +224,6 @@ static void on_imu_data(const uint32_t /*handle*/, const uint8_t /*dev_type*/, L { static std::atomic last_pkt_ts_ns{0}; static std::atomic imu_pkt_count{0}; - static std::atomic imu_gap_count{0}; - static std::atomic max_sensor_gap_us{0}; using clk = std::chrono::steady_clock; static auto last_wall = clk::now(); auto now_wall = clk::now(); @@ -234,18 +232,11 @@ static void on_imu_data(const uint32_t /*handle*/, const uint8_t /*dev_type*/, L if (prev != 0 && pkt_ts_ns > prev) { uint64_t sensor_gap_us = (pkt_ts_ns - prev) / 1000; uint64_t wall_gap_us = std::chrono::duration_cast( now_wall - last_wall).count(); - uint64_t cur_max = max_sensor_gap_us.load(); - while (sensor_gap_us > cur_max && - !max_sensor_gap_us.compare_exchange_weak(cur_max, sensor_gap_us)) {} if (sensor_gap_us > 15000) { - imu_gap_count.fetch_add(1); fprintf(stderr, "[imu-gap] sensor_gap=%.1fms wall_gap=%.1fms pkt#%llu\n", sensor_gap_us / 1000.0, wall_gap_us / 1000.0, static_cast(pkt_count)); } } last_wall = now_wall; - if (pkt_count % 1000 == 0) { - fprintf(stderr, "[imu-stats] pkts=%llu gaps>15ms=%llu max_sensor_gap=%.1fms\n", static_cast(pkt_count), static_cast(imu_gap_count.load()), max_sensor_gap_us.load() / 1000.0); - } } double ts = static_cast(pkt_ts_ns) / 1e9; diff --git a/dimos/navigation/basic_path_follower/module.py b/dimos/navigation/basic_path_follower/module.py index e4c591e1da..d219677e4d 100644 --- a/dimos/navigation/basic_path_follower/module.py +++ b/dimos/navigation/basic_path_follower/module.py @@ -101,9 +101,12 @@ def _on_odometry(self, msg: Odometry) -> None: self._current_odom = msg.to_pose_stamped() def _on_path(self, path: Path) -> None: - # An empty path is the planner clearing its plan, not a stop. Keep the - # last path. + # The planner owns path safety: it sends the route as far as it is safe, + # or an empty path when nothing ahead is traversable. Follow what we get. if len(path.poses) == 0: + with self._lock: + self._waypoints = None + self.nav_cmd_vel.publish(Twist()) return waypoints = np.array([[p.position.x, p.position.y] for p in path.poses]) with self._lock: diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs index 70a1f21358..dd919ee413 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs @@ -319,15 +319,13 @@ impl Worker { }; let plan_start = Instant::now(); - let waypoints = tokio::task::block_in_place(|| planner.plan(start, goal, &self.config)); - let waypoints = match waypoints { - Some(wp) => wp, - None => { - tracing::warn!(?start, ?goal, "no path between start and goal"); - publish_path(&self.path, &empty_path(&self.config.world_frame, now())).await; - return; - } - }; + let waypoints = + tokio::task::block_in_place(|| planner.plan_or_truncate(start, goal, &self.config)); + if waypoints.is_empty() { + // No full path and nothing safe ahead on the cached path: stop. + publish_path(&self.path, &empty_path(&self.config.world_frame, now())).await; + return; + } let plan_ms = plan_start.elapsed().as_secs_f64() * 1e3; let produced = Instant::now(); let since_last_ms = last_path_at.map_or(-1.0, |t| (produced - t).as_secs_f64() * 1e3); diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs index 7b72b705a8..9251ab847c 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs @@ -112,6 +112,9 @@ pub struct Planner { graph: PlannerGraph, voxel_map: AHashSet, by_col: ColumnIz, + // Last successful path with the goal it served, for safe truncation when a + // later replan finds no full path. + last_path: Option<((f32, f32, f32), Vec)>, } impl Planner { @@ -406,7 +409,35 @@ impl Planner { if self.graph.nodes.is_empty() { return None; } - planner::plan(&self.graph, start, goal, config) + planner::plan(&self.graph, start, goal, config).map(|(wp, _)| wp) + } + + /// Plan to the goal, or follow the cached path as far as it is still safe. + /// Returns the waypoints, empty when nothing ahead is traversable (stop). + pub fn plan_or_truncate( + &mut self, + start: (f32, f32, f32), + goal: (f32, f32, f32), + config: &Config, + ) -> Vec<(f32, f32, f32)> { + if !self.graph.nodes.is_empty() { + if let Some((waypoints, cells)) = planner::plan(&self.graph, start, goal, config) { + self.last_path = Some((goal, cells)); + return waypoints; + } + } + match &self.last_path { + Some((cached_goal, cells)) if *cached_goal == goal => { + let safe = planner::truncate_to_safe(&self.graph, cells, start, config); + tracing::warn!( + ?goal, + safe_waypoints = safe.len(), + "no full path to goal, validating and following the cached path only while safe" + ); + safe + } + _ => Vec::new(), + } } pub fn graph(&self) -> &PlannerGraph { diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs index a43f3157ad..065b57eca4 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs @@ -22,6 +22,10 @@ const SNAP_SEARCH_RADIUS_M: f32 = 1.5; /// Max snap candidates tried when connecting the start. const MAX_SNAP_ATTEMPTS: usize = 64; +/// World-frame waypoints paired with the string-pulled cell path that produced +/// them. The cell path is cached for later safe truncation. +type PlannedPath = (Vec<(f32, f32, f32)>, Vec); + /// Surface cells near the pose, nearest first in xy. pub fn snap_candidates( surface_lookup: &SurfaceLookup, @@ -84,14 +88,14 @@ fn best_iz_in_column( } /// Plan path from start pose to goal pose using the node graph. -/// Returns none if either of the poses can't be snapped to surface or if -/// there is no valid path. +/// Returns the waypoints and the string-pulled cell path, or none if either +/// pose can't be snapped to surface or there is no valid path. pub fn plan( plg: &PlannerGraph, start_pose: (f32, f32, f32), goal_pose: (f32, f32, f32), config: &Config, -) -> Option> { +) -> Option { let voxel_size = config.voxel_size; let z_tolerance_m = config.robot_height; let start_candidates = @@ -180,9 +184,52 @@ pub fn plan( }; let cells = assemble_cells(plg, &node_seq, &lead_in, &goal_segment); let cells = string_pull(plg, &cells, step_cells, &wall_cost); - Some(cells_to_waypoints( - plg, &cells, start_pose, goal_pose, voxel_size, - )) + let waypoints = cells_to_waypoints(plg, &cells, start_pose, goal_pose, voxel_size); + let path_cells: Vec = cells.iter().map(|&id| plg.cells.coord(id)).collect(); + Some((waypoints, path_cells)) +} + +/// Re-validate a cached cell path against the current surface, returning the +/// waypoints up to the first segment no longer traversable from the robot. +/// Empty when nothing ahead is safe, which the follower reads as a stop. +pub fn truncate_to_safe( + plg: &PlannerGraph, + cached: &[VoxelKey], + start_pose: (f32, f32, f32), + config: &Config, +) -> Vec<(f32, f32, f32)> { + if cached.is_empty() { + return Vec::new(); + } + let voxel_size = config.voxel_size; + let step_cells = (config.step_threshold_m / voxel_size).floor() as i32; + let wall_cost = WallCost { + clearance_m: config.wall_clearance_m, + buffer_m: config.wall_buffer_m, + buffer_weight: config.wall_buffer_weight, + voxel_size, + }; + + // The cached path runs start -> goal. Validate from its start toward the + // goal and keep the contiguous prefix that is still traversable. + let mut safe_end = 0usize; + for j in 0..cached.len() - 1 { + if segment_metrics(plg, cached[j], cached[j + 1], step_cells, &wall_cost).is_some() { + safe_end = j + 1; + } else { + break; + } + } + if safe_end == 0 { + return Vec::new(); + } + + let mut waypoints = Vec::with_capacity(safe_end + 2); + waypoints.push(start_pose); + for &(ix, iy, iz) in &cached[..=safe_end] { + waypoints.push(surface_point_xyz(ix, iy, iz, voxel_size)); + } + waypoints } /// Pick the entry node by connect cost plus cost-to-go, with its on-surface @@ -654,7 +701,62 @@ mod tests { goal_tolerance: 0.3, viz_publish_hz: 2.0, }; - plan(plg, start, goal, &config) + plan(plg, start, goal, &config).map(|(wp, _)| wp) + } + + fn surface_graph(cells: &[VoxelKey]) -> PlannerGraph { + let mut plg = PlannerGraph::new(); + build_surface_lookup(cells, &mut plg.surface_lookup); + build_surface_cells(&mut plg.cells, &plg.surface_lookup, VOXEL, 2); + plg + } + + fn truncate_config() -> Config { + Config { + world_frame: "world".into(), + voxel_size: VOXEL, + robot_height: Z_TOL, + surface_closing_radius: 0.0, + node_spacing_m: 1.0, + wall_clearance_m: 0.2, + wall_buffer_m: 0.5, + wall_buffer_weight: 4.0, + step_threshold_m: 0.25, + step_penalty_weight: 4.0, + goal_tolerance: 0.3, + viz_publish_hz: 2.0, + } + } + + #[test] + fn truncate_validates_from_start_and_cuts_at_break() { + let cfg = truncate_config(); + // Cached path runs start (x=0) -> goal (x=9) along the strip. + let cached: Vec = (0..10).map(|x| (x, 0, 0)).collect(); + let start = surface_point_xyz(0, 0, 0, VOXEL); + + // Intact surface: the entire cached path is still traversable. + let full = truncate_to_safe(&surface_graph(&cached), &cached, start, &cfg); + assert_eq!( + full.len(), + cached.len() + 1, + "start pose + every cached cell" + ); + + // Surface gone at x=5: validate from the start, cut before the gap, and + // keep the start-side prefix (cells 0..=4). + let gap: Vec = (0..10).filter(|&x| x != 5).map(|x| (x, 0, 0)).collect(); + let cut = truncate_to_safe(&surface_graph(&gap), &cached, start, &cfg); + assert_eq!(cut.len(), 6, "start pose + cells 0..=4"); + assert_eq!(*cut.last().unwrap(), surface_point_xyz(4, 0, 0, VOXEL)); + + // Surface gone at x=1: the first step is blocked, so nothing ahead is + // safe and the result is empty (a stop). + let early: Vec = (0..10).filter(|&x| x != 1).map(|x| (x, 0, 0)).collect(); + assert!( + truncate_to_safe(&surface_graph(&early), &cached, start, &cfg).is_empty(), + "first step blocked -> empty (stop)" + ); } #[test] diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs index 3b6e0ade0e..d320eff253 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs @@ -205,7 +205,7 @@ impl MLSPlanner { .into_pyarray(py) } - /// Returns `(W, 3)` float32 waypoints or `None` if no path exists. + /// Returns `(W, 3)` float32 waypoints or `None` if no full path exists. fn plan<'py>( &self, py: Python<'py>, diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 18aacbd17d..1fae6b4ad3 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -39,8 +39,8 @@ def _render_global_map(msg: Any) -> Any: def _render_path(msg: Any) -> Any: - # The planner clears its plan with an empty path on every start-pose change. - # Logging those would blank the line. Drop them so the last path stays shown. + # The planner emits an empty path when it finds no route to the goal. + # Logging those would blank the line, so drop them and keep the last path. if len(msg.poses) == 0: return None return msg From 120583624be95cb36f6a54b51bf5e5d664d22911 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 23 Jun 2026 15:24:52 -0700 Subject: [PATCH 49/71] Escape the well --- .../nav_3d/mls_planner/rust/src/planner.rs | 110 ++++++++++++++---- 1 file changed, 90 insertions(+), 20 deletions(-) diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs index 065b57eca4..d67cb91fb0 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs @@ -190,7 +190,7 @@ pub fn plan( } /// Re-validate a cached cell path against the current surface, returning the -/// waypoints up to the first segment no longer traversable from the robot. +/// route ahead of the robot up to the first segment no longer traversable. /// Empty when nothing ahead is safe, which the follower reads as a stop. pub fn truncate_to_safe( plg: &PlannerGraph, @@ -198,7 +198,7 @@ pub fn truncate_to_safe( start_pose: (f32, f32, f32), config: &Config, ) -> Vec<(f32, f32, f32)> { - if cached.is_empty() { + if cached.len() < 2 { return Vec::new(); } let voxel_size = config.voxel_size; @@ -210,28 +210,68 @@ pub fn truncate_to_safe( voxel_size, }; - // The cached path runs start -> goal. Validate from its start toward the - // goal and keep the contiguous prefix that is still traversable. - let mut safe_end = 0usize; - for j in 0..cached.len() - 1 { + // The cached path runs start -> goal, but its head is the stale original + // start. Resume from where the robot sits on it so the follower is pulled + // forward toward the goal, never back to that start. Validate each segment + // ahead and keep the contiguous traversable run. + let resume = resume_segment(cached, start_pose, voxel_size); + let mut safe_end = resume; + for j in resume..cached.len() - 1 { if segment_metrics(plg, cached[j], cached[j + 1], step_cells, &wall_cost).is_some() { safe_end = j + 1; } else { break; } } - if safe_end == 0 { + if safe_end == resume { return Vec::new(); } - let mut waypoints = Vec::with_capacity(safe_end + 2); + let mut waypoints = Vec::with_capacity(safe_end - resume + 1); waypoints.push(start_pose); - for &(ix, iy, iz) in &cached[..=safe_end] { + for &(ix, iy, iz) in &cached[resume + 1..=safe_end] { waypoints.push(surface_point_xyz(ix, iy, iz, voxel_size)); } waypoints } +/// Index of the cached segment the robot is on, by nearest-point projection in +/// the ground plane. The route ahead resumes at the following cell. +fn resume_segment(cached: &[VoxelKey], start: (f32, f32, f32), voxel_size: f32) -> usize { + let p = (start.0, start.1); + let mut best = 0usize; + let mut best_d2 = f32::INFINITY; + for i in 0..cached.len() - 1 { + let a = surface_point_xyz(cached[i].0, cached[i].1, cached[i].2, voxel_size); + let b = surface_point_xyz( + cached[i + 1].0, + cached[i + 1].1, + cached[i + 1].2, + voxel_size, + ); + let d2 = point_segment_dist2((a.0, a.1), (b.0, b.1), p); + if d2 < best_d2 { + best_d2 = d2; + best = i; + } + } + best +} + +/// Squared distance from point p to segment a-b in the plane. +fn point_segment_dist2(a: (f32, f32), b: (f32, f32), p: (f32, f32)) -> f32 { + let (abx, aby) = (b.0 - a.0, b.1 - a.1); + let denom = abx * abx + aby * aby; + let t = if denom == 0.0 { + 0.0 + } else { + (((p.0 - a.0) * abx + (p.1 - a.1) * aby) / denom).clamp(0.0, 1.0) + }; + let (cx, cy) = (a.0 + t * abx, a.1 + t * aby); + let (dx, dy) = (p.0 - cx, p.1 - cy); + dx * dx + dy * dy +} + /// Pick the entry node by connect cost plus cost-to-go, with its on-surface /// lead-in and the node sequence to the goal. fn select_entry( @@ -729,28 +769,26 @@ mod tests { } #[test] - fn truncate_validates_from_start_and_cuts_at_break() { + fn truncate_resumes_ahead_of_robot_and_cuts_at_break() { let cfg = truncate_config(); // Cached path runs start (x=0) -> goal (x=9) along the strip. let cached: Vec = (0..10).map(|x| (x, 0, 0)).collect(); let start = surface_point_xyz(0, 0, 0, VOXEL); - // Intact surface: the entire cached path is still traversable. + // Intact surface, robot at the start: keep the whole route ahead, which + // is the start pose plus cells 1..=9 (the robot's own cell is dropped). let full = truncate_to_safe(&surface_graph(&cached), &cached, start, &cfg); - assert_eq!( - full.len(), - cached.len() + 1, - "start pose + every cached cell" - ); + assert_eq!(full.len(), cached.len(), "start pose + cells ahead"); + assert_eq!(full[1], surface_point_xyz(1, 0, 0, VOXEL)); + assert_eq!(*full.last().unwrap(), surface_point_xyz(9, 0, 0, VOXEL)); - // Surface gone at x=5: validate from the start, cut before the gap, and - // keep the start-side prefix (cells 0..=4). + // Surface gone at x=5: cut before the gap, keep cells 1..=4 ahead. let gap: Vec = (0..10).filter(|&x| x != 5).map(|x| (x, 0, 0)).collect(); let cut = truncate_to_safe(&surface_graph(&gap), &cached, start, &cfg); - assert_eq!(cut.len(), 6, "start pose + cells 0..=4"); + assert_eq!(cut.len(), 5, "start pose + cells 1..=4"); assert_eq!(*cut.last().unwrap(), surface_point_xyz(4, 0, 0, VOXEL)); - // Surface gone at x=1: the first step is blocked, so nothing ahead is + // Surface gone at x=1: the first step ahead is blocked, so nothing is // safe and the result is empty (a stop). let early: Vec = (0..10).filter(|&x| x != 1).map(|x| (x, 0, 0)).collect(); assert!( @@ -759,6 +797,38 @@ mod tests { ); } + #[test] + fn truncate_does_not_backtrack_when_robot_has_advanced() { + let cfg = truncate_config(); + // Cached route start (x=0) -> goal (x=9). The robot has already walked + // to x=5, and the surface is now gone at x=8 (a door closed ahead). + let cached: Vec = (0..10).map(|x| (x, 0, 0)).collect(); + let robot = surface_point_xyz(5, 0, 0, VOXEL); + let blocked: Vec = (0..10).filter(|&x| x != 8).map(|x| (x, 0, 0)).collect(); + + let wp = truncate_to_safe(&surface_graph(&blocked), &cached, robot, &cfg); + + // Resumes ahead of the robot: no cell behind x=5, up to the x=8 break. + assert_eq!(wp[0], robot); + let xs: Vec = wp[1..] + .iter() + .map(|w| (w.0 / VOXEL).floor() as i32) + .collect(); + assert!( + xs.iter().all(|&x| x >= 5), + "path backtracked behind the robot: {xs:?}" + ); + assert_eq!( + *xs.last().unwrap(), + 7, + "should stop just before the x=8 break" + ); + assert!( + xs.windows(2).all(|p| p[1] > p[0]), + "path not monotonic forward: {xs:?}" + ); + } + #[test] fn snap_picks_in_column_cell() { let mut lookup = SurfaceLookup::new(); From 75f2335f82e5413cababe998d3c25e01fe2b2e18 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 23 Jun 2026 15:40:00 -0700 Subject: [PATCH 50/71] Best effort distance on truncating --- .../nav_3d/mls_planner/rust/src/planner.rs | 150 +++++++++++++----- 1 file changed, 109 insertions(+), 41 deletions(-) diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs index d67cb91fb0..e620658a94 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs @@ -22,6 +22,11 @@ const SNAP_SEARCH_RADIUS_M: f32 = 1.5; /// Max snap candidates tried when connecting the start. const MAX_SNAP_ATTEMPTS: usize = 64; +/// Standoff held from a blockage on a truncated path. The robot follows the +/// cached route to best effort and stops this far short of the last +/// traversable point, in meters. +const BEST_EFFORT_DISTANCE_M: f32 = 1.0; + /// World-frame waypoints paired with the string-pulled cell path that produced /// them. The cell path is cached for later safe truncation. type PlannedPath = (Vec<(f32, f32, f32)>, Vec); @@ -191,7 +196,9 @@ pub fn plan( /// Re-validate a cached cell path against the current surface, returning the /// route ahead of the robot up to the first segment no longer traversable. -/// Empty when nothing ahead is safe, which the follower reads as a stop. +/// When the route is cut by a blockage, it holds a BEST_EFFORT_DISTANCE_M +/// standoff short of it. Empty when nothing ahead is safe, which the follower +/// reads as a stop. pub fn truncate_to_safe( plg: &PlannerGraph, cached: &[VoxelKey], @@ -232,9 +239,46 @@ pub fn truncate_to_safe( for &(ix, iy, iz) in &cached[resume + 1..=safe_end] { waypoints.push(surface_point_xyz(ix, iy, iz, voxel_size)); } + + // A break short of the goal end means the route runs into a blockage, so + // hold the standoff. A run to the goal end has no blockage to stand off. + if safe_end + 1 < cached.len() { + return back_off_tail(&waypoints, BEST_EFFORT_DISTANCE_M); + } waypoints } +/// Drop `distance` meters off the goal end of the path, measured in the ground +/// plane. Empty when trimming consumes the whole path, which is the robot +/// already sitting inside the standoff. +fn back_off_tail(waypoints: &[(f32, f32, f32)], distance: f32) -> Vec<(f32, f32, f32)> { + let mut remaining = distance; + for i in (1..waypoints.len()).rev() { + let (b, a) = (waypoints[i], waypoints[i - 1]); + let seg = (b.0 - a.0).hypot(b.1 - a.1); + if seg < remaining { + remaining -= seg; + continue; + } + let t = if seg == 0.0 { + 0.0 + } else { + (seg - remaining) / seg + }; + let cut = ( + a.0 + (b.0 - a.0) * t, + a.1 + (b.1 - a.1) * t, + a.2 + (b.2 - a.2) * t, + ); + let mut out = waypoints[..i].to_vec(); + if out.last() != Some(&cut) { + out.push(cut); + } + return if out.len() >= 2 { out } else { Vec::new() }; + } + Vec::new() +} + /// Index of the cached segment the robot is on, by nearest-point projection in /// the ground plane. The route ahead resumes at the following cell. fn resume_segment(cached: &[VoxelKey], start: (f32, f32, f32), voxel_size: f32) -> usize { @@ -769,66 +813,90 @@ mod tests { } #[test] - fn truncate_resumes_ahead_of_robot_and_cuts_at_break() { + fn truncate_keeps_the_full_clear_route_ahead() { let cfg = truncate_config(); - // Cached path runs start (x=0) -> goal (x=9) along the strip. - let cached: Vec = (0..10).map(|x| (x, 0, 0)).collect(); + // Cached route 0 -> 39 (cells), still fully traversable, robot at start. + let cached: Vec = (0..40).map(|x| (x, 0, 0)).collect(); let start = surface_point_xyz(0, 0, 0, VOXEL); - // Intact surface, robot at the start: keep the whole route ahead, which - // is the start pose plus cells 1..=9 (the robot's own cell is dropped). + // No blockage, so no standoff: start pose plus every cell ahead, where + // the robot's own cell 0 is dropped. let full = truncate_to_safe(&surface_graph(&cached), &cached, start, &cfg); - assert_eq!(full.len(), cached.len(), "start pose + cells ahead"); + assert_eq!(full.len(), cached.len(), "start pose + cells 1..=39"); assert_eq!(full[1], surface_point_xyz(1, 0, 0, VOXEL)); - assert_eq!(*full.last().unwrap(), surface_point_xyz(9, 0, 0, VOXEL)); + assert_eq!(*full.last().unwrap(), surface_point_xyz(39, 0, 0, VOXEL)); + } - // Surface gone at x=5: cut before the gap, keep cells 1..=4 ahead. - let gap: Vec = (0..10).filter(|&x| x != 5).map(|x| (x, 0, 0)).collect(); - let cut = truncate_to_safe(&surface_graph(&gap), &cached, start, &cfg); - assert_eq!(cut.len(), 5, "start pose + cells 1..=4"); - assert_eq!(*cut.last().unwrap(), surface_point_xyz(4, 0, 0, VOXEL)); + #[test] + fn truncate_holds_standoff_ahead_of_advanced_robot() { + let cfg = truncate_config(); + // Robot has advanced to x=20 along a 0 -> 39 route, and the surface is + // now gone at x=35 (a door closed ahead). + let cached: Vec = (0..40).map(|x| (x, 0, 0)).collect(); + let robot = surface_point_xyz(20, 0, 0, VOXEL); + let blocked: Vec = (0..40).filter(|&x| x != 35).map(|x| (x, 0, 0)).collect(); - // Surface gone at x=1: the first step ahead is blocked, so nothing is - // safe and the result is empty (a stop). - let early: Vec = (0..10).filter(|&x| x != 1).map(|x| (x, 0, 0)).collect(); + let wp = truncate_to_safe(&surface_graph(&blocked), &cached, robot, &cfg); + assert_eq!(wp[0], robot); + + // Never behind the robot, always forward toward the goal. + let xs: Vec = wp.iter().map(|w| w.0).collect(); + assert!( + xs.iter().all(|&x| x >= robot.0 - 1e-4), + "backtracked: {xs:?}" + ); + assert!(xs.windows(2).all(|p| p[1] >= p[0]), "not forward: {xs:?}"); + + // Stops a standoff short of the last traversable cell (x=34). + let last_safe = surface_point_xyz(34, 0, 0, VOXEL); + let last = *wp.last().unwrap(); + let gap = (last_safe.0 - last.0).hypot(last_safe.1 - last.1); assert!( - truncate_to_safe(&surface_graph(&early), &cached, start, &cfg).is_empty(), - "first step blocked -> empty (stop)" + (gap - BEST_EFFORT_DISTANCE_M).abs() < VOXEL, + "standoff is {gap} m, expected ~{BEST_EFFORT_DISTANCE_M}" ); } #[test] - fn truncate_does_not_backtrack_when_robot_has_advanced() { + fn truncate_stops_inside_standoff_or_at_blockage() { let cfg = truncate_config(); - // Cached route start (x=0) -> goal (x=9). The robot has already walked - // to x=5, and the surface is now gone at x=8 (a door closed ahead). - let cached: Vec = (0..10).map(|x| (x, 0, 0)).collect(); - let robot = surface_point_xyz(5, 0, 0, VOXEL); - let blocked: Vec = (0..10).filter(|&x| x != 8).map(|x| (x, 0, 0)).collect(); + let cached: Vec = (0..40).map(|x| (x, 0, 0)).collect(); + let robot = surface_point_xyz(0, 0, 0, VOXEL); - let wp = truncate_to_safe(&surface_graph(&blocked), &cached, robot, &cfg); - - // Resumes ahead of the robot: no cell behind x=5, up to the x=8 break. - assert_eq!(wp[0], robot); - let xs: Vec = wp[1..] - .iter() - .map(|w| (w.0 / VOXEL).floor() as i32) - .collect(); + // Blockage at the next step: nothing safe ahead, stop. + let at_robot: Vec = (0..40).filter(|&x| x != 1).map(|x| (x, 0, 0)).collect(); assert!( - xs.iter().all(|&x| x >= 5), - "path backtracked behind the robot: {xs:?}" - ); - assert_eq!( - *xs.last().unwrap(), - 7, - "should stop just before the x=8 break" + truncate_to_safe(&surface_graph(&at_robot), &cached, robot, &cfg).is_empty(), + "blockage at the next step -> stop" ); + + // Blockage only ~0.4 m ahead, inside the standoff: the best-effort point + // is behind the robot, so stop. + let near: Vec = (0..40).filter(|&x| x != 5).map(|x| (x, 0, 0)).collect(); assert!( - xs.windows(2).all(|p| p[1] > p[0]), - "path not monotonic forward: {xs:?}" + truncate_to_safe(&surface_graph(&near), &cached, robot, &cfg).is_empty(), + "inside the standoff -> stop" ); } + #[test] + fn back_off_tail_trims_from_the_goal_end() { + let path = vec![ + (0.0, 0.0, 0.0), + (1.0, 0.0, 0.0), + (2.0, 0.0, 0.0), + (3.0, 0.0, 0.0), + ]; + // Trim within the last segment. + assert_eq!(*back_off_tail(&path, 0.5).last().unwrap(), (2.5, 0.0, 0.0)); + // Trim exactly to a vertex without leaving a duplicate point. + let to_vertex = back_off_tail(&path, 1.0); + assert_eq!(*to_vertex.last().unwrap(), (2.0, 0.0, 0.0)); + assert_eq!(to_vertex.len(), 3); + // Trimming more than the path length stops (empty). + assert!(back_off_tail(&path, 5.0).is_empty()); + } + #[test] fn snap_picks_in_column_cell() { let mut lookup = SurfaceLookup::new(); From d8aad1f6102683f35fab6126f76ad85908f6052b Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 23 Jun 2026 16:02:39 -0700 Subject: [PATCH 51/71] Bug fix on best effort distance --- .../nav_3d/mls_planner/rust/src/planner.rs | 152 +++++++++++++++--- 1 file changed, 132 insertions(+), 20 deletions(-) diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs index e620658a94..8048801b31 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs @@ -195,10 +195,9 @@ pub fn plan( } /// Re-validate a cached cell path against the current surface, returning the -/// route ahead of the robot up to the first segment no longer traversable. -/// When the route is cut by a blockage, it holds a BEST_EFFORT_DISTANCE_M -/// standoff short of it. Empty when nothing ahead is safe, which the follower -/// reads as a stop. +/// route ahead of the robot up to the last traversable surface cell before a +/// blockage. It holds a BEST_EFFORT_DISTANCE_M standoff short of the blockage. +/// Empty when nothing ahead is safe, which the follower reads as a stop. pub fn truncate_to_safe( plg: &PlannerGraph, cached: &[VoxelKey], @@ -219,30 +218,43 @@ pub fn truncate_to_safe( // The cached path runs start -> goal, but its head is the stale original // start. Resume from where the robot sits on it so the follower is pulled - // forward toward the goal, never back to that start. Validate each segment - // ahead and keep the contiguous traversable run. + // forward toward the goal, never back to that start. let resume = resume_segment(cached, start_pose, voxel_size); - let mut safe_end = resume; + + // Walk each chord ahead at surface resolution. On a blockage, keep the + // chord up to its last traversable cell so the path ends right at the + // obstacle, not back at the previous smoothing anchor. + let mut waypoints = vec![start_pose]; + let mut blocked = false; for j in resume..cached.len() - 1 { - if segment_metrics(plg, cached[j], cached[j + 1], step_cells, &wall_cost).is_some() { - safe_end = j + 1; - } else { + let (last_safe, cut) = last_safe_on_chord( + plg, + cached[j], + cached[j + 1], + step_cells, + &wall_cost, + voxel_size, + ); + if cut { + if let Some(p) = last_safe { + if waypoints.last() != Some(&p) { + waypoints.push(p); + } + } + blocked = true; break; } - } - if safe_end == resume { - return Vec::new(); + let (ix, iy, iz) = cached[j + 1]; + waypoints.push(surface_point_xyz(ix, iy, iz, voxel_size)); } - let mut waypoints = Vec::with_capacity(safe_end - resume + 1); - waypoints.push(start_pose); - for &(ix, iy, iz) in &cached[resume + 1..=safe_end] { - waypoints.push(surface_point_xyz(ix, iy, iz, voxel_size)); + if waypoints.len() < 2 { + return Vec::new(); } - // A break short of the goal end means the route runs into a blockage, so - // hold the standoff. A run to the goal end has no blockage to stand off. - if safe_end + 1 < cached.len() { + // A blockage ahead means the route runs into an obstacle, so hold the + // standoff. A clean run to the goal end has nothing to stand off from. + if blocked { return back_off_tail(&waypoints, BEST_EFFORT_DISTANCE_M); } waypoints @@ -279,6 +291,82 @@ fn back_off_tail(waypoints: &[(f32, f32, f32)], distance: f32) -> Vec<(f32, f32, Vec::new() } +/// Walk the straight chord a -> b at surface resolution, the same sampling the +/// segment validator uses. Returns the world point of the last traversable +/// surface cell reached and whether the chord was cut short of b. The point is +/// None only when the chord's start column is already off the surface. +fn last_safe_on_chord( + plg: &PlannerGraph, + a: VoxelKey, + b: VoxelKey, + step_cells: i32, + wc: &WallCost, + voxel_size: f32, +) -> (Option<(f32, f32, f32)>, bool) { + let (dx, dy, dz) = (b.0 - a.0, b.1 - a.1, b.2 - a.2); + let samples = dx.abs().max(dy.abs()) * 2; + if samples == 0 { + // Same column: traversable only if it is not a pure vertical move. + return if dz == 0 { + (Some(surface_point_xyz(a.0, a.1, a.2, voxel_size)), false) + } else { + (None, true) + }; + } + let (mut last_ix, mut last_iy) = (i32::MIN, i32::MIN); + let mut prev_iz: Option = None; + let mut last_safe: Option<(f32, f32, f32)> = None; + for k in 0..=samples { + let t = k as f32 / samples as f32; + let ix = (a.0 as f32 + t * dx as f32).round() as i32; + let iy = (a.1 as f32 + t * dy as f32).round() as i32; + if ix == last_ix && iy == last_iy { + continue; + } + last_ix = ix; + last_iy = iy; + let iz_line = a.2 as f32 + t * dz as f32; + let Some(zs) = plg.surface_lookup.get(&(ix, iy)) else { + return (last_safe, true); + }; + // Surface cell in this column nearest the interpolated segment height. + let mut nearest: Option<(f32, i32)> = None; + for &iz in zs { + let d = (iz as f32 - iz_line).abs(); + if nearest.is_none_or(|(bd, _)| d < bd) { + nearest = Some((d, iz)); + } + } + let Some((d, iz)) = nearest else { + return (last_safe, true); + }; + if d > step_cells as f32 { + return (last_safe, true); + } + if prev_iz.is_some_and(|p| (iz - p).abs() > step_cells) { + return (last_safe, true); + } + let pen = match plg.cells.id((ix, iy, iz)) { + Some(id) => { + let wall_dist = plg + .wall_state + .dist + .get(id as usize) + .copied() + .unwrap_or(f32::INFINITY); + penalty_of(wall_dist, wc.clearance_m, wc.buffer_m, wc.buffer_weight) + } + None => 1.0, + }; + if !pen.is_finite() { + return (last_safe, true); + } + prev_iz = Some(iz); + last_safe = Some(surface_point_xyz(ix, iy, iz, voxel_size)); + } + (last_safe, false) +} + /// Index of the cached segment the robot is on, by nearest-point projection in /// the ground plane. The route ahead resumes at the following cell. fn resume_segment(cached: &[VoxelKey], start: (f32, f32, f32), voxel_size: f32) -> usize { @@ -857,6 +945,30 @@ mod tests { ); } + #[test] + fn truncate_walks_into_a_sparse_chord_to_the_blockage() { + let cfg = truncate_config(); + // Sparse cached route: one 0 -> 39 chord, as string_pull leaves a + // straight approach. The surface is gone at x=20, mid-chord, where the + // old anchor-level cut would have discarded the whole chord and stopped + // back at x=0. + let cached: Vec = vec![(0, 0, 0), (39, 0, 0)]; + let surface: Vec = (0..40).filter(|&x| x != 20).map(|x| (x, 0, 0)).collect(); + let robot = surface_point_xyz(0, 0, 0, VOXEL); + + let wp = truncate_to_safe(&surface_graph(&surface), &cached, robot, &cfg); + + // It advances well into the chord and stops a standoff short of the gap. + let last = *wp.last().unwrap(); + assert!(last.0 > 0.5, "did not walk into the chord: {last:?}"); + let last_safe = surface_point_xyz(19, 0, 0, VOXEL); + let gap = (last_safe.0 - last.0).hypot(last_safe.1 - last.1); + assert!( + (gap - BEST_EFFORT_DISTANCE_M).abs() < VOXEL, + "standoff is {gap} m from the last safe cell, expected ~{BEST_EFFORT_DISTANCE_M}" + ); + } + #[test] fn truncate_stops_inside_standoff_or_at_blockage() { let cfg = truncate_config(); From 11d7035766a63992be7c13d4f8c3e796e16a6230 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 23 Jun 2026 16:32:44 -0700 Subject: [PATCH 52/71] Clean up --- .../navigation/basic_path_follower/module.py | 19 ++++++------- .../nav_3d/mls_planner/rust/src/planner.rs | 27 +++++++++++++------ dimos/robot/unitree/connection.py | 7 +++-- 3 files changed, 34 insertions(+), 19 deletions(-) diff --git a/dimos/navigation/basic_path_follower/module.py b/dimos/navigation/basic_path_follower/module.py index d219677e4d..5b0b4c491f 100644 --- a/dimos/navigation/basic_path_follower/module.py +++ b/dimos/navigation/basic_path_follower/module.py @@ -21,6 +21,7 @@ from dimos_lcm.std_msgs import Bool # type: ignore[import-untyped] import numpy as np +from numpy.typing import NDArray from reactivex.disposable import Disposable from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT @@ -74,7 +75,7 @@ def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._lock = RLock() self._current_odom: PoseStamped | None = None - self._waypoints: np.ndarray | None = None + self._waypoints: NDArray[np.float32] | None = None self._stop_event = Event() self._thread: Thread | None = None @@ -108,7 +109,7 @@ def _on_path(self, path: Path) -> None: self._waypoints = None self.nav_cmd_vel.publish(Twist()) return - waypoints = np.array([[p.position.x, p.position.y] for p in path.poses]) + waypoints = np.array([[p.position.x, p.position.y] for p in path.poses], dtype=np.float32) with self._lock: self._waypoints = waypoints @@ -130,8 +131,8 @@ def _follow(self) -> None: elapsed = time.perf_counter() - start_time self._stop_event.wait(max(0.0, period - elapsed)) - def _step(self, odom: PoseStamped, waypoints: np.ndarray) -> None: - position = np.array([odom.position.x, odom.position.y]) + def _step(self, odom: PoseStamped, waypoints: NDArray[np.float32]) -> None: + position = np.array([odom.position.x, odom.position.y], dtype=np.float32) if float(np.linalg.norm(waypoints[-1] - position)) < self.config.goal_tolerance: self.nav_cmd_vel.publish(Twist()) with self._lock: @@ -151,12 +152,12 @@ def _step(self, odom: PoseStamped, waypoints: np.ndarray) -> None: -self.config.max_angular, min(self.config.max_angular, self.config.heading_gain * yaw_error), ) - # Taper forward speed by cos(yaw_error): decelerate into turns and pivot - # in place near 90 degrees, rather than stop and lurch. linear = self.config.speed * max(0.0, math.cos(yaw_error)) self.nav_cmd_vel.publish(Twist(Vector3(linear, 0, 0), Vector3(0, 0, angular))) - def _lookahead_point(self, waypoints: np.ndarray, position: np.ndarray) -> np.ndarray: + def _lookahead_point( + self, waypoints: NDArray[np.float32], position: NDArray[np.float32] + ) -> NDArray[np.float32]: if len(waypoints) == 1: return np.asarray(waypoints[0]) @@ -180,8 +181,8 @@ def _lookahead_point(self, waypoints: np.ndarray, position: np.ndarray) -> np.nd return np.asarray(waypoints[-1]) def _project_onto_path( - self, waypoints: np.ndarray, position: np.ndarray - ) -> tuple[int, np.ndarray]: + self, waypoints: NDArray[np.float32], position: NDArray[np.float32] + ) -> tuple[int, NDArray[np.float32]]: best_idx = 0 best_point = np.asarray(waypoints[0]) best_dist = math.inf diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs index 8048801b31..3475ed0db3 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs @@ -459,30 +459,34 @@ fn select_entry( )) } -/// Bounded Dijkstra from the robot cell, visiting cells within the radius. -/// Returns per-cell distance and predecessor maps. +/// Bounded Dijkstra from the robot cell. Cost is wall-penalized for steering, +/// but the radius bounds metric distance, not penalized cost. fn robot_search( cells: &SurfaceCells, source: CellId, radius_m: f32, ) -> (AHashMap, AHashMap) { let mut dist: AHashMap = AHashMap::new(); + let mut geo: AHashMap = AHashMap::new(); let mut pred: AHashMap = AHashMap::new(); let mut heap: BinaryHeap = BinaryHeap::new(); dist.insert(source, 0.0); + geo.insert(source, 0.0); heap.push(Scored(0.0, source)); while let Some(Scored(d, u)) = heap.pop() { - if d > radius_m { - break; - } if d > dist.get(&u).copied().unwrap_or(f32::INFINITY) { continue; } + // Stop expanding past the metric radius. + if geo.get(&u).copied().unwrap_or(f32::INFINITY) > radius_m { + continue; + } for edge in cells.neighbors(u) { let nd = d + edge.cost; if nd < dist.get(&edge.dest).copied().unwrap_or(f32::INFINITY) { dist.insert(edge.dest, nd); + geo.insert(edge.dest, geo[&u] + edge.base_cost); pred.insert(edge.dest, u); heap.push(Scored(nd, edge.dest)); } @@ -696,9 +700,16 @@ fn string_pull( let mut rough_rise = 0.0_f32; let mut j = anchor + 1; while j < cells.len() { - if let Some((pen, rise)) = metrics(cells[j - 1], cells[j]) { - rough_pen = rough_pen.max(pen); - rough_rise += rise; + match metrics(cells[j - 1], cells[j]) { + Some((pen, rise)) => { + rough_pen = rough_pen.max(pen); + rough_rise += rise; + } + // Infeasible step. Raise the baseline so the chord breaks. + None => { + rough_pen = f32::INFINITY; + rough_rise = f32::INFINITY; + } } match metrics(cells[anchor], cells[j]) { Some((pen, rise)) if pen <= rough_pen + 1e-3 && rise <= rough_rise + 1e-3 => { diff --git a/dimos/robot/unitree/connection.py b/dimos/robot/unitree/connection.py index d9ffb08c22..a818c92d3a 100644 --- a/dimos/robot/unitree/connection.py +++ b/dimos/robot/unitree/connection.py @@ -51,10 +51,13 @@ from dimos.robot.unitree.type.odometry import Odometry from dimos.types.timestamped import Timestamped from dimos.utils.decorators.decorators import simple_mcache +from dimos.utils.logging_config import setup_logger from dimos.utils.reactive import backpressure, callback_to_observable VideoMessage: TypeAlias = NDArray[np.uint8] # Shape: (height, width, 3) +logger = setup_logger() + _T = TypeVar("_T", bound=Timestamped) @@ -226,7 +229,7 @@ async def async_move_duration() -> None: future.result() return True except Exception as e: - print(f"Failed to send movement command: {e}") + logger.warning("Failed to send movement command: %s", e) return False # Generic conversion of unitree subscription to Subject (used for all subs) @@ -340,7 +343,7 @@ def set_motion_mode(self, name: str) -> None: resp = self.publish_request(RTC_TOPIC["MOTION_SWITCHER"], {"api_id": 1001}) current = json.loads(resp["data"]["data"]).get("name") except (KeyError, TypeError, ValueError) as e: - print(f"Motion mode check failed: {e}") + logger.warning("Motion mode check failed: %s", e) if current == name: return self.publish_request( From 4bde8834a4e138687ef55bea5e323a129067c502 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 23 Jun 2026 16:54:05 -0700 Subject: [PATCH 53/71] Remove debug logging --- .../sensors/lidar/pointlio/cpp/main.cpp | 20 ------------------- 1 file changed, 20 deletions(-) diff --git a/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp b/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp index 0a8603fd97..807eae074c 100644 --- a/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp +++ b/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp @@ -219,26 +219,6 @@ static void on_imu_data(const uint32_t /*handle*/, const uint8_t /*dev_type*/, L if (!g_running.load() || data == nullptr || !g_point_lio) { return; } uint64_t pkt_ts_ns = get_timestamp_ns(data); - // Live IMU-drop instrumentation: a dropped datagram shows as a sensor-ts - // jump; wall gaps exceeding sensor gaps mean callback starvation. - { - static std::atomic last_pkt_ts_ns{0}; - static std::atomic imu_pkt_count{0}; - using clk = std::chrono::steady_clock; - static auto last_wall = clk::now(); - auto now_wall = clk::now(); - uint64_t prev = last_pkt_ts_ns.exchange(pkt_ts_ns); - uint64_t pkt_count = imu_pkt_count.fetch_add(1) + 1; - if (prev != 0 && pkt_ts_ns > prev) { - uint64_t sensor_gap_us = (pkt_ts_ns - prev) / 1000; - uint64_t wall_gap_us = std::chrono::duration_cast( now_wall - last_wall).count(); - if (sensor_gap_us > 15000) { - fprintf(stderr, "[imu-gap] sensor_gap=%.1fms wall_gap=%.1fms pkt#%llu\n", sensor_gap_us / 1000.0, wall_gap_us / 1000.0, static_cast(pkt_count)); - } - } - last_wall = now_wall; - } - double ts = static_cast(pkt_ts_ns) / 1e9; auto* imu_pts = reinterpret_cast(data->data); uint16_t dot_num = data->dot_num; From c25585bf950591d882bbc64b75e196c7f886b100 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 23 Jun 2026 17:43:36 -0700 Subject: [PATCH 54/71] Remove trivial things --- dimos/mapping/ray_tracing/rust/src/main.rs | 16 +++-- .../ray_tracing/rust/src/voxel_ray_tracer.rs | 35 ++++++----- dimos/mapping/ray_tracing/test_transformer.py | 15 ----- .../basic_path_follower/test_module.py | 4 -- .../nav_3d/mls_planner/rust/src/adjacency.rs | 35 +++++------ .../nav_3d/mls_planner/rust/src/dijkstra.rs | 12 ++-- .../nav_3d/mls_planner/rust/src/edges.rs | 32 +++++----- .../nav_3d/mls_planner/rust/src/main.rs | 47 +++++---------- .../mls_planner/rust/src/mls_planner.rs | 33 ++++------- .../nav_3d/mls_planner/rust/src/nodes.rs | 58 +++---------------- .../nav_3d/mls_planner/rust/src/planner.rs | 29 ++++------ .../nav_3d/mls_planner/rust/src/python.rs | 5 +- .../nav_3d/mls_planner/rust/src/surfaces.rs | 26 +++------ 13 files changed, 112 insertions(+), 235 deletions(-) diff --git a/dimos/mapping/ray_tracing/rust/src/main.rs b/dimos/mapping/ray_tracing/rust/src/main.rs index 60bae5ad73..fd15858f2e 100644 --- a/dimos/mapping/ray_tracing/rust/src/main.rs +++ b/dimos/mapping/ray_tracing/rust/src/main.rs @@ -28,8 +28,8 @@ struct RayTracingVoxelMap { #[output(encode = PointCloud2::encode)] local_map: Output, - // Cylinder bounds of the local map. Position holds the center, orientation - // holds radius, z_min, z_max. Stamped like local_map so consumers pair them. + // 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, @@ -57,7 +57,7 @@ impl RayTracingVoxelMap { async fn on_lidar(&mut self, msg: PointCloud2) { let Some((translation, rotation)) = self.last_pose else { - // Need at least one odometry sample before we can raycast. + // Need an odometry sample before we can raycast. return; }; let origin = (translation.x, translation.y, translation.z); @@ -79,7 +79,7 @@ impl RayTracingVoxelMap { return; } - // Register sensor-frame clouds into the world by the odom pose. + // Transform sensor-frame points into the world by the odom pose. let rot = rotation.to_rotation_matrix(); let points: Vec<(f32, f32, f32)> = points .iter() @@ -89,13 +89,12 @@ impl RayTracingVoxelMap { }) .collect(); - // The integrated points are world-frame either way. let out_frame_id = "world"; let live = update_map(&mut self.map, origin, &points, &self.config); - // The batch only feeds the local region bounds, so don't let it grow - // when the local map is disabled. + // The batch only feeds the local region bounds, so skip it when the local + // map is disabled. if self.config.emit_every > 0 { self.batch_points.extend_from_slice(&points); self.batch_origins.push(origin); @@ -344,8 +343,7 @@ fn build_local_cloud( make_cloud(data, n, frame_id, stamp) } -/// Build the global and local clouds in one pass over the map, so a frame that -/// emits both does not scan the voxel map twice. +/// Build the global and local clouds in one pass over the map. fn build_global_and_local( map: &VoxelMap, live: &AHashSet, diff --git a/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs b/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs index da28006f82..05269c313a 100644 --- a/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs +++ b/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs @@ -27,12 +27,12 @@ pub struct Config { pub min_health: i32, #[validate(range(min = 1))] pub max_health: i32, - /// Don't clear a miss when abs of ray dot normal is below this, clear it when above. - /// Higher clears only on direct hits, lower clears on slight grazes too. + /// Spare a miss when abs of ray dot normal is below this. Higher clears only + /// on direct hits, lower clears on slight grazes too. #[validate(range(min = 0.0, max = 1.0))] pub graze_cos: f32, /// Only spare a voxel whose neighborhood was hit within this many frames. - /// A stale voxel can be cleared, even if it's a grazing hit. Large disables it. + /// Large disables it. pub recency_window: u32, /// Publish the accumulated local map and region bounds every Nth frame. Zero disables them. #[validate(range(min = 0))] @@ -40,8 +40,8 @@ pub struct Config { /// Publish the global map every Nth frame. Zero disables it. #[validate(range(min = 0))] pub global_emit_every: u32, - /// Size the local region to this percentile of batch point distances, - /// so a stray far hit cannot inflate the region the planner recomputes. + /// Size the local region to this percentile of batch point distances, so a + /// stray far hit cannot inflate it. #[validate(range(min = 0.0, max = 100.0))] pub region_percentile: f32, } @@ -109,7 +109,7 @@ const NEIGHBORHOOD_CAP: usize = (2 * NORMAL_NEIGHBOR_RADIUS as usize + 1).pow(3) const NORMAL_REWEIGHT_ITERS: u32 = 3; /// Neighbor weight falloff with plane distance, as a fraction of voxel size. const NORMAL_PLANE_SIGMA_FRAC: f32 = 0.5; -/// Fraction of points that must be kept after the IRLS to count as a real plane +/// Fraction of points that must survive the IRLS to count as a real plane. const NORMAL_MIN_SUPPORT: f32 = 0.5; /// Occupancy health, accumulated point moments about the voxel center, and the @@ -273,7 +273,7 @@ fn pooled_normal_and_recency( *w = (-(dist * dist) / two_sig2).exp(); } } - // get rid of planes if we had to discard too many points to get a plane + // Reject the plane if too many points had to be discarded to fit it. let kept: f32 = nbs.iter().zip(&weights).map(|(nb, &w)| w * nb.n).sum(); if kept < NORMAL_MIN_SUPPORT * n_raw as f32 { return (None, recency); @@ -317,7 +317,7 @@ fn refresh_voxels( } /// Spare a clearing miss only when a grazing ray skims a recently hit planar -/// surface. Stale or voxels with no normal are left to the health checks. +/// surface. Stale voxels or those with no normal are left to the health checks. fn should_spare( c: &Voxel, ray_unit: Vector3, @@ -410,8 +410,8 @@ pub fn iter_global_points( }) } -/// Healthy voxel centers paired with their surface normal. -/// If no normal, it's just the null vector +/// Healthy voxel centers paired with their surface normal, the zero vector where +/// there is no plane. pub fn iter_global_normals( map: &VoxelMap, voxel_size: f32, @@ -498,7 +498,6 @@ pub fn update_map( a }); - // add new hits for v in &hits { let c = map.voxels.entry(*v).or_insert_with(|| Voxel { health: cfg.min_health, @@ -523,7 +522,6 @@ pub fn update_map( } } - // refresh cached normals and recency wherever the neighborhood changed refresh_voxels(map, &hits, &removed, cfg.voxel_size); hits @@ -538,8 +536,9 @@ fn world_to_voxel(x: f32, y: f32, z: f32, inv: f32) -> VoxelKey { ) } -/// Amanatides & Woo 3d DDA. Records voxels on ray in between the end of the shadow region -/// and origin if it is in the map. Voxels within grace region of the endpoint are spared from being marked as misses. +/// Amanatides and Woo 3d DDA. Records in-map voxels along the ray between the +/// origin and the end of the shadow region. Voxels within the grace region of +/// the endpoint are spared from being marked as misses. #[allow(clippy::too_many_arguments)] fn find_misses_along_ray( misses: &mut AHashSet, @@ -655,12 +654,12 @@ fn find_misses_along_ray( let dist_sq = ddx * ddx + ddy * ddy + ddz * ddz; if past_endpoint { - // continue past the endpoint and in to the shadow realm + // Past the endpoint, keep going until we leave the shadow region. if dist_sq > shadow_sq { return; } } else if dist_sq < grace_sq { - // too close to the endpoint to safely mark as miss because we might be clipping other voxel's rays + // Too close to the endpoint to safely mark a miss, we might be clipping another voxel's ray. continue; } @@ -785,7 +784,7 @@ mod tests { let mut map = VoxelMap::default(); map.set((3, 0, 0), 1); update_map(&mut map, (0.0, 0.0, 0.0), &[(5.5, 0.5, 0.5)], &cfg); - // make sure the initial point got cleared by the new update + // The voxel on the ray should be cleared. assert!(!map.voxels.contains_key(&(3, 0, 0))); assert_eq!(map.health((5, 0, 0)), Some(1)); } @@ -806,7 +805,7 @@ mod tests { let mut map = VoxelMap::default(); map.set((6, 0, 0), 1); update_map(&mut map, (0.0, 0.0, 0.0), &[(5.5, 0.5, 0.5)], &cfg); - // point within the shadow is no longer included, new point is included + // The voxel inside the shadow region should be cleared. assert!(!map.voxels.contains_key(&(6, 0, 0))); assert_eq!(map.health((5, 0, 0)), Some(1)); } diff --git a/dimos/mapping/ray_tracing/test_transformer.py b/dimos/mapping/ray_tracing/test_transformer.py index 2d6c91a83b..c633c519e9 100644 --- a/dimos/mapping/ray_tracing/test_transformer.py +++ b/dimos/mapping/ray_tracing/test_transformer.py @@ -50,21 +50,6 @@ def test_emit_every_n_yields_on_cadence_and_flushes_remainder() -> None: assert [r.tags["frame_count"] for r in results] == [3, 6, 7] -def test_negative_emit_every_is_rejected() -> None: - with pytest.raises(ValueError, match="emit_every"): - RayTraceMap(emit_every=-1) - - -def test_pose_propagates_to_emitted_obs() -> None: - pose = (1.5, -2.0, 0.5) - obs = _obs(_cube(), ts=1.0, pose=pose) - - [emitted] = list(RayTraceMap()(iter([obs]))) - - assert emitted.pose_tuple is not None - assert emitted.pose_tuple[:3] == pose - - def test_poseless_obs_are_skipped() -> None: points = _cube() poseless = Observation(id=1, ts=0.0, pose=None, _data=PointCloud2.from_numpy(points)) diff --git a/dimos/navigation/basic_path_follower/test_module.py b/dimos/navigation/basic_path_follower/test_module.py index a5ad8d5782..3ee85b7525 100644 --- a/dimos/navigation/basic_path_follower/test_module.py +++ b/dimos/navigation/basic_path_follower/test_module.py @@ -25,7 +25,3 @@ def test_lookahead_scales_in_linear_region(): def test_lookahead_clamped_at_ceiling(): assert lookahead_distance(2.0, 1.5, 0.4, 1.5) == 1.5 - - -def test_lookahead_monotonic_in_speed(): - assert lookahead_distance(0.8, 1.5, 0.4, 1.5) > lookahead_distance(0.45, 1.5, 0.4, 1.5) diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/adjacency.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/adjacency.rs index baf81839cc..8653e072a7 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/adjacency.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/adjacency.rs @@ -3,8 +3,7 @@ //! Surface cells indexed by dense CellId. //! -//! Uses a "slot map" to store cells. When inserting, either expand the map -//! or reuse a freed location marked with a tombstone. +//! A slot map: inserts reuse freed slots marked with a tombstone, or grow. use ahash::{AHashMap, AHashSet}; use rayon::prelude::*; @@ -13,15 +12,15 @@ use crate::voxel::VoxelKey; pub type SurfaceLookup = AHashMap<(i32, i32), Vec>; -/// Index of surface voxel +/// Index of a surface voxel. pub type CellId = u32; pub const NO_CELL: CellId = u32::MAX; -/// Represent a deleted cell that can be reincarnated on an insertion. +/// Marks a freed slot reusable by the next insert. const TOMBSTONE: VoxelKey = (i32::MIN, i32::MIN, i32::MIN); const NEIGHBORS_4: [(i32, i32); 4] = [(-1, 0), (1, 0), (0, -1), (0, 1)]; -/// Vertical extent of a `dz`-cell change in meters, the step-penalty input. +/// Vertical extent of a dz-cell change in meters, the step-penalty input. #[inline] pub fn rise(dz: i32, voxel_size: f32) -> f32 { dz.unsigned_abs() as f32 * voxel_size @@ -32,7 +31,7 @@ pub struct Edge { pub dest: CellId, /// Geometric cost, set at build time and never mutated. pub base_cost: f32, - /// Vertical change of the edge in meters, for the step penalty. + /// Vertical change of the edge in meters. pub rise: f32, /// base_cost scaled by the wall penalty plus the step penalty. pub cost: f32, @@ -56,7 +55,7 @@ impl SurfaceCells { self.coord.len() } - /// Clear all vecs but keeps space allocated. + /// Clear all vecs but keep allocated capacity. pub fn clear(&mut self) { self.coord.clear(); self.by_coord.clear(); @@ -71,9 +70,7 @@ impl SurfaceCells { self.coord[id as usize] != TOMBSTONE } - /// Get or insert a new cell. - /// - /// Only expand the list if there are no available dead cells. + /// Get or insert a cell, reusing a freed slot before growing. pub fn insert(&mut self, k: VoxelKey) -> CellId { debug_assert_ne!(k, TOMBSTONE, "voxel coord collides with tombstone sentinel"); if let Some(&id) = self.by_coord.get(&k) { @@ -92,9 +89,7 @@ impl SurfaceCells { id } - /// Remove a cell. - /// - /// Mark the cell as available with a tombstone and remove the output edges. + /// Remove a cell, tombstoning its slot and dropping its edges. #[allow(dead_code)] pub fn remove(&mut self, k: VoxelKey) -> Option { let id = self.by_coord.remove(&k)?; @@ -108,13 +103,13 @@ impl SurfaceCells { Some(id) } - /// XYZ coord to cell ID. + /// Coord to cell ID. #[inline] pub fn id(&self, k: VoxelKey) -> Option { self.by_coord.get(&k).copied() } - /// Cell ID to XYZ coord. + /// Cell ID to coord. #[inline] pub fn coord(&self, id: CellId) -> VoxelKey { self.coord[id as usize] @@ -189,8 +184,7 @@ pub fn build_surface_lookup(cells: &[VoxelKey], out: &mut SurfaceLookup) { } } -/// Populate cells with surface adjacency from the lookup. Deletes any -/// existing contents +/// Build surface adjacency from the lookup, replacing any existing contents. pub fn build_surface_cells( cells: &mut SurfaceCells, surface_lookup: &SurfaceLookup, @@ -241,9 +235,10 @@ pub fn build_surface_cells( }); } -/// Recompute outgoing edges for the seed cells and their surface neighbors, -/// matching build_surface_cells. Call after an incremental insert or remove so -/// the affected region matches a full rebuild. +/// Recompute outgoing edges for the seed cells and their surface neighbors. +/// +/// Call after an incremental insert or remove so the affected region matches a +/// full rebuild. pub fn rebuild_edges_around( cells: &mut SurfaceCells, surface_lookup: &SurfaceLookup, diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/dijkstra.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/dijkstra.rs index 25f7250528..9830fc44a3 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/dijkstra.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/dijkstra.rs @@ -21,7 +21,7 @@ pub struct DijkstraState { } impl DijkstraState { - /// Reset all vecs to the specified capacity. + /// Reset all vecs to n slots. pub fn reset(&mut self, n: usize) { self.dist.clear(); self.dist.resize(n, f32::INFINITY); @@ -32,8 +32,8 @@ impl DijkstraState { self.heap.clear(); } - /// Grow the vecs to hold `n` slots without disturbing existing labels. - /// New slots default to unreached. + /// Grow the vecs to n slots without disturbing existing labels. New slots + /// default to unreached. fn ensure_capacity(&mut self, n: usize) { if self.dist.len() < n { self.dist.resize(n, f32::INFINITY); @@ -62,7 +62,7 @@ impl Weight { } } -/// Multi-source dijkstra labeling each cell with its nearest source and path. +/// Multi-source Dijkstra labeling each cell with its nearest source and path. pub fn dijkstra( cells: &SurfaceCells, sources: &[CellId], @@ -101,7 +101,7 @@ pub fn dijkstra( } } -/// Multi-source dijkstra that re-labels only cells in `window`, seeded from +/// Multi-source Dijkstra that re-labels only cells in the window, seeded from /// in-window sources and the cached frontier just outside it. Correct while the /// window margin exceeds the reach of the change. pub fn dijkstra_region( @@ -200,7 +200,7 @@ impl PartialOrd for Scored { } impl Ord for Scored { fn cmp(&self, other: &Self) -> Ordering { - // Order on score, and use cell id for tie-breaker for repeatability + // Tie-break on cell id for repeatable ordering. other.0.total_cmp(&self.0).then(self.1.cmp(&other.1)) } } diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs index 22ccb8dbd2..57859cde68 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs @@ -3,10 +3,9 @@ //! Node-graph edge construction. //! -//! Build edges by running multi-source Dijkstra from all the start nodes. -//! This labels the surface with each cells closest source, also known as -//! the Voronoi region. We use the boundaries of these regions to build the -//! edges between start nodes. +//! Multi-source Dijkstra from the start nodes labels each cell with its closest +//! source, partitioning the surface into Voronoi regions. Edges between nodes +//! come from the boundaries between those regions. use ahash::{AHashMap, AHashSet}; use rayon::prelude::*; @@ -16,12 +15,12 @@ use crate::dijkstra::{dijkstra, dijkstra_region, walk_preds, DijkstraState, Weig use crate::nodes::NodeData; use crate::voxel::VoxelKey; -/// A node is identified by the CellId it sits on, stable across incremental +/// A node is identified by the CellId it sits on. Stable across incremental /// updates so cached edges and the Voronoi forest survive a regional rebuild. pub type NodeId = CellId; pub const NO_NODE: NodeId = NO_CELL; -/// Index into planner graph node edges +/// Index into the planner graph node edges. pub type NodeEdgeIdx = u32; #[derive(Clone, Copy, Debug)] @@ -45,8 +44,7 @@ pub struct PlannerGraph { /// Each cell's nearest node and the predecessor back toward it. The planner /// walks these to expand a node-to-node edge into its cell path. pub cell_state: DijkstraState, - /// Each cell's distance to the nearest wall, recomputed only in the changed - /// window and reused elsewhere. + /// Each cell's distance to the nearest wall. pub wall_state: DijkstraState, } @@ -56,10 +54,8 @@ impl PlannerGraph { } } -/// Assemble the cheapest paths between neighboring source nodes. -/// -/// Runs multi-source dijkstra from the sources, then adds the cheapest edges -/// between Voronoi region boundaries. +/// Assemble the cheapest edges between neighboring source nodes from their +/// Voronoi region boundaries. pub fn build_node_edges( cells: &SurfaceCells, nodes: &[NodeData], @@ -98,9 +94,9 @@ fn rebuild_node_adj(edges: &[NodeEdge], out_adj: &mut AHashMap = Arc>>; /// A map input handed from the handle loop to the worker. Only the newest is @@ -67,24 +66,21 @@ struct MlsPlanner { #[config] config: Config, - // These live only on the handle loop. We hold a local map and its bounds - // here until their stamps match, then hand the paired snapshot to the worker. + // Held on the handle loop until stamps match, then handed off paired. pending_local: Option, pending_bounds: Option, - // Shared with the worker task. The handle loop only writes these and wakes - // the worker, so it never blocks on the heavy map processing. + // Written by the handle loop, read by the worker, so the loop never blocks + // on map processing. pending: Shared, latest_start: Shared, active_goal: Shared, wake: Arc, - // Handle to the background worker, aborted on teardown. worker: Option>, } impl MlsPlanner { - /// Spawn the background task that handles map updates and planning. async fn spawn_worker(&mut self) { let worker = Worker { pending: Arc::clone(&self.pending), @@ -100,7 +96,6 @@ impl MlsPlanner { self.worker = Some(tokio::spawn(worker.run())); } - /// Abort the background worker on teardown. async fn stop_worker(&mut self) { if let Some(handle) = self.worker.take() { handle.abort(); @@ -121,8 +116,7 @@ impl MlsPlanner { self.try_pair(); } - /// Hand a local map and its stamp-matching bounds to the worker once both - /// are in hand. + /// Hand off the local map and bounds once their stamps match. fn try_pair(&mut self) { if !stamps_paired(self.pending_bounds.as_ref(), self.pending_local.as_ref()) { return; @@ -132,14 +126,13 @@ impl MlsPlanner { self.hand_off(MapUpdate::Region { cloud, bounds }); } - /// Store the newest map input and wake the worker. fn hand_off(&self, update: MapUpdate) { *self.pending.lock().expect("pending mutex") = Some(update); self.wake.notify_one(); } - /// Record the latest start pose. The worker reads it when it replans. No - /// wake here, so odometry never drives replanning. + /// Record the latest start pose. No wake here, so odometry never drives + /// replanning. async fn on_start_pose(&mut self, msg: PoseStamped) { let p = &msg.pose.position; *self.latest_start.lock().expect("start mutex") = @@ -153,8 +146,7 @@ impl MlsPlanner { } } -/// True when bounds and a local cloud are both present with matching stamps, -/// so they form one paired region update. +/// True when bounds and a local cloud are both present with matching stamps. fn stamps_paired(bounds: Option<&PoseStamped>, cloud: Option<&PointCloud2>) -> bool { match (bounds, cloud) { (Some(b), Some(c)) => same_stamp(&b.header.stamp, &c.header.stamp), @@ -169,8 +161,8 @@ fn goal_position(p: &Point) -> Option { (goal.0.is_finite() && goal.1.is_finite() && goal.2.is_finite()).then_some(goal) } -/// Owns the planner graph and does every map mutation, graph publish, and -/// replan off the handle loop. Woken by the handlers via `wake`. +/// Owns the planner graph and does map mutation, publishing, and replanning +/// off the handle loop. Woken by the handlers. struct Worker { pending: Shared, latest_start: Shared, @@ -190,7 +182,6 @@ impl Worker { let mut last_viz_at: Option = None; loop { self.wake.notified().await; - // Take the newest map input, dropping any intermediates. let update = self.pending.lock().expect("pending mutex").take(); if let Some(update) = update { self.apply_update(&mut planner, update, &mut last_viz_at) @@ -200,9 +191,8 @@ impl Worker { } } - /// Update the graph every cycle so planning sees fresh surfaces, then publish - /// the viz clouds rate-capped to viz_publish_hz, since building them is costly - /// and nothing on the planning path reads them. + /// Update the graph every cycle, but rate-cap viz publishing to + /// viz_publish_hz since building those clouds is costly and unread by planning. async fn apply_update( &self, planner: &mut Planner, @@ -300,8 +290,7 @@ impl Worker { (surface, node_cloud, edges) } - /// Replan from the latest start to the active goal. This gates and does the - /// IO. The planning itself lives in Planner::plan. + /// Gate and publish a replan. The planning itself lives in Planner::plan. async fn maybe_replan(&self, planner: &mut Planner, last_path_at: &mut Option) { let Some(start) = *self.latest_start.lock().expect("start mutex") else { return; @@ -322,7 +311,7 @@ impl Worker { let waypoints = tokio::task::block_in_place(|| planner.plan_or_truncate(start, goal, &self.config)); if waypoints.is_empty() { - // No full path and nothing safe ahead on the cached path: stop. + // No full path and nothing safe ahead on the cached path, so stop. publish_path(&self.path, &empty_path(&self.config.world_frame, now())).await; return; } @@ -559,14 +548,6 @@ mod tests { assert!(!is_at_goal((0.0, 0.0, 0.0), (0.2, 0.0, 0.0), 0.1)); } - #[test] - fn same_stamp_compares_sec_and_nsec() { - let a = Time { sec: 5, nsec: 7 }; - assert!(same_stamp(&a, &Time { sec: 5, nsec: 7 })); - assert!(!same_stamp(&a, &Time { sec: 5, nsec: 8 })); - assert!(!same_stamp(&a, &Time { sec: 6, nsec: 7 })); - } - fn stamped(stamp: Time) -> Header { Header { stamp, diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs index 9251ab847c..26cc80a0aa 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs @@ -69,7 +69,7 @@ fn validate_wall_buffer(config: &Config) -> Result<(), ValidationError> { } impl Config { - /// Compute number of dilations and erosions to do based on closing radius + /// Number of dilation and erosion passes for the closing radius. pub fn closing_passes(&self) -> u32 { (self.surface_closing_radius / self.voxel_size).ceil() as u32 } @@ -112,8 +112,8 @@ pub struct Planner { graph: PlannerGraph, voxel_map: AHashSet, by_col: ColumnIz, - // Last successful path with the goal it served, for safe truncation when a - // later replan finds no full path. + // Last successful path and its goal, for safe truncation when a later + // replan finds no full path. last_path: Option<((f32, f32, f32), Vec)>, } @@ -140,8 +140,8 @@ impl Planner { self.rebuild_graph(config); } - /// Update planner artifacts within a local region instead of recomputing - /// the entire planner on the entire map. + /// Update planner artifacts within a local region instead of rebuilding + /// from the whole map. pub fn update_region( &mut self, local_points: &[(f32, f32, f32)], @@ -159,8 +159,7 @@ impl Planner { return; }; - // A changed voxel column shifts surfaces only within pad of it, so the - // write-back box is the changed-column bbox grown by pad. + // A changed column shifts surfaces only within pad of it. let write = (bx0 - pad, bx1 + pad, by0 - pad, by1 + pad); let new_cells = extract_surfaces_region(&self.by_col, clearance, config.closing_passes(), write); @@ -169,8 +168,8 @@ impl Planner { self.rebuild_region_graph(added, removed, config); } - /// Patch cells for the changed surface, then re-place nodes and edges over - /// the change window. A no-op when no surface cell changed. + /// Patch changed cells, then re-place nodes and edges over the change + /// window. A no-op when no surface cell changed. fn rebuild_region_graph( &mut self, added: Vec, @@ -332,7 +331,7 @@ impl Planner { /// Live cells within the changed-cell bbox grown by the node-graph margin, /// which covers the reach of any node, edge, or Voronoi change. fn node_window(&self, changed: &[VoxelKey], config: &Config) -> AHashSet { - // A few extra cells beyond the morphology, wall-buffer, and spacing reach. + // Slack beyond the morphology, wall-buffer, and spacing reach. const SLACK_CELLS: i32 = 2; let voxel_size = config.voxel_size; let pad = (2 * config.closing_passes()) as i32; @@ -740,8 +739,7 @@ mod region_tests { } /// Re-observing the same geometry must change nothing: no voxel, surface, - /// cell, node, or edge moves. This is the anti-jitter guarantee, far nodes - /// stay put when their region is re-seen, matching a full rebuild. + /// cell, node, or edge moves. This is the anti-jitter guarantee. #[test] fn region_reobserve_leaves_graph_bit_identical() { let cfg = test_config(); @@ -1078,15 +1076,4 @@ mod region_tests { let last = *wp.last().expect("path has waypoints"); assert!((last.0 - goal.0).abs() < 1e-3 && (last.1 - goal.1).abs() < 1e-3); } - - #[test] - fn wall_buffer_weight_requires_buffer() { - let mut cfg = test_config(); - cfg.wall_buffer_weight = 1.0; - cfg.wall_buffer_m = 0.0; - assert!(validate_wall_buffer(&cfg).is_err()); - - cfg.wall_buffer_m = 0.3; - assert!(validate_wall_buffer(&cfg).is_ok()); - } } diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/nodes.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/nodes.rs index 1854a7986c..c326935f40 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/nodes.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/nodes.rs @@ -20,10 +20,7 @@ pub struct NodeData { pub pos: (f32, f32, f32), } -/// Distribute nodes on the surfaces. -/// -/// Runs multi source dijkstra using edges as sources, then distribute nodes -/// using a grid based NMS. +/// Place graph nodes across the surface, spaced out and biased away from walls. #[allow(clippy::too_many_arguments)] pub fn place_nodes( cells: &mut SurfaceCells, @@ -74,8 +71,7 @@ pub fn place_nodes( ); } -/// Sort candidates by descending wall distance, thin them with NMS against the -/// seed nodes, and append the survivors as nodes. +/// Thin candidates with NMS, clearest-first, against the seed nodes. fn place_from_candidates( cells: &SurfaceCells, mut candidates: Vec, @@ -172,9 +168,7 @@ fn collect_wall_adjacent_in_window( } } -/// A cell is wall-adjacent when it is missing at least one of its 4 xy-direction -/// neighbors. Membership is tracked with a 4-bit mask to avoid per-cell -/// allocation on the 349k-cell case. +/// A cell is wall-adjacent when it lacks at least one of its 4 xy neighbors. fn is_wall_adjacent(cells: &SurfaceCells, id: CellId) -> bool { let (cx, cy, _) = cells.coord(id); let mut mask: u8 = 0; @@ -238,10 +232,9 @@ fn collect_wall_adjacent_cells(cells: &SurfaceCells, out: &mut Vec) { } } -/// Space out nodes based on minimum distance. -/// -/// The seed nodes suppress nearby candidates without being emitted, keeping a -/// regional re-placement consistent with cached nodes outside the window. +/// Keep nodes at least node_spacing_m apart. Seeds suppress nearby candidates +/// without being emitted, so regional re-placement respects cached nodes +/// outside the window. fn nms_grid( cells: &SurfaceCells, candidates_sorted: &[CellId], @@ -343,9 +336,8 @@ fn scale_edges( } } -/// Lateral wall multiplier at wall distance d. Infinite inside the clearance, -/// then 1 + weight at the clearance edge decaying convexly to 1 at -/// clearance_m + buffer_m, and 1 beyond. +/// Lateral wall multiplier: infinite inside clearance, ramping convexly from +/// 1 + weight at the clearance edge down to 1 at clearance_m + buffer_m. #[inline] pub(crate) fn penalty_of(d: f32, clearance_m: f32, buffer_m: f32, weight: f32) -> f32 { if d < clearance_m { @@ -514,23 +506,6 @@ mod tests { } } - #[test] - fn sloped_patch_places_interior_nodes() { - let mut cells_in = Vec::new(); - for ix in 0..10 { - for iy in 0..10 { - cells_in.push((ix, iy, ix)); - } - } - let mut sc = build_cells(&cells_in, 2); - let mut state = DijkstraState::default(); - let mut nodes = Vec::new(); - place_nodes( - &mut sc, VOXEL, 1.0, 0.0, 0.3, 1.0, 0.0, &mut state, &mut nodes, - ); - assert!(!nodes.is_empty()); - } - #[test] fn each_disconnected_component_gets_a_node() { // Two 1-wide strips far apart: every cell is wall-adjacent so none @@ -635,21 +610,4 @@ mod tests { "step penalty must add weight * rise" ); } - - #[test] - fn wall_cells_scale_outbound_cost() { - let cells_in: Vec = (0..10).map(|ix| (ix, 0, 0)).collect(); - let mut sc = build_cells(&cells_in, 2); - let mut state = DijkstraState::default(); - let mut nodes = Vec::new(); - place_nodes( - &mut sc, VOXEL, 1.0, 0.0, 0.3, 1.0, 0.0, &mut state, &mut nodes, - ); - let id0 = sc.id((0, 0, 0)).unwrap(); - let outbound = sc.neighbors(id0); - assert!(!outbound.is_empty()); - for edge in outbound { - assert!(edge.cost >= 1.5 * VOXEL - 1e-5); - } - } } diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs index 3475ed0db3..4861fe01a9 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs @@ -22,9 +22,7 @@ const SNAP_SEARCH_RADIUS_M: f32 = 1.5; /// Max snap candidates tried when connecting the start. const MAX_SNAP_ATTEMPTS: usize = 64; -/// Standoff held from a blockage on a truncated path. The robot follows the -/// cached route to best effort and stops this far short of the last -/// traversable point, in meters. +/// On a blocked path, stop this far short of the last traversable point. const BEST_EFFORT_DISTANCE_M: f32 = 1.0; /// World-frame waypoints paired with the string-pulled cell path that produced @@ -194,10 +192,9 @@ pub fn plan( Some((waypoints, path_cells)) } -/// Re-validate a cached cell path against the current surface, returning the -/// route ahead of the robot up to the last traversable surface cell before a -/// blockage. It holds a BEST_EFFORT_DISTANCE_M standoff short of the blockage. -/// Empty when nothing ahead is safe, which the follower reads as a stop. +/// Re-validate a cached path against the current surface. Returns the route +/// ahead up to a standoff short of the first blockage, or empty when nothing +/// ahead is safe, which the follower reads as a stop. pub fn truncate_to_safe( plg: &PlannerGraph, cached: &[VoxelKey], @@ -216,14 +213,12 @@ pub fn truncate_to_safe( voxel_size, }; - // The cached path runs start -> goal, but its head is the stale original - // start. Resume from where the robot sits on it so the follower is pulled - // forward toward the goal, never back to that start. + // The cached path's head is the stale original start. Resume from where the + // robot sits on it so the follower is pulled forward, never back to it. let resume = resume_segment(cached, start_pose, voxel_size); - // Walk each chord ahead at surface resolution. On a blockage, keep the - // chord up to its last traversable cell so the path ends right at the - // obstacle, not back at the previous smoothing anchor. + // Walk each chord ahead at surface resolution. On a blockage, keep up to the + // last traversable cell so the path ends at the obstacle, not the prior anchor. let mut waypoints = vec![start_pose]; let mut blocked = false; for j in resume..cached.len() - 1 { @@ -252,17 +247,15 @@ pub fn truncate_to_safe( return Vec::new(); } - // A blockage ahead means the route runs into an obstacle, so hold the - // standoff. A clean run to the goal end has nothing to stand off from. + // Hold the standoff only when blocked; a clean run to the goal has nothing + // to stand off from. if blocked { return back_off_tail(&waypoints, BEST_EFFORT_DISTANCE_M); } waypoints } -/// Drop `distance` meters off the goal end of the path, measured in the ground -/// plane. Empty when trimming consumes the whole path, which is the robot -/// already sitting inside the standoff. +/// Trim `distance` off the goal end, measured in the ground plane. fn back_off_tail(waypoints: &[(f32, f32, f32)], distance: f32) -> Vec<(f32, f32, f32)> { let mut remaining = distance; for i in (1..waypoints.len()).rev() { diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs index d320eff253..ff360d08f5 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs @@ -80,10 +80,9 @@ impl MLSPlanner { wall_buffer_weight, step_threshold_m, step_penalty_weight, - // Only the binary's replan loop reads goal_tolerance. This - // in-process binding plans on demand and never consults it. + // Unused here. Only the binary's replan loop reads goal_tolerance. goal_tolerance: 1.0, - // Only the binary's worker publishes viz artifacts. Unused here. + // Unused here. Only the binary's worker publishes viz artifacts. viz_publish_hz: 1.0, }; config diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/surfaces.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/surfaces.rs index 19e72de8c2..c2e6f86dd5 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/surfaces.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/surfaces.rs @@ -1,9 +1,9 @@ // Copyright 2026 Dimensional Inc. // SPDX-License-Identifier: Apache-2.0 -//! Surface extraction: from a voxel map, mark cells with robot-height -//! clearance above as standable, then morphologically close per-z-level -//! holes without letting closing bridge across walls. +//! Surface extraction: mark cells with robot-height clearance above as +//! standable, then morphologically close per-z-level holes without bridging +//! across walls. use ahash::{AHashMap, AHashSet}; use image::{GrayImage, Luma}; @@ -18,8 +18,8 @@ const OFF: Luma = Luma([0]); pub type ColumnIz = AHashMap<(i32, i32), Vec>; -/// A cell is standable if it has at least the robot's height of clear -/// space above it. +/// A cell is standable if it has at least the robot's height of clear space +/// above it. fn is_standable(ix: i32, iy: i32, iz: i32, by_col: &ColumnIz, clearance_cells: i32) -> bool { let Some(zs) = by_col.get(&(ix, iy)) else { return true; @@ -108,10 +108,9 @@ pub fn remove_from_by_col(by_col: &mut ColumnIz, (ix, iy, iz): VoxelKey) { } } -/// Re-extract surface cells whose columns fall in the inclusive write box. -/// Reads by_col over the box plus the morphology halo so closing at the -/// boundary matches a full rebuild, then filters back to the box. by_col must -/// already be current. +/// Re-extract surface cells in the inclusive write box. Reads a morphology +/// halo around the box so boundary closing matches a full rebuild, then +/// filters back to the box. by_col must already be current. pub fn extract_surfaces_region( by_col: &ColumnIz, clearance_cells: i32, @@ -148,8 +147,7 @@ pub fn extract_surfaces_region( .collect() } -/// Dilation and erosion on all xy slices of the extracted surfaces -/// to fill in small holes. +/// Dilate then erode every xy slice to fill small holes. fn close_surface_holes( standable: Vec, by_col: &ColumnIz, @@ -248,12 +246,6 @@ mod tests { assert!(run(&[], 5, 0).is_empty()); } - #[test] - fn single_cell_is_topmost_surface() { - let s = run(&[(0, 0, 0)], 5, 0); - assert_eq!(s, vec![(0, 0, 0)]); - } - #[test] fn stacked_cells_within_headroom_only_topmost_is_surface() { let cells: Vec = (0..5).map(|z| (0, 0, z)).collect(); From 5e2b87b48ec8f55c2270c86ce4baa476bf1799a4 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Wed, 24 Jun 2026 14:08:13 -0700 Subject: [PATCH 55/71] Refactor transform main --- .../nav_3d/mls_planner/utils/plan_rrd.py | 199 +++++++++--------- 1 file changed, 103 insertions(+), 96 deletions(-) diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index bc955ca998..b0cc2dec59 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -28,9 +28,7 @@ import typer from dimos.mapping.ray_tracing.transformer import RayTraceMap -from dimos.memory2.store.memory import MemoryStore from dimos.memory2.store.sqlite import SqliteStore -from dimos.memory2.stream import Stream from dimos.memory2.transform import FnTransformer from dimos.memory2.type.observation import Observation from dimos.msgs.nav_msgs.Odometry import Odometry @@ -40,9 +38,6 @@ TIMELINE = "ts" -TIMING_KEYS = ["update_ms", "plan_ms", "total_ms"] -SIZE_KEYS = ["voxels", "surface_cells", "nodes", "edges"] - # Distinct path colors for overlaid configurations, config 0 first. PATH_PALETTE = [ [0, 255, 0], @@ -75,23 +70,6 @@ def _parse_configs( return out -def _print_summary(streams: dict[str, dict[str, Stream[float]]]) -> None: - print("\nper-frame summary (mean / p50 / p95 / max):") - for kind, by_key in streams.items(): - for key, stream in by_key.items(): - values = [obs.data for obs in stream] - if not values: - continue - arr = np.asarray(values, dtype=np.float64) - mean, p50, p95, peak = ( - arr.mean(), - np.percentile(arr, 50), - np.percentile(arr, 95), - arr.max(), - ) - print(f" {kind}/{key:<14} {mean:9.2f} {p50:9.2f} {p95:9.2f} {peak:9.2f}") - - PairObs = Observation[tuple[Observation[PointCloud2], Observation[Odometry]]] @@ -177,6 +155,95 @@ def _log_shared( return surface, nodes, edges +def _init_recording(db_path: FsPath, out: FsPath | None, live: bool) -> None: + rr.init("plan_rrd", recording_id=db_path.stem) + if out is not None and live: + # Generous viewer memory so the gRPC sink never backpressures the writer. + rr.spawn(connect=False, memory_limit="16GB", server_memory_limit="16GB") + rr.set_sinks(rr.GrpcSink(), rr.FileSink(str(out))) + elif out is not None: + rr.save(str(out)) + else: + rr.spawn() + register_colormap_annotation("turbo") + + +def _build_planners( + configs: list[tuple[float, float, float]], + voxel_size: float, + robot_height: float, + surface_closing_radius: float, + node_spacing: float, + step_height: float, + step_penalty_weight: float, +) -> list[tuple[str, list[int], MLSPlanner]]: + planners: list[tuple[str, list[int], MLSPlanner]] = [] + for i, (clr, buf, wgt) in enumerate(configs): + planner = MLSPlanner( + voxel_size=voxel_size, + robot_height=robot_height, + surface_closing_radius=surface_closing_radius, + node_spacing_m=node_spacing, + wall_clearance_m=clr, + wall_buffer_m=buf, + wall_buffer_weight=wgt, + step_threshold_m=step_height, + step_penalty_weight=step_penalty_weight, + ) + color = PATH_PALETTE[i % len(PATH_PALETTE)] + label = f"cfg{i}_c{clr:g}_b{buf:g}_w{wgt:g}" + planners.append((label, color, planner)) + print(f"config {i}: clearance={clr} buffer={buf} weight={wgt} color={color} -> {label}") + return planners + + +def _process_frame( + ray_obs: Observation[PointCloud2], + planners: list[tuple[str, list[int], MLSPlanner]], + goal: tuple[float, float, float], + robot_height: float, + render_voxel: float, + clearance_clamp: float, +) -> dict[str, float]: + """Plan every config for one frame, log paths/map/metrics, return the ref timing.""" + assert ray_obs.pose_tuple is not None + bounds = ray_obs.tags["region_bounds"] + px, py, pz, *_ = ray_obs.pose_tuple + start = (float(px), float(py), float(pz) - robot_height) + ox, oy, radius, z_min, z_max = bounds + pts = ray_obs.data.points_f32() + rr.set_time(TIMELINE, timestamp=ray_obs.ts) + + ref_timing: dict[str, float] = {} + surface = nodes = edges = np.empty((0,), dtype=np.float32) + for j, (label, color, planner) in enumerate(planners): + t0 = perf_counter() + planner.update_region(pts, (ox, oy), radius, z_min, z_max) + t1 = perf_counter() + waypoints = planner.plan(start, goal) + t2 = perf_counter() + _log_path_wp(waypoints, f"world/paths/{label}", color) + if j == 0: + ref_timing = { + "update_ms": (t1 - t0) * 1000, + "plan_ms": (t2 - t1) * 1000, + "total_ms": (t2 - t0) * 1000, + } + surface, nodes, edges = _log_shared(start, planner, render_voxel, clearance_clamp) + + for key, value in ref_timing.items(): + rr.log(f"metrics/timing/{key}", rr.Scalars(value)) + sizes = { + "voxels": planners[0][2].voxel_count(), + "surface_cells": len(surface), + "nodes": len(nodes), + "edges": len(edges), + } + for key, value in sizes.items(): + rr.log(f"metrics/size/{key}", rr.Scalars(value)) + return ref_timing + + def main( dataset: str = typer.Argument(..., help="Dataset .db: bare name (cwd or data/) or path"), out: FsPath | None = typer.Option( @@ -243,17 +310,7 @@ def main( ), ) -> None: db_path = resolve_named_path(dataset, ".db") - - rr.init("plan_rrd", recording_id=db_path.stem) - if out is not None and live: - # Generous viewer memory so the gRPC sink never backpressures the writer. - rr.spawn(connect=False, memory_limit="16GB", server_memory_limit="16GB") - rr.set_sinks(rr.GrpcSink(), rr.FileSink(str(out))) - elif out is not None: - rr.save(str(out)) - else: - rr.spawn() - register_colormap_annotation("turbo") + _init_recording(db_path, out, live) store = SqliteStore(path=str(db_path)) with store: @@ -277,74 +334,26 @@ def main( ) configs = _parse_configs(config, wall_clearance, wall_buffer, wall_buffer_weight) - planners: list[tuple[str, list[int], MLSPlanner]] = [] - for i, (clr, buf, wgt) in enumerate(configs): - planner = MLSPlanner( - voxel_size=voxel_size, - robot_height=robot_height, - surface_closing_radius=surface_closing_radius, - node_spacing_m=node_spacing, - wall_clearance_m=clr, - wall_buffer_m=buf, - wall_buffer_weight=wgt, - step_threshold_m=step_height, - step_penalty_weight=step_penalty_weight, - ) - color = PATH_PALETTE[i % len(PATH_PALETTE)] - label = f"cfg{i}_c{clr:g}_b{buf:g}_w{wgt:g}" - planners.append((label, color, planner)) - print(f"config {i}: clearance={clr} buffer={buf} weight={wgt} color={color} -> {label}") + planners = _build_planners( + configs, + voxel_size, + robot_height, + surface_closing_radius, + node_spacing, + step_height, + step_penalty_weight, + ) rr.log("world/goal", rr.Points3D([goal], colors=[[255, 0, 0]], radii=0.1), static=True) - metrics = MemoryStore() - timing_streams = {k: metrics.stream(f"timing_{k}", float) for k in TIMING_KEYS} - size_streams = {k: metrics.stream(f"size_{k}", float) for k in SIZE_KEYS} - try: frame = 0 for ray_obs in ray_pipeline: if ray_obs.pose_tuple is None: continue - bounds = ray_obs.tags["region_bounds"] - px, py, pz, *_ = ray_obs.pose_tuple - start = (float(px), float(py), float(pz) - robot_height) - ox, oy, radius, z_min, z_max = bounds - pts = ray_obs.data.points_f32() - rr.set_time(TIMELINE, timestamp=ray_obs.ts) - - ref_timing: dict[str, float] = {} - surface = nodes = edges = np.empty((0,), dtype=np.float32) - for j, (label, color, planner) in enumerate(planners): - t0 = perf_counter() - planner.update_region(pts, (ox, oy), radius, z_min, z_max) - t1 = perf_counter() - waypoints = planner.plan(start, goal) - t2 = perf_counter() - _log_path_wp(waypoints, f"world/paths/{label}", color) - if j == 0: - ref_timing = { - "update_ms": (t1 - t0) * 1000, - "plan_ms": (t2 - t1) * 1000, - "total_ms": (t2 - t0) * 1000, - } - surface, nodes, edges = _log_shared( - start, planner, render_voxel, clearance_clamp - ) - - for key, value in ref_timing.items(): - timing_streams[key].append(float(value), ts=ray_obs.ts) - rr.log(f"metrics/timing/{key}", rr.Scalars(value)) - sizes = { - "voxels": planners[0][2].voxel_count(), - "surface_cells": len(surface), - "nodes": len(nodes), - "edges": len(edges), - } - for key, value in sizes.items(): - size_streams[key].append(float(value), ts=ray_obs.ts) - rr.log(f"metrics/size/{key}", rr.Scalars(value)) - + ref_timing = _process_frame( + ray_obs, planners, goal, robot_height, render_voxel, clearance_clamp + ) frame += 1 print( f"frame={frame} configs={len(planners)} " @@ -354,9 +363,7 @@ def main( flush=True, ) except KeyboardInterrupt: - print("\ninterrupted; reporting metrics for completed frames") - finally: - _print_summary({"timing": timing_streams, "size": size_streams}) + print("\ninterrupted") if out is not None: print(f"wrote {out}") From ff34a566ae1ffcfca7cf70fbdefbaca36c8e6afa Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Wed, 24 Jun 2026 15:02:41 -0700 Subject: [PATCH 56/71] Bug fix --- dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs index 57859cde68..cb5837b3fb 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/edges.rs @@ -128,7 +128,9 @@ pub fn build_node_edges_region( } let win_cells: Vec = window.iter().copied().collect(); - merge_min(&mut merged, boundary_edge_map(cells, state, &win_cells)); + let mut new_edges = boundary_edge_map(cells, state, &win_cells); + new_edges.retain(|_, e| live_node.contains(&e.a) && live_node.contains(&e.b)); + merge_min(&mut merged, new_edges); out_edges.clear(); out_edges.extend(merged.into_values()); From f890182f78db4f1d7e5f258413955c9b515a2958 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Fri, 26 Jun 2026 10:25:29 +0800 Subject: [PATCH 57/71] new hack --- .../pointlio/config/no_gravity_align.yaml | 62 ++++++++++ dimos/hardware/sensors/lidar/pointlio/hack.py | 110 ++++++++++++++++++ dimos/robot/assembly/mid360_realsense_30.urdf | 108 +++++++++++++++++ .../navigation/go2_mid360_normal.urdf | 22 ++++ .../navigation/go2_mid360_rotated.urdf | 24 ++++ .../navigation/mid360_realsense_30.urdf | 108 +++++++++++++++++ .../navigation/unitree_go2_nav_3d.py | 60 +++++++++- 7 files changed, 490 insertions(+), 4 deletions(-) create mode 100644 dimos/hardware/sensors/lidar/pointlio/config/no_gravity_align.yaml create mode 100644 dimos/hardware/sensors/lidar/pointlio/hack.py create mode 100644 dimos/robot/assembly/mid360_realsense_30.urdf create mode 100644 dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_normal.urdf create mode 100644 dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_rotated.urdf create mode 100644 dimos/robot/unitree/go2/blueprints/navigation/mid360_realsense_30.urdf diff --git a/dimos/hardware/sensors/lidar/pointlio/config/no_gravity_align.yaml b/dimos/hardware/sensors/lidar/pointlio/config/no_gravity_align.yaml new file mode 100644 index 0000000000..ed9c057083 --- /dev/null +++ b/dimos/hardware/sensors/lidar/pointlio/config/no_gravity_align.yaml @@ -0,0 +1,62 @@ +# default.yaml with gravity self-leveling turned off (mapping.gravity_align). +# Used when the mount tilt is corrected by a known static transform instead of +# letting Point-LIO level the odometry to gravity. Keep the rest in sync with +# default.yaml. +common: + con_frame: false + con_frame_num: 1 + cut_frame: false + cut_frame_time_interval: 0.1 + time_lag_imu_to_lidar: 0.0 + +preprocess: + lidar_type: 1 + scan_line: 4 + scan_rate: 10 + timestamp_unit: 3 # 3 = nanosecond + blind: 0.5 + point_filter_num: 3 + +mapping: + use_imu_as_input: false + prop_at_freq_of_imu: true + check_satu: true + init_map_size: 10 + space_down_sample: true + satu_acc: 3.0 + satu_gyro: 35.0 + acc_norm: 1.0 + plane_thr: 0.1 + filter_size_surf: 0.2 + filter_size_map: 0.5 + ivox_grid_resolution: 2.0 + ivox_nearby_type: 6 + cube_side_length: 1000.0 + det_range: 100.0 + fov_degree: 360.0 + imu_en: true + start_in_aggressive_motion: false + extrinsic_est_en: false + imu_time_inte: 0.005 + lidar_meas_cov: 0.01 + acc_cov_input: 0.1 + vel_cov: 20.0 + gyr_cov_input: 0.01 + gyr_cov_output: 1000.0 + acc_cov_output: 500.0 + b_gyr_cov: 0.0001 + b_acc_cov: 0.0001 + imu_meas_acc_cov: 0.01 + imu_meas_omg_cov: 0.01 + match_s: 81.0 + gravity_align: false + gravity: [0.0, 0.0, -9.810] + gravity_init: [0.0, 0.0, -9.810] + extrinsic_T: [-0.011, -0.02329, 0.04412] # Mid-360 IMU->lidar offset (m) + extrinsic_R: [1.0, 0.0, 0.0, + 0.0, 1.0, 0.0, + 0.0, 0.0, 1.0] + +odometry: + publish_odometry_without_downsample: false + odom_only: false diff --git a/dimos/hardware/sensors/lidar/pointlio/hack.py b/dimos/hardware/sensors/lidar/pointlio/hack.py new file mode 100644 index 0000000000..8ffb359354 --- /dev/null +++ b/dimos/hardware/sensors/lidar/pointlio/hack.py @@ -0,0 +1,110 @@ +# 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. + +"""Rewrite the lidar cloud from a physical mount into a pretend one. + +The mid-360 is bolted to the robot at an angle (``rotated_urdf``), but the rest +of the stack is wired as if it sat "forward = forward" (``normal_urdf``). This +module consumes ``rotated_lidar`` and emits ``lidar`` with each point moved from +the rotated sensor frame into the normal one, so nothing downstream needs to know +the sensor is tilted. + +The correction is a single rigid transform derived from the two URDFs: +``inv(normal_base_to_sensor) @ rotated_base_to_sensor`` evaluated at the shared +``sensor_frame``. Only the cloud is rewritten — the odometry is left untouched. +""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np + +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.robot.urdf_loader import UrdfLoader +from dimos.spec import perception + + +class PointLioHackConfig(ModuleConfig): + rotated_urdf: Path + normal_urdf: Path + sensor_frame: str = "mid360_link" + + +def _transform_to_matrix(transform: Transform) -> np.ndarray: + matrix = np.eye(4) + matrix[:3, :3] = transform.rotation.to_rotation_matrix() + matrix[:3, 3] = ( + transform.translation.x, + transform.translation.y, + transform.translation.z, + ) + return matrix + + +def _base_to_frame_matrix(loader: UrdfLoader, leaf_frame: str) -> np.ndarray: + """Compose the fixed-joint chain from the model root down to ``leaf_frame``.""" + static_transforms = loader.static_transforms + chain: list[Transform] = [] + frame = leaf_frame + while frame in static_transforms: + transform = static_transforms[frame] + chain.append(transform) + frame = transform.frame_id + matrix = np.eye(4) + for transform in reversed(chain): + matrix = matrix @ _transform_to_matrix(transform) + return matrix + + +class PointLioHack(Module, perception.Lidar): + config: PointLioHackConfig + + rotated_lidar: In[PointCloud2] + lidar: Out[PointCloud2] + + async def main(self) -> AsyncIterator[None]: + rotated = _base_to_frame_matrix( + UrdfLoader(name="rotated", model_path=self.config.rotated_urdf), + self.config.sensor_frame, + ) + normal = _base_to_frame_matrix( + UrdfLoader(name="normal", model_path=self.config.normal_urdf), + self.config.sensor_frame, + ) + correction = np.linalg.inv(normal) @ rotated + self._rotation = correction[:3, :3].astype(np.float32) + self._translation = correction[:3, 3].astype(np.float32) + yield + + async def handle_rotated_lidar(self, value: PointCloud2) -> None: + points = value.points_f32() @ self._rotation.T + self._translation + self.lidar.publish( + PointCloud2.from_numpy( + points=points, + frame_id=value.frame_id, + timestamp=value.ts, + intensities=value.intensities_f32(), + ) + ) + + +# Verify protocol port compliance (mypy will flag missing ports) +if TYPE_CHECKING: + PointLioHack() diff --git a/dimos/robot/assembly/mid360_realsense_30.urdf b/dimos/robot/assembly/mid360_realsense_30.urdf new file mode 100644 index 0000000000..a132b2930f --- /dev/null +++ b/dimos/robot/assembly/mid360_realsense_30.urdf @@ -0,0 +1,108 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_normal.urdf b/dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_normal.urdf new file mode 100644 index 0000000000..aa02f520cb --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_normal.urdf @@ -0,0 +1,22 @@ + + + + + + + + + + + + + + + + + + + + + + diff --git a/dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_rotated.urdf b/dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_rotated.urdf new file mode 100644 index 0000000000..c17d084ee1 --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_rotated.urdf @@ -0,0 +1,24 @@ + + + + + + + + + + + + + + + + + + + + + + diff --git a/dimos/robot/unitree/go2/blueprints/navigation/mid360_realsense_30.urdf b/dimos/robot/unitree/go2/blueprints/navigation/mid360_realsense_30.urdf new file mode 100644 index 0000000000..a132b2930f --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/navigation/mid360_realsense_30.urdf @@ -0,0 +1,108 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 1fae6b4ad3..094c23ccd8 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -15,24 +15,55 @@ """3d navigation on Go2 with ray tracing and MLS planning""" +import math +from pathlib import Path from typing import Any from dimos.core.coordination.blueprints import autoconnect from dimos.core.global_config import global_config +from dimos.hardware.sensors.lidar.pointlio.hack import PointLioHack from dimos.hardware.sensors.lidar.pointlio.module import PointLio from dimos.mapping.ray_tracing.module import RayTracingVoxelMap +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.navigation.basic_path_follower.module import BasicPathFollower from dimos.navigation.movement_manager.movement_manager import MovementManager from dimos.navigation.nav_3d.mls_planner.goal_relay import GoalRelay from dimos.navigation.nav_3d.mls_planner.mls_planner_native import MLSPlannerNative from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import rerun_config from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.robot.urdf_loader import UrdfLoader from dimos.visualization.vis_module import vis_module +_navigation_dir = Path(__file__).parent +# The physical mid-360 mount (rotated_urdf) and the "forward = forward" mount the +# rest of the stack pretends it has (normal_urdf). PointLioHack rewrites the cloud +# from the first into the second, so everything downstream uses the normal mount. +_rotated_urdf = _navigation_dir / "go2_mid360_rotated.urdf" +_normal_urdf = _navigation_dir / "go2_mid360_normal.urdf" + +# GO2Connection is a TfModule; feeding it the normal mount publishes those frames +# on its static interval, matching the cloud the hack emits. +go2_mid360_model = UrdfLoader(name="go2_mid360", model_path=_normal_urdf) + voxel_size = 0.08 # Height of the head-mounted lidar above the ground while standing. go2_lidar_height = 0.5 +# The lidar is physically pitched down 45 deg and rolled 45 deg, so pointlio's +# body frame is tilted. Counter-rotate the robot body box by the inverse of that +# mount so it renders level (flip a sign here if the box leans the wrong way). +_body_box_roll_deg = -45.0 +_body_box_pitch_deg = -45.0 +_body_box_yaw_deg = 0.0 +_body_box_rotation = Quaternion.from_euler( + Vector3( + math.radians(_body_box_roll_deg), + math.radians(_body_box_pitch_deg), + math.radians(_body_box_yaw_deg), + ) +) + def _render_global_map(msg: Any) -> Any: return msg.to_rerun() @@ -47,12 +78,19 @@ def _render_path(msg: Any) -> Any: def _static_robot_body(rr: Any) -> list[Any]: - """Go2-shaped box on pointlio's body frame, counter-rotated for the lidar pitch.""" + """Go2-shaped box + forward arrow on pointlio's body frame, counter-rotated + so the physically pitched/rolled lidar mount renders level.""" return [ rr.Boxes3D(half_sizes=[0.35, 0.155, 0.2], colors=[(0, 255, 127)]), + # Red arrow out the nose marks the robot's forward (+x) direction. + rr.Arrows3D( + origins=[(0.35, 0.0, 0.0)], + vectors=[(0.3, 0.0, 0.0)], + colors=[(255, 40, 40)], + ), rr.Transform3D( parent_frame="tf#/body", - rotation=rr.RotationAxisAngle(axis=(0, 1, 0), degrees=-45.0), + rotation=_body_box_rotation.to_rerun(), ), ] @@ -85,13 +123,27 @@ def _static_robot_body(rr: Any) -> list[Any]: unitree_go2_nav_3d = autoconnect( vis_module(viewer_backend=global_config.viewer, rerun_config=_nav_rerun_config), # "mcf" for stair traversal - GO2Connection.blueprint(lidar=False, camera=False, motion_mode="mcf").remappings( + GO2Connection.blueprint( + static_transforms=dict(go2_mid360_model.static_transforms), + lidar=False, + camera=False, + motion_mode="mcf", + ).remappings( [ (GO2Connection, "lidar", "lidar_l1"), (GO2Connection, "odom", "odom_go2"), ] ), - PointLio.blueprint(body_frame_id="body"), + # gravity_align is off (no_gravity_align.yaml) so pointlio leaves the odometry + # in the raw mount frame; the hack applies the known rotated->normal transform + # instead. Only the lidar is rerouted through the hack; odometry flows straight + # to consumers. + PointLio.blueprint(body_frame_id="body", config="no_gravity_align.yaml").remappings( + [ + (PointLio, "lidar", "rotated_lidar"), + ] + ), + PointLioHack.blueprint(rotated_urdf=_rotated_urdf, normal_urdf=_normal_urdf), RayTracingVoxelMap.blueprint( voxel_size=voxel_size, emit_every=1, From 8640fbef7187007fdef3b8994c84281f1a245b62 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Fri, 26 Jun 2026 12:31:48 +0800 Subject: [PATCH 58/71] render time --- dimos/hardware/sensors/lidar/pointlio/hack.py | 67 +++++++++++++++---- .../navigation/go2_mid360_normal.urdf | 2 +- .../navigation/unitree_go2_nav_3d.py | 58 ++++++++-------- 3 files changed, 83 insertions(+), 44 deletions(-) diff --git a/dimos/hardware/sensors/lidar/pointlio/hack.py b/dimos/hardware/sensors/lidar/pointlio/hack.py index 8ffb359354..b2c638825f 100644 --- a/dimos/hardware/sensors/lidar/pointlio/hack.py +++ b/dimos/hardware/sensors/lidar/pointlio/hack.py @@ -12,17 +12,22 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Rewrite the lidar cloud from a physical mount into a pretend one. +"""Fake a normally-mounted mid-360 from the physically-tilted one. -The mid-360 is bolted to the robot at an angle (``rotated_urdf``), but the rest -of the stack is wired as if it sat "forward = forward" (``normal_urdf``). This -module consumes ``rotated_lidar`` and emits ``lidar`` with each point moved from -the rotated sensor frame into the normal one, so nothing downstream needs to know +The sensor is bolted to the robot at an angle (``rotated_urdf``), but the rest of +the stack is wired as if it sat in the nominal mount (``normal_urdf``). This +module rewrites both the cloud and the odometry so nothing downstream can tell the sensor is tilted. -The correction is a single rigid transform derived from the two URDFs: -``inv(normal_base_to_sensor) @ rotated_base_to_sensor`` evaluated at the shared -``sensor_frame``. Only the cloud is rewritten — the odometry is left untouched. +The correction is a single rigid transform ``C`` derived from the two URDFs: +``C = inv(normal_base_to_sensor) @ rotated_base_to_sensor`` evaluated at the +shared ``sensor_frame``. + +* cloud: each point is moved ``p' = C @ p`` (rotated sensor frame -> normal one). +* odometry: the body pose is conjugated ``T' = C @ T @ inv(C)`` and the twist + rotated by ``C``'s rotation. That is exactly the odometry a normally-mounted + sensor would have produced, so a forward walk stays forward instead of drifting + sideways the way a bare body-frame relabel would. """ from __future__ import annotations @@ -35,7 +40,12 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 from dimos.robot.urdf_loader import UrdfLoader from dimos.spec import perception @@ -73,11 +83,17 @@ def _base_to_frame_matrix(loader: UrdfLoader, leaf_frame: str) -> np.ndarray: return matrix -class PointLioHack(Module, perception.Lidar): +def _vector_to_numpy(vector: Vector3) -> np.ndarray: + return np.array([vector.x, vector.y, vector.z]) + + +class PointLioHack(Module, perception.Lidar, perception.Odometry): config: PointLioHackConfig rotated_lidar: In[PointCloud2] lidar: Out[PointCloud2] + rotated_odometry: In[Odometry] + odometry: Out[Odometry] async def main(self) -> AsyncIterator[None]: rotated = _base_to_frame_matrix( @@ -88,9 +104,10 @@ async def main(self) -> AsyncIterator[None]: UrdfLoader(name="normal", model_path=self.config.normal_urdf), self.config.sensor_frame, ) - correction = np.linalg.inv(normal) @ rotated - self._rotation = correction[:3, :3].astype(np.float32) - self._translation = correction[:3, 3].astype(np.float32) + self._correction = np.linalg.inv(normal) @ rotated + self._correction_inv = np.linalg.inv(self._correction) + self._rotation = self._correction[:3, :3].astype(np.float32) + self._translation = self._correction[:3, 3].astype(np.float32) yield async def handle_rotated_lidar(self, value: PointCloud2) -> None: @@ -104,6 +121,32 @@ async def handle_rotated_lidar(self, value: PointCloud2) -> None: ) ) + async def handle_rotated_odometry(self, value: Odometry) -> None: + pose = np.eye(4) + pose[:3, :3] = value.orientation.to_rotation_matrix() + pose[:3, 3] = _vector_to_numpy(value.position) + faked = self._correction @ pose @ self._correction_inv + + rotation = self._correction[:3, :3] + linear = rotation @ _vector_to_numpy(value.linear_velocity) + angular = rotation @ _vector_to_numpy(value.angular_velocity) + + self.odometry.publish( + Odometry( + ts=value.ts, + frame_id=value.frame_id, + child_frame_id=value.child_frame_id, + pose=Pose( + Vector3(*(float(component) for component in faked[:3, 3])), + Quaternion.from_rotation_matrix(faked[:3, :3]), + ), + twist=Twist( + Vector3(*(float(component) for component in linear)), + Vector3(*(float(component) for component in angular)), + ), + ) + ) + # Verify protocol port compliance (mypy will flag missing ports) if TYPE_CHECKING: diff --git a/dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_normal.urdf b/dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_normal.urdf index aa02f520cb..3a4ea01189 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_normal.urdf +++ b/dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_normal.urdf @@ -11,7 +11,7 @@ - + diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 094c23ccd8..bbb3b6dcf9 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -15,7 +15,6 @@ """3d navigation on Go2 with ray tracing and MLS planning""" -import math from pathlib import Path from typing import Any @@ -24,8 +23,6 @@ from dimos.hardware.sensors.lidar.pointlio.hack import PointLioHack from dimos.hardware.sensors.lidar.pointlio.module import PointLio from dimos.mapping.ray_tracing.module import RayTracingVoxelMap -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.navigation.basic_path_follower.module import BasicPathFollower from dimos.navigation.movement_manager.movement_manager import MovementManager from dimos.navigation.nav_3d.mls_planner.goal_relay import GoalRelay @@ -50,20 +47,6 @@ # Height of the head-mounted lidar above the ground while standing. go2_lidar_height = 0.5 -# The lidar is physically pitched down 45 deg and rolled 45 deg, so pointlio's -# body frame is tilted. Counter-rotate the robot body box by the inverse of that -# mount so it renders level (flip a sign here if the box leans the wrong way). -_body_box_roll_deg = -45.0 -_body_box_pitch_deg = -45.0 -_body_box_yaw_deg = 0.0 -_body_box_rotation = Quaternion.from_euler( - Vector3( - math.radians(_body_box_roll_deg), - math.radians(_body_box_pitch_deg), - math.radians(_body_box_yaw_deg), - ) -) - def _render_global_map(msg: Any) -> Any: return msg.to_rerun() @@ -77,9 +60,25 @@ def _render_path(msg: Any) -> Any: return msg +def _render_surface_map(msg: Any) -> Any: + # Walkable surface the planner solved on. Cyan voxel boxes so it reads as a + # distinct carpet over the raw global_map. + return msg.to_rerun(voxel_size=voxel_size, colors=[80, 200, 255], mode="boxes") + + +def _render_nodes(msg: Any) -> Any: + # Planning-graph nodes as fat magenta spheres so they pop above the surface. + return msg.to_rerun(voxel_size=0.25, colors=[255, 0, 200], mode="spheres") + + +def _render_node_edges(msg: Any) -> Any: + # LineSegments3D already colors edges by traversability (green/yellow/red). + return msg.to_rerun() + + def _static_robot_body(rr: Any) -> list[Any]: - """Go2-shaped box + forward arrow on pointlio's body frame, counter-rotated - so the physically pitched/rolled lidar mount renders level.""" + """Go2-shaped box + forward arrow on pointlio's body frame. The hack fakes a + level, forward-facing mount, so the box needs no counter-rotation.""" return [ rr.Boxes3D(half_sizes=[0.35, 0.155, 0.2], colors=[(0, 255, 127)]), # Red arrow out the nose marks the robot's forward (+x) direction. @@ -88,10 +87,7 @@ def _static_robot_body(rr: Any) -> list[Any]: vectors=[(0.3, 0.0, 0.0)], colors=[(255, 40, 40)], ), - rr.Transform3D( - parent_frame="tf#/body", - rotation=_body_box_rotation.to_rerun(), - ), + rr.Transform3D(parent_frame="tf#/body"), ] @@ -114,9 +110,9 @@ def _static_robot_body(rr: Any) -> list[Any]: "world/camera_info": None, "world/color_image": None, "world/lidar": None, - "world/surface_map": None, - "world/nodes": None, - "world/node_edges": None, + "world/surface_map": _render_surface_map, + "world/nodes": _render_nodes, + "world/node_edges": _render_node_edges, }, } @@ -134,13 +130,13 @@ def _static_robot_body(rr: Any) -> list[Any]: (GO2Connection, "odom", "odom_go2"), ] ), - # gravity_align is off (no_gravity_align.yaml) so pointlio leaves the odometry - # in the raw mount frame; the hack applies the known rotated->normal transform - # instead. Only the lidar is rerouted through the hack; odometry flows straight - # to consumers. + # gravity_align is off (no_gravity_align.yaml) so pointlio leaves both the cloud + # and the odometry in the raw mount frame. The hack rewrites both into the normal + # mount, so everything downstream sees a normally-mounted sensor. PointLio.blueprint(body_frame_id="body", config="no_gravity_align.yaml").remappings( [ (PointLio, "lidar", "rotated_lidar"), + (PointLio, "odometry", "rotated_odometry"), ] ), PointLioHack.blueprint(rotated_urdf=_rotated_urdf, normal_urdf=_normal_urdf), @@ -162,7 +158,7 @@ def _static_robot_body(rr: Any) -> list[Any]: wall_buffer_weight=100.0, step_threshold_m=0.16, step_penalty_weight=1.0, - viz_publish_hz=0.0, + viz_publish_hz=2.0, ).remappings([(MLSPlannerNative, "global_map", "global_map_unused")]), GoalRelay.blueprint(), BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6), From 4440728f5675ca6d955b8d915a13f54bd9c8918b Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Fri, 26 Jun 2026 14:39:55 +0800 Subject: [PATCH 59/71] - --- dimos/hardware/sensors/lidar/pointlio/hack.py | 22 +++++++++++++++---- .../hardware/sensors/lidar/pointlio/module.py | 12 ++++++---- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/dimos/hardware/sensors/lidar/pointlio/hack.py b/dimos/hardware/sensors/lidar/pointlio/hack.py index b2c638825f..a2efe422de 100644 --- a/dimos/hardware/sensors/lidar/pointlio/hack.py +++ b/dimos/hardware/sensors/lidar/pointlio/hack.py @@ -55,6 +55,10 @@ class PointLioHackConfig(ModuleConfig): rotated_urdf: Path normal_urdf: Path sensor_frame: str = "mid360_link" + # odom->body TF, faked here instead of by PointLio so the TF tree matches the + # shimmed cloud/odometry. Must match PointLio's body_start_frame_id/body_frame_id. + odom_frame: str = "odom" + body_frame: str = "body" def _transform_to_matrix(transform: Transform) -> np.ndarray: @@ -131,15 +135,15 @@ async def handle_rotated_odometry(self, value: Odometry) -> None: linear = rotation @ _vector_to_numpy(value.linear_velocity) angular = rotation @ _vector_to_numpy(value.angular_velocity) + faked_position = Vector3(*(float(component) for component in faked[:3, 3])) + faked_orientation = Quaternion.from_rotation_matrix(faked[:3, :3]) + self.odometry.publish( Odometry( ts=value.ts, frame_id=value.frame_id, child_frame_id=value.child_frame_id, - pose=Pose( - Vector3(*(float(component) for component in faked[:3, 3])), - Quaternion.from_rotation_matrix(faked[:3, :3]), - ), + pose=Pose(faked_position, faked_orientation), twist=Twist( Vector3(*(float(component) for component in linear)), Vector3(*(float(component) for component in angular)), @@ -147,6 +151,16 @@ async def handle_rotated_odometry(self, value: Odometry) -> None: ) ) + self.tf.publish( + Transform( + frame_id=self.config.odom_frame, + child_frame_id=self.config.body_frame, + translation=faked_position, + rotation=faked_orientation, + ts=value.ts, + ) + ) + # Verify protocol port compliance (mypy will flag missing ports) if TYPE_CHECKING: diff --git a/dimos/hardware/sensors/lidar/pointlio/module.py b/dimos/hardware/sensors/lidar/pointlio/module.py index 4a13625ab8..92496ba225 100644 --- a/dimos/hardware/sensors/lidar/pointlio/module.py +++ b/dimos/hardware/sensors/lidar/pointlio/module.py @@ -34,8 +34,9 @@ from pydantic import Field from pydantic.experimental.pipeline import validate_as -from reactivex.disposable import Disposable +# HACK(kronk-nav): re-enable with the _on_odom_for_tf subscription in start(). +# from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import Out @@ -129,9 +130,12 @@ class PointLio(NativeModule, perception.Lidar, perception.Odometry): def start(self) -> None: self._validate_network() super().start() - self.register_disposable( - Disposable(self.odometry.transport.subscribe(self._on_odom_for_tf, self.odometry)) - ) + # HACK(kronk-nav): odom->body TF is published by PointLioHack from the + # faked odometry instead, so the TF tree matches the shimmed map. Leaving + # this on would broadcast the raw tilted pose and fight the hack. + # self.register_disposable( + # Disposable(self.odometry.transport.subscribe(self._on_odom_for_tf, self.odometry)) + # ) def _on_odom_for_tf(self, msg: Odometry) -> None: self.tf.publish( From b610dc5d38d6f50c8eef978d925f33cccf76eb38 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Fri, 26 Jun 2026 14:58:21 +0800 Subject: [PATCH 60/71] - --- .../unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index bbb3b6dcf9..011fc1fa05 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -101,8 +101,10 @@ def _static_robot_body(rr: Any) -> list[Any]: "memory_limit": "256MB", # base_link tf comes from the go2 internal odometry, which is not the map # frame. Anchor the robot box to pointlio's body frame instead and hide the - # camera frustum that rides base_link. - "static": {"world/tf/body": _static_robot_body}, + # camera frustum that rides base_link. Use a dedicated entity path (not + # world/tf/body) so the box's Transform3D doesn't collide with the live + # odom->body TF that PointLioHack logs onto world/tf/body. + "static": {"world/robot_body": _static_robot_body}, "visual_override": { **rerun_config["visual_override"], "world/global_map": _render_global_map, From 3bf64f8b49e39fb47e0e699ef6160c4b546cc56a Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Fri, 26 Jun 2026 16:32:32 +0800 Subject: [PATCH 61/71] - --- .../unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 011fc1fa05..e15086a01d 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -135,7 +135,11 @@ def _static_robot_body(rr: Any) -> list[Any]: # gravity_align is off (no_gravity_align.yaml) so pointlio leaves both the cloud # and the odometry in the raw mount frame. The hack rewrites both into the normal # mount, so everything downstream sees a normally-mounted sensor. - PointLio.blueprint(body_frame_id="body", config="no_gravity_align.yaml").remappings( + PointLio.blueprint( + body_frame_id="body", + config="no_gravity_align.yaml", + space_down_sample=False, + ).remappings( [ (PointLio, "lidar", "rotated_lidar"), (PointLio, "odometry", "rotated_odometry"), From 24c5b533e243d1aafb1455d165dd90b90ba4d426 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Fri, 26 Jun 2026 17:04:07 +0800 Subject: [PATCH 62/71] - --- .../go2/blueprints/navigation/unitree_go2_nav_3d.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index e15086a01d..c40285dee4 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -22,6 +22,7 @@ from dimos.core.global_config import global_config from dimos.hardware.sensors.lidar.pointlio.hack import PointLioHack from dimos.hardware.sensors.lidar.pointlio.module import PointLio +from dimos.hardware.sensors.lidar.pointlio.recorder import PointlioRecorder from dimos.mapping.ray_tracing.module import RayTracingVoxelMap from dimos.navigation.basic_path_follower.module import BasicPathFollower from dimos.navigation.movement_manager.movement_manager import MovementManager @@ -146,6 +147,14 @@ def _static_robot_body(rr: Any) -> list[Any]: ] ), PointLioHack.blueprint(rotated_urdf=_rotated_urdf, normal_urdf=_normal_urdf), + # Record the hack's faked (normal-mount) cloud + odometry, not PointLio's raw + # tilted output, so a replay reproduces what the nav stack actually consumed. + PointlioRecorder.blueprint().remappings( + [ + (PointlioRecorder, "pointlio_lidar", "lidar"), + (PointlioRecorder, "pointlio_odometry", "odometry"), + ] + ), RayTracingVoxelMap.blueprint( voxel_size=voxel_size, emit_every=1, From 3be411f323b4ef60fd4246e8484690c10777fec6 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Fri, 26 Jun 2026 19:23:25 +0800 Subject: [PATCH 63/71] fix: sqlite-vec aarch64 wheel + pointlio nav-3d config The sqlite-vec 0.1.6 aarch64 wheel ships a 32-bit vec0.so (upstream packaging bug) which fails to load on the Jetson with "wrong ELF class: ELFCLASS32". Pin >=0.1.7 to skip it (resolves to 0.1.9). Also fix the unitree-go2-nav-3d PointLio config: pass gravity_align=False instead of the removed config="no_gravity_align.yaml", drop the redundant body_frame_id, and set auto_build=True so the native binary recompiles. --- .../blueprints/navigation/unitree_go2_nav_3d.py | 6 +++--- pyproject.toml | 2 +- uv.lock | 14 +++++++------- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index c40285dee4..f20c1db119 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -133,13 +133,13 @@ def _static_robot_body(rr: Any) -> list[Any]: (GO2Connection, "odom", "odom_go2"), ] ), - # gravity_align is off (no_gravity_align.yaml) so pointlio leaves both the cloud + # gravity_align is off so pointlio leaves both the cloud # and the odometry in the raw mount frame. The hack rewrites both into the normal # mount, so everything downstream sees a normally-mounted sensor. PointLio.blueprint( - body_frame_id="body", - config="no_gravity_align.yaml", + gravity_align=False, space_down_sample=False, + auto_build=True, ).remappings( [ (PointLio, "lidar", "rotated_lidar"), diff --git a/pyproject.toml b/pyproject.toml index a60761a753..d3861e777a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,7 +129,7 @@ dependencies = [ "toolz>=1.1.0", "protobuf>=6.33.5,<7", "psutil>=7.0.0", - "sqlite-vec>=0.1.6", + "sqlite-vec>=0.1.7", "lz4>=4.4.5", "bleak>=3.0.2", "cryptography>=46.0.5", diff --git a/uv.lock b/uv.lock index 33a33fa51b..07fec5a6c4 100644 --- a/uv.lock +++ b/uv.lock @@ -2492,7 +2492,7 @@ requires-dist = [ { name = "sortedcontainers", specifier = "==2.4.0" }, { name = "sounddevice", marker = "extra == 'agents'" }, { name = "soundfile", marker = "extra == 'web'" }, - { name = "sqlite-vec", specifier = ">=0.1.6" }, + { name = "sqlite-vec", specifier = ">=0.1.7" }, { name = "sse-starlette", marker = "extra == 'web'", specifier = ">=2.2.1" }, { name = "structlog", specifier = ">=25.5.0,<26" }, { name = "tensorboard", marker = "extra == 'misc'", specifier = "==2.20.0" }, @@ -10472,14 +10472,14 @@ wheels = [ [[package]] name = "sqlite-vec" -version = "0.1.6" +version = "0.1.9" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/88/ed/aabc328f29ee6814033d008ec43e44f2c595447d9cccd5f2aabe60df2933/sqlite_vec-0.1.6-py3-none-macosx_10_6_x86_64.whl", hash = "sha256:77491bcaa6d496f2acb5cc0d0ff0b8964434f141523c121e313f9a7d8088dee3", size = 164075, upload-time = "2024-11-20T16:40:29.847Z" }, - { url = "https://files.pythonhosted.org/packages/a7/57/05604e509a129b22e303758bfa062c19afb020557d5e19b008c64016704e/sqlite_vec-0.1.6-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fdca35f7ee3243668a055255d4dee4dea7eed5a06da8cad409f89facf4595361", size = 165242, upload-time = "2024-11-20T16:40:31.206Z" }, - { url = "https://files.pythonhosted.org/packages/f2/48/dbb2cc4e5bad88c89c7bb296e2d0a8df58aab9edc75853728c361eefc24f/sqlite_vec-0.1.6-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7b0519d9cd96164cd2e08e8eed225197f9cd2f0be82cb04567692a0a4be02da3", size = 103704, upload-time = "2024-11-20T16:40:33.729Z" }, - { url = "https://files.pythonhosted.org/packages/80/76/97f33b1a2446f6ae55e59b33869bed4eafaf59b7f4c662c8d9491b6a714a/sqlite_vec-0.1.6-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux1_x86_64.whl", hash = "sha256:823b0493add80d7fe82ab0fe25df7c0703f4752941aee1c7b2b02cec9656cb24", size = 151556, upload-time = "2024-11-20T16:40:35.387Z" }, - { url = "https://files.pythonhosted.org/packages/6a/98/e8bc58b178266eae2fcf4c9c7a8303a8d41164d781b32d71097924a6bebe/sqlite_vec-0.1.6-py3-none-win_amd64.whl", hash = "sha256:c65bcfd90fa2f41f9000052bcb8bb75d38240b2dae49225389eca6c3136d3f0c", size = 281540, upload-time = "2024-11-20T16:40:37.296Z" }, + { url = "https://files.pythonhosted.org/packages/68/85/9fad0045d8e7c8df3e0fa5a56c630e8e15ad6e5ca2e6106fceb666aa6638/sqlite_vec-0.1.9-py3-none-macosx_10_6_x86_64.whl", hash = "sha256:1b62a7f0a060d9475575d4e599bbf94a13d85af896bc1ce86ee80d1b5b48e5fb", size = 131171, upload-time = "2026-03-31T08:02:31.717Z" }, + { url = "https://files.pythonhosted.org/packages/a4/3d/3677e0cd2f92e5ebc43cd29fbf565b75582bff1ccfa0b8327c7508e1084f/sqlite_vec-0.1.9-py3-none-macosx_11_0_arm64.whl", hash = "sha256:1d52e30513bae4cc9778ddbf6145610434081be4c3afe57cd877893bad9f6b6c", size = 165434, upload-time = "2026-03-31T08:02:32.712Z" }, + { url = "https://files.pythonhosted.org/packages/00/d4/f2b936d3bdc38eadcbd2a87875815db36430fab0363182ba5d12cd8e0b51/sqlite_vec-0.1.9-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4e921e592f24a5f9a18f590b6ddd530eb637e2d474e3b1972f9bbeb773aa3cb9", size = 160076, upload-time = "2026-03-31T08:02:33.796Z" }, + { url = "https://files.pythonhosted.org/packages/6f/ad/6afd073b0f817b3e03f9e37ad626ae341805891f23c74b5292818f49ac63/sqlite_vec-0.1.9-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux1_x86_64.whl", hash = "sha256:1515727990b49e79bcaf75fdee2ffc7d461f8b66905013231251f1c8938e7786", size = 163388, upload-time = "2026-03-31T08:02:34.888Z" }, + { url = "https://files.pythonhosted.org/packages/42/89/81b2907cda14e566b9bf215e2ad82fc9b349edf07d2010756ffdb902f328/sqlite_vec-0.1.9-py3-none-win_amd64.whl", hash = "sha256:4a28dc12fa4b53d7b1dced22da2488fade444e96b5d16fd2d698cd670675cf32", size = 292804, upload-time = "2026-03-31T08:02:36.035Z" }, ] [[package]] From 2e0398319d2944617c0b6e9b8235ba826568ffb3 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 30 Jun 2026 17:29:29 -0500 Subject: [PATCH 64/71] feat(pointlio): mount-correction transform arg + frame_mapping-driven frames Add a row-major 4x4 transform to PointLioConfig so the C++ rewrites the cloud, odometry, and odom->body TF into a corrected mount frame, replacing the StaticTfCorrection hack. Drop frame_id/child_frame_id/sensor_frame_id config fields in favor of frame_mapping, injected into the C++ args at start. --- .../sensors/lidar/pointlio/cpp/main.cpp | 89 +++++++++++++++---- .../hardware/sensors/lidar/pointlio/module.py | 65 ++++++++++---- .../lidar/pointlio/mount_correction.py | 81 +++++++++++++++++ dimos/robot/assembly/mid360_realsense_30.py | 3 +- .../basic/unitree_go2_mid360_record.py | 3 +- 5 files changed, 204 insertions(+), 37 deletions(-) create mode 100644 dimos/hardware/sensors/lidar/pointlio/mount_correction.py diff --git a/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp b/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp index 9c7bcb3714..d2aef4a4cf 100644 --- a/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp +++ b/dimos/hardware/sensors/lidar/pointlio/cpp/main.cpp @@ -16,6 +16,8 @@ #include +#include + #include #include #include @@ -81,6 +83,17 @@ static std::string g_child_frame_id; // required via --child_frame_id static std::string g_sensor_frame_id; // required via --sensor_frame_id static float g_frequency = 10.0f; +// Optional rigid transform applied to the published cloud + odometry (and, via +// the Python TF republisher that reads the odometry, the odom->body TF). Lets a +// physically-tilted mount be reported as a different mount without touching the +// SLAM core. Identity unless --transform supplies a row-major 4x4 (16 doubles). +// Cloud points are moved p' = C @ p; the body pose is conjugated T' = C @ T @ +// inv(C) and the twist rotated by C's rotation, so cloud and odometry stay +// mutually consistent. +static bool g_has_transform = false; +static Eigen::Matrix4d g_transform = Eigen::Matrix4d::Identity(); +static Eigen::Matrix4d g_transform_inv = Eigen::Matrix4d::Identity(); + // Frame accumulator (Livox SDK raw → CustomMsg) static std::mutex g_pc_mutex; // Serializes all Point-LIO EKF access. The SDK delivers IMU on its own callback @@ -142,11 +155,22 @@ static void publish_lidar(PointCloudXYZI::Ptr cloud, double timestamp, const std pc.data_length = pc.row_step; pc.data.resize(pc.data_length); + const Eigen::Matrix3d transform_rotation = g_transform.block<3, 3>(0, 0); + const Eigen::Vector3d transform_translation = g_transform.block<3, 1>(0, 3); for (int point_idx = 0; point_idx < num_points; ++point_idx) { float* dst = reinterpret_cast(pc.data.data() + point_idx * 16); - dst[0] = cloud->points[point_idx].x; - dst[1] = cloud->points[point_idx].y; - dst[2] = cloud->points[point_idx].z; + double x = cloud->points[point_idx].x; + double y = cloud->points[point_idx].y; + double z = cloud->points[point_idx].z; + if (g_has_transform) { + const Eigen::Vector3d moved = transform_rotation * Eigen::Vector3d(x, y, z) + transform_translation; + x = moved.x(); + y = moved.y(); + z = moved.z(); + } + dst[0] = static_cast(x); + dst[1] = static_cast(y); + dst[2] = static_cast(z); dst[3] = cloud->points[point_idx].intensity; } @@ -161,25 +185,48 @@ static void publish_odometry(const custom_messages::Odometry& odom, double times msg.child_frame_id = g_child_frame_id; // Pose in the SLAM/sensor frame. - msg.pose.pose.position.x = odom.pose.pose.position.x; - msg.pose.pose.position.y = odom.pose.pose.position.y; - msg.pose.pose.position.z = odom.pose.pose.position.z; - msg.pose.pose.orientation.x = odom.pose.pose.orientation.x; - msg.pose.pose.orientation.y = odom.pose.pose.orientation.y; - msg.pose.pose.orientation.z = odom.pose.pose.orientation.z; - msg.pose.pose.orientation.w = odom.pose.pose.orientation.w; + Eigen::Vector3d position( + odom.pose.pose.position.x, odom.pose.pose.position.y, odom.pose.pose.position.z); + Eigen::Quaterniond orientation( + odom.pose.pose.orientation.w, odom.pose.pose.orientation.x, + odom.pose.pose.orientation.y, odom.pose.pose.orientation.z); + Eigen::Vector3d linear( + odom.twist.twist.linear.x, odom.twist.twist.linear.y, odom.twist.twist.linear.z); + Eigen::Vector3d angular( + odom.twist.twist.angular.x, odom.twist.twist.angular.y, odom.twist.twist.angular.z); + + if (g_has_transform) { + Eigen::Matrix4d pose = Eigen::Matrix4d::Identity(); + pose.block<3, 3>(0, 0) = orientation.toRotationMatrix(); + pose.block<3, 1>(0, 3) = position; + const Eigen::Matrix4d faked = g_transform * pose * g_transform_inv; + const Eigen::Matrix3d transform_rotation = g_transform.block<3, 3>(0, 0); + position = faked.block<3, 1>(0, 3); + orientation = Eigen::Quaterniond(faked.block<3, 3>(0, 0)); + orientation.normalize(); + linear = transform_rotation * linear; + angular = transform_rotation * angular; + } + + msg.pose.pose.position.x = position.x(); + msg.pose.pose.position.y = position.y(); + msg.pose.pose.position.z = position.z(); + msg.pose.pose.orientation.x = orientation.x(); + msg.pose.pose.orientation.y = orientation.y(); + msg.pose.pose.orientation.z = orientation.z(); + msg.pose.pose.orientation.w = orientation.w(); for (int idx = 0; idx < 36; ++idx) { msg.pose.covariance[idx] = odom.pose.covariance[idx]; } // Velocity from Point-LIO's IESKF state (its key output over FAST-LIO). - msg.twist.twist.linear.x = odom.twist.twist.linear.x; - msg.twist.twist.linear.y = odom.twist.twist.linear.y; - msg.twist.twist.linear.z = odom.twist.twist.linear.z; - msg.twist.twist.angular.x = odom.twist.twist.angular.x; - msg.twist.twist.angular.y = odom.twist.twist.angular.y; - msg.twist.twist.angular.z = odom.twist.twist.angular.z; + msg.twist.twist.linear.x = linear.x(); + msg.twist.twist.linear.y = linear.y(); + msg.twist.twist.linear.z = linear.z(); + msg.twist.twist.angular.x = angular.x(); + msg.twist.twist.angular.y = angular.y(); + msg.twist.twist.angular.z = angular.z(); std::memset(msg.twist.covariance, 0, sizeof(msg.twist.covariance)); g_lcm->publish(g_odometry_topic, &msg); @@ -368,6 +415,16 @@ int main(int argc, char** argv) { if (auto gi = parse_doubles(mod.arg("gravity_init", "")); !gi.empty()) params.gravity_init = gi; if (auto et = parse_doubles(mod.arg("extrinsic_t", "")); !et.empty()) params.extrinsic_T = et; if (auto er = parse_doubles(mod.arg("extrinsic_r", "")); !er.empty()) params.extrinsic_R = er; + // Optional post-SLAM mount correction (row-major 4x4); identity when absent. + if (auto transform = parse_doubles(mod.arg("transform", "")); transform.size() == 16) { + g_has_transform = true; + for (int row = 0; row < 4; ++row) { + for (int col = 0; col < 4; ++col) { + g_transform(row, col) = transform[row * 4 + col]; + } + } + g_transform_inv = g_transform.inverse(); + } // odometry params.publish_odometry_without_downsample = mod.arg_bool("publish_odometry_without_downsample", params.publish_odometry_without_downsample); diff --git a/dimos/hardware/sensors/lidar/pointlio/module.py b/dimos/hardware/sensors/lidar/pointlio/module.py index 057da597f7..25d2659e4c 100644 --- a/dimos/hardware/sensors/lidar/pointlio/module.py +++ b/dimos/hardware/sensors/lidar/pointlio/module.py @@ -34,10 +34,9 @@ import os from typing import TYPE_CHECKING, Literal -from pydantic import Field +from pydantic import Field, field_validator +from reactivex.disposable import Disposable -# HACK(kronk-nav): re-enable with the _on_odom_for_tf subscription in start(). -# from reactivex.disposable import Disposable from dimos.core.core import rpc from dimos.core.native_module import NativeModule, NativeModuleConfig from dimos.core.stream import Out @@ -59,7 +58,7 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.navigation.cmu_nav.frames import FRAME_BODY, FRAME_ODOM +from dimos.navigation.cmu_nav.frames import FRAME_BODY, FRAME_ODOM, FRAME_SENSOR from dimos.spec import perception # Human-readable enums; the C++ binary (main.cpp) maps these strings to @@ -82,13 +81,32 @@ class PointLioConfig(NativeModuleConfig): lidar_ip: str | None = Field(default_factory=lambda: os.environ.get("DIMOS_POINTLIO_LIDAR_IP")) frequency: float = 10.0 - # Odometry is published as frame_id (fixed) -> child_frame_id (moving body), - # and also broadcast on TF. The point cloud is stamped with sensor_frame_id - # (the lidar's own frame — get_body_cloud is the undistorted scan, not yet - # transformed into the body frame). - frame_id: str = FRAME_ODOM - child_frame_id: str = FRAME_BODY - sensor_frame_id: str = "mid360_link" + # Optional rigid mount correction applied to the published cloud, odometry, + # and TF before they leave the binary: row-major 4x4 (16 floats), None = + # identity. The cloud is moved p' = C @ p; the body pose is conjugated + # T' = C @ T @ inv(C) and the twist rotated by C's rotation, so the cloud and + # odometry stay mutually consistent and a forward walk stays forward. + transform: list[float] | None = None + + @field_validator("transform") + @classmethod + def _validate_transform(cls, value: list[float] | None) -> list[float] | None: + if value is not None and len(value) != 16: + raise ValueError(f"transform must be a row-major 4x4 (16 floats), got {len(value)}") + return value + + # Common-name -> real-frame map. Odometry is published as odom (fixed) -> + # body (moving) and broadcast on TF; the cloud is stamped with the sensor + # frame. A blueprint can rename any of these (e.g. send body to a dead leaf + # when a downstream shim republishes the real odom->body) without touching + # this code. These names are also handed to the C++ binary as CLI args. + frame_mapping: dict[str, str] = Field( + default_factory=lambda: { + FRAME_ODOM: FRAME_ODOM, + FRAME_BODY: FRAME_BODY, + FRAME_SENSOR: "mid360_link", + } + ) # Point-LIO internal processing rates (Hz) msr_freq: float = 50.0 @@ -178,19 +196,28 @@ class PointLio(NativeModule, perception.Lidar, perception.Odometry): @rpc def start(self) -> None: self._validate_network() + # The C++ binary stamps its cloud/odometry/TF with these frames. They live + # in frame_mapping (so a blueprint can rename any of them) and are not + # auto-emitted by to_cli_args, so hand them over explicitly. + self.config.extra_args = [ + *self.config.extra_args, + "--frame_id", + self.frame_mapping[FRAME_ODOM], + "--child_frame_id", + self.frame_mapping[FRAME_BODY], + "--sensor_frame_id", + self.frame_mapping[FRAME_SENSOR], + ] super().start() - # HACK(kronk-nav): odom->body TF is published by PointLioHack from the - # faked odometry instead, so the TF tree matches the shimmed map. Leaving - # this on would broadcast the raw tilted pose and fight the hack. - # self.register_disposable( - # Disposable(self.odometry.transport.subscribe(self._on_odom_for_tf, self.odometry)) - # ) + self.register_disposable( + Disposable(self.odometry.transport.subscribe(self._on_odom_for_tf, self.odometry)) + ) def _on_odom_for_tf(self, msg: Odometry) -> None: self.tf.publish( Transform( - frame_id=self.frame_id, - child_frame_id=self.config.child_frame_id, + frame_id=self.frame_mapping[FRAME_ODOM], + child_frame_id=self.frame_mapping[FRAME_BODY], translation=Vector3( msg.pose.position.x, msg.pose.position.y, diff --git a/dimos/hardware/sensors/lidar/pointlio/mount_correction.py b/dimos/hardware/sensors/lidar/pointlio/mount_correction.py new file mode 100644 index 0000000000..6f99398faf --- /dev/null +++ b/dimos/hardware/sensors/lidar/pointlio/mount_correction.py @@ -0,0 +1,81 @@ +# 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. + +"""Compute the rigid transform that makes a physically-tilted sensor look like a +differently-mounted one. + +``mount_correction_matrix(original_static_tf, new_static_tf)`` returns the +row-major 4x4 (flattened to 16 floats) to hand to ``PointLio``'s ``transform`` +arg. PointLio then moves every cloud point ``p' = C @ p`` and conjugates the body +pose ``T' = C @ T @ inv(C)`` before publishing, so the whole stack downstream sees +a sensor mounted per ``new_static_tf`` even though it is physically mounted per +``original_static_tf``. Usage:: + + PointLio.blueprint( + transform=mount_correction_matrix(rotated_urdf, normal_urdf), + config="no_gravity_align.yaml", + ) + +``C = inv(new_base_to_sensor) @ original_base_to_sensor`` at the shared +``sensor_frame``. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.robot.urdf_loader import UrdfLoader + + +def _transform_to_matrix(transform: Transform) -> np.ndarray: + matrix = np.eye(4) + matrix[:3, :3] = transform.rotation.to_rotation_matrix() + matrix[:3, 3] = ( + transform.translation.x, + transform.translation.y, + transform.translation.z, + ) + return matrix + + +def _base_to_frame_matrix(loader: UrdfLoader, leaf_frame: str) -> np.ndarray: + """Compose the fixed-joint chain from the model root down to ``leaf_frame``.""" + static_transforms = loader.static_transforms + chain: list[Transform] = [] + frame = leaf_frame + while frame in static_transforms: + transform = static_transforms[frame] + chain.append(transform) + frame = transform.frame_id + matrix = np.eye(4) + for transform in reversed(chain): + matrix = matrix @ _transform_to_matrix(transform) + return matrix + + +def mount_correction_matrix( + original_static_tf: Path, + new_static_tf: Path, + sensor_frame: str = "mid360_link", +) -> list[float]: + """Row-major 4x4 (16 floats) mapping the ``original`` mount to the ``new`` one.""" + original = _base_to_frame_matrix( + UrdfLoader(name="original", model_path=original_static_tf), sensor_frame + ) + new = _base_to_frame_matrix(UrdfLoader(name="new", model_path=new_static_tf), sensor_frame) + correction = np.linalg.inv(new) @ original + return [float(value) for value in correction.flatten()] diff --git a/dimos/robot/assembly/mid360_realsense_30.py b/dimos/robot/assembly/mid360_realsense_30.py index bc884a04a5..95823ae10b 100644 --- a/dimos/robot/assembly/mid360_realsense_30.py +++ b/dimos/robot/assembly/mid360_realsense_30.py @@ -58,6 +58,7 @@ 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.navigation.cmu_nav.frames import FRAME_ODOM from dimos.protocol.tf.static_tf_publisher import ( FrameSpec, StaticTfPublisher, @@ -157,7 +158,7 @@ class Mid360RealsenseRecorder(PointlioRecorder): (Mid360, "imu", "livox_imu"), ] ), - PointLio.blueprint(frame_id="world").remappings( + PointLio.blueprint(frame_mapping={FRAME_ODOM: "world"}).remappings( [ (PointLio, "lidar", "pointlio_lidar"), (PointLio, "odometry", "pointlio_odometry"), diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py index be041b4788..40322b1b11 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py @@ -38,6 +38,7 @@ from dimos.hardware.sensors.lidar.livox.module import Mid360 from dimos.hardware.sensors.lidar.pointlio.module import PointLio from dimos.hardware.sensors.lidar.virtual_mid360.recorder import Mid360PcapRecorder +from dimos.navigation.cmu_nav.frames import FRAME_ODOM from dimos.navigation.movement_manager.movement_manager import MovementManager from dimos.robot.unitree.go2.connection import GO2Connection from dimos.robot.unitree.go2.go2_mid360_recorder import Go2Mid360Recorder @@ -77,7 +78,7 @@ def _default_recording_dir() -> Path: (Mid360, "imu", "livox_imu"), ] ), - PointLio.blueprint(frame_id="world").remappings( + PointLio.blueprint(frame_mapping={FRAME_ODOM: "world"}).remappings( [ (PointLio, "lidar", "pointlio_lidar"), (PointLio, "odometry", "pointlio_odometry"), From bfaf897ce7529e05e8f7f2abe60b29888fd96c84 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 30 Jun 2026 17:29:44 -0500 Subject: [PATCH 65/71] refactor(go2): relocate mid360 URDFs, wire nav_3d mount shim + pcap recording Move the mid360 mount URDFs up to the go2 package and export their paths from config. nav_3d now builds the mount correction from those URDFs (replacing the deleted PointLioHack), records the corrected pointlio streams, and captures the raw Livox pcap. --- dimos/hardware/sensors/lidar/pointlio/hack.py | 167 ------------------ .../navigation/mid360_realsense_30.urdf | 108 ----------- .../navigation/unitree_go2_nav_3d.py | 40 ++--- dimos/robot/unitree/go2/config.py | 7 +- .../navigation => }/go2_mid360_normal.urdf | 0 .../navigation => }/go2_mid360_rotated.urdf | 0 6 files changed, 25 insertions(+), 297 deletions(-) delete mode 100644 dimos/hardware/sensors/lidar/pointlio/hack.py delete mode 100644 dimos/robot/unitree/go2/blueprints/navigation/mid360_realsense_30.urdf rename dimos/robot/unitree/go2/{blueprints/navigation => }/go2_mid360_normal.urdf (100%) rename dimos/robot/unitree/go2/{blueprints/navigation => }/go2_mid360_rotated.urdf (100%) diff --git a/dimos/hardware/sensors/lidar/pointlio/hack.py b/dimos/hardware/sensors/lidar/pointlio/hack.py deleted file mode 100644 index a2efe422de..0000000000 --- a/dimos/hardware/sensors/lidar/pointlio/hack.py +++ /dev/null @@ -1,167 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Fake a normally-mounted mid-360 from the physically-tilted one. - -The sensor is bolted to the robot at an angle (``rotated_urdf``), but the rest of -the stack is wired as if it sat in the nominal mount (``normal_urdf``). This -module rewrites both the cloud and the odometry so nothing downstream can tell -the sensor is tilted. - -The correction is a single rigid transform ``C`` derived from the two URDFs: -``C = inv(normal_base_to_sensor) @ rotated_base_to_sensor`` evaluated at the -shared ``sensor_frame``. - -* cloud: each point is moved ``p' = C @ p`` (rotated sensor frame -> normal one). -* odometry: the body pose is conjugated ``T' = C @ T @ inv(C)`` and the twist - rotated by ``C``'s rotation. That is exactly the odometry a normally-mounted - sensor would have produced, so a forward walk stays forward instead of drifting - sideways the way a bare body-frame relabel would. -""" - -from __future__ import annotations - -from collections.abc import AsyncIterator -from pathlib import Path -from typing import TYPE_CHECKING - -import numpy as np - -from dimos.core.module import Module, ModuleConfig -from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Transform import Transform -from dimos.msgs.geometry_msgs.Twist import Twist -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 -from dimos.robot.urdf_loader import UrdfLoader -from dimos.spec import perception - - -class PointLioHackConfig(ModuleConfig): - rotated_urdf: Path - normal_urdf: Path - sensor_frame: str = "mid360_link" - # odom->body TF, faked here instead of by PointLio so the TF tree matches the - # shimmed cloud/odometry. Must match PointLio's body_start_frame_id/body_frame_id. - odom_frame: str = "odom" - body_frame: str = "body" - - -def _transform_to_matrix(transform: Transform) -> np.ndarray: - matrix = np.eye(4) - matrix[:3, :3] = transform.rotation.to_rotation_matrix() - matrix[:3, 3] = ( - transform.translation.x, - transform.translation.y, - transform.translation.z, - ) - return matrix - - -def _base_to_frame_matrix(loader: UrdfLoader, leaf_frame: str) -> np.ndarray: - """Compose the fixed-joint chain from the model root down to ``leaf_frame``.""" - static_transforms = loader.static_transforms - chain: list[Transform] = [] - frame = leaf_frame - while frame in static_transforms: - transform = static_transforms[frame] - chain.append(transform) - frame = transform.frame_id - matrix = np.eye(4) - for transform in reversed(chain): - matrix = matrix @ _transform_to_matrix(transform) - return matrix - - -def _vector_to_numpy(vector: Vector3) -> np.ndarray: - return np.array([vector.x, vector.y, vector.z]) - - -class PointLioHack(Module, perception.Lidar, perception.Odometry): - config: PointLioHackConfig - - rotated_lidar: In[PointCloud2] - lidar: Out[PointCloud2] - rotated_odometry: In[Odometry] - odometry: Out[Odometry] - - async def main(self) -> AsyncIterator[None]: - rotated = _base_to_frame_matrix( - UrdfLoader(name="rotated", model_path=self.config.rotated_urdf), - self.config.sensor_frame, - ) - normal = _base_to_frame_matrix( - UrdfLoader(name="normal", model_path=self.config.normal_urdf), - self.config.sensor_frame, - ) - self._correction = np.linalg.inv(normal) @ rotated - self._correction_inv = np.linalg.inv(self._correction) - self._rotation = self._correction[:3, :3].astype(np.float32) - self._translation = self._correction[:3, 3].astype(np.float32) - yield - - async def handle_rotated_lidar(self, value: PointCloud2) -> None: - points = value.points_f32() @ self._rotation.T + self._translation - self.lidar.publish( - PointCloud2.from_numpy( - points=points, - frame_id=value.frame_id, - timestamp=value.ts, - intensities=value.intensities_f32(), - ) - ) - - async def handle_rotated_odometry(self, value: Odometry) -> None: - pose = np.eye(4) - pose[:3, :3] = value.orientation.to_rotation_matrix() - pose[:3, 3] = _vector_to_numpy(value.position) - faked = self._correction @ pose @ self._correction_inv - - rotation = self._correction[:3, :3] - linear = rotation @ _vector_to_numpy(value.linear_velocity) - angular = rotation @ _vector_to_numpy(value.angular_velocity) - - faked_position = Vector3(*(float(component) for component in faked[:3, 3])) - faked_orientation = Quaternion.from_rotation_matrix(faked[:3, :3]) - - self.odometry.publish( - Odometry( - ts=value.ts, - frame_id=value.frame_id, - child_frame_id=value.child_frame_id, - pose=Pose(faked_position, faked_orientation), - twist=Twist( - Vector3(*(float(component) for component in linear)), - Vector3(*(float(component) for component in angular)), - ), - ) - ) - - self.tf.publish( - Transform( - frame_id=self.config.odom_frame, - child_frame_id=self.config.body_frame, - translation=faked_position, - rotation=faked_orientation, - ts=value.ts, - ) - ) - - -# Verify protocol port compliance (mypy will flag missing ports) -if TYPE_CHECKING: - PointLioHack() diff --git a/dimos/robot/unitree/go2/blueprints/navigation/mid360_realsense_30.urdf b/dimos/robot/unitree/go2/blueprints/navigation/mid360_realsense_30.urdf deleted file mode 100644 index a132b2930f..0000000000 --- a/dimos/robot/unitree/go2/blueprints/navigation/mid360_realsense_30.urdf +++ /dev/null @@ -1,108 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index c40285dee4..2d1efb30f6 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -15,34 +15,32 @@ """3d navigation on Go2 with ray tracing and MLS planning""" -from pathlib import Path from typing import Any from dimos.core.coordination.blueprints import autoconnect from dimos.core.global_config import global_config -from dimos.hardware.sensors.lidar.pointlio.hack import PointLioHack from dimos.hardware.sensors.lidar.pointlio.module import PointLio +from dimos.hardware.sensors.lidar.pointlio.mount_correction import mount_correction_matrix from dimos.hardware.sensors.lidar.pointlio.recorder import PointlioRecorder +from dimos.hardware.sensors.lidar.virtual_mid360.recorder import Mid360PcapRecorder from dimos.mapping.ray_tracing.module import RayTracingVoxelMap from dimos.navigation.basic_path_follower.module import BasicPathFollower from dimos.navigation.movement_manager.movement_manager import MovementManager from dimos.navigation.nav_3d.mls_planner.goal_relay import GoalRelay from dimos.navigation.nav_3d.mls_planner.mls_planner_native import MLSPlannerNative from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import rerun_config +from dimos.robot.unitree.go2.config import mid360_rotated_urdf_path, mid360_urdf_path from dimos.robot.unitree.go2.connection import GO2Connection from dimos.robot.urdf_loader import UrdfLoader from dimos.visualization.vis_module import vis_module -_navigation_dir = Path(__file__).parent -# The physical mid-360 mount (rotated_urdf) and the "forward = forward" mount the -# rest of the stack pretends it has (normal_urdf). PointLioHack rewrites the cloud -# from the first into the second, so everything downstream uses the normal mount. -_rotated_urdf = _navigation_dir / "go2_mid360_rotated.urdf" -_normal_urdf = _navigation_dir / "go2_mid360_normal.urdf" +_mount_correction = mount_correction_matrix( + original_static_tf=mid360_rotated_urdf_path, new_static_tf=mid360_urdf_path +) # GO2Connection is a TfModule; feeding it the normal mount publishes those frames # on its static interval, matching the cloud the hack emits. -go2_mid360_model = UrdfLoader(name="go2_mid360", model_path=_normal_urdf) +go2_mid360_model = UrdfLoader(name="go2_mid360", model_path=mid360_urdf_path) voxel_size = 0.08 # Height of the head-mounted lidar above the ground while standing. @@ -104,7 +102,7 @@ def _static_robot_body(rr: Any) -> list[Any]: # frame. Anchor the robot box to pointlio's body frame instead and hide the # camera frustum that rides base_link. Use a dedicated entity path (not # world/tf/body) so the box's Transform3D doesn't collide with the live - # odom->body TF that PointLioHack logs onto world/tf/body. + # odom->body TF that PointLio logs onto world/tf/body. "static": {"world/robot_body": _static_robot_body}, "visual_override": { **rerun_config["visual_override"], @@ -133,28 +131,28 @@ def _static_robot_body(rr: Any) -> list[Any]: (GO2Connection, "odom", "odom_go2"), ] ), - # gravity_align is off (no_gravity_align.yaml) so pointlio leaves both the cloud - # and the odometry in the raw mount frame. The hack rewrites both into the normal - # mount, so everything downstream sees a normally-mounted sensor. + # gravity_align is off (no_gravity_align.yaml) so pointlio runs in the raw mount + # frame; the transform rewrites its cloud + odometry (and odom->body TF) into the + # normal mount before publishing. PointLio.blueprint( - body_frame_id="body", config="no_gravity_align.yaml", space_down_sample=False, - ).remappings( - [ - (PointLio, "lidar", "rotated_lidar"), - (PointLio, "odometry", "rotated_odometry"), - ] + transform=_mount_correction, ), - PointLioHack.blueprint(rotated_urdf=_rotated_urdf, normal_urdf=_normal_urdf), - # Record the hack's faked (normal-mount) cloud + odometry, not PointLio's raw + # Record the corrected (normal-mount) cloud + odometry, not PointLio's raw # tilted output, so a replay reproduces what the nav stack actually consumed. + # The remap only renames the wire topics to PointLio's lidar/odometry; the + # recorder's ports stay pointlio_lidar/pointlio_odometry, so that's what the + # db streams are named. PointlioRecorder.blueprint().remappings( [ (PointlioRecorder, "pointlio_lidar", "lidar"), (PointlioRecorder, "pointlio_odometry", "odometry"), ] ), + # Raw Livox UDP capture (tcpdump). Pulls the lidar IP + iface from env + # (DIMOS_MID360_LIDAR_IP / DIMOS_MID360_PCAP_IFACE); pcap lands in recordings/. + Mid360PcapRecorder.blueprint(), RayTracingVoxelMap.blueprint( voxel_size=voxel_size, emit_every=1, diff --git a/dimos/robot/unitree/go2/config.py b/dimos/robot/unitree/go2/config.py index c4a2f5be59..b1426e4604 100644 --- a/dimos/robot/unitree/go2/config.py +++ b/dimos/robot/unitree/go2/config.py @@ -34,9 +34,14 @@ def camera_info_static() -> CameraInfo: return CameraInfo.from_yaml(str(yaml_path)) +go2_dir = Path(__file__).parent +mid360_urdf_path = go2_dir / "go2_mid360_normal.urdf" +mid360_rotated_urdf_path = go2_dir / "go2_mid360_rotated.urdf" + + Go2Config = UrdfLoader( name="unitree_go2", - model_path=Path(__file__).parent / "go2.urdf", + model_path=go2_dir / "go2.urdf", ) diff --git a/dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_normal.urdf b/dimos/robot/unitree/go2/go2_mid360_normal.urdf similarity index 100% rename from dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_normal.urdf rename to dimos/robot/unitree/go2/go2_mid360_normal.urdf diff --git a/dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_rotated.urdf b/dimos/robot/unitree/go2/go2_mid360_rotated.urdf similarity index 100% rename from dimos/robot/unitree/go2/blueprints/navigation/go2_mid360_rotated.urdf rename to dimos/robot/unitree/go2/go2_mid360_rotated.urdf From 4e75fc321a0bfef791a779457bd56959d0ebe340 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 30 Jun 2026 17:44:07 -0500 Subject: [PATCH 66/71] refactor(pointlio): use DIMOS_MID360_LIDAR_IP for the lidar IP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PointLio reads the same physical Mid-360 as the rest of the rig, so collapse its DIMOS_POINTLIO_LIDAR_IP onto the shared DIMOS_MID360_LIDAR_IP — one env var for both modules instead of duplicating the same IP. --- dimos/hardware/sensors/lidar/pointlio/module.py | 6 +++--- dimos/robot/assembly/mid360_realsense_30.py | 4 ++-- .../go2/blueprints/basic/unitree_go2_mid360_record.py | 7 +++---- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/dimos/hardware/sensors/lidar/pointlio/module.py b/dimos/hardware/sensors/lidar/pointlio/module.py index 25d2659e4c..4b95470d11 100644 --- a/dimos/hardware/sensors/lidar/pointlio/module.py +++ b/dimos/hardware/sensors/lidar/pointlio/module.py @@ -76,9 +76,9 @@ class PointLioConfig(NativeModuleConfig): executable: str = "result/bin/pointlio_native" build_command: str | None = "nix build .#pointlio_native" # lidar_ip required; host_ip optional (auto-derived from lidar_ip's subnet). - # Both fall back to DIMOS_POINTLIO_LIDAR_IP / DIMOS_POINTLIO_HOST_IP. + # Fall back to DIMOS_MID360_LIDAR_IP / DIMOS_POINTLIO_HOST_IP. host_ip: str | None = Field(default_factory=lambda: os.environ.get("DIMOS_POINTLIO_HOST_IP")) - lidar_ip: str | None = Field(default_factory=lambda: os.environ.get("DIMOS_POINTLIO_LIDAR_IP")) + lidar_ip: str | None = Field(default_factory=lambda: os.environ.get("DIMOS_MID360_LIDAR_IP")) frequency: float = 10.0 # Optional rigid mount correction applied to the published cloud, odometry, @@ -244,7 +244,7 @@ def _validate_network(self) -> None: if not lidar_ip: raise RuntimeError( "PointLio: lidar_ip not set — it's network-specific. Set it in the config " - "or via the DIMOS_POINTLIO_LIDAR_IP env var." + "or via the DIMOS_MID360_LIDAR_IP env var." ) # host_ip optional: derive the local NIC on lidar_ip's /24 when unset or # not one of our IPs (shared with the Mid360 driver). diff --git a/dimos/robot/assembly/mid360_realsense_30.py b/dimos/robot/assembly/mid360_realsense_30.py index 95823ae10b..f0415687c2 100644 --- a/dimos/robot/assembly/mid360_realsense_30.py +++ b/dimos/robot/assembly/mid360_realsense_30.py @@ -25,9 +25,9 @@ captures a raw .pcap of the Mid-360 UDP stream). The lidar IPs come from each module's own config (``DIMOS_MID360_LIDAR_IP`` for the -Mid-360 / pcap capture, ``DIMOS_POINTLIO_LIDAR_IP`` for Point-LIO):: +Mid-360 / pcap capture and Point-LIO both read ``DIMOS_MID360_LIDAR_IP``):: - export DIMOS_MID360_LIDAR_IP=192.168.1.155 DIMOS_POINTLIO_LIDAR_IP=192.168.1.155 + export DIMOS_MID360_LIDAR_IP=192.168.1.155 dimos run mid360-realsense-record # db only dimos run mid360-realsense-record-with-pcap # db + raw pcap diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py index 40322b1b11..9517992504 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py @@ -20,11 +20,10 @@ published continuously onto tf so they're captured in the recording. Raw Livox capture is opt-in: set ``RECORD_PCAP=1`` to also record a .pcap of the Mid-360 UDP stream. -The lidar IPs come from each module's own config (``DIMOS_MID360_LIDAR_IP`` for the -Mid-360 / pcap capture, ``DIMOS_POINTLIO_LIDAR_IP`` for Point-LIO). Run it for a -timestamped ``recordings/`` folder:: +The lidar IP comes from ``DIMOS_MID360_LIDAR_IP`` (shared by the Mid-360 / pcap +capture and Point-LIO). Run it for a timestamped ``recordings/`` folder:: - export DIMOS_MID360_LIDAR_IP=192.168.1.171 DIMOS_POINTLIO_LIDAR_IP=192.168.1.171 + export DIMOS_MID360_LIDAR_IP=192.168.1.171 uv run python dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py """ From 27bee579daefd59f7029570d7d2bb69aa0498953 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 30 Jun 2026 17:46:52 -0500 Subject: [PATCH 67/71] feat(nav_3d): toggle the rotated-mid360 mount correction, off by default The rotated-mount un-rotation assumed the sensor was always physically rotated. Gate it behind USE_ROTATED_MID360_MOUNT (default off) so the blueprint runs the normal mount with an identity transform unless the rotated rig is in use. --- .../navigation/unitree_go2_nav_3d.py | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 24edaaae00..a816e8590d 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -34,8 +34,17 @@ from dimos.robot.urdf_loader import UrdfLoader from dimos.visualization.vis_module import vis_module -_mount_correction = mount_correction_matrix( - original_static_tf=mid360_rotated_urdf_path, new_static_tf=mid360_urdf_path +# Flip on only when the Mid-360 is physically mounted rotated; this un-rotates its +# Point-LIO cloud + odometry back into the normal mount frame. Off = sensor mounted +# normally, no correction (identity). +USE_ROTATED_MID360_MOUNT = False + +_mount_correction = ( + mount_correction_matrix( + original_static_tf=mid360_rotated_urdf_path, new_static_tf=mid360_urdf_path + ) + if USE_ROTATED_MID360_MOUNT + else None ) # GO2Connection is a TfModule; feeding it the normal mount publishes those frames @@ -131,9 +140,10 @@ def _static_robot_body(rr: Any) -> list[Any]: (GO2Connection, "odom", "odom_go2"), ] ), - # gravity_align is off so pointlio runs in the raw mount frame; the transform - # rewrites its cloud + odometry (and odom->body TF) into the normal mount before - # publishing. auto_build recompiles the native binary for the transform support. + # gravity_align is off so pointlio runs in the raw mount frame. When the rotated + # mount is enabled, transform rewrites its cloud + odometry (and odom->body TF) + # into the normal mount; otherwise transform is identity. auto_build recompiles + # the native binary for the transform support. PointLio.blueprint( gravity_align=False, space_down_sample=False, From baed84912e02dc4403fc089fb949d6ea1ecd88b1 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Tue, 30 Jun 2026 18:05:04 -0500 Subject: [PATCH 68/71] fix(go2): drop dead Thread-based camera_info publish in connection The camera block referenced onimage/Thread/publish_camera_info, none of which exist after camera_info publishing moved to the _on_static_publish hook. Subscribe the on_image closure and drop the stale thread wiring so mypy passes. --- dimos/robot/unitree/go2/connection.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 0465de0b01..20a25187a2 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -250,12 +250,7 @@ def on_image(image: Image) -> None: self.register_disposable(Disposable(self.cmd_vel.subscribe(self.move))) if self.config.camera: - self.register_disposable(self.connection.video_stream().subscribe(onimage)) - self._camera_info_thread = Thread( - target=self.publish_camera_info, - daemon=True, - ) - self._camera_info_thread.start() + self.register_disposable(self.connection.video_stream().subscribe(on_image)) if self.config.motion_mode and isinstance(self.connection, UnitreeWebRTCConnection): self.connection.set_motion_mode(self.config.motion_mode) From 6eb40df48f250b272c4752732aef36b7801ae3b0 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 1 Jul 2026 05:36:04 -0700 Subject: [PATCH 69/71] feat(nav_3d): split rotated mount into unitree-go2-nav-3d-rotated blueprint Replace the USE_ROTATED_MID360_MOUNT constant with two blueprints from a shared _nav_3d() factory: unitree-go2-nav-3d (normal mount) and unitree-go2-nav-3d-rotated (applies the mount_correction transform). --- dimos/robot/all_blueprints.py | 1 + .../navigation/unitree_go2_nav_3d.py | 160 ++++++++++-------- 2 files changed, 86 insertions(+), 75 deletions(-) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 61220a0fae..44cefcf859 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -121,6 +121,7 @@ "unitree-go2-memory": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_memory", "unitree-go2-mid360-record": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_mid360_record:unitree_go2_mid360_record", "unitree-go2-nav-3d": "dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_nav_3d:unitree_go2_nav_3d", + "unitree-go2-nav-3d-rotated": "dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_nav_3d:unitree_go2_nav_3d_rotated", "unitree-go2-relocalization": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_relocalization", "unitree-go2-ros": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2_ros:unitree_go2_ros", "unitree-go2-security": "dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_security:unitree_go2_security", diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index a816e8590d..120a6ae218 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -17,7 +17,7 @@ from typing import Any -from dimos.core.coordination.blueprints import autoconnect +from dimos.core.coordination.blueprints import Blueprint, autoconnect from dimos.core.global_config import global_config from dimos.hardware.sensors.lidar.pointlio.module import PointLio from dimos.hardware.sensors.lidar.pointlio.mount_correction import mount_correction_matrix @@ -34,21 +34,15 @@ from dimos.robot.urdf_loader import UrdfLoader from dimos.visualization.vis_module import vis_module -# Flip on only when the Mid-360 is physically mounted rotated; this un-rotates its -# Point-LIO cloud + odometry back into the normal mount frame. Off = sensor mounted -# normally, no correction (identity). -USE_ROTATED_MID360_MOUNT = False - -_mount_correction = ( - mount_correction_matrix( - original_static_tf=mid360_rotated_urdf_path, new_static_tf=mid360_urdf_path - ) - if USE_ROTATED_MID360_MOUNT - else None +# Correction that un-rotates a physically-rotated Mid-360's Point-LIO cloud + +# odometry back into the normal mount frame. Used by the `_rotated` blueprint; the +# plain blueprint passes None (identity — sensor mounted normally). +_rotated_mount_correction = mount_correction_matrix( + original_static_tf=mid360_rotated_urdf_path, new_static_tf=mid360_urdf_path ) # GO2Connection is a TfModule; feeding it the normal mount publishes those frames -# on its static interval, matching the cloud the hack emits. +# on its static interval, matching the corrected cloud both blueprints emit. go2_mid360_model = UrdfLoader(name="go2_mid360", model_path=mid360_urdf_path) voxel_size = 0.08 @@ -126,65 +120,81 @@ def _static_robot_body(rr: Any) -> list[Any]: }, } -unitree_go2_nav_3d = autoconnect( - vis_module(viewer_backend=global_config.viewer, rerun_config=_nav_rerun_config), - # "mcf" for stair traversal - GO2Connection.blueprint( - static_transforms=dict(go2_mid360_model.static_transforms), - lidar=False, - camera=False, - motion_mode="mcf", - ).remappings( - [ - (GO2Connection, "lidar", "lidar_l1"), - (GO2Connection, "odom", "odom_go2"), - ] - ), - # gravity_align is off so pointlio runs in the raw mount frame. When the rotated - # mount is enabled, transform rewrites its cloud + odometry (and odom->body TF) - # into the normal mount; otherwise transform is identity. auto_build recompiles - # the native binary for the transform support. - PointLio.blueprint( - gravity_align=False, - space_down_sample=False, - auto_build=True, - transform=_mount_correction, - ), - # Record the corrected (normal-mount) cloud + odometry, not PointLio's raw - # tilted output, so a replay reproduces what the nav stack actually consumed. - # The remap only renames the wire topics to PointLio's lidar/odometry; the - # recorder's ports stay pointlio_lidar/pointlio_odometry, so that's what the - # db streams are named. - PointlioRecorder.blueprint().remappings( - [ - (PointlioRecorder, "pointlio_lidar", "lidar"), - (PointlioRecorder, "pointlio_odometry", "odometry"), - ] - ), - # Raw Livox UDP capture (tcpdump). Pulls the lidar IP + iface from env - # (DIMOS_MID360_LIDAR_IP / DIMOS_MID360_PCAP_IFACE); pcap lands in recordings/. - Mid360PcapRecorder.blueprint(), - RayTracingVoxelMap.blueprint( - voxel_size=voxel_size, - emit_every=1, - global_emit_every=50, - max_health=10, - graze_cos=0.85, - ), - # global_map is remapped off so the planner runs purely on the - # incremental local_map + region_bounds pair. - MLSPlannerNative.blueprint( - world_frame="odom", - voxel_size=voxel_size, - robot_height=go2_lidar_height, - wall_clearance_m=0.2, - wall_buffer_m=0.75, - wall_buffer_weight=100.0, - step_threshold_m=0.16, - step_penalty_weight=1.0, - viz_publish_hz=2.0, - ).remappings([(MLSPlannerNative, "global_map", "global_map_unused")]), - GoalRelay.blueprint(), - BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6), - MovementManager.blueprint(), -).global_config(n_workers=10, robot_model="unitree_go2", obstacle_avoidance=False) + +def _nav_3d(mount_correction: list[float] | None) -> Blueprint: + """Go2 3D nav + recording stack. ``mount_correction`` un-rotates a physically + rotated Mid-360 back into the normal mount frame; None = normal mount (no + correction).""" + return autoconnect( + vis_module(viewer_backend=global_config.viewer, rerun_config=_nav_rerun_config), + # "mcf" for stair traversal + GO2Connection.blueprint( + static_transforms=dict(go2_mid360_model.static_transforms), + lidar=False, + camera=False, + motion_mode="mcf", + ).remappings( + [ + (GO2Connection, "lidar", "lidar_l1"), + (GO2Connection, "odom", "odom_go2"), + ] + ), + # gravity_align is off so pointlio runs in the raw mount frame. For the + # rotated blueprint, transform rewrites its cloud + odometry (and odom->body + # TF) into the normal mount; otherwise transform is None (identity). + # auto_build recompiles the native binary for the transform support. + PointLio.blueprint( + gravity_align=False, + space_down_sample=False, + auto_build=True, + transform=mount_correction, + ), + # Record the corrected (normal-mount) cloud + odometry, not PointLio's raw + # tilted output, so a replay reproduces what the nav stack actually consumed. + # The remap only renames the wire topics to PointLio's lidar/odometry; the + # recorder's ports stay pointlio_lidar/pointlio_odometry, so that's what the + # db streams are named. + PointlioRecorder.blueprint().remappings( + [ + (PointlioRecorder, "pointlio_lidar", "lidar"), + (PointlioRecorder, "pointlio_odometry", "odometry"), + ] + ), + # Raw Livox UDP capture (tcpdump). Pulls the lidar IP + iface from env + # (DIMOS_MID360_LIDAR_IP / DIMOS_MID360_PCAP_IFACE); pcap lands in recordings/. + Mid360PcapRecorder.blueprint(), + RayTracingVoxelMap.blueprint( + voxel_size=voxel_size, + emit_every=1, + global_emit_every=50, + max_health=10, + graze_cos=0.85, + ), + # global_map is remapped off so the planner runs purely on the + # incremental local_map + region_bounds pair. + MLSPlannerNative.blueprint( + world_frame="odom", + voxel_size=voxel_size, + robot_height=go2_lidar_height, + wall_clearance_m=0.2, + wall_buffer_m=0.75, + wall_buffer_weight=100.0, + step_threshold_m=0.16, + step_penalty_weight=1.0, + viz_publish_hz=2.0, + ).remappings([(MLSPlannerNative, "global_map", "global_map_unused")]), + GoalRelay.blueprint(), + BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6), + MovementManager.blueprint(), + ) + + +# Standard mount (level, forward-facing): no correction. +unitree_go2_nav_3d = _nav_3d(None).global_config( + n_workers=10, robot_model="unitree_go2", obstacle_avoidance=False +) + +# Rotated mount: un-rotate the cloud + odometry back into the normal mount frame. +unitree_go2_nav_3d_rotated = _nav_3d(_rotated_mount_correction).global_config( + n_workers=10, robot_model="unitree_go2", obstacle_avoidance=False +) From 5bb48a0997d8b20b3da2cb06e9dde16494feeb46 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 1 Jul 2026 06:24:57 -0700 Subject: [PATCH 70/71] =?UTF-8?q?feat(memory2):=20db=5Fto=5Frrd=20?= =?UTF-8?q?=E2=80=94=20render=20a=20recording's=20clouds,=20odoms=20+=20TF?= =?UTF-8?q?=20to=20rerun?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Standalone tool: loads the TF tree if present, then logs every PointCloud2 and Odometry stream (clouds placed via TF frame_id when TF exists). --seconds bounds the window from the recording start. --- dimos/memory2/utils/db_to_rrd.py | 154 +++++++++++++++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 dimos/memory2/utils/db_to_rrd.py diff --git a/dimos/memory2/utils/db_to_rrd.py b/dimos/memory2/utils/db_to_rrd.py new file mode 100644 index 0000000000..54fa1ec3bd --- /dev/null +++ b/dimos/memory2/utils/db_to_rrd.py @@ -0,0 +1,154 @@ +# 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. + +"""Render a memory2 recording (.db) into rerun. + +Loads the TF tree if the recording has one, then logs every PointCloud2 and +Odometry stream. Point clouds are placed via TF (by their ``frame_id``) when a +TF stream is present, otherwise logged as-is. ``--seconds`` bounds the window +from the start of the recording. +""" + +from __future__ import annotations + +from pathlib import Path +import shutil +import subprocess + +import rerun as rr +import typer + +from dimos.memory2.store.sqlite import SqliteStore +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.utils.data import resolve_named_path +from dimos.visualization.rerun.init import rerun_init + +TIMELINE = "time" + + +def _open_viewer(rrd: str) -> None: + exe = shutil.which("rerun") + if exe: + subprocess.Popen([exe, rrd]) + print(f"opening {rrd} in rerun") + else: + print(f"rerun viewer not found on PATH; open manually:\n rerun {rrd}") + + +def _classify(store: SqliteStore) -> tuple[list[str], list[str], list[str], float | None]: + """Sort streams into (clouds, odoms, tfs) by their first observation's type, + and return the earliest timestamp across them (the timeline anchor).""" + clouds: list[str] = [] + odoms: list[str] = [] + tfs: list[str] = [] + t0: float | None = None + for name in store.list_streams(): + try: + first = store.streams[name].first() + except LookupError: + continue + data = first.data + if isinstance(data, PointCloud2): + clouds.append(name) + elif isinstance(data, Odometry): + odoms.append(name) + elif isinstance(data, TFMessage): + tfs.append(name) + else: + continue + t0 = first.ts if t0 is None else min(t0, first.ts) + return clouds, odoms, tfs, t0 + + +def main( + dataset: str = typer.Argument(..., help="Recording .db: bare name (cwd or data/) or path"), + out: Path | None = typer.Option( + None, "--out", help="Output .rrd path (default: next to the .db)" + ), + seconds: float | None = typer.Option( + None, "--seconds", help="Only render the first N seconds of the recording" + ), + no_gui: bool = typer.Option(False, "--no-gui", help="Write the .rrd but don't open the viewer"), +) -> None: + db_path = resolve_named_path(dataset, ".db") + store = SqliteStore(path=str(db_path), must_exist=True) + with store: + clouds, odoms, tfs, t0 = _classify(store) + if t0 is None: + print("no PointCloud2 / Odometry / TF streams found in this recording") + raise typer.Exit(1) + has_tf = bool(tfs) + print(f"clouds={clouds}") + print(f"odoms={odoms}") + print(f"tf={tfs or '(none)'}") + + rerun_init("db_to_rrd") + rrd_path = str(out) if out is not None else str(db_path.with_suffix(".rrd")) + rr.save(rrd_path) + register_colormap_annotation("turbo") + + def in_window(ts: float) -> bool: + return seconds is None or ts - t0 <= seconds + + # TF first so the tf graph exists before framed clouds reference it. + for name in tfs: + for obs in store.streams[name]: + if not in_window(obs.ts): + break + if obs.data is None: + continue + rr.set_time(TIMELINE, duration=obs.ts - t0) + for path, archetype in obs.data.to_rerun(): + rr.log(path, archetype) + + for name in clouds: + entity = f"world/clouds/{name}" + for obs in store.streams[name]: + if not in_window(obs.ts): + break + cloud = obs.data + if cloud is None: + continue + rr.set_time(TIMELINE, duration=obs.ts - t0) + # With a TF tree, place the cloud in its own frame; rerun resolves + # it to world. Without TF, log it in whatever frame it was stored. + if has_tf and cloud.frame_id: + rr.log(entity, rr.Transform3D(parent_frame=f"tf#/{cloud.frame_id}")) + rr.log(entity, cloud.to_rerun(mode="points")) + + for name in odoms: + entity = f"world/odom/{name}" + trail: list[list[float]] = [] + for obs in store.streams[name]: + if not in_window(obs.ts): + break + odom = obs.data + if odom is None: + continue + rr.set_time(TIMELINE, duration=obs.ts - t0) + rr.log(entity, odom.to_rerun()) + trail.append([odom.x, odom.y, odom.z]) + if trail: + rr.log(f"{entity}/trail", rr.LineStrips3D([trail]), static=True) + + rr.rerun_shutdown() + print(f"wrote {rrd_path}") + if not no_gui: + _open_viewer(rrd_path) + + +if __name__ == "__main__": + typer.run(main) From d0da4befa660f6c26327ac1cf27b68ff01be0aa9 Mon Sep 17 00:00:00 2001 From: Jeff Hykin Date: Wed, 1 Jul 2026 06:34:54 -0700 Subject: [PATCH 71/71] fix(db_to_rrd): draw clouds in their stored frame, don't reframe via TF Point-LIO clouds are already registered in their world frame, and the recorded TF tree usually lacks the sensor frames, so referencing tf#/ left clouds unplaced. Log clouds raw. Also move each odom trail off the posed entity (was double-transformed by the pose) onto world/trails/. --- dimos/memory2/utils/db_to_rrd.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/dimos/memory2/utils/db_to_rrd.py b/dimos/memory2/utils/db_to_rrd.py index 54fa1ec3bd..5845a6250d 100644 --- a/dimos/memory2/utils/db_to_rrd.py +++ b/dimos/memory2/utils/db_to_rrd.py @@ -15,9 +15,10 @@ """Render a memory2 recording (.db) into rerun. Loads the TF tree if the recording has one, then logs every PointCloud2 and -Odometry stream. Point clouds are placed via TF (by their ``frame_id``) when a -TF stream is present, otherwise logged as-is. ``--seconds`` bounds the window -from the start of the recording. +Odometry stream. Clouds are drawn in their stored coordinates (Point-LIO's cloud +is already registered into its world frame); each odom stream gets a moving pose +frame plus its full path as a trail. ``--seconds`` bounds the window from the +start of the recording. """ from __future__ import annotations @@ -90,7 +91,6 @@ def main( if t0 is None: print("no PointCloud2 / Odometry / TF streams found in this recording") raise typer.Exit(1) - has_tf = bool(tfs) print(f"clouds={clouds}") print(f"odoms={odoms}") print(f"tf={tfs or '(none)'}") @@ -114,6 +114,9 @@ def in_window(ts: float) -> bool: for path, archetype in obs.data.to_rerun(): rr.log(path, archetype) + # Clouds are logged in their stored coordinates (Point-LIO's cloud is + # already registered into its world frame). Each stream is its own entity + # so you can toggle them independently in the viewer. for name in clouds: entity = f"world/clouds/{name}" for obs in store.streams[name]: @@ -123,14 +126,9 @@ def in_window(ts: float) -> bool: if cloud is None: continue rr.set_time(TIMELINE, duration=obs.ts - t0) - # With a TF tree, place the cloud in its own frame; rerun resolves - # it to world. Without TF, log it in whatever frame it was stored. - if has_tf and cloud.frame_id: - rr.log(entity, rr.Transform3D(parent_frame=f"tf#/{cloud.frame_id}")) rr.log(entity, cloud.to_rerun(mode="points")) for name in odoms: - entity = f"world/odom/{name}" trail: list[list[float]] = [] for obs in store.streams[name]: if not in_window(obs.ts): @@ -139,10 +137,13 @@ def in_window(ts: float) -> bool: if odom is None: continue rr.set_time(TIMELINE, duration=obs.ts - t0) - rr.log(entity, odom.to_rerun()) + # Moving pose frame (has its own transform)... + rr.log(f"world/poses/{name}", odom.to_rerun()) trail.append([odom.x, odom.y, odom.z]) + # ...and the full path as a static line in world coords (kept off the + # posed entity so it isn't transformed by the pose). if trail: - rr.log(f"{entity}/trail", rr.LineStrips3D([trail]), static=True) + rr.log(f"world/trails/{name}", rr.LineStrips3D([trail]), static=True) rr.rerun_shutdown() print(f"wrote {rrd_path}")