From d180a840816b2e46dc7f260f740f73847f9fbe9c Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 12:41:01 -0700 Subject: [PATCH 1/9] 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 d64f53fca82bab3786cdb4b50ff707d2fe8f7677 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 12:53:29 -0700 Subject: [PATCH 2/9] 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 157ee74003f2a595278dc3378242f18add477dc0 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 12:59:26 -0700 Subject: [PATCH 3/9] 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 35e82a14304f049f77138f101aa151ca59ca4c24 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 13:45:36 -0700 Subject: [PATCH 4/9] 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 e158d6c5179d29c2c68caec7a12d206d871dce19 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 15:40:33 -0700 Subject: [PATCH 5/9] 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 e78c81a49a53a70816e2f4e34b58deccee626815 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 17:17:55 -0700 Subject: [PATCH 6/9] 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 ff50001c45a9a385543cbb1c4ea55ca961a3d28e Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 18:12:39 -0700 Subject: [PATCH 7/9] 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 8afc1f16bc69440a32a77542f21c6c4a29bfb568 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 19:57:47 -0700 Subject: [PATCH 8/9] 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 eedf4e3ddcf49d8532ca9b0d53ef70f0ca07748d Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Mon, 22 Jun 2026 20:32:12 -0700 Subject: [PATCH 9/9] 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 {