diff --git a/dimos/core/global_config.py b/dimos/core/global_config.py index 6ea946ebea..d52434d4fd 100644 --- a/dimos/core/global_config.py +++ b/dimos/core/global_config.py @@ -73,6 +73,9 @@ class GlobalConfig(BaseSettings): env_file=".env", env_file_encoding="utf-8", extra="ignore", + # Coerce on assignment so string overrides from the CLI (-o g.field=val) + # and env become real typed values instead of raw strings. + validate_assignment=True, ) def update(self, **kwargs: object) -> None: diff --git a/dimos/hardware/sensors/lidar/fastlio2/cpp/config/mid360.json b/dimos/hardware/sensors/lidar/fastlio2/cpp/config/mid360.json index ff6cc6dbf6..806d110ee2 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/cpp/config/mid360.json +++ b/dimos/hardware/sensors/lidar/fastlio2/cpp/config/mid360.json @@ -11,7 +11,7 @@ "blind": 1 }, "mapping": { - "acc_cov": 0.1, + "acc_cov": 1.0, "gyr_cov": 0.1, "b_acc_cov": 0.0001, "b_gyr_cov": 0.0001, 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.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..ba2ca2a0fe 100644 --- a/dimos/mapping/ray_tracing/rust/src/main.rs +++ b/dimos/mapping/ray_tracing/rust/src/main.rs @@ -6,8 +6,9 @@ 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}; @@ -26,11 +27,20 @@ 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)>, + frame_count: u32, + batch_points: Vec<(f32, f32, f32)>, + batch_origins: Vec<(f32, f32, f32)>, } impl RayTracingVoxelMap { @@ -67,29 +77,84 @@ impl RayTracingVoxelMap { 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, + &msg.header.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: msg.header.frame_id.clone(), + }, + 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, @@ -97,13 +162,6 @@ impl RayTracingVoxelMap { &msg.header.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 +238,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, - }; + } +} - (global_cloud, local_cloud) +/// 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) +} + +/// 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 +383,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 +401,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 +420,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 9487c85bbd..96d4346577 100644 --- a/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs +++ b/dimos/mapping/ray_tracing/rust/src/voxel_ray_tracer.rs @@ -2,7 +2,9 @@ // SPDX-License-Identifier: Apache-2.0 use ahash::{AHashMap, AHashSet}; +use arrayvec::ArrayVec; use nalgebra::{Matrix3, Vector3}; +use rayon::prelude::*; use serde::Deserialize; use validator::{Validate, ValidationError}; @@ -33,6 +35,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> { @@ -84,7 +97,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; @@ -94,6 +107,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; @@ -189,15 +203,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 { @@ -205,6 +221,7 @@ fn pooled_normal( let Some(v) = voxels.get(&nk) else { continue; }; + recency = recency.max(v.last_hit); if v.num_pts == 0 { continue; } @@ -225,12 +242,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()); @@ -261,25 +278,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 @@ -302,14 +303,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 { @@ -323,16 +321,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 @@ -360,6 +354,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, @@ -424,35 +460,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 { @@ -620,17 +667,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)); + } } } } @@ -650,6 +690,9 @@ mod tests { max_health: 1, graze_cos: 0.5, recency_window: 60, + emit_every: 1, + global_emit_every: 1, + region_percentile: 95.0, } } @@ -697,6 +740,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(); @@ -809,6 +879,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; @@ -961,6 +1034,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 @@ -1032,6 +1108,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. @@ -1091,6 +1170,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. @@ -1219,6 +1301,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..44ccb82ad2 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,33 @@ 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_origin() -> None: + margin = 0.2 + 0.1 + obs = _obs(_ring((2.0, 3.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.0 - margin) + assert z_max == pytest.approx(0.0 + 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..e4afbd9287 100644 --- a/dimos/mapping/ray_tracing/transformer.py +++ b/dimos/mapping/ray_tracing/transformer.py @@ -16,18 +16,24 @@ 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.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.""" @@ -38,6 +44,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 +53,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 +107,26 @@ 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)) + pts = obs.data.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 e362c6697c..ea5154f54e 100644 --- a/dimos/mapping/ray_tracing/voxel_map.py +++ b/dimos/mapping/ray_tracing/voxel_map.py @@ -17,11 +17,11 @@ from __future__ import annotations try: - from dimos_voxel_ray_tracing import VoxelRayMapper + from dimos_voxel_ray_tracing import VoxelRayMapper, local_bounds except ImportError as e: raise ImportError( "dimos_voxel_ray_tracing is not built. Run: " "uv run maturin develop --uv -m dimos/mapping/ray_tracing/rust/Cargo.toml" ) from e -__all__ = ["VoxelRayMapper"] +__all__ = ["VoxelRayMapper", "local_bounds"] 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/dimos/navigation/basic_path_follower/module.py b/dimos/navigation/basic_path_follower/module.py new file mode 100644 index 0000000000..c0cbb44ed9 --- /dev/null +++ b/dimos/navigation/basic_path_follower/module.py @@ -0,0 +1,256 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import math +from threading import Event, RLock, Thread +import time +from typing import Any + +from dimos_lcm.std_msgs import Bool # type: ignore[import-untyped] +import numpy as np +from reactivex.disposable import Disposable + +from dimos.constants import DEFAULT_THREAD_JOIN_TIMEOUT +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.geometry_msgs.Twist import Twist +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.msgs.nav_msgs.Path import Path +from dimos.utils.logging_config import setup_logger +from dimos.utils.trigonometry import angle_diff + +logger = setup_logger() + + +class BasicPathFollowerConfig(ModuleConfig): + speed: float = 0.5 + control_frequency: float = 10.0 + goal_tolerance: float = 0.3 + # Lookahead grows with speed (time_s * speed) within these bounds, so the + # robot tracks tightly when slow and stays stable when fast. + lookahead_time_s: float = 1.5 + min_lookahead_m: float = 0.4 + max_lookahead_m: float = 1.5 + heading_gain: float = 1.5 + max_angular: float = 1.0 + + +def lookahead_distance(speed: float, time_s: float, lo: float, hi: float) -> float: + return min(hi, max(lo, time_s * speed)) + + +class BasicPathFollower(Module): + """Follow a planned path by chasing a lookahead point with P-controlled heading. + + Consumes world-frame paths and odometry. Publishes nav_cmd_vel until the + last waypoint is within goal tolerance, then publishes goal_reached. + A stop_movement message cancels the current path. Empty paths are ignored + so continuous replanning does not stutter the follower. + """ + + config: BasicPathFollowerConfig + + path: In[Path] + odometry: In[Odometry] + stop_movement: In[Bool] + + nav_cmd_vel: Out[Twist] + goal_reached: Out[Bool] + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._lock = RLock() + self._current_odom: PoseStamped | None = None + self._waypoints: np.ndarray | None = None + self._stop_event = Event() + self._thread: Thread | None = None + self._path_count = 0 + self._stats_last = 0.0 + self._last_path_t = 0.0 + self._max_gap = 0.0 + self._cmd_log_last = 0.0 + + @rpc + def start(self) -> None: + super().start() + self.register_disposable(Disposable(self.odometry.subscribe(self._on_odometry))) + self.register_disposable(Disposable(self.path.subscribe(self._on_path))) + if self.stop_movement.transport is not None: + self.register_disposable(Disposable(self.stop_movement.subscribe(self._on_stop))) + self._stats_last = time.perf_counter() + self._thread = Thread(target=self._follow, daemon=True) + self._thread.start() + + @rpc + def stop(self) -> None: + self._stop_event.set() + if self._thread is not None: + self._thread.join(timeout=DEFAULT_THREAD_JOIN_TIMEOUT) + self.nav_cmd_vel.publish(Twist()) + super().stop() + + def _on_odometry(self, msg: Odometry) -> None: + with self._lock: + self._current_odom = msg.to_pose_stamped() + + def _on_path(self, path: Path) -> None: + # The planner clears its plan with an empty path on every start-pose + # change and on a failed replan. Under continuous replanning that is not + # a stop, so keep the last path. Stops arrive via stop_movement. + if len(path.poses) == 0: + return + waypoints = np.array([[p.position.x, p.position.y] for p in path.poses]) + now = time.perf_counter() + with self._lock: + if self._last_path_t: + self._max_gap = max(self._max_gap, now - self._last_path_t) + self._last_path_t = now + self._waypoints = waypoints + self._path_count += 1 + + def _on_stop(self, msg: Bool) -> None: + if msg.data: + with self._lock: + self._waypoints = None + self.nav_cmd_vel.publish(Twist()) + + def _follow(self) -> None: + period = 1.0 / self.config.control_frequency + while not self._stop_event.is_set(): + start_time = time.perf_counter() + with self._lock: + odom = self._current_odom + waypoints = self._waypoints + if odom is not None and waypoints is not None: + self._step(odom, waypoints) + self._maybe_log_stats(odom, waypoints) + elapsed = time.perf_counter() - start_time + self._stop_event.wait(max(0.0, period - elapsed)) + + # DEBUG: replanning telemetry, remove before merge + def _maybe_log_stats(self, odom: PoseStamped | None, waypoints: np.ndarray | None) -> None: + now = time.perf_counter() + elapsed = now - self._stats_last + if elapsed < 2.0: + return + self._stats_last = now + with self._lock: + count = self._path_count + gap = self._max_gap + self._path_count = 0 + self._max_gap = 0.0 + if count == 0: + return + rate = count / elapsed + lag = float("nan") + if odom is not None and waypoints is not None: + position = np.array([odom.position.x, odom.position.y]) + lag = float(np.linalg.norm(waypoints[0] - position)) + logger.debug( + "path follower stats", + replan_hz=round(rate, 1), + max_gap_ms=round(gap * 1000), + path_lag_m=round(lag, 2), + ) + + def _step(self, odom: PoseStamped, waypoints: np.ndarray) -> None: + position = np.array([odom.position.x, odom.position.y]) + if float(np.linalg.norm(waypoints[-1] - position)) < self.config.goal_tolerance: + self.nav_cmd_vel.publish(Twist()) + with self._lock: + if self._waypoints is waypoints: + self._waypoints = None + self.goal_reached.publish(Bool(True)) + logger.info("Goal reached") + return + + target = self._lookahead_point(waypoints, position) + yaw_error = angle_diff( + math.atan2(target[1] - position[1], target[0] - position[0]), + odom.orientation.euler[2], + ) + + angular = max( + -self.config.max_angular, + min(self.config.max_angular, self.config.heading_gain * yaw_error), + ) + # Taper forward speed with heading error instead of a hard stop: the + # robot decelerates into a turn and pivots in place only near 90 degrees + # (cos -> 0). Slower speed at sharp corners means a tighter turn radius, + # so it tracks the safe path instead of stopping and lurching. + linear = self.config.speed * max(0.0, math.cos(yaw_error)) + + # DEBUG: command trace, remove before merge + now = time.perf_counter() + if now - self._cmd_log_last >= 0.2: + self._cmd_log_last = now + logger.debug( + "follower cmd", + lin=round(linear, 3), + ang=round(angular, 3), + yaw_err_deg=round(math.degrees(yaw_error), 1), + target_dx=round(float(target[0] - position[0]), 2), + target_dy=round(float(target[1] - position[1]), 2), + goal_dist=round(float(np.linalg.norm(waypoints[-1] - position)), 2), + ) + + self.nav_cmd_vel.publish(Twist(Vector3(linear, 0, 0), Vector3(0, 0, angular))) + + def _lookahead_point(self, waypoints: np.ndarray, position: np.ndarray) -> np.ndarray: + if len(waypoints) == 1: + return np.asarray(waypoints[0]) + + # Project onto the path, then march the lookahead along it and + # interpolate. Returning an actual waypoint would let the target jump to + # a distant vertex and cut the corner. + seg_idx, start = self._project_onto_path(waypoints, position) + remaining = lookahead_distance( + self.config.speed, + self.config.lookahead_time_s, + self.config.min_lookahead_m, + self.config.max_lookahead_m, + ) + for i in range(seg_idx, len(waypoints) - 1): + end = waypoints[i + 1] + seg = end - start + seg_len = float(np.linalg.norm(seg)) + if seg_len >= remaining: + return np.asarray(start + (remaining / seg_len) * seg) + remaining -= seg_len + start = end + return np.asarray(waypoints[-1]) + + def _project_onto_path( + self, waypoints: np.ndarray, position: np.ndarray + ) -> tuple[int, np.ndarray]: + best_idx = 0 + best_point = np.asarray(waypoints[0]) + best_dist = math.inf + for i in range(len(waypoints) - 1): + a = waypoints[i] + ab = waypoints[i + 1] - a + denom = float(np.dot(ab, ab)) + t = 0.0 if denom == 0 else float(np.clip(np.dot(position - a, ab) / denom, 0.0, 1.0)) + proj = a + t * ab + dist = float(np.linalg.norm(position - proj)) + if dist < best_dist: + best_dist = dist + best_idx = i + best_point = proj + return best_idx, best_point diff --git a/dimos/navigation/basic_path_follower/test_module.py b/dimos/navigation/basic_path_follower/test_module.py new file mode 100644 index 0000000000..a5ad8d5782 --- /dev/null +++ b/dimos/navigation/basic_path_follower/test_module.py @@ -0,0 +1,31 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dimos.navigation.basic_path_follower.module import lookahead_distance + + +def test_lookahead_floor_at_low_speed(): + assert lookahead_distance(0.1, 1.5, 0.4, 1.5) == 0.4 + + +def test_lookahead_scales_in_linear_region(): + assert lookahead_distance(0.5, 1.5, 0.4, 1.5) == 0.75 + + +def test_lookahead_clamped_at_ceiling(): + assert lookahead_distance(2.0, 1.5, 0.4, 1.5) == 1.5 + + +def test_lookahead_monotonic_in_speed(): + assert lookahead_distance(0.8, 1.5, 0.4, 1.5) > lookahead_distance(0.45, 1.5, 0.4, 1.5) diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py new file mode 100644 index 0000000000..5731e1abfe --- /dev/null +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -0,0 +1,58 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from reactivex.disposable import Disposable + +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.stream import In, Out +from dimos.msgs.geometry_msgs.PointStamped import PointStamped +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.msgs.nav_msgs.Odometry import Odometry + + +class GoalRelayConfig(ModuleConfig): + pass + + +class GoalRelay(Module): + """Adapt odometry and goal points to the planner's PoseStamped inputs. + + The MLS planner consumes start and goal as PoseStamped and replans itself + on map updates. This stateless adapter just relays: each odometry becomes a + start_pose, each goal becomes a goal_pose. A non-finite goal is relayed + through so the planner can clear its active goal. + """ + + 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..65897c06b6 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,76 @@ 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 every live cell within one step of a changed +/// cell, matching build_surface_cells edge logic. Call after surgical +/// insert or remove so the affected region matches a full rebuild. The seed +/// coordinates are the changed cells, and their surface neighbors are included. +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..eb0859a5d3 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)); + } + } + } +} + +/// Bounded multi-source dijkstra that re-labels only cells in `window`, seeding +/// the wavefront 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 { @@ -96,15 +200,104 @@ impl PartialOrd for Scored { } impl Ord for Scored { fn cmp(&self, other: &Self) -> Ordering { - // Order on score, and use cell id for tie-breaker for repeatability - other.0.total_cmp(&self.0).then(self.1.cmp(&other.1)) + // Min-distance pops first. Coordinate tie-break keeps the result + // independent of CellId assignment, so incremental matches full rebuild. + other.0.total_cmp(&self.0).then(other.1.cmp(&self.1)) } } #[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 +313,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 +333,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 +359,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 +379,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 +388,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..f294c241bb 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,11 @@ pub struct PlannerGraph { pub surface_lookup: SurfaceLookup, pub nodes: Vec, pub node_edges: Vec, - pub node_adj: Vec>, + pub node_adj: AHashMap>, + /// Voronoi forest read by the planner to expand node edges into cell paths. pub cell_state: DijkstraState, + /// Persistent wall-distance field, reseeded regionally from cached values. + pub wall_state: DijkstraState, } impl PlannerGraph { @@ -59,7 +63,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 +74,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); } - 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); +} + +/// Regional counterpart to build_node_edges: recompute the Voronoi only inside +/// the window, keep cached edges whose boundary is outside it, and rescan the +/// window for fresh 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; } + 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 +172,11 @@ fn best_boundary_edges(cells: &SurfaceCells, state: &DijkstraState, out: &mut Ve continue; } let cost = du + edge.cost + dv; - 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 +193,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 +277,25 @@ 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] 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..be90e808e0 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,53 @@ // 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 coalesced 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, + // Incremental path: a local map slice paired by stamp with the region + // bounds it covers, published by the ray tracer alongside local_map. + #[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 +69,244 @@ struct MlsPlanner { #[config] config: Config, - planner: Planner, - latest_start: Option<(f32, f32, f32)>, + // Handle-loop-local: pair a local map with its stamp-matching bounds before + // handing the matched 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 worker that owns the planner graph and does all heavy map + /// processing, so the handle loop stays free to drain its inputs. + 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. Cheap stamp compare. Runs on the handle loop. + 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 (coalescing over any unprocessed one) and + /// wake the worker. + fn hand_off(&self, update: MapUpdate) { + *self.pending.lock().expect("pending mutex") = Some(update); + self.wake.notify_one(); + } + + /// Store-only: 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)); + } + + /// Arm the active goal, or clear it on a non-finite goal (the cancel + /// signal), then wake the worker so a fresh click replans immediately + /// against the current graph. A goal arrives once per click, so this is not + /// the odometry-rate external trigger the refactor removed. + 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; + // Coalesced: 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; + } + } + + /// Mutate the graph from a map update, then publish the viz artifacts at no + /// more than `viz_publish_hz`. The graph itself is always updated; only the + /// surface_map / nodes / node_edges rebuild + publish is rate-capped, since + /// rebuilding the whole graph every cycle is the dominant per-cycle cost and + /// nothing on the planning path consumes those outputs. The CPU-bound + /// section runs under `block_in_place` so the runtime can still schedule the + /// handle loop on another thread. + 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); + } + } + + /// Pure-CPU graph mutation from a map update. Returns whether the graph was + /// updated (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. Pure glue: it gates and + /// does IO, all planning 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 +314,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 15321a087a..cd23d9913a 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 @@ -4,17 +4,21 @@ //! Config and the owned-state Planner that builds and queries the MLS graph. use ahash::AHashSet; +use rayon::prelude::*; use serde::Deserialize; -use validator::Validate; +use validator::{Validate, 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}; -#[derive(Debug, Deserialize, Validate)] +#[derive(Debug, Clone, Deserialize, Validate)] +#[validate(schema(function = "validate_wall_buffer"))] #[serde(deny_unknown_fields)] pub struct Config { pub world_frame: String, @@ -22,16 +26,85 @@ pub struct Config { 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)] @@ -39,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, ); @@ -97,24 +406,610 @@ 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) + ); + } } 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..dad28dd1a6 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,11 @@ //! nodes at local maxima via NMS, and rescale cell-edge costs to push paths //! toward corridor centers. -use ahash::AHashMap; +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 +22,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 +41,51 @@ 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 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 +94,128 @@ 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); + + 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, + ); - apply_wall_safe_penalty(cells, &state.dist, node_wall_buffer_m); + 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. -fn collect_wall_adjacent_cells(cells: &SurfaceCells, out: &mut Vec) { +/// 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, 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, - }; + 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); } - if mask != 0b1111 { + } + 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 in cells.ids() { + if is_wall_adjacent(cells, id) { out.push(id); } } @@ -95,9 +227,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 +249,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 +283,69 @@ 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 } #[cfg(test)] @@ -193,7 +378,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,7 +399,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.is_empty()); } @@ -223,7 +412,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 +429,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..2d603bec48 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs @@ -4,48 +4,61 @@ use std::cmp::Ordering; use std::collections::BinaryHeap; -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,258 @@ 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(); - 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, + let Some(&goal_node) = goal_segment.last() else { + tracing::warn!(?goal_coord, "plan failed: goal has no predecessor chain"); + return None; + }; + if !node_cells.contains(&goal_node) { + tracing::warn!( + ?goal_coord, + "plan failed: goal region does not reach a graph node" + ); + return None; + } + + // 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: no start candidate connects to the goal component", + ); + 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 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; } - 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; + 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) +} + +/// 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 +359,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 +389,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 +583,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 +633,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 +661,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 +670,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!( - wp.iter() - .any(|w| (w.0 - nx).abs() < 1e-5 && (w.1 - ny).abs() < 1e-5), - "node ({nx}, {ny}) should appear among waypoints" + 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!( + 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..ff3301f466 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( @@ -92,17 +193,54 @@ def main( max_range: float = typer.Option(30.0, "--max-range", help="Max ray cast distance (m)"), 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 +258,108 @@ 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, ) - ).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), - ) + 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)) - 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") - - 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/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 24d90dd7fe..6f969ca951 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -111,6 +111,7 @@ "unitree-go2-keyboard-teleop": "dimos.robot.unitree.go2.blueprints.basic.unitree_go2_keyboard_teleop:unitree_go2_keyboard_teleop", "unitree-go2-markers": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_markers", "unitree-go2-memory": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_memory", + "unitree-go2-nav-3d": "dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_nav_3d:unitree_go2_nav_3d", "unitree-go2-relocalization": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2:unitree_go2_relocalization", "unitree-go2-ros": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2_ros:unitree_go2_ros", "unitree-go2-security": "dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_security:unitree_go2_security", @@ -135,6 +136,7 @@ "arm-teleop-module": "dimos.teleop.quest.quest_extensions.ArmTeleopModule", "b-box-navigation-module": "dimos.navigation.bbox_navigation.BBoxNavigationModule", "b1-connection-module": "dimos.robot.unitree.b1.connection.B1ConnectionModule", + "basic-path-follower": "dimos.navigation.basic_path_follower.module.BasicPathFollower", "camera-module": "dimos.hardware.sensors.camera.module.CameraModule", "cartesian-motion-controller": "dimos.manipulation.control.servo_control.cartesian_motion_controller.CartesianMotionController", "click-start-goal-router": "dimos.navigation.nav_stack.modules.click_start_goal_router.click_start_goal_router.ClickStartGoalRouter", @@ -165,6 +167,7 @@ "go2-fleet-connection": "dimos.robot.unitree.go2.fleet_connection.Go2FleetConnection", "go2-memory": "dimos.robot.unitree.go2.blueprints.smart.unitree_go2.Go2Memory", "go2-teleop-module": "dimos.teleop.quest.quest_extensions.Go2TeleopModule", + "goal-relay": "dimos.navigation.nav_3d.mls_planner.goal_relay.GoalRelay", "google-maps-skill-container": "dimos.agents.skills.google_maps_skill_container.GoogleMapsSkillContainer", "gps-nav-skill-container": "dimos.agents.skills.gps_nav_skill.GpsNavSkillContainer", "grasping-module": "dimos.manipulation.grasping.grasping.GraspingModule", diff --git a/dimos/robot/unitree/connection.py b/dimos/robot/unitree/connection.py index 44101cc19d..1afe2638e5 100644 --- a/dimos/robot/unitree/connection.py +++ b/dimos/robot/unitree/connection.py @@ -15,6 +15,7 @@ import asyncio from dataclasses import dataclass import functools +import json import threading import time from typing import Any, TypeAlias, TypeVar @@ -320,6 +321,30 @@ def set_obstacle_avoidance(self, enabled: bool = True) -> None: {"api_id": 1001, "parameter": {"enable": int(enabled)}}, ) + def set_motion_mode(self, name: str) -> None: + """Select the robot's top-level motion mode via the motion switcher. + + 'mcf' is the AI/sport controller (firmware >= 1.1.7) that handles + terrain, including walking up and down stairs. 'normal' is the basic + controller. Stair traversal lives in mcf, so the nav stack selects it. + """ + # api_id 1001 = CheckMode, 1002 = SelectMode, param {"name": }. + current = None + try: + resp = self.publish_request(RTC_TOPIC["MOTION_SWITCHER"], {"api_id": 1001}) + current = json.loads(resp["data"]["data"]).get("name") + except (KeyError, TypeError, ValueError) as e: + print(f"Motion mode check failed: {e}") + print(f"Motion mode: current={current!r} requested={name!r}") + if current == name: + return + self.publish_request( + RTC_TOPIC["MOTION_SWITCHER"], + {"api_id": 1002, "parameter": {"name": name}}, + ) + # Switching controllers makes the robot re-stand; give it a moment. + time.sleep(5) + def free_walk(self) -> bool: """Activate FreeWalk locomotion mode — enables walking and velocity commands.""" return bool(self.publish_request(RTC_TOPIC["SPORT_MOD"], {"api_id": SPORT_CMD["FreeWalk"]})) diff --git a/dimos/robot/unitree/dimsim_connection.py b/dimos/robot/unitree/dimsim_connection.py index 5afcd1fda7..8f5d222f5a 100644 --- a/dimos/robot/unitree/dimsim_connection.py +++ b/dimos/robot/unitree/dimsim_connection.py @@ -94,6 +94,9 @@ def balance_stand(self) -> bool: def set_obstacle_avoidance(self, enabled: bool = True) -> None: pass + def set_motion_mode(self, name: str) -> None: + pass + def enable_rage_mode(self) -> bool: return True diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py index 928b2dc23a..2fb7d09094 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_basic.py @@ -87,7 +87,7 @@ def _go2_rerun_blueprint() -> Any: ) -rerun_config = { +rerun_config: dict[str, Any] = { "blueprint": _go2_rerun_blueprint, # Custom converters for specific rerun entity paths # Normally all these would be specified in their respectative modules diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py new file mode 100644 index 0000000000..d9d0ee2bac --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Go2 3D nav stack: fastlio odometry, raytraced voxel map, MLS planning, path following. + +Expects a mid360 + fastlio running against a lidar mounted on the robot. +The go2's built-in lidar and odometry are remapped aside and unused for +navigation. The map, paths, and robot pose all live in fastlio's world frame. +""" + +import os +from typing import Any + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.global_config import global_config +from dimos.hardware.sensors.lidar.fastlio2.module import FastLio2 +from dimos.mapping.ray_tracing.module import RayTracingVoxelMap +from dimos.navigation.basic_path_follower.module import BasicPathFollower +from dimos.navigation.movement_manager.movement_manager import MovementManager +from dimos.navigation.nav_3d.mls_planner.goal_relay import GoalRelay +from dimos.navigation.nav_3d.mls_planner.mls_planner_native import MLSPlannerNative +from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import rerun_config +from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.visualization.vis_module import vis_module + +voxel_size = 0.1 +# Height of the head-mounted lidar above the ground while standing. +go2_lidar_height = 0.5 + + +def _render_global_map(msg: Any) -> Any: + # don't use a min z height because why would we want that + return msg.to_rerun() + + +def _render_path(msg: Any) -> Any: + # The planner clears its plan with an empty path on every start-pose change. + # Logging those would blank the line. Drop them so the last path stays shown. + if len(msg.poses) == 0: + return None + return msg + + +def _static_robot_body(rr: Any) -> list[Any]: + """Go2-shaped box on fastlio's body frame, counter-rotated for the lidar pitch.""" + return [ + rr.Boxes3D(half_sizes=[0.35, 0.155, 0.2], colors=[(0, 255, 127)]), + rr.Transform3D( + parent_frame="tf#/body", + rotation=rr.RotationAxisAngle(axis=(0, 1, 0), degrees=-45.0), + ), + ] + + +_nav_rerun_config = { + **rerun_config, + "max_hz": { + **rerun_config["max_hz"], + "world/global_map": 1.0, + "world/local_map": 1.0, + }, + "memory_limit": "256MB", + # base_link tf comes from the go2 internal odometry, which is not the map + # frame. Anchor the robot box to fastlio's body frame instead and hide the + # camera frustum that rides base_link. + "static": {"world/tf/body": _static_robot_body}, + "visual_override": { + **rerun_config["visual_override"], + "world/global_map": _render_global_map, + "world/path": _render_path, + "world/camera_info": None, + "world/color_image": None, + "world/lidar": None, + "world/surface_map": None, + "world/nodes": None, + "world/node_edges": None, + }, +} + +unitree_go2_nav_3d = autoconnect( + vis_module(viewer_backend=global_config.viewer, rerun_config=_nav_rerun_config), + # motion_mode="mcf" selects the AI/sport controller that walks up and down + # stairs. With it, obstacle_avoidance is turned off below, matching the app + # recipe for stair traversal. + GO2Connection.blueprint(lidar=False, camera=False, motion_mode="mcf").remappings( + [ + (GO2Connection, "lidar", "lidar_l1"), + (GO2Connection, "odom", "odom_go2"), + ] + ), + FastLio2.blueprint( + host_ip=os.getenv("LIDAR_HOST_IP", "192.168.1.5"), + lidar_ip=os.getenv("LIDAR_IP", "192.168.1.155"), + map_freq=-1.0, + ).remappings([(FastLio2, "global_map", "global_map_fastlio")]), + RayTracingVoxelMap.blueprint( + voxel_size=voxel_size, emit_every=2, global_emit_every=50, max_health=5 + ), + # global_map is remapped off so the planner runs purely on the + # incremental local_map + region_bounds pair. + MLSPlannerNative.blueprint( + world_frame="odom", + voxel_size=voxel_size, + robot_height=go2_lidar_height, + wall_clearance_m=0.3, + wall_buffer_m=0.75, + wall_buffer_weight=100.0, + step_threshold_m=0.25, + step_penalty_weight=1.0, + viz_publish_hz=0.0, + ).remappings([(MLSPlannerNative, "global_map", "global_map_unused")]), + GoalRelay.blueprint(), + BasicPathFollower.blueprint(speed=0.5, heading_gain=0.8, max_angular=0.6), + MovementManager.blueprint(), +).global_config(n_workers=10, robot_model="unitree_go2", obstacle_avoidance=False) + +__all__ = ["unitree_go2_nav_3d"] diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index 5568a473ef..63ed750150 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -67,6 +67,11 @@ class Go2Mode(str, Enum): class ConnectionConfig(ModuleConfig): ip: str = Field(default_factory=lambda m: m["g"].robot_ip) mode: Go2Mode = Go2Mode.DEFAULT + lidar: bool = True + camera: bool = True + # Top-level motion controller: "mcf" is the AI/sport mode that traverses + # terrain (stairs); "normal" is basic. None leaves the current mode as-is. + motion_mode: str | None = None class Go2ConnectionProtocol(Protocol): @@ -82,6 +87,7 @@ def standup(self) -> bool: ... def liedown(self) -> bool: ... def balance_stand(self) -> bool: ... def set_obstacle_avoidance(self, enabled: bool = True) -> None: ... + def set_motion_mode(self, name: str) -> None: ... def enable_rage_mode(self) -> bool: ... def publish_request(self, topic: str, data: dict) -> dict: ... # type: ignore[type-arg] @@ -171,6 +177,9 @@ def balance_stand(self) -> bool: def set_obstacle_avoidance(self, enabled: bool = True) -> None: pass + def set_motion_mode(self, name: str) -> None: + pass + def enable_rage_mode(self) -> bool: return True @@ -241,16 +250,23 @@ def onimage(image: Image) -> None: self.color_image.publish(image) self._latest_video_frame = image - self.register_disposable(self.connection.lidar_stream().subscribe(self.lidar.publish)) + if self.config.lidar: + self.register_disposable(self.connection.lidar_stream().subscribe(self.lidar.publish)) self.register_disposable(self.connection.odom_stream().subscribe(self._publish_tf)) - self.register_disposable(self.connection.video_stream().subscribe(onimage)) self.register_disposable(Disposable(self.cmd_vel.subscribe(self.move))) - self._camera_info_thread = Thread( - target=self.publish_camera_info, - daemon=True, - ) - self._camera_info_thread.start() + if self.config.camera: + self.register_disposable(self.connection.video_stream().subscribe(onimage)) + self._camera_info_thread = Thread( + target=self.publish_camera_info, + daemon=True, + ) + self._camera_info_thread.start() + + # Select the terrain-capable controller (mcf) before standing, so the + # robot can walk up and down stairs under our velocity commands. + if self.config.motion_mode: + self.connection.set_motion_mode(self.config.motion_mode) self.standup() time.sleep(3) diff --git a/dimos/robot/unitree/mujoco_connection.py b/dimos/robot/unitree/mujoco_connection.py index 4c455899e8..9ce7d41de5 100644 --- a/dimos/robot/unitree/mujoco_connection.py +++ b/dimos/robot/unitree/mujoco_connection.py @@ -237,6 +237,9 @@ def balance_stand(self) -> bool: def set_obstacle_avoidance(self, enabled: bool = True) -> None: pass + def set_motion_mode(self, name: str) -> None: + pass + def enable_rage_mode(self) -> bool: return True diff --git a/dimos/visualization/rerun/bridge.py b/dimos/visualization/rerun/bridge.py index 2f5fb1efa9..2d43767c94 100644 --- a/dimos/visualization/rerun/bridge.py +++ b/dimos/visualization/rerun/bridge.py @@ -289,6 +289,14 @@ def _on_message(self, msg: Any, topic: Any) -> None: if not rerun_data: return + # Place data on the timeline by capture time, not by when the bridge got + # to it. Otherwise heavy messages (clouds) log late and desync from light + # ones (path, pose) captured at the same instant. View on the "capture" + # timeline for synchronized playback. + ts = getattr(msg, "ts", None) + if ts is not None: + rr.set_time("capture", timestamp=ts) + # TFMessage for example returns list of (entity_path, archetype) tuples if is_rerun_multi(rerun_data): for path, archetype in rerun_data: diff --git a/pyproject.toml b/pyproject.toml index 6806733e06..e7f6fa8864 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -523,6 +523,8 @@ module = [ "cyclonedds", "cyclonedds.*", "dimos_lcm.*", + "dimos_mls_planner", + "dimos_voxel_ray_tracing", "etils", "faster_whisper", "geometry_msgs.*",