From b4fce2bab980a55c12f8fb8cf3477047557a53bc Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 30 Jun 2026 16:05:03 -0700 Subject: [PATCH 01/26] Transform buffer --- .../sensors/lidar/virtual_mid360/Cargo.lock | 187 +++++++ dimos/mapping/ray_tracing/rust/Cargo.lock | 2 + .../nav_3d/mls_planner/rust/Cargo.lock | 96 +++- examples/native-modules/rust/Cargo.lock | 179 ++++++ native/rust/Cargo.lock | 178 ++++++ native/rust/README.md | 26 + native/rust/dimos-module-macros/src/lib.rs | 18 +- native/rust/dimos-module/Cargo.toml | 3 +- native/rust/dimos-module/src/lib.rs | 2 + native/rust/dimos-module/src/module.rs | 18 + native/rust/dimos-module/src/tf.rs | 526 ++++++++++++++++++ 11 files changed, 1226 insertions(+), 9 deletions(-) create mode 100644 native/rust/dimos-module/src/tf.rs diff --git a/dimos/hardware/sensors/lidar/virtual_mid360/Cargo.lock b/dimos/hardware/sensors/lidar/virtual_mid360/Cargo.lock index 53993a985c..fe9507498f 100644 --- a/dimos/hardware/sensors/lidar/virtual_mid360/Cargo.lock +++ b/dimos/hardware/sensors/lidar/virtual_mid360/Cargo.lock @@ -11,6 +11,27 @@ dependencies = [ "memchr", ] +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" + [[package]] name = "byteorder" version = "1.5.0" @@ -80,6 +101,8 @@ version = "0.1.0" dependencies = [ "dimos-lcm", "dimos-module-macros", + "lcm-msgs", + "nalgebra", "serde", "serde_json", "tokio", @@ -133,6 +156,30 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "glam" +version = "0.30.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19fc433e8437a212d1b6f1e68c7824af3aed907da60afa994e7f542d18d12aa9" + +[[package]] +name = "glam" +version = "0.31.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556f6b2ea90b8d15a74e0e7bb41671c9bdf38cd9f78c284d750b9ce58a2b5be7" + +[[package]] +name = "glam" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f70749695b063ecbf6b62949ccccde2e733ec3ecbbd71d467dca4e5c6c97cca0" + +[[package]] +name = "glam" +version = "0.33.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f22fb22f065b308be0d8724e3706c7fa3fc2a6c7d6899df4cad7860e7a75436" + [[package]] name = "icu_collections" version = "2.2.0" @@ -254,6 +301,14 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" +[[package]] +name = "lcm-msgs" +version = "0.1.0" +source = "git+https://github.com/dimensionalOS/dimos-lcm.git?branch=rust-codegen#e7c9428b7201cdfeadecd181c77c9e2d60a14503" +dependencies = [ + "byteorder", +] + [[package]] name = "libc" version = "0.2.186" @@ -281,6 +336,16 @@ dependencies = [ "regex-automata", ] +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.2" @@ -298,6 +363,37 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "nalgebra" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adc43a60c217b0c6ff46e47f26911015ad8d2e5a8be1af668c67e370d99a4346" +dependencies = [ + "approx", + "glam 0.30.10", + "glam 0.31.1", + "glam 0.32.1", + "glam 0.33.2", + "matrixmultiply", + "nalgebra-macros", + "num-complex", + "num-rational", + "num-traits", + "simba", + "typenum", +] + +[[package]] +name = "nalgebra-macros" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "973e7178a678cfd059ccec50887658d482ce16b0aa9da3888ddeab5cd5eb4889" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -307,6 +403,54 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num-bigint" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -374,6 +518,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "regex" version = "1.12.4" @@ -403,6 +553,15 @@ version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" +[[package]] +name = "safe_arch" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f7caad094bd561859bcd467734a720c3c1f5d1f338995351fefe2190c45efed" +dependencies = [ + "bytemuck", +] + [[package]] name = "serde" version = "1.0.228" @@ -465,6 +624,18 @@ dependencies = [ "libc", ] +[[package]] +name = "simba" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f45c644a9f3a386f9288625d9f0c1e999e1acf07a37df35d0516c7f199d9cb2" +dependencies = [ + "approx", + "num-complex", + "num-traits", + "wide", +] + [[package]] name = "smallvec" version = "1.15.2" @@ -645,6 +816,12 @@ dependencies = [ "tracing-serde", ] +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + [[package]] name = "unicode-ident" version = "1.0.24" @@ -723,6 +900,16 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wide" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfdfe6a32973f2d1b268b8895845a8a96cac2f0191e72c27cc929036060dbf89" +dependencies = [ + "bytemuck", + "safe_arch", +] + [[package]] name = "windows-link" version = "0.2.1" diff --git a/dimos/mapping/ray_tracing/rust/Cargo.lock b/dimos/mapping/ray_tracing/rust/Cargo.lock index 122775cdb5..fa86a50997 100644 --- a/dimos/mapping/ray_tracing/rust/Cargo.lock +++ b/dimos/mapping/ray_tracing/rust/Cargo.lock @@ -145,6 +145,8 @@ version = "0.1.0" dependencies = [ "dimos-lcm", "dimos-module-macros", + "lcm-msgs", + "nalgebra", "serde", "serde_json", "tokio", diff --git a/dimos/navigation/nav_3d/mls_planner/rust/Cargo.lock b/dimos/navigation/nav_3d/mls_planner/rust/Cargo.lock index 296ebd114b..929b0c041e 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/Cargo.lock +++ b/dimos/navigation/nav_3d/mls_planner/rust/Cargo.lock @@ -185,6 +185,8 @@ version = "0.1.0" dependencies = [ "dimos-lcm", "dimos-module-macros", + "lcm-msgs", + "nalgebra 0.35.0", "serde", "serde_json", "tokio", @@ -293,6 +295,30 @@ dependencies = [ "wasip2", ] +[[package]] +name = "glam" +version = "0.30.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19fc433e8437a212d1b6f1e68c7824af3aed907da60afa994e7f542d18d12aa9" + +[[package]] +name = "glam" +version = "0.31.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556f6b2ea90b8d15a74e0e7bb41671c9bdf38cd9f78c284d750b9ce58a2b5be7" + +[[package]] +name = "glam" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f70749695b063ecbf6b62949ccccde2e733ec3ecbbd71d467dca4e5c6c97cca0" + +[[package]] +name = "glam" +version = "0.33.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f22fb22f065b308be0d8724e3706c7fa3fc2a6c7d6899df4cad7860e7a75436" + [[package]] name = "heck" version = "0.5.0" @@ -431,7 +457,7 @@ dependencies = [ "getrandom 0.2.17", "image", "itertools", - "nalgebra", + "nalgebra 0.32.6", "num", "rand", "rand_distr", @@ -577,10 +603,41 @@ dependencies = [ "num-complex", "num-rational", "num-traits", - "simba", + "simba 0.8.1", "typenum", ] +[[package]] +name = "nalgebra" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adc43a60c217b0c6ff46e47f26911015ad8d2e5a8be1af668c67e370d99a4346" +dependencies = [ + "approx", + "glam 0.30.10", + "glam 0.31.1", + "glam 0.32.1", + "glam 0.33.2", + "matrixmultiply", + "nalgebra-macros", + "num-complex", + "num-rational", + "num-traits", + "simba 0.10.0", + "typenum", +] + +[[package]] +name = "nalgebra-macros" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "973e7178a678cfd059ccec50887658d482ce16b0aa9da3888ddeab5cd5eb4889" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "ndarray" version = "0.16.1" @@ -991,6 +1048,15 @@ dependencies = [ "bytemuck", ] +[[package]] +name = "safe_arch" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f7caad094bd561859bcd467734a720c3c1f5d1f338995351fefe2190c45efed" +dependencies = [ + "bytemuck", +] + [[package]] name = "serde" version = "1.0.228" @@ -1063,7 +1129,19 @@ dependencies = [ "num-complex", "num-traits", "paste", - "wide", + "wide 0.7.33", +] + +[[package]] +name = "simba" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f45c644a9f3a386f9288625d9f0c1e999e1acf07a37df35d0516c7f199d9cb2" +dependencies = [ + "approx", + "num-complex", + "num-traits", + "wide 1.5.0", ] [[package]] @@ -1409,7 +1487,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ce5da8ecb62bcd8ec8b7ea19f69a51275e91299be594ea5cc6ef7819e16cd03" dependencies = [ "bytemuck", - "safe_arch", + "safe_arch 0.7.4", +] + +[[package]] +name = "wide" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfdfe6a32973f2d1b268b8895845a8a96cac2f0191e72c27cc929036060dbf89" +dependencies = [ + "bytemuck", + "safe_arch 1.0.0", ] [[package]] diff --git a/examples/native-modules/rust/Cargo.lock b/examples/native-modules/rust/Cargo.lock index 8a7fda6bba..c7a37d778a 100644 --- a/examples/native-modules/rust/Cargo.lock +++ b/examples/native-modules/rust/Cargo.lock @@ -11,6 +11,27 @@ dependencies = [ "memchr", ] +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" + [[package]] name = "byteorder" version = "1.5.0" @@ -80,6 +101,8 @@ version = "0.1.0" dependencies = [ "dimos-lcm", "dimos-module-macros", + "lcm-msgs", + "nalgebra", "serde", "serde_json", "tokio", @@ -145,6 +168,30 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "glam" +version = "0.30.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19fc433e8437a212d1b6f1e68c7824af3aed907da60afa994e7f542d18d12aa9" + +[[package]] +name = "glam" +version = "0.31.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556f6b2ea90b8d15a74e0e7bb41671c9bdf38cd9f78c284d750b9ce58a2b5be7" + +[[package]] +name = "glam" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f70749695b063ecbf6b62949ccccde2e733ec3ecbbd71d467dca4e5c6c97cca0" + +[[package]] +name = "glam" +version = "0.33.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f22fb22f065b308be0d8724e3706c7fa3fc2a6c7d6899df4cad7860e7a75436" + [[package]] name = "icu_collections" version = "2.2.0" @@ -301,6 +348,16 @@ dependencies = [ "regex-automata", ] +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.0" @@ -318,6 +375,37 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "nalgebra" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adc43a60c217b0c6ff46e47f26911015ad8d2e5a8be1af668c67e370d99a4346" +dependencies = [ + "approx", + "glam 0.30.10", + "glam 0.31.1", + "glam 0.32.1", + "glam 0.33.2", + "matrixmultiply", + "nalgebra-macros", + "num-complex", + "num-rational", + "num-traits", + "simba", + "typenum", +] + +[[package]] +name = "nalgebra-macros" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "973e7178a678cfd059ccec50887658d482ce16b0aa9da3888ddeab5cd5eb4889" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -327,6 +415,54 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num-bigint" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -394,6 +530,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "regex" version = "1.12.3" @@ -423,6 +565,15 @@ version = "0.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" +[[package]] +name = "safe_arch" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f7caad094bd561859bcd467734a720c3c1f5d1f338995351fefe2190c45efed" +dependencies = [ + "bytemuck", +] + [[package]] name = "serde" version = "1.0.228" @@ -485,6 +636,18 @@ dependencies = [ "libc", ] +[[package]] +name = "simba" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f45c644a9f3a386f9288625d9f0c1e999e1acf07a37df35d0516c7f199d9cb2" +dependencies = [ + "approx", + "num-complex", + "num-traits", + "wide", +] + [[package]] name = "smallvec" version = "1.15.1" @@ -665,6 +828,12 @@ dependencies = [ "tracing-serde", ] +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + [[package]] name = "unicode-ident" version = "1.0.24" @@ -731,6 +900,16 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wide" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfdfe6a32973f2d1b268b8895845a8a96cac2f0191e72c27cc929036060dbf89" +dependencies = [ + "bytemuck", + "safe_arch", +] + [[package]] name = "windows-link" version = "0.2.1" diff --git a/native/rust/Cargo.lock b/native/rust/Cargo.lock index 165d4b923f..dcd6cb8d0e 100644 --- a/native/rust/Cargo.lock +++ b/native/rust/Cargo.lock @@ -11,6 +11,27 @@ dependencies = [ "memchr", ] +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "bytemuck" +version = "1.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8efb64bd706a16a1bdde310ae86b351e4d21550d98d056f22f8a7f7a2183fec" + [[package]] name = "byteorder" version = "1.5.0" @@ -81,6 +102,7 @@ dependencies = [ "dimos-lcm", "dimos-module-macros", "lcm-msgs", + "nalgebra", "serde", "serde_json", "tokio", @@ -135,6 +157,30 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "glam" +version = "0.30.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19fc433e8437a212d1b6f1e68c7824af3aed907da60afa994e7f542d18d12aa9" + +[[package]] +name = "glam" +version = "0.31.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556f6b2ea90b8d15a74e0e7bb41671c9bdf38cd9f78c284d750b9ce58a2b5be7" + +[[package]] +name = "glam" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f70749695b063ecbf6b62949ccccde2e733ec3ecbbd71d467dca4e5c6c97cca0" + +[[package]] +name = "glam" +version = "0.33.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f22fb22f065b308be0d8724e3706c7fa3fc2a6c7d6899df4cad7860e7a75436" + [[package]] name = "icu_collections" version = "2.2.0" @@ -291,6 +337,16 @@ dependencies = [ "regex-automata", ] +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.0" @@ -308,6 +364,37 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "nalgebra" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adc43a60c217b0c6ff46e47f26911015ad8d2e5a8be1af668c67e370d99a4346" +dependencies = [ + "approx", + "glam 0.30.10", + "glam 0.31.1", + "glam 0.32.1", + "glam 0.33.2", + "matrixmultiply", + "nalgebra-macros", + "num-complex", + "num-rational", + "num-traits", + "simba", + "typenum", +] + +[[package]] +name = "nalgebra-macros" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "973e7178a678cfd059ccec50887658d482ce16b0aa9da3888ddeab5cd5eb4889" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -317,6 +404,54 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num-bigint" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -384,6 +519,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "regex" version = "1.12.3" @@ -413,6 +554,15 @@ version = "0.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" +[[package]] +name = "safe_arch" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f7caad094bd561859bcd467734a720c3c1f5d1f338995351fefe2190c45efed" +dependencies = [ + "bytemuck", +] + [[package]] name = "serde" version = "1.0.228" @@ -475,6 +625,18 @@ dependencies = [ "libc", ] +[[package]] +name = "simba" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f45c644a9f3a386f9288625d9f0c1e999e1acf07a37df35d0516c7f199d9cb2" +dependencies = [ + "approx", + "num-complex", + "num-traits", + "wide", +] + [[package]] name = "smallvec" version = "1.15.1" @@ -676,6 +838,12 @@ dependencies = [ "syn", ] +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + [[package]] name = "unicode-ident" version = "1.0.24" @@ -742,6 +910,16 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wide" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfdfe6a32973f2d1b268b8895845a8a96cac2f0191e72c27cc929036060dbf89" +dependencies = [ + "bytemuck", + "safe_arch", +] + [[package]] name = "windows-link" version = "0.2.1" diff --git a/native/rust/README.md b/native/rust/README.md index 2b79d7d840..0e776fb9ee 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -55,6 +55,7 @@ async fn main() { - `#[input(decode = fn, handler = fn)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. - `#[output(encode = fn)]`: on a field of type `Output`. `encode` is required. - `#[config]`: on one field. The type must be defined with `#[native_config]` (see [Config](#config)). At most one per struct. If absent, `Config` defaults to `dimos_module::NoConfig`. +- `#[tf]`: on a field of type `Tf`. Subscribes to the `tf` topic and answers transform queries (see [Transforms](#transforms)). No arguments. - Unattributed fields are initialized via `Default::default()` and treated as module state. ## Config @@ -97,6 +98,31 @@ At runtime `run()` enforces the mapping on the Python payload: deserialization r Field name = port name. Ports map to topics via the stdin JSON; unmapped ports fall back to `/{port}`. +## Transforms + +A `#[tf]` field gives a module a consumer-side view of the transform graph, the Rust counterpart to Python's `tf.get()`. It subscribes to the `tf` topic (mapped like any other port, default `/tf`), buffers each `parent -> child` edge it sees, and answers queries by composing transforms along the shortest path through the graph. + +```rust +#[derive(Module)] +struct VoxelMap { + #[input(decode = PointCloud2::decode)] + lidar: Input, + #[tf] + tf: Tf, +} + +impl VoxelMap { + async fn handle_lidar(&mut self, cloud: PointCloud2) { + // De-rotate a scan from the lidar's mount frame into the robot base frame. + if let Some(t) = self.tf.get_latest("base_link", "mid360_link") { + let point_in_base = t.isometry() * point_in_lidar; + } + } +} +``` + +`Tf` is a cheap-to-clone handle; the graph fills in the background as `tf` messages arrive. `get(parent, child, time, tolerance)` selects the sample nearest `time` (latest when `None`) and returns `None` when no path connects the frames or no sample falls within `tolerance` seconds. `get_latest` is the no-time shorthand. The result exposes an `nalgebra` `Isometry3` via `isometry()`, ready to apply to a point. Lookups are nearest-in-time, not interpolated. This is consumer-only; modules do not publish transforms. + ## What `#[derive(Module)]` generates Just for reference, in the example above the macro expands to: diff --git a/native/rust/dimos-module-macros/src/lib.rs b/native/rust/dimos-module-macros/src/lib.rs index 53658ed8f7..e9b4618675 100644 --- a/native/rust/dimos-module-macros/src/lib.rs +++ b/native/rust/dimos-module-macros/src/lib.rs @@ -3,7 +3,7 @@ use proc_macro2::TokenStream as TokenStream2; use quote::{format_ident, quote}; use syn::{parse_macro_input, Data, DeriveInput, Field, Fields, Ident, Path, Type}; -#[proc_macro_derive(Module, attributes(input, output, config, module))] +#[proc_macro_derive(Module, attributes(input, output, config, tf, module))] pub fn derive_module(input: TokenStream) -> TokenStream { let input = parse_macro_input!(input as DeriveInput); match expand(input) { @@ -164,6 +164,7 @@ enum FieldKind { Input { decode: Path, handler: Ident }, Output { encode: Path }, Config, + Tf, State, } @@ -259,6 +260,7 @@ fn expand(input: DeriveInput) -> syn::Result { quote!(#name: builder.output(#name_str, #encode)) } FieldKind::Config => quote!(#name: config), + FieldKind::Tf => quote!(#name: builder.tf()), FieldKind::State => quote!(#name: ::core::default::Default::default()), } }); @@ -341,7 +343,7 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { if found.is_some() { return Err(syn::Error::new_spanned( attr, - "field has multiple module attributes; only one of #[input], #[output], #[config] is allowed", + "field has multiple module attributes; only one of #[input], #[output], #[config], #[tf] is allowed", )); } let mut decode: Option = None; @@ -366,7 +368,7 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { if found.is_some() { return Err(syn::Error::new_spanned( attr, - "field has multiple module attributes; only one of #[input], #[output], #[config] is allowed", + "field has multiple module attributes; only one of #[input], #[output], #[config], #[tf] is allowed", )); } let mut encode: Option = None; @@ -388,10 +390,18 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { if found.is_some() { return Err(syn::Error::new_spanned( attr, - "field has multiple module attributes; only one of #[input], #[output], #[config] is allowed", + "field has multiple module attributes; only one of #[input], #[output], #[config], #[tf] is allowed", )); } found = Some(FieldKind::Config); + } else if path.is_ident("tf") { + if found.is_some() { + return Err(syn::Error::new_spanned( + attr, + "field has multiple module attributes; only one of #[input], #[output], #[config], #[tf] is allowed", + )); + } + found = Some(FieldKind::Tf); } } diff --git a/native/rust/dimos-module/Cargo.toml b/native/rust/dimos-module/Cargo.toml index b695a077a8..429ffb0a08 100644 --- a/native/rust/dimos-module/Cargo.toml +++ b/native/rust/dimos-module/Cargo.toml @@ -8,6 +8,8 @@ license = "Apache-2.0" [dependencies] dimos-lcm = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } dimos-module-macros = { version = "=0.1.0", path = "../dimos-module-macros" } +lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } +nalgebra = "0.35.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time", "signal", "io-std", "io-util"] } serde = { version = "1", features = ["derive"] } serde_json = "1" @@ -16,5 +18,4 @@ tracing-subscriber = { version = "0.3", features = ["json", "env-filter"] } validator = { version = "0.20", features = ["derive"] } [dev-dependencies] -lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } tracing-test = "0.2" diff --git a/native/rust/dimos-module/src/lib.rs b/native/rust/dimos-module/src/lib.rs index a3399b1ea1..39dfac044c 100644 --- a/native/rust/dimos-module/src/lib.rs +++ b/native/rust/dimos-module/src/lib.rs @@ -1,11 +1,13 @@ pub mod lcm; pub mod log; pub mod module; +pub mod tf; pub mod transport; pub use dimos_module_macros::{native_config, Module}; pub use lcm::LcmTransport; pub use module::{run, Builder, Input, Module, ModuleConfig, NativeConfig, NoConfig, Output}; +pub use tf::Tf; pub use transport::Transport; // Re-export LcmOptions so callers don't need to depend on dimos-lcm directly. diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index 740fa45af4..82075081f4 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -266,6 +266,7 @@ pub struct Builder { topics: HashMap, routes: HashMap>>, publish_tx: mpsc::Sender<(String, Vec)>, + tf: Option, } impl Builder { @@ -277,6 +278,7 @@ impl Builder { topics, routes: HashMap::new(), publish_tx, + tf: None, } } @@ -317,6 +319,22 @@ impl Builder { sender: self.publish_tx.clone(), } } + + /// A handle that answers transform queries from the `/tf` topic. + /// + /// The first call subscribes to the resolved `tf` topic and starts filling + /// the transform graph in the background. Repeated calls share one graph. + pub fn tf(&mut self) -> crate::tf::Tf { + if let Some(tf) = &self.tf { + return tf.clone(); + } + let topic = self.topic_for("tf"); + let (tf, route) = + crate::tf::tf_subscription(topic.clone(), crate::tf::DEFAULT_TF_BUFFER_SIZE); + self.routes.entry(topic).or_default().push(route); + self.tf = Some(tf.clone()); + tf + } } pub(crate) fn spawn_pubsub_tasks( diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs new file mode 100644 index 0000000000..e5614af248 --- /dev/null +++ b/native/rust/dimos-module/src/tf.rs @@ -0,0 +1,526 @@ +//! A consumer-side transform client for native modules. +//! +//! Mirrors the Python `dimos.protocol.tf` semantics: a module subscribes to the +//! `/tf` topic, every `TFMessage` edge is buffered per `(parent, child)` pair, +//! and [`Tf::get`] answers a query by composing transforms along the shortest +//! path through the frame graph. Lookups are nearest-in-time within a tolerance, +//! not interpolated, matching the Python buffer. + +use std::collections::{HashMap, HashSet, VecDeque}; +use std::sync::{Arc, RwLock}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use nalgebra::{Isometry3, Quaternion, Translation3, UnitQuaternion, Vector3}; + +use crate::module::Route; + +/// How many seconds of history each edge keeps. +pub const DEFAULT_TF_BUFFER_SIZE: f64 = 10.0; + +fn now_secs() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs_f64()) + .unwrap_or(0.0) +} + +/// A rigid transform from `parent` to `child` at a point in time. +/// +/// The isometry maps a point expressed in `child` coordinates into `parent` +/// coordinates: `p_parent = transform.isometry() * p_child`. +#[derive(Clone, Debug)] +pub struct Transform { + pub parent: String, + pub child: String, + pub ts: f64, + iso: Isometry3, +} + +impl Transform { + /// The transform as an `nalgebra` isometry, ready to apply to a point. + pub fn isometry(&self) -> Isometry3 { + self.iso + } + + /// Translation component (`parent`-frame position of the `child` origin). + pub fn translation(&self) -> Vector3 { + self.iso.translation.vector + } + + /// Rotation component. + pub fn rotation(&self) -> UnitQuaternion { + self.iso.rotation + } + + fn inverse(&self) -> Transform { + Transform { + parent: self.child.clone(), + child: self.parent.clone(), + ts: self.ts, + iso: self.iso.inverse(), + } + } + + // self (a -> b) followed by other (b -> c) gives a -> c. + fn compose(&self, other: &Transform) -> Transform { + Transform { + parent: self.parent.clone(), + child: other.child.clone(), + ts: self.ts, + iso: self.iso * other.iso, + } + } +} + +struct Sample { + ts: f64, + iso: Isometry3, +} + +// One edge's time-sorted history, pruned to a fixed-duration window. +struct TBuffer { + buffer_size: f64, + samples: Vec, +} + +impl TBuffer { + fn new(buffer_size: f64) -> Self { + Self { + buffer_size, + samples: Vec::new(), + } + } + + fn add(&mut self, ts: f64, iso: Isometry3) { + let pos = self.samples.partition_point(|s| s.ts <= ts); + self.samples.insert(pos, Sample { ts, iso }); + self.prune(ts - self.buffer_size); + } + + fn prune(&mut self, min_ts: f64) { + let drop_to = self.samples.partition_point(|s| s.ts < min_ts); + if drop_to > 0 { + self.samples.drain(0..drop_to); + } + } + + fn last(&self) -> Option<&Sample> { + self.samples.last() + } + + // Nearest sample in time. On a tie, prefer the later sample. Returns None + // when the closest sample is further than `tolerance` from `ts`. + fn find_closest(&self, ts: f64, tolerance: Option) -> Option<&Sample> { + let pos = self.samples.partition_point(|s| s.ts < ts); + let prev = pos.checked_sub(1).and_then(|i| self.samples.get(i)); + let next = self.samples.get(pos); + let best = match (prev, next) { + (Some(p), Some(n)) => { + if (n.ts - ts).abs() <= (ts - p.ts).abs() { + n + } else { + p + } + } + (Some(p), None) => p, + (None, Some(n)) => n, + (None, None) => return None, + }; + match tolerance { + Some(tol) if (best.ts - ts).abs() > tol => None, + _ => Some(best), + } + } +} + +/// The transform graph: one [`TBuffer`] per `(parent, child)` edge. +struct MultiTBuffer { + buffer_size: f64, + buffers: HashMap<(String, String), TBuffer>, +} + +impl MultiTBuffer { + fn new(buffer_size: f64) -> Self { + Self { + buffer_size, + buffers: HashMap::new(), + } + } + + fn receive(&mut self, parent: &str, child: &str, ts: f64, iso: Isometry3) { + let buffer_size = self.buffer_size; + self.buffers + .entry((parent.to_string(), child.to_string())) + .or_insert_with(|| TBuffer::new(buffer_size)) + .add(ts, iso); + } + + fn connections(&self, frame: &str) -> Vec { + let mut out = Vec::new(); + for (parent, child) in self.buffers.keys() { + if parent == frame { + out.push(child.clone()); + } + if child == frame { + out.push(parent.clone()); + } + } + out + } + + fn sample( + &self, + buf: &TBuffer, + parent: &str, + child: &str, + time: Option, + tolerance: Option, + ) -> Option { + let s = match time { + None => buf.last()?, + Some(t) => buf.find_closest(t, tolerance)?, + }; + Some(Transform { + parent: parent.to_string(), + child: child.to_string(), + ts: s.ts, + iso: s.iso, + }) + } + + // A single forward or reverse edge (reverse returns the inverse). + fn edge( + &self, + parent: &str, + child: &str, + time: Option, + tolerance: Option, + ) -> Option { + if parent == child { + return Some(Transform { + parent: parent.to_string(), + child: child.to_string(), + ts: time.unwrap_or_else(now_secs), + iso: Isometry3::identity(), + }); + } + if let Some(buf) = self.buffers.get(&(parent.to_string(), child.to_string())) { + return self.sample(buf, parent, child, time, tolerance); + } + if let Some(buf) = self.buffers.get(&(child.to_string(), parent.to_string())) { + return self + .sample(buf, child, parent, time, tolerance) + .map(|t| t.inverse()); + } + None + } + + fn get( + &self, + parent: &str, + child: &str, + time: Option, + tolerance: Option, + ) -> Option { + if let Some(direct) = self.edge(parent, child, time, tolerance) { + return Some(direct); + } + let path = self.bfs(parent, child, time, tolerance)?; + let mut steps = path.into_iter(); + let first = steps.next()?; + Some(steps.fold(first, |acc, step| acc.compose(&step))) + } + + // Shortest path of edges from parent to child (Python's BFS over the graph). + fn bfs( + &self, + parent: &str, + child: &str, + time: Option, + tolerance: Option, + ) -> Option> { + let mut queue: VecDeque<(String, Vec)> = VecDeque::new(); + queue.push_back((parent.to_string(), Vec::new())); + let mut visited: HashSet = HashSet::new(); + visited.insert(parent.to_string()); + + while let Some((frame, path)) = queue.pop_front() { + if frame == child { + return Some(path); + } + for next in self.connections(&frame) { + if visited.insert(next.clone()) { + if let Some(edge) = self.edge(&frame, &next, time, tolerance) { + let mut extended = path.clone(); + extended.push(edge); + queue.push_back((next, extended)); + } + } + } + } + None + } +} + +/// A cheap-to-clone handle for querying the transform graph. +/// +/// Obtain one from `Builder::tf` (or a `#[tf]` field on a `#[derive(Module)]` +/// struct). The graph is filled in the background as `/tf` messages arrive. +#[derive(Clone)] +pub struct Tf { + buffer: Arc>, +} + +impl Tf { + /// The transform from `parent` to `child`. + /// + /// `time` selects the sample nearest that stamp (latest sample when `None`), + /// and `tolerance` bounds how far that sample may be in seconds. Returns + /// `None` when no path connects the frames or no sample is within tolerance. + pub fn get( + &self, + parent: &str, + child: &str, + time: Option, + tolerance: Option, + ) -> Option { + self.buffer + .read() + .expect("tf buffer lock poisoned") + .get(parent, child, time, tolerance) + } + + /// The latest available transform from `parent` to `child`. + pub fn get_latest(&self, parent: &str, child: &str) -> Option { + self.get(parent, child, None, None) + } +} + +// Decodes /tf messages into the shared graph. Registered as a Route so the +// module's existing recv loop dispatches tf traffic to it. +struct TfRoute { + topic: String, + buffer: Arc>, +} + +impl Route for TfRoute { + fn try_dispatch(&self, data: &[u8]) { + let msg = match lcm_msgs::tf2_msgs::TFMessage::decode(data) { + Ok(msg) => msg, + Err(e) => { + crate::error_throttled!( + Duration::from_secs(1), + topic = %self.topic, + error = %e, + "tf decode error" + ); + return; + } + }; + let mut buffer = self.buffer.write().expect("tf buffer lock poisoned"); + for st in &msg.transforms { + let t = &st.transform.translation; + let q = &st.transform.rotation; + let iso = Isometry3::from_parts( + Translation3::new(t.x, t.y, t.z), + UnitQuaternion::from_quaternion(Quaternion::new(q.w, q.x, q.y, q.z)), + ); + let ts = st.header.stamp.sec as f64 + st.header.stamp.nsec as f64 * 1e-9; + buffer.receive(&st.header.frame_id, &st.child_frame_id, ts, iso); + } + } +} + +// Builds the shared graph plus the handle and the route that feeds it. +pub(crate) fn tf_subscription(topic: String, buffer_size: f64) -> (Tf, Box) { + let buffer = Arc::new(RwLock::new(MultiTBuffer::new(buffer_size))); + let tf = Tf { + buffer: Arc::clone(&buffer), + }; + let route = Box::new(TfRoute { topic, buffer }); + (tf, route) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::f64::consts::PI; + + fn tf_with(buffer_size: f64) -> (Tf, MultiHandle) { + let buffer = Arc::new(RwLock::new(MultiTBuffer::new(buffer_size))); + ( + Tf { + buffer: Arc::clone(&buffer), + }, + MultiHandle { buffer }, + ) + } + + // Test-only writer that bypasses LCM and pushes edges straight into the graph. + struct MultiHandle { + buffer: Arc>, + } + + impl MultiHandle { + fn add(&self, parent: &str, child: &str, ts: f64, xyz: (f64, f64, f64), yaw: f64) { + let iso = Isometry3::from_parts( + Translation3::new(xyz.0, xyz.1, xyz.2), + UnitQuaternion::from_euler_angles(0.0, 0.0, yaw), + ); + self.buffer.write().unwrap().receive(parent, child, ts, iso); + } + } + + #[test] + fn direct_edge() { + let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + h.add("base_link", "arm", 1.0, (1.0, -1.0, 0.0), 0.0); + let t = tf.get_latest("base_link", "arm").unwrap(); + assert!((t.translation().x - 1.0).abs() < 1e-9); + assert!((t.translation().y + 1.0).abs() < 1e-9); + assert_eq!(t.parent, "base_link"); + assert_eq!(t.child, "arm"); + } + + #[test] + fn reverse_edge_returns_inverse() { + let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + h.add("base_link", "arm", 1.0, (1.0, 2.0, 3.0), 0.0); + let inv = tf.get_latest("arm", "base_link").unwrap(); + assert!((inv.translation().x + 1.0).abs() < 1e-9); + assert!((inv.translation().y + 2.0).abs() < 1e-9); + assert!((inv.translation().z + 3.0).abs() < 1e-9); + assert_eq!(inv.parent, "arm"); + assert_eq!(inv.child, "base_link"); + } + + // Mirrors test_tf_ros_example: a 30-degree yaw then a pure translation. + #[test] + fn composes_ros_example_chain() { + let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + h.add("base_link", "arm", 1.0, (1.0, -1.0, 0.0), PI / 6.0); + h.add("arm", "end_effector", 1.0, (1.0, 1.0, 0.0), 0.0); + let t = tf.get_latest("base_link", "end_effector").unwrap(); + assert!( + (t.translation().x - 1.366).abs() < 1e-3, + "{}", + t.translation().x + ); + assert!( + (t.translation().y - 0.366).abs() < 1e-3, + "{}", + t.translation().y + ); + assert_eq!(t.parent, "base_link"); + assert_eq!(t.child, "end_effector"); + } + + // Mirrors test_tf_main: world->robot->sensor multi-hop composition. + #[test] + fn composes_multi_hop_chain() { + let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + h.add("world", "robot", 1.0, (1.0, 2.0, 3.0), 0.0); + h.add("robot", "sensor", 1.0, (0.5, 0.0, 0.2), PI / 2.0); + let t = tf.get_latest("world", "sensor").unwrap(); + assert!((t.translation().x - 1.5).abs() < 1e-3); + assert!((t.translation().y - 2.0).abs() < 1e-3); + assert!((t.translation().z - 3.2).abs() < 1e-3); + } + + #[test] + fn missing_path_returns_none() { + let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + h.add("world", "robot", 1.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.get_latest("world", "unconnected").is_none()); + } + + #[test] + fn identity_for_same_frame() { + let (tf, _h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + // No query time: identity is stamped now, not the epoch. + let t = tf.get_latest("base_link", "base_link").unwrap(); + assert!((t.translation().norm()).abs() < 1e-12); + assert!( + t.ts > 0.0, + "identity ts should be a fresh stamp, got {}", + t.ts + ); + // Explicit query time is echoed back. + let at = tf.get("base_link", "base_link", Some(42.0), None).unwrap(); + assert!((at.ts - 42.0).abs() < 1e-9); + } + + #[test] + fn time_query_picks_nearest_sample() { + let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + h.add("a", "b", 10.0, (1.0, 0.0, 0.0), 0.0); + h.add("a", "b", 20.0, (2.0, 0.0, 0.0), 0.0); + let near_10 = tf.get("a", "b", Some(11.0), None).unwrap(); + assert!((near_10.translation().x - 1.0).abs() < 1e-9); + let near_20 = tf.get("a", "b", Some(18.0), None).unwrap(); + assert!((near_20.translation().x - 2.0).abs() < 1e-9); + } + + #[test] + fn time_query_outside_tolerance_returns_none() { + let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + h.add("a", "b", 10.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.get("a", "b", Some(50.0), Some(1.0)).is_none()); + assert!(tf.get("a", "b", Some(10.5), Some(1.0)).is_some()); + } + + #[test] + fn prunes_samples_outside_window() { + let mut buf = TBuffer::new(5.0); + buf.add(1.0, Isometry3::identity()); + buf.add(2.0, Isometry3::identity()); + buf.add(10.0, Isometry3::identity()); + // The window is [5.0, 10.0]; the 1.0 and 2.0 samples are dropped. + assert_eq!(buf.samples.len(), 1); + assert!((buf.last().unwrap().ts - 10.0).abs() < 1e-9); + } + + #[test] + fn tf_route_decodes_into_graph() { + use lcm_msgs::geometry_msgs::{ + Quaternion as LQuat, Transform as LTransform, Vector3 as LVec3, + }; + use lcm_msgs::std_msgs::{Header, Time}; + use lcm_msgs::tf2_msgs::TFMessage; + + let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_BUFFER_SIZE); + let msg = TFMessage { + transforms: vec![lcm_msgs::geometry_msgs::TransformStamped { + header: Header { + seq: 0, + stamp: Time { + sec: 5, + nsec: 500_000_000, + }, + frame_id: "base_link".to_string(), + }, + child_frame_id: "mid360_link".to_string(), + transform: LTransform { + translation: LVec3 { + x: 0.1, + y: 0.2, + z: 0.3, + }, + rotation: LQuat { + x: 0.0, + y: 0.0, + z: 0.0, + w: 1.0, + }, + }, + }], + }; + route.try_dispatch(&msg.encode()); + + let t = tf.get_latest("base_link", "mid360_link").unwrap(); + assert!((t.translation().x - 0.1).abs() < 1e-9); + assert!((t.translation().y - 0.2).abs() < 1e-9); + assert!((t.translation().z - 0.3).abs() < 1e-9); + assert!((t.ts - 5.5).abs() < 1e-9); + } +} From 6a5564e9fa18879390b910230a6cd76bdce05081 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Thu, 2 Jul 2026 11:20:26 -0700 Subject: [PATCH 02/26] Use base to sensor transform --- dimos/hardware/sensors/lidar/fastlio2/module.py | 16 ++++++++++++++-- dimos/hardware/sensors/lidar/pointlio/module.py | 13 ++++++++++++- dimos/navigation/basic_path_follower/module.py | 7 ++++++- .../navigation/nav_3d/mls_planner/goal_relay.py | 12 +++++++++--- .../blueprints/navigation/unitree_go2_nav_3d.py | 5 ++++- 5 files changed, 45 insertions(+), 8 deletions(-) diff --git a/dimos/hardware/sensors/lidar/fastlio2/module.py b/dimos/hardware/sensors/lidar/fastlio2/module.py index 6f0694ad4a..0c6cd5fc2d 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/module.py +++ b/dimos/hardware/sensors/lidar/fastlio2/module.py @@ -28,6 +28,7 @@ import time from typing import TYPE_CHECKING, Literal +import numpy as np from pydantic import Field from reactivex.disposable import Disposable @@ -147,6 +148,7 @@ def start(self) -> None: ) def _on_odom_for_tf(self, msg: Odometry) -> None: + ts = msg.ts or time.time() self.tf.publish( Transform( frame_id=self.frame_id, @@ -162,8 +164,18 @@ def _on_odom_for_tf(self, msg: Odometry) -> None: msg.pose.orientation.z, msg.pose.orientation.w, ), - ts=msg.ts or time.time(), - ) + ts=ts, + ), + # The IMU-to-lidar extrinsic. + Transform( + frame_id=self.config.child_frame_id, + child_frame_id=self.config.sensor_frame_id, + translation=Vector3(*self.config.extrinsic_t), + rotation=Quaternion.from_rotation_matrix( + np.asarray(self.config.extrinsic_r).reshape(3, 3) + ), + ts=ts, + ), ) @rpc diff --git a/dimos/hardware/sensors/lidar/pointlio/module.py b/dimos/hardware/sensors/lidar/pointlio/module.py index 9db92e9f6d..bca739029e 100644 --- a/dimos/hardware/sensors/lidar/pointlio/module.py +++ b/dimos/hardware/sensors/lidar/pointlio/module.py @@ -34,6 +34,7 @@ import os from typing import TYPE_CHECKING, Literal +import numpy as np from pydantic import Field from reactivex.disposable import Disposable @@ -201,7 +202,17 @@ def _on_odom_for_tf(self, msg: Odometry) -> None: # Match the odometry ts exactly; no `or time.time()` fallback (a # real ts of 0.0 must not become wall time). ts=msg.ts, - ) + ), + # The IMU-to-lidar extrinsic. + Transform( + frame_id=self.config.child_frame_id, + child_frame_id=self.config.sensor_frame_id, + translation=Vector3(*self.config.extrinsic_t), + rotation=Quaternion.from_rotation_matrix( + np.asarray(self.config.extrinsic_r).reshape(3, 3) + ), + ts=msg.ts, + ), ) @rpc diff --git a/dimos/navigation/basic_path_follower/module.py b/dimos/navigation/basic_path_follower/module.py index 5b0b4c491f..601599a590 100644 --- a/dimos/navigation/basic_path_follower/module.py +++ b/dimos/navigation/basic_path_follower/module.py @@ -40,6 +40,7 @@ class BasicPathFollowerConfig(ModuleConfig): + base_frame: str = "base_link" speed: float = 0.5 control_frequency: float = 10.0 goal_tolerance: float = 0.3 @@ -98,8 +99,12 @@ def stop(self) -> None: super().stop() def _on_odometry(self, msg: Odometry) -> None: + # Steer from the robot base pose, not the LIO body frame. + base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, 1.0) + if base is None: + return with self._lock: - self._current_odom = msg.to_pose_stamped() + self._current_odom = base.to_pose(ts=msg.ts) def _on_path(self, path: Path) -> None: # The planner owns path safety: it sends the route as far as it is safe, diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py index 905c4f8973..e0fcfde8b5 100644 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -25,11 +25,14 @@ class GoalRelayConfig(ModuleConfig): - pass + base_frame: str = "base_link" class GoalRelay(Module): - """Adapt odometry and goal points to the planner's PoseStamped inputs.""" + """Adapt odometry and goal points to the planner's PoseStamped inputs. + + Odometry is corrected to the robot base frame via tf. + """ config: GoalRelayConfig @@ -46,7 +49,10 @@ def start(self) -> None: self.register_disposable(Disposable(self.goal.subscribe(self._on_goal))) def _on_odometry(self, msg: Odometry) -> None: - self.start_pose.publish(msg.to_pose_stamped()) + base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, 1.0) + if base is None: + return + self.start_pose.publish(base.to_pose(ts=msg.ts)) def _on_goal(self, point: PointStamped) -> None: self.goal_pose.publish(point.to_pose_stamped()) diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 1fae6b4ad3..3c73dbc892 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -27,6 +27,7 @@ from dimos.navigation.nav_3d.mls_planner.mls_planner_native import MLSPlannerNative from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import rerun_config from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.robot.unitree.go2.go2_mid360_static_transforms import Go2Mid360StaticTf from dimos.visualization.vis_module import vis_module voxel_size = 0.08 @@ -91,7 +92,9 @@ def _static_robot_body(rr: Any) -> list[Any]: (GO2Connection, "odom", "odom_go2"), ] ), - PointLio.blueprint(body_frame_id="body"), + PointLio.blueprint(), + # Publishes the rig's mount frames onto tf. + Go2Mid360StaticTf.blueprint(), RayTracingVoxelMap.blueprint( voxel_size=voxel_size, emit_every=1, From 1adf95b87a5212d381797b6d860bc91379a7b366 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Thu, 2 Jul 2026 11:42:51 -0700 Subject: [PATCH 03/26] Show robot and sensor axis --- .../nav_3d/mls_planner/utils/plan_rrd.py | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) 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 b0cc2dec59..26d91ac778 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -31,13 +31,25 @@ from dimos.memory2.store.sqlite import SqliteStore from dimos.memory2.transform import FnTransformer from dimos.memory2.type.observation import Observation +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation +from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.navigation.nav_3d.mls_planner.mls_planner import MLSPlanner +from dimos.protocol.tf.tf import MultiTBuffer from dimos.utils.data import resolve_named_path TIMELINE = "ts" +# Body-frame axis-triad length for the odometry transform (m). +ODOM_AXIS_LEN = 0.5 + +# Mount frames as recorded on the tf stream. +BASE_FRAME = "base_link" +SENSOR_FRAME = "mid360_link" + # Distinct path colors for overlaid configurations, config 0 first. PATH_PALETTE = [ [0, 255, 0], @@ -107,6 +119,63 @@ def _log_path_wp(waypoints: NDArray[np.float32] | None, entity: str, color: list rr.log(entity, rr.LineStrips3D([points], colors=[color], radii=0.05)) +def _base_from_body(store: SqliteStore) -> Transform | None: + """The body -> base_link transform from the recording's static mount frames. + + Treats the LIO body frame as the sensor frame. The few-cm IMU-to-lidar + extrinsic is invisible at triad scale. + """ + buffer = MultiTBuffer() + try: + for i, obs in enumerate(store.stream("tf", TFMessage).order_by("ts")): + buffer.receive_transform(*obs.data.transforms) + if i >= 20: + break + except Exception as e: + print(f"no usable tf stream in the recording ({e}); skipping the base_link triad") + return None + return buffer.get(SENSOR_FRAME, BASE_FRAME) + + +def _log_odometry( + pose: tuple[float, ...], + ts: float, + trail: list[tuple[float, float, float]], + base_from_body: Transform | None, +) -> None: + """Log the odometry pose as a moving body-frame transform with an XYZ axis + triad, plus the trajectory trail growing over time. The triad is a static + child of world/odom, so it inherits this transform and sweeps along the path.""" + px, py, pz, qx, qy, qz, qw = pose + rr.set_time(TIMELINE, timestamp=ts) + rr.log( + "world/odom", + rr.Transform3D(translation=[px, py, pz], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw])), + ) + trail.append((px, py, pz)) + if len(trail) > 1: + rr.log("world/odom_path", rr.LineStrips3D([trail], colors=[[255, 255, 255]], radii=0.015)) + if base_from_body is None: + return + body = Transform( + translation=Vector3(px, py, pz), + rotation=Quaternion(qx, qy, qz, qw), + frame_id="world", + child_frame_id=base_from_body.frame_id, + ts=ts, + ) + base = body + base_from_body + rr.log( + "world/base_link", + rr.Transform3D( + translation=[base.translation.x, base.translation.y, base.translation.z], + quaternion=rr.Quaternion( + xyzw=[base.rotation.x, base.rotation.y, base.rotation.z, base.rotation.w] + ), + ), + ) + + def _clearance_colors(clearance: NDArray[np.float32], clamp_m: float) -> NDArray[np.uint8]: """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) @@ -346,6 +415,27 @@ def main( rr.log("world/goal", rr.Points3D([goal], colors=[[255, 0, 0]], radii=0.1), static=True) + # Static XYZ axis triads in the odometry body frame and the derived + # robot base frame. + base_from_body = _base_from_body(store) + entities = ["world/odom/axes"] + (["world/base_link/axes"] if base_from_body else []) + for entity in entities: + rr.log( + entity, + rr.Arrows3D( + origins=[[0.0, 0.0, 0.0]] * 3, + vectors=[ + [ODOM_AXIS_LEN, 0.0, 0.0], + [0.0, ODOM_AXIS_LEN, 0.0], + [0.0, 0.0, ODOM_AXIS_LEN], + ], + colors=[[255, 0, 0], [0, 255, 0], [0, 0, 255]], + radii=ODOM_AXIS_LEN / 25, + ), + static=True, + ) + odom_trail: list[tuple[float, float, float]] = [] + try: frame = 0 for ray_obs in ray_pipeline: @@ -354,6 +444,7 @@ def main( ref_timing = _process_frame( ray_obs, planners, goal, robot_height, render_voxel, clearance_clamp ) + _log_odometry(ray_obs.pose_tuple, ray_obs.ts, odom_trail, base_from_body) frame += 1 print( f"frame={frame} configs={len(planners)} " From ee6187a41a43493a80eeb0e0d36ce40a2d8fb627 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 7 Jul 2026 17:16:07 -0700 Subject: [PATCH 04/26] Plan from below base link and visual changes --- .../nav_3d/mls_planner/goal_relay.py | 12 +- .../nav_3d/mls_planner/rust/src/main.rs | 2 - .../nav_3d/mls_planner/utils/plan_rrd.py | 103 +++++++++++++----- .../navigation/unitree_go2_nav_3d.py | 15 ++- 4 files changed, 97 insertions(+), 35 deletions(-) diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py index e0fcfde8b5..4ac77140c9 100644 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -26,6 +26,10 @@ class GoalRelayConfig(ModuleConfig): base_frame: str = "base_link" + sensor_frame: str = "mid360_link" + # The lidar's height above the ground. base_link's height is derived from the + # base -> sensor mount transform, so this is the only fixed measurement. + lidar_height: float = 0.0 class GoalRelay(Module): @@ -50,9 +54,13 @@ def start(self) -> None: def _on_odometry(self, msg: Odometry) -> None: base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, 1.0) - if base is None: + mount = self.tf.get(self.config.base_frame, self.config.sensor_frame, msg.ts, 1.0) + if base is None or mount is None: return - self.start_pose.publish(base.to_pose(ts=msg.ts)) + base_height = self.config.lidar_height - mount.translation.z + start = base.to_pose(ts=msg.ts) + start.position.z -= base_height + self.start_pose.publish(start) def _on_goal(self, point: PointStamped) -> None: self.goal_pose.publish(point.to_pose_stamped()) diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs index 66c3afa4b4..85025f6e95 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs @@ -314,8 +314,6 @@ impl Worker { let Some(start) = *self.latest_start.lock().expect("start mutex") else { return; }; - // Ground-project the sensor pose so the start snaps to the supporting surface. - let start = (start.0, start.1, start.2 - self.config.robot_height); let goal = { let mut guard = self.active_goal.lock().expect("goal mutex"); let Some(goal) = *guard else { diff --git a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py index 47436da675..af26cb17fb 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -43,7 +43,7 @@ TIMELINE = "ts" -# Body-frame axis-triad length for the odometry transform (m). +# Axis-triad length for the odometry transform (m). ODOM_AXIS_LEN = 0.5 # Arrow radius as a fraction of the triad length. AXIS_RADIUS_RATIO = 25 @@ -52,6 +52,10 @@ BASE_FRAME = "base_link" SENSOR_FRAME = "mid360_link" +# Robot footprint (m): length is forward (x), width is left (y). +ROBOT_LENGTH = 0.6858 +ROBOT_WIDTH = 0.3175 + # Distinct path colors for overlaid configurations, config 0 first. PATH_PALETTE = [ [0, 255, 0], @@ -121,12 +125,8 @@ def _log_path_wp(waypoints: NDArray[np.float32] | None, entity: str, color: list rr.log(entity, rr.LineStrips3D([points], colors=[color], radii=0.05)) -def _base_from_body(store: SqliteStore) -> Transform | None: - """The body -> base_link transform from the recording's static mount frames. - - Treats the LIO body frame as the sensor frame. The few-cm IMU-to-lidar - extrinsic is invisible at triad scale. - """ +def _base_from_sensor(store: SqliteStore) -> Transform | None: + """The sensor -> base_link transform from the recording's static mount frames.""" buffer = MultiTBuffer() try: for i, obs in enumerate(store.stream("tf", TFMessage).order_by("ts")): @@ -139,34 +139,40 @@ def _base_from_body(store: SqliteStore) -> Transform | None: return buffer.get(SENSOR_FRAME, BASE_FRAME) +def _base_pose(pose: tuple[float, ...], ts: float, base_from_sensor: Transform) -> Transform: + """Compose the odometry pose with the recorded mount frames into world -> base_link.""" + px, py, pz, qx, qy, qz, qw = pose + sensor = Transform( + translation=Vector3(px, py, pz), + rotation=Quaternion(qx, qy, qz, qw), + frame_id="world", + child_frame_id=base_from_sensor.frame_id, + ts=ts, + ) + return sensor + base_from_sensor + + def _log_odometry( pose: tuple[float, ...], ts: float, trail: list[tuple[float, float, float]], - base_from_body: Transform | None, + base_from_sensor: Transform | None, ) -> None: - """Log the odometry pose as a moving body-frame transform with an XYZ axis + """Log the sensor pose as a moving mid360_link transform with an XYZ axis triad, plus the trajectory trail growing over time. The triad is a static - child of world/odom, so it inherits this transform and sweeps along the path.""" + child of world/mid360_link, so it inherits this transform and sweeps along the path.""" px, py, pz, qx, qy, qz, qw = pose rr.set_time(TIMELINE, timestamp=ts) rr.log( - "world/odom", + "world/mid360_link", rr.Transform3D(translation=[px, py, pz], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw])), ) trail.append((px, py, pz)) if len(trail) > 1: - rr.log("world/odom_path", rr.LineStrips3D([trail], colors=[[255, 255, 255]], radii=0.015)) - if base_from_body is None: + rr.log("world/mid360_path", rr.LineStrips3D([trail], colors=[[255, 255, 255]], radii=0.015)) + if base_from_sensor is None: return - body = Transform( - translation=Vector3(px, py, pz), - rotation=Quaternion(qx, qy, qz, qw), - frame_id="world", - child_frame_id=base_from_body.frame_id, - ts=ts, - ) - base = body + base_from_body + base = _base_pose(pose, ts, base_from_sensor) rr.log( "world/base_link", rr.Transform3D( @@ -279,6 +285,7 @@ def _process_frame( ray_obs: Observation[PointCloud2], planners: list[tuple[str, list[int], MLSPlanner]], goal: tuple[float, float, float], + base_from_sensor: Transform | None, robot_height: float, render_voxel: float, clearance_clamp: float, @@ -288,7 +295,20 @@ def _process_frame( assert ray_obs.pose_tuple is not None bounds = ray_obs.tags["region_bounds"] px, py, pz, *_ = ray_obs.pose_tuple - start = (float(px), float(py), float(pz) - robot_height) + # Plan from the robot base, ground-projected to the supporting surface. Without + # a tf stream fall back to the sensor pose dropped by the robot height. + if base_from_sensor is not None: + base = _base_pose(ray_obs.pose_tuple, ray_obs.ts, base_from_sensor) + # The mount transform gives the lidar's height above base_link, so the + # base sits this far below the lidar's known ground height. + base_height = robot_height - base_from_sensor.inverse().translation.z + start = ( + float(base.translation.x), + float(base.translation.y), + float(base.translation.z) - base_height, + ) + else: + 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) @@ -359,7 +379,9 @@ def main( help="Min occupied neighbors a surface voxel needs to be emitted; " "0 emits all, higher drops isolated returns", ), - robot_height: float = typer.Option(0.3, "--robot-height", help="Robot height (m)"), + robot_height: float = typer.Option( + 0.45, "--robot-height", help="Robot height, ground to tallest point / lidar (m)" + ), max_overhead: float = typer.Option( 2.0, "--max-overhead", help="Ignore surface more than this far above the sensor (m)" ), @@ -455,10 +477,12 @@ def main( rr.log("world/goal", rr.Points3D([goal], colors=[[255, 0, 0]], radii=0.1), static=True) - # Static XYZ axis triads in the odometry body frame and the derived + # Static XYZ axis triads in the odometry sensor frame and the derived # robot base frame. - base_from_body = _base_from_body(store) - entities = ["world/odom/axes"] + (["world/base_link/axes"] if base_from_body else []) + base_from_sensor = _base_from_sensor(store) + entities = ["world/mid360_link/axes"] + ( + ["world/base_link/axes"] if base_from_sensor else [] + ) for entity in entities: rr.log( entity, @@ -474,7 +498,29 @@ def main( ), static=True, ) - odom_trail: list[tuple[float, float, float]] = [] + # also show the outline of the robot + if base_from_sensor is not None: + rr.log( + "world/base_link/outline", + rr.Boxes3D( + half_sizes=[ROBOT_LENGTH / 2, ROBOT_WIDTH / 2, robot_height / 2], + colors=[(0, 255, 127)], + ), + static=True, + ) + # Light red clearance cylinder centered on the robot base. + # wall_clearance is the planner's proxy for the robot radius. + rr.log( + "world/base_link/clearance", + rr.Cylinders3D( + lengths=[robot_height], + radii=[wall_clearance], + colors=[(255, 120, 120, 80)], + fill_mode="solid", + ), + static=True, + ) + sensor_trail: list[tuple[float, float, float]] = [] try: frame = 0 @@ -485,12 +531,13 @@ def main( ray_obs, planners, goal, + base_from_sensor, robot_height, render_voxel, clearance_clamp, ref_clearance, ) - _log_odometry(ray_obs.pose_tuple, ray_obs.ts, odom_trail, base_from_body) + _log_odometry(ray_obs.pose_tuple, ray_obs.ts, sensor_trail, base_from_sensor) frame += 1 print( f"frame={frame} configs={len(planners)} " diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 17fc4c54d2..ad7c9aa59f 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -44,6 +44,12 @@ voxel_size = 0.08 +# Robot footprint (m): length is forward (x), width is left (y). +ROBOT_LENGTH = 0.6858 +ROBOT_WIDTH = 0.3175 +# Ground to the tallest point (the lidar). Drives clearance and the ground projection. +ROBOT_HEIGHT = 0.45 + # Body-frame axis-triad length (m). _axis_len = 0.5 # Arrow radius as a fraction of the triad length. @@ -87,7 +93,10 @@ def _render_path(msg: Any) -> Any: def _static_robot_body(rr: Any) -> list[Any]: """Go2-shaped box on pointlio's sensor frame, counter-rotated for the lidar pitch.""" return [ - rr.Boxes3D(half_sizes=[0.35, 0.155, 0.2], colors=[(0, 255, 127)]), + rr.Boxes3D( + half_sizes=[ROBOT_LENGTH / 2, ROBOT_WIDTH / 2, ROBOT_HEIGHT / 2], + colors=[(0, 255, 127)], + ), rr.Transform3D( parent_frame="tf#/mid360_link", rotation=rr.RotationAxisAngle(axis=(0, 1, 0), degrees=-math.degrees(MID360_PITCH_DOWN)), @@ -178,7 +187,7 @@ def _static_sensor_axes(rr: Any) -> list[Any]: MLSPlannerNative.blueprint( world_frame="odom", voxel_size=voxel_size, - robot_height=0.3, + robot_height=ROBOT_HEIGHT, surface_closing_radius=0.3, wall_clearance_m=0.1, wall_buffer_m=0.75, @@ -187,7 +196,7 @@ def _static_sensor_axes(rr: Any) -> list[Any]: step_penalty_weight=4.0, viz_publish_hz=0.0, ).remappings([(MLSPlannerNative, "global_map", "global_map_unused")]), - GoalRelay.blueprint(), + GoalRelay.blueprint(lidar_height=ROBOT_HEIGHT), BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6).remappings( [(BasicPathFollower, "odometry", "body_odometry")] ), From ae3aa87428c2c5c60de7323d20c97fe113e1ad29 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 7 Jul 2026 18:26:43 -0700 Subject: [PATCH 05/26] Save in path --- .../unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index ad7c9aa59f..ffcc52c7c8 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -18,9 +18,9 @@ from datetime import datetime import math import os +from pathlib import Path from typing import Any -from dimos.constants import STATE_DIR from dimos.core.coordination.blueprints import autoconnect from dimos.core.global_config import global_config from dimos.core.stream import In @@ -75,7 +75,7 @@ def _recording_db_path() -> str: stamp = ( now.strftime("%Y-%m-%d") + "_" + now.strftime("%I-%M%p").lower() + "-" + now.strftime("%Z") ) - return str(STATE_DIR / "recordings" / stamp / "mem2.db") + return str(Path("recordings") / stamp / "mem2.db") def _render_global_map(msg: Any) -> Any: From 15688160388f13abd5a53c3c72f1e73332e864ec Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Wed, 8 Jul 2026 15:04:30 -0700 Subject: [PATCH 06/26] Remove body --- dimos/hardware/sensors/lidar/fastlio2/module.py | 17 ++--------------- dimos/hardware/sensors/lidar/pointlio/module.py | 13 +------------ .../blueprints/navigation/unitree_go2_nav_3d.py | 2 +- 3 files changed, 4 insertions(+), 28 deletions(-) diff --git a/dimos/hardware/sensors/lidar/fastlio2/module.py b/dimos/hardware/sensors/lidar/fastlio2/module.py index 7e072c9cb0..37b691d135 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/module.py +++ b/dimos/hardware/sensors/lidar/fastlio2/module.py @@ -25,10 +25,8 @@ from __future__ import annotations import os -import time from typing import TYPE_CHECKING, Literal -import numpy as np from pydantic import Field from reactivex.disposable import Disposable @@ -145,7 +143,6 @@ def start(self) -> None: ) def _on_odom_for_tf(self, msg: Odometry) -> None: - ts = msg.ts or time.time() self.tf.publish( Transform( frame_id=self.frame_id, @@ -161,18 +158,8 @@ def _on_odom_for_tf(self, msg: Odometry) -> None: msg.pose.orientation.z, msg.pose.orientation.w, ), - ts=ts, - ), - # The IMU-to-lidar extrinsic. - Transform( - frame_id=self.config.child_frame_id, - child_frame_id=self.config.sensor_frame_id, - translation=Vector3(*self.config.extrinsic_t), - rotation=Quaternion.from_rotation_matrix( - np.asarray(self.config.extrinsic_r).reshape(3, 3) - ), - ts=ts, - ), + ts=msg.ts, + ) ) @rpc diff --git a/dimos/hardware/sensors/lidar/pointlio/module.py b/dimos/hardware/sensors/lidar/pointlio/module.py index 31c7463343..352768254b 100644 --- a/dimos/hardware/sensors/lidar/pointlio/module.py +++ b/dimos/hardware/sensors/lidar/pointlio/module.py @@ -34,7 +34,6 @@ import os from typing import TYPE_CHECKING, Literal -import numpy as np from pydantic import Field from reactivex.disposable import Disposable @@ -199,17 +198,7 @@ def _on_odom_for_tf(self, msg: Odometry) -> None: # Match the odometry ts exactly; no `or time.time()` fallback (a # real ts of 0.0 must not become wall time). ts=msg.ts, - ), - # The IMU-to-lidar extrinsic. - Transform( - frame_id=self.config.child_frame_id, - child_frame_id=self.config.sensor_frame_id, - translation=Vector3(*self.config.extrinsic_t), - rotation=Quaternion.from_rotation_matrix( - np.asarray(self.config.extrinsic_r).reshape(3, 3) - ), - ts=msg.ts, - ), + ) ) @rpc diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index ffcc52c7c8..1e4e99f442 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -171,7 +171,7 @@ def _static_sensor_axes(rr: Any) -> list[Any]: (GO2Connection, "odom", "odom_go2"), ] ), - PointLio.blueprint(child_frame_id="body"), + PointLio.blueprint(), # Publishes the rig's mount frames onto tf. Go2Mid360StaticTf.blueprint(), RayTracingVoxelMap.blueprint( From f42fa107371623e37ad973f35dcee7021703f95a Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Wed, 8 Jul 2026 17:29:37 -0700 Subject: [PATCH 07/26] Refactor tolerance --- .../navigation/basic_path_follower/module.py | 5 ++- .../nav_3d/mls_planner/goal_relay.py | 9 +++-- .../nav_3d/mls_planner/utils/plan_rrd.py | 7 +--- .../navigation/unitree_go2_nav_3d.py | 4 +-- native/rust/dimos-module/src/tf.rs | 34 +++++++++++++------ 5 files changed, 36 insertions(+), 23 deletions(-) diff --git a/dimos/navigation/basic_path_follower/module.py b/dimos/navigation/basic_path_follower/module.py index 601599a590..e3a3c4847f 100644 --- a/dimos/navigation/basic_path_follower/module.py +++ b/dimos/navigation/basic_path_follower/module.py @@ -38,6 +38,9 @@ logger = setup_logger() +# Max staleness (s) for tf lookups against a live odometry stamp. +TF_LOOKUP_TOLERANCE_S = 0.1 + class BasicPathFollowerConfig(ModuleConfig): base_frame: str = "base_link" @@ -100,7 +103,7 @@ def stop(self) -> None: def _on_odometry(self, msg: Odometry) -> None: # Steer from the robot base pose, not the LIO body frame. - base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, 1.0) + base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, TF_LOOKUP_TOLERANCE_S) if base is None: return with self._lock: diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py index 4ac77140c9..619c947a5b 100644 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -23,6 +23,9 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.nav_msgs.Odometry import Odometry +# Max staleness (s) for tf lookups against a live odometry stamp. +TF_LOOKUP_TOLERANCE_S = 0.1 + class GoalRelayConfig(ModuleConfig): base_frame: str = "base_link" @@ -53,8 +56,10 @@ def start(self) -> None: self.register_disposable(Disposable(self.goal.subscribe(self._on_goal))) def _on_odometry(self, msg: Odometry) -> None: - base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, 1.0) - mount = self.tf.get(self.config.base_frame, self.config.sensor_frame, msg.ts, 1.0) + base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, TF_LOOKUP_TOLERANCE_S) + mount = self.tf.get( + self.config.base_frame, self.config.sensor_frame, msg.ts, TF_LOOKUP_TOLERANCE_S + ) if base is None or mount is None: return base_height = self.config.lidar_height - mount.translation.z 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 af26cb17fb..8fec170b4f 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -158,9 +158,7 @@ def _log_odometry( trail: list[tuple[float, float, float]], base_from_sensor: Transform | None, ) -> None: - """Log the sensor pose as a moving mid360_link transform with an XYZ axis - triad, plus the trajectory trail growing over time. The triad is a static - child of world/mid360_link, so it inherits this transform and sweeps along the path.""" + """Log the moving sensor pose and its growing trajectory trail.""" px, py, pz, qx, qy, qz, qw = pose rr.set_time(TIMELINE, timestamp=ts) rr.log( @@ -477,8 +475,6 @@ def main( rr.log("world/goal", rr.Points3D([goal], colors=[[255, 0, 0]], radii=0.1), static=True) - # Static XYZ axis triads in the odometry sensor frame and the derived - # robot base frame. base_from_sensor = _base_from_sensor(store) entities = ["world/mid360_link/axes"] + ( ["world/base_link/axes"] if base_from_sensor else [] @@ -498,7 +494,6 @@ def main( ), static=True, ) - # also show the outline of the robot if base_from_sensor is not None: rr.log( "world/base_link/outline", diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 1e4e99f442..5f571d63a8 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -197,9 +197,7 @@ def _static_sensor_axes(rr: Any) -> list[Any]: viz_publish_hz=0.0, ).remappings([(MLSPlannerNative, "global_map", "global_map_unused")]), GoalRelay.blueprint(lidar_height=ROBOT_HEIGHT), - BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6).remappings( - [(BasicPathFollower, "odometry", "body_odometry")] - ), + BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6), MovementManager.blueprint(), ).global_config(n_workers=10, robot_model="unitree_go2", obstacle_avoidance=False) diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index e5614af248..0a401fa8be 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -1,10 +1,22 @@ -//! A consumer-side transform client for native modules. +// 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. + +//! Consumer-side transform client for native modules. //! -//! Mirrors the Python `dimos.protocol.tf` semantics: a module subscribes to the -//! `/tf` topic, every `TFMessage` edge is buffered per `(parent, child)` pair, -//! and [`Tf::get`] answers a query by composing transforms along the shortest -//! path through the frame graph. Lookups are nearest-in-time within a tolerance, -//! not interpolated, matching the Python buffer. +//! Each `/tf` edge is buffered per `(parent, child)`, and [`Tf::get`] composes +//! the shortest path through the frame graph. Lookups are nearest-in-time within +//! a tolerance, not interpolated. use std::collections::{HashMap, HashSet, VecDeque}; use std::sync::{Arc, RwLock}; @@ -37,7 +49,7 @@ pub struct Transform { } impl Transform { - /// The transform as an `nalgebra` isometry, ready to apply to a point. + /// The transform as an isometry. pub fn isometry(&self) -> Isometry3 { self.iso } @@ -77,7 +89,7 @@ struct Sample { iso: Isometry3, } -// One edge's time-sorted history, pruned to a fixed-duration window. +// One edge's time-sorted history, capped to a fixed-duration window. struct TBuffer { buffer_size: f64, samples: Vec, @@ -231,7 +243,7 @@ impl MultiTBuffer { Some(steps.fold(first, |acc, step| acc.compose(&step))) } - // Shortest path of edges from parent to child (Python's BFS over the graph). + // Shortest path of edges from parent to child. fn bfs( &self, parent: &str, @@ -394,7 +406,7 @@ mod tests { assert_eq!(inv.child, "base_link"); } - // Mirrors test_tf_ros_example: a 30-degree yaw then a pure translation. + // A 30-degree yaw then a pure translation. #[test] fn composes_ros_example_chain() { let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); @@ -415,7 +427,7 @@ mod tests { assert_eq!(t.child, "end_effector"); } - // Mirrors test_tf_main: world->robot->sensor multi-hop composition. + // world->robot->sensor multi-hop composition. #[test] fn composes_multi_hop_chain() { let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); From 92d0fc8f382a65d3d4de701b2afb67dfbed59384 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Thu, 9 Jul 2026 12:58:06 -0700 Subject: [PATCH 08/26] Rust tf example --- dimos/core/native_module.py | 24 ++++- examples/native-modules/rust/Cargo.toml | 4 + .../native-modules/rust/src/tf_listener.rs | 54 ++++++++++ examples/native-modules/rust_tf.py | 101 ++++++++++++++++++ 4 files changed, 182 insertions(+), 1 deletion(-) create mode 100644 examples/native-modules/rust/src/tf_listener.rs create mode 100644 examples/native-modules/rust_tf.py diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index abf43f70fc..7110207b8b 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -112,6 +112,27 @@ class LogFormat(enum.Enum): } +def _tf_channel() -> str: + """Wire channel for the tf stream on the active transport. + + tf can't be a declared port (every Module reserves the name), so the + coordinator advertises its channel to native modules here instead. Match + to the correct transport. + """ + from dimos.core.global_config import global_config + from dimos.core.transport_factory import transport_topic + from dimos.msgs.tf2_msgs.TFMessage import TFMessage + + name = transport_topic("/tf") + if global_config.transport == "zenoh": + from dimos.protocol.pubsub.impl.zenohpubsub import Topic + + return Topic(name, TFMessage).key_expr + from dimos.protocol.pubsub.impl.lcmpubsub import Topic + + return str(Topic(name, TFMessage)) + + class NativeModuleConfig(ModuleConfig): """Configuration for a native (C/C++) subprocess module.""" @@ -255,7 +276,8 @@ def start(self) -> None: assert self._process.stdin is not None if self.config.stdin_config: config_dict = self.config.to_config_dict() - blob: dict[str, Any] = {"topics": topics, "config": config_dict or None} + stdin_topics = {**topics, "tf": _tf_channel()} + blob: dict[str, Any] = {"topics": stdin_topics, "config": config_dict or None} qos = self._collect_output_qos() if qos: blob["qos"] = qos diff --git a/examples/native-modules/rust/Cargo.toml b/examples/native-modules/rust/Cargo.toml index 6bfee2b74f..15398ca349 100644 --- a/examples/native-modules/rust/Cargo.toml +++ b/examples/native-modules/rust/Cargo.toml @@ -11,6 +11,10 @@ path = "src/ping.rs" name = "pong" path = "src/pong.rs" +[[bin]] +name = "tf_listener" +path = "src/tf_listener.rs" + [dependencies] dimos-module = { path = "../../../native/rust/dimos-module" } lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } diff --git a/examples/native-modules/rust/src/tf_listener.rs b/examples/native-modules/rust/src/tf_listener.rs new file mode 100644 index 0000000000..b611710bb4 --- /dev/null +++ b/examples/native-modules/rust/src/tf_listener.rs @@ -0,0 +1,54 @@ +// Copyright 2026 Dimensional Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use dimos_module::{run_with_transport, Module, Tf}; +use tokio::time::{interval, Duration}; + +#[derive(Module)] +#[module(setup = start_lookup)] +struct TfListener { + #[tf] + tf: Tf, +} + +impl TfListener { + async fn start_lookup(&mut self) { + let tf = self.tf.clone(); + tokio::spawn(async move { + let mut ticker = interval(Duration::from_millis(500)); + loop { + ticker.tick().await; + match tf.get_latest("a", "c") { + Some(t) => { + let p = t.translation(); + tracing::info!( + parent = %t.parent, + child = %t.child, + x = p.x, + y = p.y, + z = p.z, + "transform:", + ); + } + None => tracing::info!("a -> c not available yet"), + } + } + }); + } +} + +#[tokio::main] +async fn main() { + run_with_transport::().await; +} diff --git a/examples/native-modules/rust_tf.py b/examples/native-modules/rust_tf.py new file mode 100644 index 0000000000..cee2511e0f --- /dev/null +++ b/examples/native-modules/rust_tf.py @@ -0,0 +1,101 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Python publishes a tf chain, a Rust #[tf] module reads it back. + +Run with: + python examples/native-modules/rust_tf.py +""" + +from __future__ import annotations + +import asyncio +import math +from pathlib import Path +import time + +from dimos.core.coordination.blueprints import autoconnect +from dimos.core.coordination.module_coordinator import ModuleCoordinator +from dimos.core.core import rpc +from dimos.core.module import Module +from dimos.core.native_module import NativeModule, NativeModuleConfig +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 + +_RUST_DIR = Path(__file__).parent / "rust" +_EXAMPLES = _RUST_DIR / "target" / "release" +_BUILD = "cargo build --release" + + +class TfProducer(Module): + """Publishes a time-varying a -> b -> c transform chain onto /tf.""" + + _running: bool = False + + @rpc + def start(self) -> None: + super().start() + self._running = True + self.spawn(self._publish_loop()) + + async def _publish_loop(self) -> None: + start = time.time() + while self._running: + t = time.time() - start + now = time.time() + self.tf.publish( + Transform( + translation=Vector3(0.0, math.cos(t), math.sin(t)), + frame_id="a", + child_frame_id="b", + ts=now, + ), + Transform( + translation=Vector3(1.0, 0.0, 0.0), + frame_id="b", + child_frame_id="c", + ts=now, + ), + ) + await asyncio.sleep(0.1) + + @rpc + def stop(self) -> None: + self._running = False + super().stop() + + +class TfListenerConfig(NativeModuleConfig): + executable: str = str(_EXAMPLES / "tf_listener") + build_command: str = _BUILD + cwd: str = str(_RUST_DIR) + stdin_config: bool = True + + +class TfListenerModule(NativeModule): + """Rust module that looks up a -> c and logs it. + + Expect to see (1.0, cos(t), sin(t)) + """ + + config: TfListenerConfig + + +def blueprint(): + return autoconnect(TfProducer.blueprint(), TfListenerModule.blueprint()) + + +if __name__ == "__main__": + bp = blueprint().global_config(viewer="none") + ModuleCoordinator.build(bp).loop() From 67f9662b74512a5e0600a5fe42da0028f489871d Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Thu, 9 Jul 2026 13:43:23 -0700 Subject: [PATCH 09/26] Clean up --- dimos/core/native_module.py | 24 ++----------------- dimos/core/transport_factory.py | 17 +++++++++++++ .../navigation/basic_path_follower/module.py | 4 +--- dimos/navigation/constants.py | 20 ++++++++++++++++ .../nav_3d/mls_planner/goal_relay.py | 4 +--- .../nav_3d/mls_planner/utils/plan_rrd.py | 5 +--- .../navigation/unitree_go2_nav_3d.py | 7 +----- dimos/robot/unitree/go2/constants.py | 21 ++++++++++++++++ examples/native-modules/rust_tf.py | 8 +++---- native/rust/dimos-module/src/tf.rs | 8 +++---- 10 files changed, 71 insertions(+), 47 deletions(-) create mode 100644 dimos/navigation/constants.py create mode 100644 dimos/robot/unitree/go2/constants.py diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index 7110207b8b..b22e4f7761 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -60,6 +60,7 @@ class MyCppModule(NativeModule): from dimos.core.core import rpc from dimos.core.global_config import global_config from dimos.core.module import Module, ModuleConfig +from dimos.core.transport_factory import tf_channel from dimos.utils.logging_config import setup_logger if sys.platform.startswith("linux"): @@ -112,27 +113,6 @@ class LogFormat(enum.Enum): } -def _tf_channel() -> str: - """Wire channel for the tf stream on the active transport. - - tf can't be a declared port (every Module reserves the name), so the - coordinator advertises its channel to native modules here instead. Match - to the correct transport. - """ - from dimos.core.global_config import global_config - from dimos.core.transport_factory import transport_topic - from dimos.msgs.tf2_msgs.TFMessage import TFMessage - - name = transport_topic("/tf") - if global_config.transport == "zenoh": - from dimos.protocol.pubsub.impl.zenohpubsub import Topic - - return Topic(name, TFMessage).key_expr - from dimos.protocol.pubsub.impl.lcmpubsub import Topic - - return str(Topic(name, TFMessage)) - - class NativeModuleConfig(ModuleConfig): """Configuration for a native (C/C++) subprocess module.""" @@ -276,7 +256,7 @@ def start(self) -> None: assert self._process.stdin is not None if self.config.stdin_config: config_dict = self.config.to_config_dict() - stdin_topics = {**topics, "tf": _tf_channel()} + stdin_topics = {**topics, "tf": tf_channel()} blob: dict[str, Any] = {"topics": stdin_topics, "config": config_dict or None} qos = self._collect_output_qos() if qos: diff --git a/dimos/core/transport_factory.py b/dimos/core/transport_factory.py index f8dac21a8c..7bc9435e19 100644 --- a/dimos/core/transport_factory.py +++ b/dimos/core/transport_factory.py @@ -54,6 +54,23 @@ def transport_topic(name: str, g: GlobalConfig = global_config) -> str: return name if name.startswith("/") else "/" + name +def tf_channel(g: GlobalConfig = global_config) -> str: + """Wire channel for the tf stream on the active backend. + + tf can't be a declared port because every Module reserves the name, so + native modules are handed this channel directly. Backends format the type + suffix differently, so match the active one. + """ + from dimos.msgs.tf2_msgs.TFMessage import TFMessage + + name = transport_topic("/tf", g) + if g.transport == "zenoh": + return ZenohTopic(name, TFMessage).key_expr + from dimos.protocol.pubsub.impl.lcmpubsub import Topic + + return str(Topic(name, TFMessage)) + + # High-rate sensor streams: drop stale frames under congestion, never stall the # publisher. Matched by message type since that is what makes them high-rate. _LATEST_WINS_TYPES = ("sensor_msgs.Image", "sensor_msgs.PointCloud2") diff --git a/dimos/navigation/basic_path_follower/module.py b/dimos/navigation/basic_path_follower/module.py index e3a3c4847f..103e1c960a 100644 --- a/dimos/navigation/basic_path_follower/module.py +++ b/dimos/navigation/basic_path_follower/module.py @@ -33,14 +33,12 @@ 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.navigation.constants import TF_LOOKUP_TOLERANCE_S from dimos.utils.logging_config import setup_logger from dimos.utils.trigonometry import angle_diff logger = setup_logger() -# Max staleness (s) for tf lookups against a live odometry stamp. -TF_LOOKUP_TOLERANCE_S = 0.1 - class BasicPathFollowerConfig(ModuleConfig): base_frame: str = "base_link" diff --git a/dimos/navigation/constants.py b/dimos/navigation/constants.py new file mode 100644 index 0000000000..151691f783 --- /dev/null +++ b/dimos/navigation/constants.py @@ -0,0 +1,20 @@ +# 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. + +"""Shared navigation constants.""" + +# Max staleness in seconds for tf lookups against a live odometry stamp. Must +# comfortably exceed the coarsest static-transform publish period so leading-edge +# phase or modest clock skew does not silently drop frames. +TF_LOOKUP_TOLERANCE_S = 0.25 diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py index 619c947a5b..23e035b63b 100644 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -22,9 +22,7 @@ from dimos.msgs.geometry_msgs.PointStamped import PointStamped from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.nav_msgs.Odometry import Odometry - -# Max staleness (s) for tf lookups against a live odometry stamp. -TF_LOOKUP_TOLERANCE_S = 0.1 +from dimos.navigation.constants import TF_LOOKUP_TOLERANCE_S class GoalRelayConfig(ModuleConfig): 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 8fec170b4f..1ee91f45c1 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -39,6 +39,7 @@ from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.navigation.nav_3d.mls_planner.mls_planner import MLSPlanner from dimos.protocol.tf.tf import MultiTBuffer +from dimos.robot.unitree.go2.constants import ROBOT_LENGTH, ROBOT_WIDTH from dimos.utils.data import resolve_named_path TIMELINE = "ts" @@ -52,10 +53,6 @@ BASE_FRAME = "base_link" SENSOR_FRAME = "mid360_link" -# Robot footprint (m): length is forward (x), width is left (y). -ROBOT_LENGTH = 0.6858 -ROBOT_WIDTH = 0.3175 - # Distinct path colors for overlaid configurations, config 0 first. PATH_PALETTE = [ [0, 255, 0], diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index 940057bb55..469ba00b5f 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -38,6 +38,7 @@ from dimos.navigation.nav_3d.mls_planner.mls_planner_native import MLSPlannerNative from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import rerun_config from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT, ROBOT_LENGTH, ROBOT_WIDTH from dimos.robot.unitree.go2.go2_mid360_static_transforms import ( MID360_PITCH_DOWN, Go2Mid360StaticTf, @@ -46,12 +47,6 @@ voxel_size = 0.08 -# Robot footprint (m): length is forward (x), width is left (y). -ROBOT_LENGTH = 0.6858 -ROBOT_WIDTH = 0.3175 -# Ground to the tallest point (the lidar). Drives clearance and the ground projection. -ROBOT_HEIGHT = 0.45 - # Body-frame axis-triad length (m). _axis_len = 0.5 # Arrow radius as a fraction of the triad length. diff --git a/dimos/robot/unitree/go2/constants.py b/dimos/robot/unitree/go2/constants.py new file mode 100644 index 0000000000..92126bb334 --- /dev/null +++ b/dimos/robot/unitree/go2/constants.py @@ -0,0 +1,21 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Physical constants for the Unitree Go2.""" + +# Robot footprint in meters. Length is forward x, width is left y. +ROBOT_LENGTH = 0.6858 +ROBOT_WIDTH = 0.3175 +# Ground to the tallest point. Used for clearance and the ground projection. +ROBOT_HEIGHT = 0.45 diff --git a/examples/native-modules/rust_tf.py b/examples/native-modules/rust_tf.py index cee2511e0f..5a7bcfd2d0 100644 --- a/examples/native-modules/rust_tf.py +++ b/examples/native-modules/rust_tf.py @@ -92,10 +92,8 @@ class TfListenerModule(NativeModule): config: TfListenerConfig -def blueprint(): - return autoconnect(TfProducer.blueprint(), TfListenerModule.blueprint()) - - if __name__ == "__main__": - bp = blueprint().global_config(viewer="none") + bp = autoconnect(TfProducer.blueprint(), TfListenerModule.blueprint()).global_config( + viewer="none" + ) ModuleCoordinator.build(bp).loop() diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index 0a401fa8be..10c3479b45 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -73,7 +73,7 @@ impl Transform { } } - // self (a -> b) followed by other (b -> c) gives a -> c. + // self a -> b, then other b -> c, gives a -> c. fn compose(&self, other: &Transform) -> Transform { Transform { parent: self.parent.clone(), @@ -121,7 +121,7 @@ impl TBuffer { } // Nearest sample in time. On a tie, prefer the later sample. Returns None - // when the closest sample is further than `tolerance` from `ts`. + // when the closest sample is further than tolerance from ts. fn find_closest(&self, ts: f64, tolerance: Option) -> Option<&Sample> { let pos = self.samples.partition_point(|s| s.ts < ts); let prev = pos.checked_sub(1).and_then(|i| self.samples.get(i)); @@ -200,7 +200,7 @@ impl MultiTBuffer { }) } - // A single forward or reverse edge (reverse returns the inverse). + // A single forward or reverse edge. Reverse returns the inverse. fn edge( &self, parent: &str, @@ -487,7 +487,7 @@ mod tests { buf.add(1.0, Isometry3::identity()); buf.add(2.0, Isometry3::identity()); buf.add(10.0, Isometry3::identity()); - // The window is [5.0, 10.0]; the 1.0 and 2.0 samples are dropped. + // The window is [5.0, 10.0]. The 1.0 and 2.0 samples are dropped. assert_eq!(buf.samples.len(), 1); assert!((buf.last().unwrap().ts - 10.0).abs() < 1e-9); } From 03965d1232729e558b00024160bee8d9e9e25797 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 14 Jul 2026 10:40:53 -0700 Subject: [PATCH 10/26] BFS fix --- native/rust/dimos-module/src/tf.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index 10c3479b45..0c4bdc470d 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -261,8 +261,9 @@ impl MultiTBuffer { return Some(path); } for next in self.connections(&frame) { - if visited.insert(next.clone()) { + if !visited.contains(&next) { if let Some(edge) = self.edge(&frame, &next, time, tolerance) { + visited.insert(next.clone()); let mut extended = path.clone(); extended.push(edge); queue.push_back((next, extended)); From c4ce81c6b6d36e68c34c5dee36b434dc86e09e97 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Tue, 14 Jul 2026 11:05:45 -0700 Subject: [PATCH 11/26] Remove unused code --- .../nav_3d/mls_planner/goal_relay.py | 21 ++++-- .../nav_3d/mls_planner/odom_body_frame.py | 62 ----------------- .../mls_planner/test_odom_body_frame.py | 69 ------------------- dimos/robot/all_blueprints.py | 1 - 4 files changed, 16 insertions(+), 137 deletions(-) delete mode 100644 dimos/navigation/nav_3d/mls_planner/odom_body_frame.py delete mode 100644 dimos/navigation/nav_3d/mls_planner/test_odom_body_frame.py diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py index 23e035b63b..a7ec5e9fec 100644 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -14,6 +14,8 @@ from __future__ import annotations +from typing import Any + from reactivex.disposable import Disposable from dimos.core.core import rpc @@ -47,6 +49,10 @@ class GoalRelay(Module): start_pose: Out[PoseStamped] goal_pose: Out[PoseStamped] + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._base_height: float | None = None + @rpc def start(self) -> None: super().start() @@ -55,15 +61,20 @@ def start(self) -> None: def _on_odometry(self, msg: Odometry) -> None: base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, TF_LOOKUP_TOLERANCE_S) - mount = self.tf.get( - self.config.base_frame, self.config.sensor_frame, msg.ts, TF_LOOKUP_TOLERANCE_S - ) - if base is None or mount is None: + base_height = self._base_from_ground() + if base is None or base_height is None: return - base_height = self.config.lidar_height - mount.translation.z start = base.to_pose(ts=msg.ts) start.position.z -= base_height self.start_pose.publish(start) + def _base_from_ground(self) -> float | None: + # The base -> sensor mount is static rig geometry, so resolve it once. + if self._base_height is None: + mount = self.tf.get(self.config.base_frame, self.config.sensor_frame) + if mount is not None: + self._base_height = self.config.lidar_height - mount.translation.z + return self._base_height + def _on_goal(self, point: PointStamped) -> None: self.goal_pose.publish(point.to_pose_stamped()) diff --git a/dimos/navigation/nav_3d/mls_planner/odom_body_frame.py b/dimos/navigation/nav_3d/mls_planner/odom_body_frame.py deleted file mode 100644 index 712665f73c..0000000000 --- a/dimos/navigation/nav_3d/mls_planner/odom_body_frame.py +++ /dev/null @@ -1,62 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from __future__ import annotations - -from pydantic import Field -from reactivex.disposable import Disposable - -from dimos.core.core import rpc -from dimos.core.module import Module, ModuleConfig -from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.nav_msgs.Odometry import Odometry - - -class OdomBodyFrameConfig(ModuleConfig): - # base_link from sensor mount rotation, xyzw. - mount_rotation: list[float] = Field(default_factory=lambda: [0.0, 0.0, 0.0, 1.0]) - body_frame_id: str = "base_link" - - -class OdomBodyFrame(Module): - """Re-express tilted-sensor LIO odometry in the level robot body frame. - - Composes out the fixed mount rotation from the orientation. Position and - twist pass through. - """ - - config: OdomBodyFrameConfig - - odometry: In[Odometry] - body_odometry: Out[Odometry] - - @rpc - def start(self) -> None: - super().start() - self._mount_inv = Quaternion(*self.config.mount_rotation).inverse() - self.register_disposable(Disposable(self.odometry.subscribe(self._on_odometry))) - - def _on_odometry(self, msg: Odometry) -> None: - leveled = msg.orientation * self._mount_inv - self.body_odometry.publish( - Odometry( - ts=msg.ts, - frame_id=msg.frame_id, - child_frame_id=self.config.body_frame_id, - pose=Pose(msg.position, leveled), - twist=msg.twist, - ) - ) diff --git a/dimos/navigation/nav_3d/mls_planner/test_odom_body_frame.py b/dimos/navigation/nav_3d/mls_planner/test_odom_body_frame.py deleted file mode 100644 index 4b483eec71..0000000000 --- a/dimos/navigation/nav_3d/mls_planner/test_odom_body_frame.py +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from types import SimpleNamespace - -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.navigation.nav_3d.mls_planner.odom_body_frame import ( - OdomBodyFrame, - OdomBodyFrameConfig, -) - - -def _level(mount_rotation, orientation): - """Run one odometry message through the handler and return the output. - - Builds the module without its transport so no runtime threads spawn. - """ - module = object.__new__(OdomBodyFrame) - module.config = OdomBodyFrameConfig( - mount_rotation=list(mount_rotation), body_frame_id="base_link" - ) - module._mount_inv = Quaternion(*module.config.mount_rotation).inverse() - captured = [] - module.body_odometry = SimpleNamespace(publish=captured.append) - module._on_odometry( - Odometry( - ts=1.0, - frame_id="odom", - child_frame_id="mid360_link", - pose=Pose(Vector3(1.0, 2.0, 3.0), orientation), - ) - ) - return captured[0] - - -def test_composes_out_the_mount_pitch(): - # A level body reads its own mount tilt as the sensor's world orientation, so - # composing the mount out returns identity. - mount = Quaternion.from_euler(Vector3(0.0, 0.3, 0.0)) - out = _level(mount.to_tuple(), mount) - assert out.orientation.angle_to(Quaternion(0.0, 0.0, 0.0, 1.0)) < 1e-5 - - -def test_preserves_body_yaw_under_mount_tilt(): - # A body yawed by a known angle keeps that yaw after the mount is composed out. - mount = Quaternion.from_euler(Vector3(0.0, 0.3, 0.0)) - body = Quaternion.from_euler(Vector3(0.0, 0.0, 0.7)) - out = _level(mount.to_tuple(), body * mount) - assert out.orientation.angle_to(body) < 1e-5 - - -def test_relabels_child_frame_and_passes_position_through(): - out = _level([0.0, 0.0, 0.0, 1.0], Quaternion(0.0, 0.0, 0.0, 1.0)) - assert out.child_frame_id == "base_link" - assert out.position.to_tuple() == (1.0, 2.0, 3.0) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 643830a719..fe94be24e8 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -222,7 +222,6 @@ "object-tracker2-d": "dimos.perception.object_tracker_2d.ObjectTracker2D", "object-tracker3-d": "dimos.perception.object_tracker_3d.ObjectTracker3D", "object-tracking": "dimos.perception.object_tracker.ObjectTracking", - "odom-body-frame": "dimos.navigation.nav_3d.mls_planner.odom_body_frame.OdomBodyFrame", "osm-skill": "dimos.agents.skills.osm.OsmSkill", "path-follower": "dimos.navigation.cmu_nav.modules.path_follower.path_follower.PathFollower", "patrolling-module": "dimos.navigation.patrolling.module.PatrollingModule", From b3b54388da9e7bcb1a30f123e50a777b2ed187c8 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Thu, 23 Jul 2026 11:41:55 -0700 Subject: [PATCH 12/26] Tf publish --- examples/native-modules/rust/Cargo.toml | 4 + .../native-modules/rust/src/tf_broadcaster.rs | 51 ++++++ .../native-modules/rust/src/tf_listener.rs | 4 +- examples/native-modules/rust_tf.py | 28 +++- native/rust/README.md | 16 +- native/rust/dimos-module/src/lib.rs | 4 +- native/rust/dimos-module/src/module.rs | 11 +- native/rust/dimos-module/src/tf.rs | 157 +++++++++++++++++- 8 files changed, 253 insertions(+), 22 deletions(-) create mode 100644 examples/native-modules/rust/src/tf_broadcaster.rs diff --git a/examples/native-modules/rust/Cargo.toml b/examples/native-modules/rust/Cargo.toml index 15398ca349..712e4fef45 100644 --- a/examples/native-modules/rust/Cargo.toml +++ b/examples/native-modules/rust/Cargo.toml @@ -15,6 +15,10 @@ path = "src/pong.rs" name = "tf_listener" path = "src/tf_listener.rs" +[[bin]] +name = "tf_broadcaster" +path = "src/tf_broadcaster.rs" + [dependencies] dimos-module = { path = "../../../native/rust/dimos-module" } lcm-msgs = { git = "https://github.com/dimensionalOS/dimos-lcm.git", branch = "rust-codegen" } diff --git a/examples/native-modules/rust/src/tf_broadcaster.rs b/examples/native-modules/rust/src/tf_broadcaster.rs new file mode 100644 index 0000000000..da9c18f75f --- /dev/null +++ b/examples/native-modules/rust/src/tf_broadcaster.rs @@ -0,0 +1,51 @@ +// Copyright 2026 Dimensional Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::time::{SystemTime, UNIX_EPOCH}; + +use dimos_module::nalgebra::Isometry3; +use dimos_module::{run_with_transport, Module, Tf, Transform}; +use tokio::time::{interval, Duration}; + +#[derive(Module)] +#[module(setup = start_broadcast)] +struct TfBroadcaster { + #[tf] + tf: Tf, +} + +impl TfBroadcaster { + async fn start_broadcast(&mut self) { + let tf = self.tf.clone(); + tokio::spawn(async move { + let mut ticker = interval(Duration::from_millis(100)); + loop { + ticker.tick().await; + let ts = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("system clock before epoch") + .as_secs_f64(); + let t = Transform::new("c", "d", ts, Isometry3::translation(0.5, 0.0, 0.0)); + if tf.publish(&[t]).await.is_err() { + break; + } + } + }); + } +} + +#[tokio::main] +async fn main() { + run_with_transport::().await; +} diff --git a/examples/native-modules/rust/src/tf_listener.rs b/examples/native-modules/rust/src/tf_listener.rs index b611710bb4..c4227fc35a 100644 --- a/examples/native-modules/rust/src/tf_listener.rs +++ b/examples/native-modules/rust/src/tf_listener.rs @@ -29,7 +29,7 @@ impl TfListener { let mut ticker = interval(Duration::from_millis(500)); loop { ticker.tick().await; - match tf.get_latest("a", "c") { + match tf.get_latest("a", "d") { Some(t) => { let p = t.translation(); tracing::info!( @@ -41,7 +41,7 @@ impl TfListener { "transform:", ); } - None => tracing::info!("a -> c not available yet"), + None => tracing::info!("a -> d not available yet"), } } }); diff --git a/examples/native-modules/rust_tf.py b/examples/native-modules/rust_tf.py index 5a7bcfd2d0..0e142d4af5 100644 --- a/examples/native-modules/rust_tf.py +++ b/examples/native-modules/rust_tf.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Python publishes a tf chain, a Rust #[tf] module reads it back. +"""Python and Rust each publish part of a tf chain, a Rust #[tf] module reads it back. + +Python publishes a -> b -> c, a Rust broadcaster publishes c -> d, and a Rust +listener composes the full a -> d transform. Run with: python examples/native-modules/rust_tf.py @@ -84,16 +87,29 @@ class TfListenerConfig(NativeModuleConfig): class TfListenerModule(NativeModule): - """Rust module that looks up a -> c and logs it. + """Rust module that looks up a -> d and logs it. - Expect to see (1.0, cos(t), sin(t)) + Expect to see (1.5, cos(t), sin(t)) """ config: TfListenerConfig +class TfBroadcasterConfig(NativeModuleConfig): + executable: str = str(_EXAMPLES / "tf_broadcaster") + build_command: str = _BUILD + cwd: str = str(_RUST_DIR) + stdin_config: bool = True + + +class TfBroadcasterModule(NativeModule): + """Rust module that publishes the c -> d transform.""" + + config: TfBroadcasterConfig + + if __name__ == "__main__": - bp = autoconnect(TfProducer.blueprint(), TfListenerModule.blueprint()).global_config( - viewer="none" - ) + bp = autoconnect( + TfProducer.blueprint(), TfBroadcasterModule.blueprint(), TfListenerModule.blueprint() + ).global_config(viewer="none") ModuleCoordinator.build(bp).loop() diff --git a/native/rust/README.md b/native/rust/README.md index 96c076e7e3..f0241ed932 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -58,7 +58,7 @@ Every transport is compiled into the binary. `run_with_transport` opens the one - `#[input(decode = fn, handler = fn)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. - `#[output(encode = fn)]`: on a field of type `Output`. `encode` is required. - `#[config]`: on one field. The type must be defined with `#[native_config]` (see [Config](#config)). At most one per struct. If absent, `Config` defaults to `dimos_module::NoConfig`. -- `#[tf]`: on a field of type `Tf`. Subscribes to the `tf` topic and answers transform queries (see [Transforms](#transforms)). No arguments. +- `#[tf]`: on a field of type `Tf`. Subscribes to the `tf` topic, answers transform queries, and publishes transforms (see [Transforms](#transforms)). No arguments. - Unattributed fields are initialized via `Default::default()` and treated as module state. ## Config @@ -103,7 +103,7 @@ Field name = port name. Ports map to topics via the stdin JSON; unmapped ports f ## Transforms -A `#[tf]` field gives a module a consumer-side view of the transform graph, the Rust counterpart to Python's `tf.get()`. It subscribes to the `tf` topic (mapped like any other port, default `/tf`), buffers each `parent -> child` edge it sees, and answers queries by composing transforms along the shortest path through the graph. +A `#[tf]` field gives a module a view of the transform graph, the Rust counterpart to Python's `tf.get()` and `tf.publish()`. It subscribes to the `tf` topic (mapped like any other port, default `/tf`), buffers each `parent -> child` edge it sees, and answers queries by composing transforms along the shortest path through the graph. ```rust #[derive(Module)] @@ -124,7 +124,17 @@ impl VoxelMap { } ``` -`Tf` is a cheap-to-clone handle; the graph fills in the background as `tf` messages arrive. `get(parent, child, time, tolerance)` selects the sample nearest `time` (latest when `None`) and returns `None` when no path connects the frames or no sample falls within `tolerance` seconds. `get_latest` is the no-time shorthand. The result exposes an `nalgebra` `Isometry3` via `isometry()`, ready to apply to a point. Lookups are nearest-in-time, not interpolated. This is consumer-only; modules do not publish transforms. +`Tf` is a cheap-to-clone handle; the graph fills in the background as `tf` messages arrive. `get(parent, child, time, tolerance)` selects the sample nearest `time` (latest when `None`) and returns `None` when no path connects the frames or no sample falls within `tolerance` seconds. `get_latest` is the no-time shorthand. The result exposes an `nalgebra` `Isometry3` via `isometry()`, ready to apply to a point. Lookups are nearest-in-time, not interpolated. + +`publish` sends transforms onto the same `tf` topic, the counterpart to Python's `tf.publish()`. Published transforms also feed the module's own graph, so a `get` right after the publish sees them. Build the isometry from `dimos_module::nalgebra`, re-exported so the version matches the SDK's types: + +```rust +use dimos_module::nalgebra::Isometry3; +use dimos_module::Transform; + +let iso = Isometry3::translation(0.5, 0.0, 0.0); +self.tf.publish(&[Transform::new("base_link", "gripper", ts, iso)]).await?; +``` ## What `#[derive(Module)]` generates diff --git a/native/rust/dimos-module/src/lib.rs b/native/rust/dimos-module/src/lib.rs index 52ae7edcc9..b5efc805be 100644 --- a/native/rust/dimos-module/src/lib.rs +++ b/native/rust/dimos-module/src/lib.rs @@ -22,10 +22,12 @@ pub mod zenoh; pub use dimos_module_macros::{native_config, Module}; pub use lcm::LcmTransport; pub use module::{run, Builder, Input, Module, ModuleConfig, NativeConfig, NoConfig, Output}; -pub use tf::Tf; +pub use tf::{Tf, Transform}; pub use transport::Transport; pub use zenoh::ZenohTransport; +pub use nalgebra; + // Re-export LcmOptions so callers don't need to depend on dimos-lcm directly. pub use dimos_lcm::LcmOptions; diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index 8ab73061ed..a48f91e109 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -332,17 +332,20 @@ impl Builder { } } - /// A handle that answers transform queries from the `/tf` topic. + /// A handle that answers transform queries and publishes on the `tf` topic. /// - /// The first call subscribes to the resolved `tf` topic and starts filling - /// the transform graph in the background. Repeated calls share one graph. + /// The first call subscribes to the resolved `tf` topic, starts filling the + /// transform graph in the background, and wires a publish queue for the same + /// topic. Repeated calls share one graph. pub fn tf(&mut self) -> crate::tf::Tf { if let Some(tf) = &self.tf { return tf.clone(); } let topic = self.topic_for("tf"); + let (tx, rx) = mpsc::channel(PUBLISH_CHANNEL_CAPACITY); + self.outputs.push((topic.clone(), rx)); let (tf, route) = - crate::tf::tf_subscription(topic.clone(), crate::tf::DEFAULT_TF_BUFFER_SIZE); + crate::tf::tf_subscription(topic.clone(), crate::tf::DEFAULT_TF_BUFFER_SIZE, tx); self.routes.entry(topic).or_default().push(route); self.tf = Some(tf.clone()); tf diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index 0c4bdc470d..e41a2d3190 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -12,17 +12,20 @@ // See the License for the specific language governing permissions and // limitations under the License. -//! Consumer-side transform client for native modules. +//! Transform client for native modules. //! //! Each `/tf` edge is buffered per `(parent, child)`, and [`Tf::get`] composes //! the shortest path through the frame graph. Lookups are nearest-in-time within -//! a tolerance, not interpolated. +//! a tolerance, not interpolated. [`Tf::publish`] sends transforms onto the same +//! topic and feeds the local graph. use std::collections::{HashMap, HashSet, VecDeque}; +use std::io; use std::sync::{Arc, RwLock}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use nalgebra::{Isometry3, Quaternion, Translation3, UnitQuaternion, Vector3}; +use tokio::sync::mpsc; use crate::module::Route; @@ -49,6 +52,20 @@ pub struct Transform { } impl Transform { + pub fn new( + parent: impl Into, + child: impl Into, + ts: f64, + iso: Isometry3, + ) -> Self { + Self { + parent: parent.into(), + child: child.into(), + ts, + iso, + } + } + /// The transform as an isometry. pub fn isometry(&self) -> Isometry3 { self.iso @@ -275,13 +292,14 @@ impl MultiTBuffer { } } -/// A cheap-to-clone handle for querying the transform graph. +/// A cheap-to-clone handle for querying and publishing transforms. /// /// Obtain one from `Builder::tf` (or a `#[tf]` field on a `#[derive(Module)]` /// struct). The graph is filled in the background as `/tf` messages arrive. #[derive(Clone)] pub struct Tf { buffer: Arc>, + sender: mpsc::Sender>, } impl Tf { @@ -307,6 +325,61 @@ impl Tf { pub fn get_latest(&self, parent: &str, child: &str) -> Option { self.get(parent, child, None, None) } + + /// Publish transforms on the `tf` topic. + /// + /// The transforms also feed the local graph, so a `get` right after sees + /// them without waiting for the transport round trip. + pub async fn publish(&self, transforms: &[Transform]) -> io::Result<()> { + { + let mut buffer = self.buffer.write().expect("tf buffer lock poisoned"); + for t in transforms { + buffer.receive(&t.parent, &t.child, t.ts, t.iso); + } + } + let msg = lcm_msgs::tf2_msgs::TFMessage { + transforms: transforms.iter().map(to_stamped).collect(), + }; + self.sender + .send(msg.encode()) + .await + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "background task gone")) + } +} + +fn to_stamped(t: &Transform) -> lcm_msgs::geometry_msgs::TransformStamped { + let mut sec = t.ts.floor(); + let mut nsec = ((t.ts - sec) * 1e9).round(); + if nsec >= 1e9 { + sec += 1.0; + nsec -= 1e9; + } + let p = t.iso.translation.vector; + let q = t.iso.rotation; + lcm_msgs::geometry_msgs::TransformStamped { + header: lcm_msgs::std_msgs::Header { + seq: 0, + stamp: lcm_msgs::std_msgs::Time { + sec: sec as i32, + nsec: nsec as i32, + }, + frame_id: t.parent.clone(), + }, + child_frame_id: t.child.clone(), + transform: lcm_msgs::geometry_msgs::Transform { + translation: lcm_msgs::geometry_msgs::Vector3 { + x: p.x, + y: p.y, + z: p.z, + }, + rotation: lcm_msgs::geometry_msgs::Quaternion { + x: q.i, + y: q.j, + z: q.k, + w: q.w, + }, + }, + } } // Decodes /tf messages into the shared graph. Registered as a Route so the @@ -344,11 +417,17 @@ impl Route for TfRoute { } } -// Builds the shared graph plus the handle and the route that feeds it. -pub(crate) fn tf_subscription(topic: String, buffer_size: f64) -> (Tf, Box) { +// Builds the shared graph plus the handle and the route that feeds it. The +// sender carries published messages to the tf topic's publish worker. +pub(crate) fn tf_subscription( + topic: String, + buffer_size: f64, + sender: mpsc::Sender>, +) -> (Tf, Box) { let buffer = Arc::new(RwLock::new(MultiTBuffer::new(buffer_size))); let tf = Tf { buffer: Arc::clone(&buffer), + sender, }; let route = Box::new(TfRoute { topic, buffer }); (tf, route) @@ -360,11 +439,19 @@ mod tests { use std::f64::consts::PI; fn tf_with(buffer_size: f64) -> (Tf, MultiHandle) { + let (tf, _rx, handle) = tf_with_publish(buffer_size); + (tf, handle) + } + + fn tf_with_publish(buffer_size: f64) -> (Tf, mpsc::Receiver>, MultiHandle) { let buffer = Arc::new(RwLock::new(MultiTBuffer::new(buffer_size))); + let (tx, rx) = mpsc::channel(8); ( Tf { buffer: Arc::clone(&buffer), + sender: tx, }, + rx, MultiHandle { buffer }, ) } @@ -501,7 +588,8 @@ mod tests { use lcm_msgs::std_msgs::{Header, Time}; use lcm_msgs::tf2_msgs::TFMessage; - let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_BUFFER_SIZE); + let (tx, _rx) = mpsc::channel(8); + let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_BUFFER_SIZE, tx); let msg = TFMessage { transforms: vec![lcm_msgs::geometry_msgs::TransformStamped { header: Header { @@ -536,4 +624,61 @@ mod tests { assert!((t.translation().z - 0.3).abs() < 1e-9); assert!((t.ts - 5.5).abs() < 1e-9); } + + #[tokio::test] + async fn publish_feeds_local_graph() { + let (tf, _rx, _h) = tf_with_publish(DEFAULT_TF_BUFFER_SIZE); + let iso = Isometry3::from_parts( + Translation3::new(1.0, 2.0, 3.0), + UnitQuaternion::from_euler_angles(0.0, 0.0, PI / 2.0), + ); + tf.publish(&[Transform::new("map", "base_link", 7.0, iso)]) + .await + .unwrap(); + + let t = tf.get_latest("map", "base_link").unwrap(); + assert!((t.translation().x - 1.0).abs() < 1e-9); + assert!((t.translation().y - 2.0).abs() < 1e-9); + assert!((t.translation().z - 3.0).abs() < 1e-9); + assert!((t.ts - 7.0).abs() < 1e-9); + } + + // Publish on one handle, dispatch the wire bytes into another graph. + #[tokio::test] + async fn publish_round_trips_through_route() { + let (tf_out, mut rx, _h) = tf_with_publish(DEFAULT_TF_BUFFER_SIZE); + let iso = Isometry3::from_parts( + Translation3::new(0.5, -0.5, 0.25), + UnitQuaternion::from_euler_angles(0.0, 0.0, PI / 6.0), + ); + tf_out + .publish(&[Transform::new("a", "b", 3.25, iso)]) + .await + .unwrap(); + let bytes = rx.recv().await.unwrap(); + + let (tx, _rx2) = mpsc::channel(8); + let (tf_in, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_BUFFER_SIZE, tx); + route.try_dispatch(&bytes); + + let t = tf_in.get_latest("a", "b").unwrap(); + assert!((t.translation().x - 0.5).abs() < 1e-9); + assert!((t.translation().y + 0.5).abs() < 1e-9); + assert!((t.translation().z - 0.25).abs() < 1e-9); + assert!((t.ts - 3.25).abs() < 1e-9); + let (_, _, yaw) = t.rotation().euler_angles(); + assert!((yaw - PI / 6.0).abs() < 1e-9); + } + + #[test] + fn stamp_rounding_does_not_overflow_nsec() { + let st = to_stamped(&Transform::new( + "a", + "b", + 1.9999999999, + Isometry3::identity(), + )); + assert_eq!(st.header.stamp.sec, 2); + assert_eq!(st.header.stamp.nsec, 0); + } } From 0de58f4ae3572703c45b20ee46619656e2fe8cbb Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Thu, 23 Jul 2026 11:49:38 -0700 Subject: [PATCH 13/26] Switch to vecdeque --- dimos/core/transport_factory.py | 8 +++----- .../nav_3d/mls_planner/utils/plan_rrd.py | 4 +++- native/rust/dimos-module/src/tf.rs | 18 ++++++------------ 3 files changed, 12 insertions(+), 18 deletions(-) diff --git a/dimos/core/transport_factory.py b/dimos/core/transport_factory.py index 9b37b7602b..3ab47587c2 100644 --- a/dimos/core/transport_factory.py +++ b/dimos/core/transport_factory.py @@ -27,6 +27,8 @@ pLCMTransport, pZenohTransport, ) +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.pubsub.impl.lcmpubsub import Topic as LCMTopic from dimos.protocol.pubsub.impl.zenohpubsub import ( QOS_LATEST_WINS, QOS_NEVER_DROP, @@ -61,14 +63,10 @@ def tf_channel(g: GlobalConfig = global_config) -> str: native modules are handed this channel directly. Backends format the type suffix differently, so match the active one. """ - from dimos.msgs.tf2_msgs.TFMessage import TFMessage - name = transport_topic("/tf", g) if g.transport == "zenoh": return ZenohTopic(name, TFMessage).key_expr - from dimos.protocol.pubsub.impl.lcmpubsub import Topic - - return str(Topic(name, TFMessage)) + return str(LCMTopic(name, TFMessage)) # High-rate sensor streams: drop stale frames under congestion, never stall the 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 6b63416213..51f286ae74 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -54,6 +54,8 @@ # Mount frames as recorded on the tf stream. BASE_FRAME = "base_link" SENSOR_FRAME = "mid360_link" +# Static mount edges settle within the first handful of tf messages. +TF_WARMUP_SAMPLES = 20 # The travelled trail. Blue, but light enough to read against the map's turbo # lows, and clear of PATH_PALETTE so it never reads as a planned path. @@ -142,7 +144,7 @@ def _base_from_sensor(store: SqliteStore) -> Transform | None: try: for i, obs in enumerate(store.stream("tf", TFMessage).order_by("ts")): buffer.receive_transform(*obs.data.transforms) - if i >= 20: + if i >= TF_WARMUP_SAMPLES: break except Exception as e: print(f"no usable tf stream in the recording ({e}); skipping the base_link triad") diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index e41a2d3190..8bad8d3254 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -41,8 +41,7 @@ fn now_secs() -> f64 { /// A rigid transform from `parent` to `child` at a point in time. /// -/// The isometry maps a point expressed in `child` coordinates into `parent` -/// coordinates: `p_parent = transform.isometry() * p_child`. +/// It maps a point expressed in `child` coordinates into `parent` coordinates. #[derive(Clone, Debug)] pub struct Transform { pub parent: String, @@ -66,11 +65,6 @@ impl Transform { } } - /// The transform as an isometry. - pub fn isometry(&self) -> Isometry3 { - self.iso - } - /// Translation component (`parent`-frame position of the `child` origin). pub fn translation(&self) -> Vector3 { self.iso.translation.vector @@ -109,14 +103,14 @@ struct Sample { // One edge's time-sorted history, capped to a fixed-duration window. struct TBuffer { buffer_size: f64, - samples: Vec, + samples: VecDeque, } impl TBuffer { fn new(buffer_size: f64) -> Self { Self { buffer_size, - samples: Vec::new(), + samples: VecDeque::new(), } } @@ -128,13 +122,13 @@ impl TBuffer { fn prune(&mut self, min_ts: f64) { let drop_to = self.samples.partition_point(|s| s.ts < min_ts); - if drop_to > 0 { - self.samples.drain(0..drop_to); + for _ in 0..drop_to { + self.samples.pop_front(); } } fn last(&self) -> Option<&Sample> { - self.samples.last() + self.samples.back() } // Nearest sample in time. On a tie, prefer the later sample. Returns None From 274b4058f3dfb91916080ebd53f702b58bfaa40b Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Thu, 23 Jul 2026 12:48:06 -0700 Subject: [PATCH 14/26] Remove odom body frame and use tf --- .../navigation/basic_path_follower/module.py | 6 ++ .../nav_3d/mls_planner/goal_relay.py | 34 +++++++--- .../nav_3d/mls_planner/odom_body_frame.py | 67 ------------------- .../mls_planner/test_odom_body_frame.py | 60 ----------------- dimos/robot/all_blueprints.py | 1 - dimos/robot/unitree/go2/zenoh/blueprints.py | 33 +++------ 6 files changed, 40 insertions(+), 161 deletions(-) delete mode 100644 dimos/navigation/nav_3d/mls_planner/odom_body_frame.py delete mode 100644 dimos/navigation/nav_3d/mls_planner/test_odom_body_frame.py diff --git a/dimos/navigation/basic_path_follower/module.py b/dimos/navigation/basic_path_follower/module.py index 103e1c960a..449c9466cd 100644 --- a/dimos/navigation/basic_path_follower/module.py +++ b/dimos/navigation/basic_path_follower/module.py @@ -103,6 +103,12 @@ def _on_odometry(self, msg: Odometry) -> None: # Steer from the robot base pose, not the LIO body frame. base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, TF_LOOKUP_TOLERANCE_S) if base is None: + logger.warning( + "No %s -> %s transform for odometry stamp %.3f, dropping frame.", + msg.frame_id, + self.config.base_frame, + msg.ts, + ) return with self._lock: self._current_odom = base.to_pose(ts=msg.ts) diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py index a7ec5e9fec..c960a932ca 100644 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -25,14 +25,17 @@ from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.navigation.constants import TF_LOOKUP_TOLERANCE_S +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() class GoalRelayConfig(ModuleConfig): base_frame: str = "base_link" sensor_frame: str = "mid360_link" - # The lidar's height above the ground. base_link's height is derived from the - # base -> sensor mount transform, so this is the only fixed measurement. - lidar_height: float = 0.0 + # Lidar height above the ground while standing. Leave as None to skip the + # ground correction and pass odometry through unshifted. + lidar_height: float | None = None class GoalRelay(Module): @@ -61,18 +64,33 @@ def start(self) -> None: def _on_odometry(self, msg: Odometry) -> None: base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, TF_LOOKUP_TOLERANCE_S) - base_height = self._base_from_ground() - if base is None or base_height is None: + if base is None: + logger.warning( + "No %s -> %s transform for odometry stamp %.3f, dropping frame.", + msg.frame_id, + self.config.base_frame, + msg.ts, + ) return start = base.to_pose(ts=msg.ts) - start.position.z -= base_height + if self.config.lidar_height is not None: + base_height = self._base_from_ground() + if base_height is None: + return + start.position.z -= base_height self.start_pose.publish(start) def _base_from_ground(self) -> float | None: # The base -> sensor mount is static rig geometry, so resolve it once. - if self._base_height is None: + if self._base_height is None and self.config.lidar_height is not None: mount = self.tf.get(self.config.base_frame, self.config.sensor_frame) - if mount is not None: + if mount is None: + logger.warning( + "No %s -> %s mount transform on tf, cannot ground-project the start pose.", + self.config.base_frame, + self.config.sensor_frame, + ) + else: self._base_height = self.config.lidar_height - mount.translation.z return self._base_height diff --git a/dimos/navigation/nav_3d/mls_planner/odom_body_frame.py b/dimos/navigation/nav_3d/mls_planner/odom_body_frame.py deleted file mode 100644 index 2afef1e2f1..0000000000 --- a/dimos/navigation/nav_3d/mls_planner/odom_body_frame.py +++ /dev/null @@ -1,67 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from __future__ import annotations - -from typing import Any - -from pydantic import Field -from reactivex.disposable import Disposable - -from dimos.core.core import rpc -from dimos.core.module import Module, ModuleConfig -from dimos.core.stream import In, Out -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.nav_msgs.Odometry import Odometry - - -class OdomBodyFrameConfig(ModuleConfig): - # base_link from sensor mount rotation, xyzw. - mount_rotation: list[float] = Field(default_factory=lambda: [0.0, 0.0, 0.0, 1.0]) - body_frame_id: str = "base_link" - - -class OdomBodyFrame(Module): - """Re-express tilted-sensor LIO odometry in the level robot body frame. - - Composes out the fixed mount rotation from the orientation. Position and - twist pass through. - """ - - config: OdomBodyFrameConfig - - odometry: In[Odometry] - body_odometry: Out[Odometry] - - def __init__(self, **kwargs: Any) -> None: - super().__init__(**kwargs) - self._mount_inv = Quaternion(*self.config.mount_rotation).inverse() - - @rpc - def start(self) -> None: - super().start() - self.register_disposable(Disposable(self.odometry.subscribe(self._on_odometry))) - - def _on_odometry(self, msg: Odometry) -> None: - leveled = msg.orientation * self._mount_inv - self.body_odometry.publish( - Odometry( - ts=msg.ts, - frame_id=msg.frame_id, - child_frame_id=self.config.body_frame_id, - pose=Pose(msg.position, leveled), - twist=msg.twist, - ) - ) diff --git a/dimos/navigation/nav_3d/mls_planner/test_odom_body_frame.py b/dimos/navigation/nav_3d/mls_planner/test_odom_body_frame.py deleted file mode 100644 index 571b0fbdd0..0000000000 --- a/dimos/navigation/nav_3d/mls_planner/test_odom_body_frame.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from dimos.msgs.geometry_msgs.Pose import Pose -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 -from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.navigation.nav_3d.mls_planner.odom_body_frame import OdomBodyFrame - - -def _level(mount_rotation, orientation): - """Run one odometry message through the handler and return the output.""" - module = OdomBodyFrame(mount_rotation=list(mount_rotation), body_frame_id="base_link") - try: - captured = [] - module.body_odometry.subscribe(captured.append) - module._on_odometry( - Odometry( - ts=1.0, - frame_id="odom", - child_frame_id="mid360_link", - pose=Pose(Vector3(1.0, 2.0, 3.0), orientation), - ) - ) - return captured[0] - finally: - module.stop() - - -def test_composes_out_the_mount_pitch(): - # A level body reads its own mount tilt as the sensor's world orientation, so - # composing the mount out returns identity. - mount = Quaternion.from_euler(Vector3(0.0, 0.3, 0.0)) - out = _level(mount.to_tuple(), mount) - assert out.orientation.angle_to(Quaternion(0.0, 0.0, 0.0, 1.0)) < 1e-5 - - -def test_preserves_body_yaw_under_mount_tilt(): - # A body yawed by a known angle keeps that yaw after the mount is composed out. - mount = Quaternion.from_euler(Vector3(0.0, 0.3, 0.0)) - body = Quaternion.from_euler(Vector3(0.0, 0.0, 0.7)) - out = _level(mount.to_tuple(), body * mount) - assert out.orientation.angle_to(body) < 1e-5 - - -def test_relabels_child_frame_and_passes_position_through(): - out = _level([0.0, 0.0, 0.0, 1.0], Quaternion(0.0, 0.0, 0.0, 1.0)) - assert out.child_frame_id == "base_link" - assert out.position.to_tuple() == (1.0, 2.0, 3.0) diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 9194f48e92..93522fb5f8 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -251,7 +251,6 @@ "object-tracker2-d": "dimos.perception.object_tracker_2d.ObjectTracker2D", "object-tracker3-d": "dimos.perception.object_tracker_3d.ObjectTracker3D", "object-tracking": "dimos.perception.object_tracker.ObjectTracking", - "odom-body-frame": "dimos.navigation.nav_3d.mls_planner.odom_body_frame.OdomBodyFrame", "osm-skill": "dimos.agents.skills.osm.OsmSkill", "path-follower": "dimos.navigation.cmu_nav.modules.path_follower.path_follower.PathFollower", "path-following-coordinator": "dimos.control.path_following_coordinator.PathFollowingCoordinator", diff --git a/dimos/robot/unitree/go2/zenoh/blueprints.py b/dimos/robot/unitree/go2/zenoh/blueprints.py index 05fb406254..6abd29d898 100644 --- a/dimos/robot/unitree/go2/zenoh/blueprints.py +++ b/dimos/robot/unitree/go2/zenoh/blueprints.py @@ -26,22 +26,19 @@ ``DanLocalPlanner`` + ``DanHolonomicTC`` pair from ``unitree-go2-mls-htc``. """ -import math from typing import Any from dimos.core.coordination.blueprints import autoconnect from dimos.core.global_config import global_config from dimos.mapping.ray_tracing.module import RayTracingVoxelMap -from dimos.msgs.geometry_msgs.Quaternion import Quaternion -from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.navigation.basic_path_follower.module import BasicPathFollower from dimos.navigation.dannav.holonomic_tc.module import DanHolonomicTC from dimos.navigation.dannav.local_planner.module import DanLocalPlanner from dimos.navigation.movement_manager.movement_manager import MovementManager from dimos.navigation.nav_3d.mls_planner.goal_relay import GoalRelay from dimos.navigation.nav_3d.mls_planner.mls_planner_native import MLSPlannerNative -from dimos.navigation.nav_3d.mls_planner.odom_body_frame import OdomBodyFrame from dimos.navigation.nav_3d.mls_planner.viz import planner_visual_override +from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT from dimos.robot.unitree.go2.zenoh.zenohconnection import GO2Zenoh from dimos.visualization.vis_module import vis_module @@ -50,20 +47,11 @@ # cost-coloured edges. Drives both its publishing and the rerun overrides. planner_viz_hz = 2.0 -# Feeds both the static tf GO2Zenoh publishes and the rotation that levels its odometry — -# they must agree or nav steers off-heading. Verified against Point-LIO's own attitude. +# The mount tree GO2Zenoh publishes onto tf; nav reads its odometry corrections from +# there. Verified against Point-LIO's own attitude. MID360_MOUNT_RPY_DEG = (-60.0, 0.0, -90.0) -def _mount_rotation() -> list[float]: - """base_link <- lidar rotation, so nav reads odometry in the level body frame. - - base_link -> front_camera carries no rotation, so this is just the mount rpy above. - """ - rpy = Vector3(*(math.radians(d) for d in MID360_MOUNT_RPY_DEG)) - return list(Quaternion.from_euler(rpy).to_tuple()) - - def _camera_info_to_pinhole(camera_info: Any) -> Any: """Log the pinhole onto the video's entity instead of camera_info's own. @@ -183,11 +171,8 @@ def _rerun_config(visual_override: dict[str, Any] | None = None) -> dict[str, An go2_zenoh_nav = autoconnect( go2_zenoh_raycaster, _mls_planner, - OdomBodyFrame.blueprint(mount_rotation=_mount_rotation()), - GoalRelay.blueprint(), - BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6).remappings( - [(BasicPathFollower, "odometry", "body_odometry")] - ), + GoalRelay.blueprint(lidar_height=ROBOT_HEIGHT), + BasicPathFollower.blueprint(speed=0.5, heading_gain=0.4, max_angular=0.6), MovementManager.blueprint(), ).global_config(transport="zenoh", n_workers=8, robot_model="unitree_go2") @@ -196,12 +181,10 @@ def _rerun_config(visual_override: dict[str, Any] | None = None) -> dict[str, An # forwards committed paths on path, so world/planner_path is muted in rerun. go2_zenoh_htc = autoconnect( go2_zenoh_raycaster, - OdomBodyFrame.blueprint(mount_rotation=_mount_rotation()), _mls_planner.remappings([(MLSPlannerNative, "path", "planner_path")]), - # Fed the leveled odometry, so its start_pose doubles as the body-frame PoseStamped - # the Dan modules consume — mirroring mls_htc, where planner start and follower odom - # are the same topic. - GoalRelay.blueprint().remappings([(GoalRelay, "odometry", "body_odometry")]), + # start_pose is base_link odometry, so the Dan modules use it as both planner + # start and follower odom. + GoalRelay.blueprint(lidar_height=ROBOT_HEIGHT), # Setting resample_spacing_m to > 0.0 will smooth out jagged paths returned by MLSP DanLocalPlanner.blueprint(resample_spacing_m=0.1).remappings( [(DanLocalPlanner, "odom", "start_pose")] From 865bd90d419ac364fafb5a7136dced5a070e7c11 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Thu, 23 Jul 2026 15:04:44 -0700 Subject: [PATCH 15/26] Cleaning --- dimos/core/transport_factory.py | 7 +------ dimos/navigation/basic_path_follower/module.py | 18 +++++++++--------- dimos/navigation/constants.py | 4 +--- .../nav_3d/mls_planner/goal_relay.py | 1 - .../nav_3d/mls_planner/utils/plan_rrd.py | 14 +++++--------- dimos/robot/unitree/go2/zenoh/blueprints.py | 11 ++--------- native/rust/dimos-module/src/tf.rs | 3 +-- 7 files changed, 19 insertions(+), 39 deletions(-) diff --git a/dimos/core/transport_factory.py b/dimos/core/transport_factory.py index 3ab47587c2..0756b31a1c 100644 --- a/dimos/core/transport_factory.py +++ b/dimos/core/transport_factory.py @@ -57,12 +57,7 @@ def transport_topic(name: str, g: GlobalConfig = global_config) -> str: def tf_channel(g: GlobalConfig = global_config) -> str: - """Wire channel for the tf stream on the active backend. - - tf can't be a declared port because every Module reserves the name, so - native modules are handed this channel directly. Backends format the type - suffix differently, so match the active one. - """ + """Wire channel for the tf stream on the active backend.""" name = transport_topic("/tf", g) if g.transport == "zenoh": return ZenohTopic(name, TFMessage).key_expr diff --git a/dimos/navigation/basic_path_follower/module.py b/dimos/navigation/basic_path_follower/module.py index 449c9466cd..6d8dce66e9 100644 --- a/dimos/navigation/basic_path_follower/module.py +++ b/dimos/navigation/basic_path_follower/module.py @@ -76,7 +76,7 @@ class BasicPathFollower(Module): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._lock = RLock() - self._current_odom: PoseStamped | None = None + self._current_pose: PoseStamped | None = None self._waypoints: NDArray[np.float32] | None = None self._stop_event = Event() self._thread: Thread | None = None @@ -100,7 +100,7 @@ def stop(self) -> None: super().stop() def _on_odometry(self, msg: Odometry) -> None: - # Steer from the robot base pose, not the LIO body frame. + # Steer from the robot base pose, not the sensor frame. base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, TF_LOOKUP_TOLERANCE_S) if base is None: logger.warning( @@ -111,7 +111,7 @@ def _on_odometry(self, msg: Odometry) -> None: ) return with self._lock: - self._current_odom = base.to_pose(ts=msg.ts) + self._current_pose = base.to_pose(ts=msg.ts) def _on_path(self, path: Path) -> None: # The planner owns path safety: it sends the route as far as it is safe, @@ -136,15 +136,15 @@ def _follow(self) -> None: while not self._stop_event.is_set(): start_time = time.perf_counter() with self._lock: - odom = self._current_odom + pose = self._current_pose waypoints = self._waypoints - if odom is not None and waypoints is not None: - self._step(odom, waypoints) + if pose is not None and waypoints is not None: + self._step(pose, waypoints) elapsed = time.perf_counter() - start_time self._stop_event.wait(max(0.0, period - elapsed)) - def _step(self, odom: PoseStamped, waypoints: NDArray[np.float32]) -> None: - position = np.array([odom.position.x, odom.position.y], dtype=np.float32) + def _step(self, pose: PoseStamped, waypoints: NDArray[np.float32]) -> None: + position = np.array([pose.position.x, pose.position.y], dtype=np.float32) if float(np.linalg.norm(waypoints[-1] - position)) < self.config.goal_tolerance: self.nav_cmd_vel.publish(Twist()) with self._lock: @@ -157,7 +157,7 @@ def _step(self, odom: PoseStamped, waypoints: NDArray[np.float32]) -> None: target = self._lookahead_point(waypoints, position) yaw_error = angle_diff( math.atan2(target[1] - position[1], target[0] - position[0]), - odom.orientation.euler[2], + pose.orientation.euler[2], ) angular = max( diff --git a/dimos/navigation/constants.py b/dimos/navigation/constants.py index 151691f783..3c67e62820 100644 --- a/dimos/navigation/constants.py +++ b/dimos/navigation/constants.py @@ -14,7 +14,5 @@ """Shared navigation constants.""" -# Max staleness in seconds for tf lookups against a live odometry stamp. Must -# comfortably exceed the coarsest static-transform publish period so leading-edge -# phase or modest clock skew does not silently drop frames. +# Maximum time window to look for transforms before refusing a tf. TF_LOOKUP_TOLERANCE_S = 0.25 diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py index c960a932ca..8250a15b7a 100644 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -81,7 +81,6 @@ def _on_odometry(self, msg: Odometry) -> None: self.start_pose.publish(start) def _base_from_ground(self) -> float | None: - # The base -> sensor mount is static rig geometry, so resolve it once. if self._base_height is None and self.config.lidar_height is not None: mount = self.tf.get(self.config.base_frame, self.config.sensor_frame) if mount is None: 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 51f286ae74..91ad9c8753 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -46,22 +46,18 @@ TIMELINE = "ts" -# Axis-triad length for the odometry transform (m). ODOM_AXIS_LEN = 0.5 -# Arrow radius as a fraction of the triad length. AXIS_RADIUS_RATIO = 25 # Mount frames as recorded on the tf stream. BASE_FRAME = "base_link" SENSOR_FRAME = "mid360_link" -# Static mount edges settle within the first handful of tf messages. +# need to wait to get tfs sometimes TF_WARMUP_SAMPLES = 20 -# The travelled trail. Blue, but light enough to read against the map's turbo -# lows, and clear of PATH_PALETTE so it never reads as a planned path. ODOM_PATH_COLOR = [80, 160, 255] -# Distinct path colors for overlaid configurations, config 0 first. +# Different colors for each path when running with multiple configs PATH_PALETTE = [ [0, 255, 0], [255, 0, 255], @@ -139,7 +135,7 @@ def _log_path_wp(waypoints: NDArray[np.float32] | None, entity: str, color: list def _base_from_sensor(store: SqliteStore) -> Transform | None: - """The sensor -> base_link transform from the recording's static mount frames.""" + """Sensor to robot base link transform using the stored tf stream.""" buffer = MultiTBuffer() try: for i, obs in enumerate(store.stream("tf", TFMessage).order_by("ts")): @@ -153,7 +149,7 @@ def _base_from_sensor(store: SqliteStore) -> Transform | None: def _base_pose(pose: tuple[float, ...], ts: float, base_from_sensor: Transform) -> Transform: - """Compose the odometry pose with the recorded mount frames into world -> base_link.""" + """World to robot base link, using the sensor to robot base link transform.""" px, py, pz, qx, qy, qz, qw = pose sensor = Transform( translation=Vector3(px, py, pz), @@ -171,7 +167,7 @@ def _log_odometry( trail: list[tuple[float, float, float]], base_from_sensor: Transform | None, ) -> None: - """Log the moving sensor pose and its growing trajectory trail.""" + """Trace the sensor moving throughout the scene.""" px, py, pz, qx, qy, qz, qw = pose rr.set_time(TIMELINE, timestamp=ts) rr.log( diff --git a/dimos/robot/unitree/go2/zenoh/blueprints.py b/dimos/robot/unitree/go2/zenoh/blueprints.py index 6abd29d898..553167206b 100644 --- a/dimos/robot/unitree/go2/zenoh/blueprints.py +++ b/dimos/robot/unitree/go2/zenoh/blueprints.py @@ -44,11 +44,10 @@ voxel_size = 0.08 # Raise above 0 (2.0 works) to draw what the planner searched over: surface, nodes and -# cost-coloured edges. Drives both its publishing and the rerun overrides. +# cost-colored edges. Drives both its publishing and the rerun overrides. planner_viz_hz = 2.0 -# The mount tree GO2Zenoh publishes onto tf; nav reads its odometry corrections from -# there. Verified against Point-LIO's own attitude. +# GO2Zenoh publishes this mount onto tf, where nav reads its odometry corrections. MID360_MOUNT_RPY_DEG = (-60.0, 0.0, -90.0) @@ -176,16 +175,10 @@ def _rerun_config(visual_override: dict[str, Any] | None = None) -> dict[str, An MovementManager.blueprint(), ).global_config(transport="zenoh", n_workers=8, robot_model="unitree_go2") -# The nav stack with BasicPathFollower swapped for the DanLocalPlanner + DanHolonomicTC -# pair from unitree-go2-mls-htc. The raw planner stream moves to planner_path; the gate -# forwards committed paths on path, so world/planner_path is muted in rerun. go2_zenoh_htc = autoconnect( go2_zenoh_raycaster, _mls_planner.remappings([(MLSPlannerNative, "path", "planner_path")]), - # start_pose is base_link odometry, so the Dan modules use it as both planner - # start and follower odom. GoalRelay.blueprint(lidar_height=ROBOT_HEIGHT), - # Setting resample_spacing_m to > 0.0 will smooth out jagged paths returned by MLSP DanLocalPlanner.blueprint(resample_spacing_m=0.1).remappings( [(DanLocalPlanner, "odom", "start_pose")] ), diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index 8bad8d3254..24e4f8f782 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -65,7 +65,7 @@ impl Transform { } } - /// Translation component (`parent`-frame position of the `child` origin). + /// Translation component. pub fn translation(&self) -> Vector3 { self.iso.translation.vector } @@ -84,7 +84,6 @@ impl Transform { } } - // self a -> b, then other b -> c, gives a -> c. fn compose(&self, other: &Transform) -> Transform { Transform { parent: self.parent.clone(), From 49d69cdf5e354842d9d34b05632e3c797a2f3758 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Thu, 23 Jul 2026 18:47:43 -0700 Subject: [PATCH 16/26] Fix tf --- dimos/core/test_transport_factory.py | 9 + .../hardware/sensors/lidar/fastlio2/module.py | 3 +- .../navigation/basic_path_follower/module.py | 18 +- .../basic_path_follower/test_module.py | 53 +++++- dimos/navigation/constants.py | 18 -- .../nav_3d/mls_planner/goal_relay.py | 46 ++--- .../nav_3d/mls_planner/mls_planner_native.py | 4 + .../nav_3d/mls_planner/rust/src/main.rs | 1 + .../mls_planner/rust/src/mls_planner.rs | 6 + .../nav_3d/mls_planner/rust/src/planner.rs | 2 + .../nav_3d/mls_planner/rust/src/python.rs | 2 + .../nav_3d/mls_planner/test_goal_relay.py | 114 ++++++++++++ .../nav_3d/mls_planner/utils/plan_rrd.py | 104 ++++++----- dimos/navigation/test_tf_pose.py | 123 ++++++++++++ dimos/navigation/tf_pose.py | 94 ++++++++++ .../navigation/unitree_go2_mls_htc.py | 2 + .../navigation/unitree_go2_nav_3d.py | 1 - dimos/robot/unitree/go2/constants.py | 2 +- .../go2/go2_mid360_static_transforms.py | 20 +- dimos/robot/unitree/go2/zenoh/blueprints.py | 2 +- .../native-modules/rust/src/tf_listener.rs | 2 +- native/rust/README.md | 4 +- native/rust/dimos-module/src/module.rs | 24 ++- native/rust/dimos-module/src/tf.rs | 176 +++++++++++++----- 24 files changed, 655 insertions(+), 175 deletions(-) delete mode 100644 dimos/navigation/constants.py create mode 100644 dimos/navigation/nav_3d/mls_planner/test_goal_relay.py create mode 100644 dimos/navigation/test_tf_pose.py create mode 100644 dimos/navigation/tf_pose.py diff --git a/dimos/core/test_transport_factory.py b/dimos/core/test_transport_factory.py index c813392c23..2e4def1681 100644 --- a/dimos/core/test_transport_factory.py +++ b/dimos/core/test_transport_factory.py @@ -29,6 +29,7 @@ make_transport, rpc_backend, tf_backend, + tf_channel, transport_topic, ) from dimos.msgs.geometry_msgs.Twist import Twist @@ -54,6 +55,14 @@ def test_transport_topic_zenoh() -> None: assert transport_topic("/coordinator/joint_state", ZENOH) == "dimos/coordinator/joint_state" +def test_tf_channel_lcm() -> None: + assert tf_channel(LCM) == "/tf#tf2_msgs.TFMessage" + + +def test_tf_channel_zenoh() -> None: + assert tf_channel(ZENOH) == "dimos/tf/tf2_msgs.TFMessage" + + def test_make_transport_lcm_typed() -> None: t = make_transport("/camera/color", Image, g=LCM) assert type(t) is LCMTransport diff --git a/dimos/hardware/sensors/lidar/fastlio2/module.py b/dimos/hardware/sensors/lidar/fastlio2/module.py index 37b691d135..21c8363023 100644 --- a/dimos/hardware/sensors/lidar/fastlio2/module.py +++ b/dimos/hardware/sensors/lidar/fastlio2/module.py @@ -25,6 +25,7 @@ from __future__ import annotations import os +import time from typing import TYPE_CHECKING, Literal from pydantic import Field @@ -158,7 +159,7 @@ def _on_odom_for_tf(self, msg: Odometry) -> None: msg.pose.orientation.z, msg.pose.orientation.w, ), - ts=msg.ts, + ts=msg.ts or time.time(), ) ) diff --git a/dimos/navigation/basic_path_follower/module.py b/dimos/navigation/basic_path_follower/module.py index 6d8dce66e9..c067c5c5eb 100644 --- a/dimos/navigation/basic_path_follower/module.py +++ b/dimos/navigation/basic_path_follower/module.py @@ -33,7 +33,7 @@ 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.navigation.constants import TF_LOOKUP_TOLERANCE_S +from dimos.navigation.tf_pose import OdomBasePose from dimos.utils.logging_config import setup_logger from dimos.utils.trigonometry import angle_diff @@ -76,6 +76,7 @@ class BasicPathFollower(Module): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._lock = RLock() + self._base_pose: OdomBasePose | None = None self._current_pose: PoseStamped | None = None self._waypoints: NDArray[np.float32] | None = None self._stop_event = Event() @@ -100,18 +101,13 @@ def stop(self) -> None: super().stop() def _on_odometry(self, msg: Odometry) -> None: - # Steer from the robot base pose, not the sensor frame. - base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, TF_LOOKUP_TOLERANCE_S) - if base is None: - logger.warning( - "No %s -> %s transform for odometry stamp %.3f, dropping frame.", - msg.frame_id, - self.config.base_frame, - msg.ts, - ) + if self._base_pose is None: + self._base_pose = OdomBasePose(self.tf, self.config.base_frame) + pose = self._base_pose.resolve(msg) + if pose is None: return with self._lock: - self._current_pose = base.to_pose(ts=msg.ts) + self._current_pose = pose def _on_path(self, path: Path) -> None: # The planner owns path safety: it sends the route as far as it is safe, diff --git a/dimos/navigation/basic_path_follower/test_module.py b/dimos/navigation/basic_path_follower/test_module.py index 3ee85b7525..09248698de 100644 --- a/dimos/navigation/basic_path_follower/test_module.py +++ b/dimos/navigation/basic_path_follower/test_module.py @@ -12,7 +12,58 @@ # See the License for the specific language governing permissions and # limitations under the License. -from dimos.navigation.basic_path_follower.module import lookahead_distance +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.navigation.basic_path_follower.module import BasicPathFollower, lookahead_distance +from dimos.protocol.tf.tf import MultiTBuffer + + +class FakeTF(MultiTBuffer): + def stop(self) -> None: + pass + + +def _odom() -> Odometry: + return Odometry( + ts=1.0, + frame_id="odom", + child_frame_id="mid360_link", + pose=Pose(Vector3(1.0, 2.0, 3.0), Quaternion(0.0, 0.0, 0.0, 1.0)), + ) + + +def test_on_odometry_steers_from_the_base_pose(): + tf = FakeTF() + tf.receive_transform( + Transform( + translation=Vector3(0.0, 0.0, 0.163), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="base_link", + child_frame_id="mid360_link", + ts=1.0, + ) + ) + module = BasicPathFollower() + module._tf = tf + try: + module._on_odometry(_odom()) + assert module._current_pose is not None + assert abs(module._current_pose.position.z - (3.0 - 0.163)) < 1e-9 + finally: + module.stop() + + +def test_on_odometry_drops_frames_without_the_mount_tf(): + module = BasicPathFollower() + module._tf = FakeTF() + try: + module._on_odometry(_odom()) + assert module._current_pose is None + finally: + module.stop() def test_lookahead_floor_at_low_speed(): diff --git a/dimos/navigation/constants.py b/dimos/navigation/constants.py deleted file mode 100644 index 3c67e62820..0000000000 --- a/dimos/navigation/constants.py +++ /dev/null @@ -1,18 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Shared navigation constants.""" - -# Maximum time window to look for transforms before refusing a tf. -TF_LOOKUP_TOLERANCE_S = 0.25 diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py index 8250a15b7a..d662d40adb 100644 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -24,7 +24,7 @@ from dimos.msgs.geometry_msgs.PointStamped import PointStamped from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.nav_msgs.Odometry import Odometry -from dimos.navigation.constants import TF_LOOKUP_TOLERANCE_S +from dimos.navigation.tf_pose import OdomBasePose, base_height_above_ground from dimos.utils.logging_config import setup_logger logger = setup_logger() @@ -32,17 +32,13 @@ class GoalRelayConfig(ModuleConfig): base_frame: str = "base_link" - sensor_frame: str = "mid360_link" - # Lidar height above the ground while standing. Leave as None to skip the - # ground correction and pass odometry through unshifted. + # Lidar height above the ground while standing. None skips the ground + # correction. lidar_height: float | None = None class GoalRelay(Module): - """Adapt odometry and goal points to the planner's PoseStamped inputs. - - Odometry is corrected to the robot base frame via tf. - """ + """Adapt odometry and goal points to the planner's PoseStamped inputs.""" config: GoalRelayConfig @@ -54,6 +50,7 @@ class GoalRelay(Module): def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) + self._base_pose: OdomBasePose | None = None self._base_height: float | None = None @rpc @@ -63,34 +60,25 @@ def start(self) -> None: self.register_disposable(Disposable(self.goal.subscribe(self._on_goal))) def _on_odometry(self, msg: Odometry) -> None: - base = self.tf.get(msg.frame_id, self.config.base_frame, msg.ts, TF_LOOKUP_TOLERANCE_S) - if base is None: - logger.warning( - "No %s -> %s transform for odometry stamp %.3f, dropping frame.", - msg.frame_id, - self.config.base_frame, - msg.ts, - ) + if self._base_pose is None: + self._base_pose = OdomBasePose(self.tf, self.config.base_frame) + start = self._base_pose.resolve(msg) + if start is None: return - start = base.to_pose(ts=msg.ts) if self.config.lidar_height is not None: - base_height = self._base_from_ground() + base_height = self._resolve_base_height(msg.child_frame_id, self.config.lidar_height) if base_height is None: return start.position.z -= base_height self.start_pose.publish(start) - def _base_from_ground(self) -> float | None: - if self._base_height is None and self.config.lidar_height is not None: - mount = self.tf.get(self.config.base_frame, self.config.sensor_frame) - if mount is None: - logger.warning( - "No %s -> %s mount transform on tf, cannot ground-project the start pose.", - self.config.base_frame, - self.config.sensor_frame, - ) - else: - self._base_height = self.config.lidar_height - mount.translation.z + def _resolve_base_height(self, sensor_frame: str, lidar_height: float) -> float | None: + if self._base_height is None: + assert self._base_pose is not None + leg = self._base_pose.sensor_to_base(sensor_frame) + if leg is None: + return None + self._base_height = base_height_above_ground(lidar_height, -leg) return self._base_height def _on_goal(self, point: PointStamped) -> None: diff --git a/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py b/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py index 739029361a..2af72f9532 100644 --- a/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py +++ b/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py @@ -33,6 +33,10 @@ class MLSPlannerNativeConfig(NativeModuleConfig): world_frame: str = "map" voxel_size: float = 0.08 robot_height: float = 0.3 + # Subtracted from the start pose z before snapping to a surface. For start + # poses stamped at the sensor instead of the ground; leave 0 when the + # publisher already ground-projects (GoalRelay does). + start_z_offset_m: float = 0.0 max_overhead_m: float = 2.0 surface_closing_radius: float = 0.3 diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs index 4118e62727..181b932f58 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/main.rs @@ -319,6 +319,7 @@ impl Worker { let Some(start) = *self.latest_start.lock().expect("start mutex") else { return; }; + let start = (start.0, start.1, start.2 - self.config.start_z_offset_m); let goal = { let mut guard = self.active_goal.lock().expect("goal mutex"); let Some(goal) = *guard else { diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs index 2f50236635..c5d02ae5ce 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs @@ -37,6 +37,11 @@ pub struct Config { pub voxel_size: f32, #[validate(range(exclusive_min = 0.0))] pub robot_height: f32, + /// Subtracted from the start pose z before snapping to a surface. For start + /// poses stamped at the sensor instead of the ground; 0 when the publisher + /// already ground-projects. + #[validate(range(min = 0.0))] + pub start_z_offset_m: f32, /// Ignore surface more than this far above the sensor. #[validate(range(min = 0.0))] pub max_overhead_m: f32, @@ -572,6 +577,7 @@ mod region_tests { world_frame: String::new(), voxel_size: 0.1, robot_height: 0.5, + start_z_offset_m: 0.0, max_overhead_m: 2.0, surface_closing_radius: 0.3, node_spacing_m: 1.0, diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs index 183231d404..8245cc1d33 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/planner.rs @@ -901,6 +901,7 @@ mod tests { world_frame: "world".into(), voxel_size: VOXEL, robot_height: Z_TOL, + start_z_offset_m: 0.0, max_overhead_m: 2.0, surface_closing_radius: 0.0, node_spacing_m: 1.0, @@ -927,6 +928,7 @@ mod tests { world_frame: "world".into(), voxel_size: VOXEL, robot_height: Z_TOL, + start_z_offset_m: 0.0, max_overhead_m: 2.0, surface_closing_radius: 0.0, node_spacing_m: 1.0, diff --git a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs index 16d25c84e5..a59d8f14a9 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/python.rs @@ -86,6 +86,8 @@ impl MLSPlanner { world_frame: String::new(), voxel_size, robot_height, + // Unused here. Only the binary's replan loop projects the start. + start_z_offset_m: 0.0, max_overhead_m, surface_closing_radius, node_spacing_m, diff --git a/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py b/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py new file mode 100644 index 0000000000..760e29f0ff --- /dev/null +++ b/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py @@ -0,0 +1,114 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.navigation.nav_3d.mls_planner.goal_relay import GoalRelay +from dimos.protocol.tf.tf import MultiTBuffer + +MOUNT_Z = 0.163 + + +class FakeTF(MultiTBuffer): + """In-memory tf with the stop() hook and call counter the module tests need.""" + + def __init__(self) -> None: + super().__init__() + self.gets = 0 + + def get(self, *args, **kwargs): # type: ignore[no-untyped-def] + self.gets += 1 + return super().get(*args, **kwargs) + + def stop(self) -> None: + pass + + +def _mount() -> Transform: + return Transform( + translation=Vector3(0.0, 0.0, MOUNT_Z), + rotation=Quaternion(0.0, 0.0, 0.0, 1.0), + frame_id="base_link", + child_frame_id="mid360_link", + ts=1.0, + ) + + +def _odom(z: float = 3.0) -> Odometry: + return Odometry( + ts=1.0, + frame_id="odom", + child_frame_id="mid360_link", + pose=Pose(Vector3(1.0, 2.0, z), Quaternion(0.0, 0.0, 0.0, 1.0)), + ) + + +def _relay(tf: FakeTF, **config) -> tuple[GoalRelay, list]: # type: ignore[no-untyped-def] + module = GoalRelay(**config) + module._tf = tf + captured: list = [] + module.start_pose.subscribe(captured.append) + return module, captured + + +def test_start_pose_is_ground_projected(): + tf = FakeTF() + tf.receive_transform(_mount()) + module, captured = _relay(tf, lidar_height=0.45) + try: + module._on_odometry(_odom()) + # Base sits MOUNT_Z below the sensor, then drops by the base's height + # above ground (0.45 - MOUNT_Z): together exactly the lidar height. + assert len(captured) == 1 + assert abs(captured[0].position.z - (3.0 - 0.45)) < 1e-9 + finally: + module.stop() + + +def test_drops_frames_without_the_mount_tf(): + module, captured = _relay(FakeTF(), lidar_height=0.45) + try: + module._on_odometry(_odom()) + assert captured == [] + finally: + module.stop() + + +def test_no_lidar_height_skips_the_ground_correction(): + tf = FakeTF() + tf.receive_transform(_mount()) + module, captured = _relay(tf) + try: + module._on_odometry(_odom()) + assert len(captured) == 1 + assert abs(captured[0].position.z - (3.0 - MOUNT_Z)) < 1e-9 + finally: + module.stop() + + +def test_mount_is_looked_up_once(): + tf = FakeTF() + tf.receive_transform(_mount()) + module, captured = _relay(tf, lidar_height=0.45) + try: + module._on_odometry(_odom()) + module._on_odometry(_odom(z=4.0)) + assert len(captured) == 2 + assert abs(captured[1].position.z - (4.0 - 0.45)) < 1e-9 + assert tf.gets == 1 + finally: + module.stop() 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 91ad9c8753..6c63c53e21 100644 --- a/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py +++ b/dimos/navigation/nav_3d/mls_planner/utils/plan_rrd.py @@ -31,6 +31,7 @@ from dimos.mapping.ray_tracing.transformer import RayTraceMap from dimos.memory2.store.sqlite import SqliteStore +from dimos.memory2.tf import StreamTF from dimos.memory2.transform import FnTransformer from dimos.memory2.type.observation import Observation from dimos.msgs.geometry_msgs.Quaternion import Quaternion @@ -38,24 +39,21 @@ from dimos.msgs.geometry_msgs.Vector3 import Vector3 from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2, register_colormap_annotation -from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.navigation.nav_3d.mls_planner.mls_planner import MLSPlanner -from dimos.protocol.tf.tf import MultiTBuffer -from dimos.robot.unitree.go2.constants import ROBOT_LENGTH, ROBOT_WIDTH +from dimos.navigation.tf_pose import base_height_above_ground +from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT, ROBOT_LENGTH, ROBOT_WIDTH from dimos.utils.data import resolve_named_path TIMELINE = "ts" -ODOM_AXIS_LEN = 0.5 +AXIS_LEN = 0.5 AXIS_RADIUS_RATIO = 25 # Mount frames as recorded on the tf stream. BASE_FRAME = "base_link" SENSOR_FRAME = "mid360_link" -# need to wait to get tfs sometimes -TF_WARMUP_SAMPLES = 20 -ODOM_PATH_COLOR = [80, 160, 255] +SENSOR_PATH_COLOR = [80, 160, 255] # Different colors for each path when running with multiple configs PATH_PALETTE = [ @@ -135,21 +133,16 @@ def _log_path_wp(waypoints: NDArray[np.float32] | None, entity: str, color: list def _base_from_sensor(store: SqliteStore) -> Transform | None: - """Sensor to robot base link transform using the stored tf stream.""" - buffer = MultiTBuffer() - try: - for i, obs in enumerate(store.stream("tf", TFMessage).order_by("ts")): - buffer.receive_transform(*obs.data.transforms) - if i >= TF_WARMUP_SAMPLES: - break - except Exception as e: - print(f"no usable tf stream in the recording ({e}); skipping the base_link triad") + """Sensor to robot base link transform from the recorded tf stream.""" + tf = StreamTF.from_store(store) + if tf is None: + print("no tf stream in the recording; skipping the base_link triad") return None - return buffer.get(SENSOR_FRAME, BASE_FRAME) + return tf.get(SENSOR_FRAME, BASE_FRAME) def _base_pose(pose: tuple[float, ...], ts: float, base_from_sensor: Transform) -> Transform: - """World to robot base link, using the sensor to robot base link transform.""" + """World to robot base link.""" px, py, pz, qx, qy, qz, qw = pose sensor = Transform( translation=Vector3(px, py, pz), @@ -161,11 +154,34 @@ def _base_pose(pose: tuple[float, ...], ts: float, base_from_sensor: Transform) return sensor + base_from_sensor +def _plan_start( + pose: tuple[float, ...], + ts: float, + base_from_sensor: Transform | None, + base_height: float, + robot_height: float, +) -> tuple[tuple[float, float, float], Transform | None]: + """Ground-projected planner start, plus the base pose when tf has the mount. + + Without a tf stream the start is the sensor pose dropped by the robot height. + """ + px, py, pz, *_ = pose + if base_from_sensor is None: + return (float(px), float(py), float(pz) - robot_height), None + base = _base_pose(pose, ts, base_from_sensor) + start = ( + float(base.translation.x), + float(base.translation.y), + float(base.translation.z) - base_height, + ) + return start, base + + def _log_odometry( pose: tuple[float, ...], ts: float, trail: list[tuple[float, float, float]], - base_from_sensor: Transform | None, + base: Transform | None, ) -> None: """Trace the sensor moving throughout the scene.""" px, py, pz, qx, qy, qz, qw = pose @@ -176,10 +192,11 @@ def _log_odometry( ) trail.append((px, py, pz)) if len(trail) > 1: - rr.log("world/mid360_path", rr.LineStrips3D([trail], colors=[ODOM_PATH_COLOR], radii=0.015)) - if base_from_sensor is None: + rr.log( + "world/mid360_path", rr.LineStrips3D([trail], colors=[SENSOR_PATH_COLOR], radii=0.015) + ) + if base is None: return - base = _base_pose(pose, ts, base_from_sensor) rr.log( "world/base_link", rr.Transform3D( @@ -359,8 +376,7 @@ def _process_frame( ray_obs: Observation[PointCloud2], planners: list[tuple[str, list[int], MLSPlanner]], goal: tuple[float, float, float], - base_from_sensor: Transform | None, - robot_height: float, + start: tuple[float, float, float], render_voxel: float, clearance_clamp: float, hard_clearance: float, @@ -369,21 +385,7 @@ def _process_frame( """Plan every config for one frame, log paths/map/metrics, return the ref timing.""" assert ray_obs.pose_tuple is not None bounds = ray_obs.tags["region_bounds"] - px, py, pz, *_ = ray_obs.pose_tuple - # Plan from the robot base, ground-projected to the supporting surface. Without - # a tf stream fall back to the sensor pose dropped by the robot height. - if base_from_sensor is not None: - base = _base_pose(ray_obs.pose_tuple, ray_obs.ts, base_from_sensor) - # The mount transform gives the lidar's height above base_link, so the - # base sits this far below the lidar's known ground height. - base_height = robot_height - base_from_sensor.inverse().translation.z - start = ( - float(base.translation.x), - float(base.translation.y), - float(base.translation.z) - base_height, - ) - else: - start = (float(px), float(py), float(pz) - robot_height) + _, _, pz, *_ = ray_obs.pose_tuple ox, oy, radius, z_min, z_max = bounds pts = ray_obs.data.points_f32() rr.set_time(TIMELINE, timestamp=ray_obs.ts) @@ -455,7 +457,7 @@ def main( "0 emits all, higher drops isolated returns", ), robot_height: float = typer.Option( - 0.45, "--robot-height", help="Robot height, ground to tallest point / lidar (m)" + ROBOT_HEIGHT, "--robot-height", help="Robot height, ground to tallest point / lidar (m)" ), max_overhead: float = typer.Option( 2.0, "--max-overhead", help="Ignore surface more than this far above the sensor (m)" @@ -566,6 +568,11 @@ def main( rr.log("world/goal", rr.Points3D([goal], colors=[[255, 0, 0]], radii=0.1), static=True) base_from_sensor = _base_from_sensor(store) + base_height = ( + base_height_above_ground(robot_height, base_from_sensor.inverse()) + if base_from_sensor is not None + else 0.0 + ) entities = ["world/mid360_link/axes"] + ( ["world/base_link/axes"] if base_from_sensor else [] ) @@ -575,12 +582,12 @@ def main( rr.Arrows3D( origins=[[0.0, 0.0, 0.0]] * 3, vectors=[ - [ODOM_AXIS_LEN, 0.0, 0.0], - [0.0, ODOM_AXIS_LEN, 0.0], - [0.0, 0.0, ODOM_AXIS_LEN], + [AXIS_LEN, 0.0, 0.0], + [0.0, AXIS_LEN, 0.0], + [0.0, 0.0, AXIS_LEN], ], colors=[[255, 0, 0], [0, 255, 0], [0, 0, 255]], - radii=ODOM_AXIS_LEN / AXIS_RADIUS_RATIO, + radii=AXIS_LEN / AXIS_RADIUS_RATIO, ), static=True, ) @@ -593,7 +600,6 @@ def main( ), static=True, ) - # Light red clearance cylinder centered on the robot base. # wall_clearance is the planner's proxy for the robot radius. rr.log( "world/base_link/clearance", @@ -612,18 +618,20 @@ def main( for ray_obs in ray_pipeline: if ray_obs.pose_tuple is None: continue + start, base = _plan_start( + ray_obs.pose_tuple, ray_obs.ts, base_from_sensor, base_height, robot_height + ) ref_timing = _process_frame( ray_obs, planners, goal, - base_from_sensor, - robot_height, + start, render_voxel, clearance_clamp, ref_clearance, crop, ) - _log_odometry(ray_obs.pose_tuple, ray_obs.ts, sensor_trail, base_from_sensor) + _log_odometry(ray_obs.pose_tuple, ray_obs.ts, sensor_trail, base) frame += 1 print( f"frame={frame} configs={len(planners)} " diff --git a/dimos/navigation/test_tf_pose.py b/dimos/navigation/test_tf_pose.py new file mode 100644 index 0000000000..0f1f487447 --- /dev/null +++ b/dimos/navigation/test_tf_pose.py @@ -0,0 +1,123 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.Quaternion import Quaternion +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.msgs.nav_msgs.Odometry import Odometry +from dimos.navigation.tf_pose import OdomBasePose, base_height_above_ground +from dimos.protocol.tf.tf import MultiTBuffer + +IDENTITY = Quaternion(0.0, 0.0, 0.0, 1.0) + + +class CountingTF(MultiTBuffer): + def __init__(self) -> None: + super().__init__() + self.gets = 0 + + def get(self, *args, **kwargs): # type: ignore[no-untyped-def] + self.gets += 1 + return super().get(*args, **kwargs) + + +def _mount(z: float = 0.163, pitch: float = 0.0) -> Transform: + return Transform( + translation=Vector3(0.0, 0.0, z), + rotation=Quaternion.from_euler(Vector3(0.0, pitch, 0.0)), + frame_id="base_link", + child_frame_id="mid360_link", + ts=1.0, + ) + + +def _odom(orientation: Quaternion = IDENTITY) -> Odometry: + return Odometry( + ts=1.0, + frame_id="odom", + child_frame_id="mid360_link", + pose=Pose(Vector3(1.0, 2.0, 3.0), orientation), + ) + + +def test_translates_to_base_frame(): + tf = MultiTBuffer() + tf.receive_transform(_mount()) + pose = OdomBasePose(tf, "base_link").resolve(_odom()) + assert pose is not None + assert pose.frame_id == "odom" + assert pose.ts == 1.0 + assert abs(pose.position.x - 1.0) < 1e-9 + assert abs(pose.position.y - 2.0) < 1e-9 + assert abs(pose.position.z - (3.0 - 0.163)) < 1e-9 + + +def test_composes_out_the_mount_pitch(): + # A level body reads its own mount tilt as the sensor's world orientation, so + # composing the mount out returns identity. + mount = _mount(pitch=0.3) + tf = MultiTBuffer() + tf.receive_transform(mount) + pose = OdomBasePose(tf, "base_link").resolve(_odom(orientation=mount.rotation)) + assert pose is not None + assert pose.orientation.angle_to(IDENTITY) < 1e-5 + + +def test_preserves_body_yaw_under_mount_tilt(): + mount = _mount(pitch=0.3) + body = Quaternion.from_euler(Vector3(0.0, 0.0, 0.7)) + tf = MultiTBuffer() + tf.receive_transform(mount) + pose = OdomBasePose(tf, "base_link").resolve(_odom(orientation=body * mount.rotation)) + assert pose is not None + assert pose.orientation.angle_to(body) < 1e-5 + + +def test_drops_frames_until_the_mount_leg_arrives(): + tf = MultiTBuffer() + resolver = OdomBasePose(tf, "base_link") + assert resolver.resolve(_odom()) is None + tf.receive_transform(_mount()) + resolver._next_lookup = 0.0 + assert resolver.resolve(_odom()) is not None + + +def test_missing_leg_lookups_are_throttled(): + tf = CountingTF() + resolver = OdomBasePose(tf, "base_link") + assert resolver.resolve(_odom()) is None + assert resolver.resolve(_odom()) is None + assert tf.gets == 1 + + +def test_mount_leg_is_looked_up_once(): + tf = CountingTF() + tf.receive_transform(_mount()) + resolver = OdomBasePose(tf, "base_link") + assert resolver.resolve(_odom()) is not None + assert resolver.resolve(_odom()) is not None + assert tf.gets == 1 + + +def test_base_frame_odometry_passes_through(): + resolver = OdomBasePose(MultiTBuffer(), "base_link") + msg = Odometry(ts=1.0, frame_id="odom", child_frame_id="base_link") + pose = resolver.resolve(msg) + assert pose is not None + assert pose.frame_id == "odom" + + +def test_base_height_above_ground(): + assert abs(base_height_above_ground(0.45, _mount()) - (0.45 - 0.163)) < 1e-9 diff --git a/dimos/navigation/tf_pose.py b/dimos/navigation/tf_pose.py new file mode 100644 index 0000000000..05549ab712 --- /dev/null +++ b/dimos/navigation/tf_pose.py @@ -0,0 +1,94 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Resolve odometry into base-frame poses using the static mount tf.""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING + +from dimos.msgs.geometry_msgs.Transform import Transform +from dimos.utils.logging_config import setup_logger + +if TYPE_CHECKING: + from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped + from dimos.msgs.nav_msgs.Odometry import Odometry + from dimos.protocol.tf.tf import TFLookup + +logger = setup_logger() + + +def base_height_above_ground(lidar_height: float, base_to_sensor: Transform) -> float: + """Height of the base frame origin above the ground while standing. + + The lidar sits ``lidar_height`` above the ground and + ``base_to_sensor.translation.z`` above the base origin. + """ + return lidar_height - base_to_sensor.translation.z + + +class OdomBasePose: + """Turn odometry messages into the base-frame pose they imply. + + The odometry message itself carries the live parent -> child edge, so only the + static leg from the message's child frame to the base frame is looked up on tf, + once, then reused for every following message. + """ + + # While the leg is missing, retry the lookup at most this often. The buffer + # warns on every miss, so per-message retries would flood the log. + RETRY_PERIOD_S = 1.0 + + def __init__(self, tf: TFLookup, base_frame: str) -> None: + self._tf = tf + self.base_frame = base_frame + self._legs: dict[str, Transform] = {} + self._waiting = False + self._next_lookup = 0.0 + + def resolve(self, msg: Odometry) -> PoseStamped | None: + """The base pose for one message. None until tf has the mount leg.""" + if msg.child_frame_id == self.base_frame: + return msg.to_pose_stamped() + leg = self.sensor_to_base(msg.child_frame_id) + if leg is None: + return None + odom = Transform.from_pose(msg.child_frame_id, msg.to_pose_stamped()) + return (odom + leg).to_pose(ts=msg.ts) + + def sensor_to_base(self, sensor_frame: str) -> Transform | None: + """The cached static sensor -> base leg. Logs once per outage, not per message.""" + if sensor_frame == self.base_frame: + return Transform.identity() + leg = self._legs.get(sensor_frame) + if leg is None: + if self._waiting and time.monotonic() < self._next_lookup: + return None + leg = self._tf.get(sensor_frame, self.base_frame) + if leg is None: + self._next_lookup = time.monotonic() + self.RETRY_PERIOD_S + if not self._waiting: + self._waiting = True + logger.warning( + "No %s -> %s transform on tf yet, dropping odometry until it arrives.", + sensor_frame, + self.base_frame, + ) + return None + if self._waiting: + self._waiting = False + logger.info("Got the %s -> %s transform, resuming.", sensor_frame, self.base_frame) + self._legs[sensor_frame] = leg + return leg diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_mls_htc.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_mls_htc.py index 1468c8c989..eb05bd5e73 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_mls_htc.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_mls_htc.py @@ -82,6 +82,8 @@ def _render_path(msg: Any) -> Any: world_frame="world", voxel_size=voxel_size, robot_height=go2_lidar_height, + # The start pose is raw go2 odometry, so the planner ground-projects it. + start_z_offset_m=go2_lidar_height, wall_clearance_m=0.2, wall_buffer_m=0.75, wall_buffer_weight=100.0, diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index afbdf10a96..e054d927c5 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -177,7 +177,6 @@ def _static_sensor_axes(rr: Any) -> list[Any]: ] ), PointLio.blueprint(), - # Publishes the rig's mount frames onto tf. Go2Mid360StaticTf.blueprint(), RayTracingVoxelMap.blueprint( voxel_size=voxel_size, diff --git a/dimos/robot/unitree/go2/constants.py b/dimos/robot/unitree/go2/constants.py index 92126bb334..51d06da0f8 100644 --- a/dimos/robot/unitree/go2/constants.py +++ b/dimos/robot/unitree/go2/constants.py @@ -17,5 +17,5 @@ # Robot footprint in meters. Length is forward x, width is left y. ROBOT_LENGTH = 0.6858 ROBOT_WIDTH = 0.3175 -# Ground to the tallest point. Used for clearance and the ground projection. +# Ground to the tallest point. ROBOT_HEIGHT = 0.45 diff --git a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py index 358bcafc9d..0d9cb26923 100644 --- a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py +++ b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py @@ -14,9 +14,9 @@ """Static mount frames for the Go2 + Mid-360 + front-camera rig. -Published continuously onto tf while recording (see :class:`Go2Mid360StaticTf`) so the -mount geometry lands in the recording's tf stream and companion streams (camera, go2 -lidar) can be anchored to ``base_link``. +Published continuously onto tf (see :class:`Go2Mid360StaticTf`) so the mount geometry +lands in the tf stream and companion streams (camera, go2 lidar) can be anchored to +``base_link``. Mount geometry (measured on the physical rig) --------------------------------------------- @@ -24,6 +24,11 @@ - front_camera -> mid360_link: lidar is 3.2cm back, 12cm up, pitched 44 deg down. - front_camera -> camera_optical: the standard ROS optical rotation (x-right, y-down, z-forward). + +The published tree is rooted at ``mid360_link``: Point-LIO owns the live +``odom -> mid360_link`` edge, and the rerun bridge keys tf entities by child frame, so +publishing the two edges above the lidar inverted keeps the static tree off the entity +the live edge writes. The tf buffer composes either direction. """ from __future__ import annotations @@ -50,14 +55,9 @@ ] -def base_link_from_mid360() -> Transform: - """Composed base_link -> mid360_link transform from the static mount tree.""" - edges = {t.child_frame_id: t for t in frames_to_edge_transforms(FRAMES)} - return edges["front_camera"] + edges["mid360_link"] - - class Go2Mid360StaticTf(StaticTfPublisher): """Publishes the Go2/Mid-360 mount tree onto tf on a fixed interval.""" def transforms(self) -> list[Transform]: - return frames_to_edge_transforms(FRAMES) + edges = {t.child_frame_id: t for t in frames_to_edge_transforms(FRAMES)} + return [-edges["mid360_link"], -edges["front_camera"], edges["camera_optical"]] diff --git a/dimos/robot/unitree/go2/zenoh/blueprints.py b/dimos/robot/unitree/go2/zenoh/blueprints.py index 553167206b..708c400023 100644 --- a/dimos/robot/unitree/go2/zenoh/blueprints.py +++ b/dimos/robot/unitree/go2/zenoh/blueprints.py @@ -136,7 +136,7 @@ def _rerun_config(visual_override: dict[str, Any] | None = None) -> dict[str, An _mls_planner = MLSPlannerNative.blueprint( world_frame="odom", voxel_size=voxel_size, - robot_height=0.3, + robot_height=ROBOT_HEIGHT, surface_closing_radius=0.3, wall_clearance_m=0.1, wall_buffer_m=0.75, diff --git a/examples/native-modules/rust/src/tf_listener.rs b/examples/native-modules/rust/src/tf_listener.rs index c4227fc35a..8d7efd1790 100644 --- a/examples/native-modules/rust/src/tf_listener.rs +++ b/examples/native-modules/rust/src/tf_listener.rs @@ -38,7 +38,7 @@ impl TfListener { x = p.x, y = p.y, z = p.z, - "transform:", + "tf lookup", ); } None => tracing::info!("a -> d not available yet"), diff --git a/native/rust/README.md b/native/rust/README.md index f0241ed932..5cc3534a15 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -118,13 +118,13 @@ impl VoxelMap { async fn handle_lidar(&mut self, cloud: PointCloud2) { // De-rotate a scan from the lidar's mount frame into the robot base frame. if let Some(t) = self.tf.get_latest("base_link", "mid360_link") { - let point_in_base = t.isometry() * point_in_lidar; + let point_in_base = t.rotation() * point_in_lidar + t.translation(); } } } ``` -`Tf` is a cheap-to-clone handle; the graph fills in the background as `tf` messages arrive. `get(parent, child, time, tolerance)` selects the sample nearest `time` (latest when `None`) and returns `None` when no path connects the frames or no sample falls within `tolerance` seconds. `get_latest` is the no-time shorthand. The result exposes an `nalgebra` `Isometry3` via `isometry()`, ready to apply to a point. Lookups are nearest-in-time, not interpolated. +`Tf` is a cheap-to-clone handle; the graph fills in the background as `tf` messages arrive. `get(parent, child, time, tolerance)` selects the sample nearest `time` (latest when `None`) and returns `None` when no path connects the frames or no sample falls within `tolerance` seconds. `get_latest` is the no-time shorthand. The result exposes its `nalgebra` parts via `translation()` (a `Vector3`) and `rotation()` (a `UnitQuaternion`). Lookups are nearest-in-time, not interpolated. `publish` sends transforms onto the same `tf` topic, the counterpart to Python's `tf.publish()`. Published transforms also feed the module's own graph, so a `get` right after the publish sees them. Build the isometry from `dimos_module::nalgebra`, re-exported so the version matches the SDK's types: diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index a48f91e109..399d8c19b9 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -334,9 +334,8 @@ impl Builder { /// A handle that answers transform queries and publishes on the `tf` topic. /// - /// The first call subscribes to the resolved `tf` topic, starts filling the - /// transform graph in the background, and wires a publish queue for the same - /// topic. Repeated calls share one graph. + /// The graph fills in the background as `tf` messages arrive. Repeated calls + /// share one graph. pub fn tf(&mut self) -> crate::tf::Tf { if let Some(tf) = &self.tf { return tf.clone(); @@ -345,7 +344,7 @@ impl Builder { let (tx, rx) = mpsc::channel(PUBLISH_CHANNEL_CAPACITY); self.outputs.push((topic.clone(), rx)); let (tf, route) = - crate::tf::tf_subscription(topic.clone(), crate::tf::DEFAULT_TF_BUFFER_SIZE, tx); + crate::tf::tf_subscription(topic.clone(), crate::tf::DEFAULT_TF_WINDOW_SECS, tx); self.routes.entry(topic).or_default().push(route); self.tf = Some(tf.clone()); tf @@ -791,6 +790,23 @@ mod tests { assert_eq!(output.topic, "/robot/cmd_vel"); } + #[test] + fn tf_uses_mapped_topic() { + let mut builder = builder_with_topics(&[("tf", "/robot/tf")]); + builder.tf(); + assert!(builder.routes.contains_key("/robot/tf")); + assert_eq!(builder.outputs[0].0, "/robot/tf"); + } + + #[test] + fn repeated_tf_calls_share_one_graph() { + let mut builder = builder_with_topics(&[("tf", "/tf")]); + builder.tf(); + builder.tf(); + assert_eq!(builder.outputs.len(), 1); + assert_eq!(builder.routes.get("/tf").map(Vec::len), Some(1)); + } + // recv/publish concurrency #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index 24e4f8f782..64d7d4e8f5 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -30,7 +30,7 @@ use tokio::sync::mpsc; use crate::module::Route; /// How many seconds of history each edge keeps. -pub const DEFAULT_TF_BUFFER_SIZE: f64 = 10.0; +pub(crate) const DEFAULT_TF_WINDOW_SECS: f64 = 10.0; fn now_secs() -> f64 { SystemTime::now() @@ -65,12 +65,10 @@ impl Transform { } } - /// Translation component. pub fn translation(&self) -> Vector3 { self.iso.translation.vector } - /// Rotation component. pub fn rotation(&self) -> UnitQuaternion { self.iso.rotation } @@ -101,14 +99,14 @@ struct Sample { // One edge's time-sorted history, capped to a fixed-duration window. struct TBuffer { - buffer_size: f64, + window_secs: f64, samples: VecDeque, } impl TBuffer { - fn new(buffer_size: f64) -> Self { + fn new(window_secs: f64) -> Self { Self { - buffer_size, + window_secs, samples: VecDeque::new(), } } @@ -116,7 +114,7 @@ impl TBuffer { fn add(&mut self, ts: f64, iso: Isometry3) { let pos = self.samples.partition_point(|s| s.ts <= ts); self.samples.insert(pos, Sample { ts, iso }); - self.prune(ts - self.buffer_size); + self.prune(ts - self.window_secs); } fn prune(&mut self, min_ts: f64) { @@ -153,27 +151,47 @@ impl TBuffer { _ => Some(best), } } + + // One transform for this edge: the latest sample, or the one nearest `time`. + fn sample( + &self, + parent: &str, + child: &str, + time: Option, + tolerance: Option, + ) -> Option { + let s = match time { + None => self.last()?, + Some(t) => self.find_closest(t, tolerance)?, + }; + Some(Transform { + parent: parent.to_string(), + child: child.to_string(), + ts: s.ts, + iso: s.iso, + }) + } } /// The transform graph: one [`TBuffer`] per `(parent, child)` edge. struct MultiTBuffer { - buffer_size: f64, + window_secs: f64, buffers: HashMap<(String, String), TBuffer>, } impl MultiTBuffer { - fn new(buffer_size: f64) -> Self { + fn new(window_secs: f64) -> Self { Self { - buffer_size, + window_secs, buffers: HashMap::new(), } } fn receive(&mut self, parent: &str, child: &str, ts: f64, iso: Isometry3) { - let buffer_size = self.buffer_size; + let window_secs = self.window_secs; self.buffers .entry((parent.to_string(), child.to_string())) - .or_insert_with(|| TBuffer::new(buffer_size)) + .or_insert_with(|| TBuffer::new(window_secs)) .add(ts, iso); } @@ -190,26 +208,6 @@ impl MultiTBuffer { out } - fn sample( - &self, - buf: &TBuffer, - parent: &str, - child: &str, - time: Option, - tolerance: Option, - ) -> Option { - let s = match time { - None => buf.last()?, - Some(t) => buf.find_closest(t, tolerance)?, - }; - Some(Transform { - parent: parent.to_string(), - child: child.to_string(), - ts: s.ts, - iso: s.iso, - }) - } - // A single forward or reverse edge. Reverse returns the inverse. fn edge( &self, @@ -227,11 +225,11 @@ impl MultiTBuffer { }); } if let Some(buf) = self.buffers.get(&(parent.to_string(), child.to_string())) { - return self.sample(buf, parent, child, time, tolerance); + return buf.sample(parent, child, time, tolerance); } if let Some(buf) = self.buffers.get(&(child.to_string(), parent.to_string())) { - return self - .sample(buf, child, parent, time, tolerance) + return buf + .sample(child, parent, time, tolerance) .map(|t| t.inverse()); } None @@ -466,7 +464,7 @@ mod tests { #[test] fn direct_edge() { - let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); h.add("base_link", "arm", 1.0, (1.0, -1.0, 0.0), 0.0); let t = tf.get_latest("base_link", "arm").unwrap(); assert!((t.translation().x - 1.0).abs() < 1e-9); @@ -477,7 +475,7 @@ mod tests { #[test] fn reverse_edge_returns_inverse() { - let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); h.add("base_link", "arm", 1.0, (1.0, 2.0, 3.0), 0.0); let inv = tf.get_latest("arm", "base_link").unwrap(); assert!((inv.translation().x + 1.0).abs() < 1e-9); @@ -490,7 +488,7 @@ mod tests { // A 30-degree yaw then a pure translation. #[test] fn composes_ros_example_chain() { - let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); h.add("base_link", "arm", 1.0, (1.0, -1.0, 0.0), PI / 6.0); h.add("arm", "end_effector", 1.0, (1.0, 1.0, 0.0), 0.0); let t = tf.get_latest("base_link", "end_effector").unwrap(); @@ -511,7 +509,7 @@ mod tests { // world->robot->sensor multi-hop composition. #[test] fn composes_multi_hop_chain() { - let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); h.add("world", "robot", 1.0, (1.0, 2.0, 3.0), 0.0); h.add("robot", "sensor", 1.0, (0.5, 0.0, 0.2), PI / 2.0); let t = tf.get_latest("world", "sensor").unwrap(); @@ -522,14 +520,14 @@ mod tests { #[test] fn missing_path_returns_none() { - let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); h.add("world", "robot", 1.0, (1.0, 0.0, 0.0), 0.0); assert!(tf.get_latest("world", "unconnected").is_none()); } #[test] fn identity_for_same_frame() { - let (tf, _h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + let (tf, _h) = tf_with(DEFAULT_TF_WINDOW_SECS); // No query time: identity is stamped now, not the epoch. let t = tf.get_latest("base_link", "base_link").unwrap(); assert!((t.translation().norm()).abs() < 1e-12); @@ -545,7 +543,7 @@ mod tests { #[test] fn time_query_picks_nearest_sample() { - let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); h.add("a", "b", 10.0, (1.0, 0.0, 0.0), 0.0); h.add("a", "b", 20.0, (2.0, 0.0, 0.0), 0.0); let near_10 = tf.get("a", "b", Some(11.0), None).unwrap(); @@ -556,7 +554,7 @@ mod tests { #[test] fn time_query_outside_tolerance_returns_none() { - let (tf, h) = tf_with(DEFAULT_TF_BUFFER_SIZE); + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); h.add("a", "b", 10.0, (1.0, 0.0, 0.0), 0.0); assert!(tf.get("a", "b", Some(50.0), Some(1.0)).is_none()); assert!(tf.get("a", "b", Some(10.5), Some(1.0)).is_some()); @@ -582,7 +580,7 @@ mod tests { use lcm_msgs::tf2_msgs::TFMessage; let (tx, _rx) = mpsc::channel(8); - let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_BUFFER_SIZE, tx); + let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_WINDOW_SECS, tx); let msg = TFMessage { transforms: vec![lcm_msgs::geometry_msgs::TransformStamped { header: Header { @@ -620,7 +618,7 @@ mod tests { #[tokio::test] async fn publish_feeds_local_graph() { - let (tf, _rx, _h) = tf_with_publish(DEFAULT_TF_BUFFER_SIZE); + let (tf, _rx, _h) = tf_with_publish(DEFAULT_TF_WINDOW_SECS); let iso = Isometry3::from_parts( Translation3::new(1.0, 2.0, 3.0), UnitQuaternion::from_euler_angles(0.0, 0.0, PI / 2.0), @@ -639,7 +637,7 @@ mod tests { // Publish on one handle, dispatch the wire bytes into another graph. #[tokio::test] async fn publish_round_trips_through_route() { - let (tf_out, mut rx, _h) = tf_with_publish(DEFAULT_TF_BUFFER_SIZE); + let (tf_out, mut rx, _h) = tf_with_publish(DEFAULT_TF_WINDOW_SECS); let iso = Isometry3::from_parts( Translation3::new(0.5, -0.5, 0.25), UnitQuaternion::from_euler_angles(0.0, 0.0, PI / 6.0), @@ -651,7 +649,7 @@ mod tests { let bytes = rx.recv().await.unwrap(); let (tx, _rx2) = mpsc::channel(8); - let (tf_in, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_BUFFER_SIZE, tx); + let (tf_in, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_WINDOW_SECS, tx); route.try_dispatch(&bytes); let t = tf_in.get_latest("a", "b").unwrap(); @@ -674,4 +672,88 @@ mod tests { assert_eq!(st.header.stamp.sec, 2); assert_eq!(st.header.stamp.nsec, 0); } + + #[test] + fn add_out_of_order_keeps_samples_sorted() { + let mut buf = TBuffer::new(DEFAULT_TF_WINDOW_SECS); + buf.add(3.0, Isometry3::identity()); + buf.add(1.0, Isometry3::identity()); + buf.add(2.0, Isometry3::identity()); + assert!((buf.last().unwrap().ts - 3.0).abs() < 1e-9); + let s = buf.find_closest(1.9, None).unwrap(); + assert!((s.ts - 2.0).abs() < 1e-9); + } + + #[test] + fn tie_prefers_the_later_sample() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("a", "b", 10.0, (1.0, 0.0, 0.0), 0.0); + h.add("a", "b", 12.0, (2.0, 0.0, 0.0), 0.0); + let t = tf.get("a", "b", Some(11.0), None).unwrap(); + assert!((t.translation().x - 2.0).abs() < 1e-9); + } + + // Two routes to d: three hops through b, c and two through x. BFS must + // compose the two-hop route. + #[test] + fn bfs_takes_the_fewest_hops_on_a_branching_graph() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("a", "b", 1.0, (1.0, 0.0, 0.0), 0.0); + h.add("b", "c", 1.0, (1.0, 0.0, 0.0), 0.0); + h.add("c", "d", 1.0, (1.0, 0.0, 0.0), 0.0); + h.add("a", "x", 1.0, (10.0, 0.0, 0.0), 0.0); + h.add("x", "d", 1.0, (1.0, 0.0, 0.0), 0.0); + let t = tf.get_latest("a", "d").unwrap(); + assert!( + (t.translation().x - 11.0).abs() < 1e-9, + "{}", + t.translation().x + ); + } + + // Inverse of a rotated edge is t' = -R^T t, the classic sign/order trap. + #[test] + fn reverse_edge_inverts_rotation_and_translation() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("base_link", "arm", 1.0, (1.0, 2.0, 3.0), PI / 2.0); + let inv = tf.get_latest("arm", "base_link").unwrap(); + assert!( + (inv.translation().x + 2.0).abs() < 1e-9, + "{:?}", + inv.translation() + ); + assert!((inv.translation().y - 1.0).abs() < 1e-9); + assert!((inv.translation().z + 3.0).abs() < 1e-9); + let (_, _, yaw) = inv.rotation().euler_angles(); + assert!((yaw + PI / 2.0).abs() < 1e-9); + } + + #[test] + fn composed_chain_accumulates_rotation() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("a", "b", 1.0, (0.0, 0.0, 0.0), PI / 6.0); + h.add("b", "c", 1.0, (0.0, 0.0, 0.0), PI / 6.0); + let t = tf.get_latest("a", "c").unwrap(); + let (_, _, yaw) = t.rotation().euler_angles(); + assert!((yaw - PI / 3.0).abs() < 1e-9); + } + + #[test] + fn dispatch_of_undecodable_bytes_leaves_the_graph_empty() { + let (tx, _rx) = mpsc::channel(8); + let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_WINDOW_SECS, tx); + route.try_dispatch(b"garbage"); + assert!(tf.get_latest("a", "b").is_none()); + } + + #[tokio::test] + async fn publish_errors_when_the_background_task_is_gone() { + let (tf, rx, _h) = tf_with_publish(DEFAULT_TF_WINDOW_SECS); + drop(rx); + let err = tf + .publish(&[Transform::new("a", "b", 1.0, Isometry3::identity())]) + .await + .unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::BrokenPipe); + } } From aae072ffdcc146885d542300a51701dbfb943923 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Fri, 31 Jul 2026 10:56:46 -0700 Subject: [PATCH 17/26] IO ports --- dimos/core/native_module.py | 6 +- dimos/core/test_native_module.py | 47 +++++- native/rust/README.md | 7 + native/rust/dimos-module-macros/src/lib.rs | 91 +++++++--- native/rust/dimos-module/src/lib.rs | 7 +- native/rust/dimos-module/src/module.rs | 188 ++++++++++++++++++--- native/rust/dimos-module/src/tf.rs | 5 +- 7 files changed, 293 insertions(+), 58 deletions(-) diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index 4fef48bcff..5aca07a363 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -473,7 +473,7 @@ def _maybe_build(self) -> None: def _collect_topics(self) -> dict[str, str]: topics: dict[str, str] = {} - for name in list(self.inputs) + list(self.outputs): + for name in list(self.inputs) + list(self.outputs) + list(self.ios): stream = getattr(self, name, None) if stream is None: continue @@ -486,9 +486,9 @@ def _collect_topics(self) -> dict[str, str]: return topics def _collect_output_qos(self) -> dict[str, dict[str, str]]: - """Publisher QoS per output channel, keyed by channel.""" + """Publisher QoS per published channel, keyed by channel.""" qos_map: dict[str, dict[str, str]] = {} - for name in self.outputs: + for name in list(self.outputs) + list(self.ios): stream = getattr(self, name, None) if stream is None: continue diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index 4a68caedc8..49ca8e8729 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -18,6 +18,7 @@ The echo script writes received CLI args to a temp file for assertions. """ +import contextlib from io import BytesIO import json from pathlib import Path @@ -33,11 +34,13 @@ from dimos.core.core import rpc from dimos.core.module import Module from dimos.core.native_module import LogFormat, NativeModule, NativeModuleConfig -from dimos.core.stream import In, Out -from dimos.core.transport import LCMTransport +from dimos.core.stream import IO, In, Out +from dimos.core.transport import LCMTransport, ZenohTransport from dimos.msgs.geometry_msgs.Twist import Twist from dimos.msgs.sensor_msgs.Imu import Imu from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 +from dimos.msgs.tf2_msgs.TFMessage import TFMessage +from dimos.protocol.pubsub.impl.zenohpubsub import QOS_NEVER_DROP, Topic as ZenohTopic _ECHO = str(Path(__file__).parent / "demos" / "native_echo.py") @@ -87,6 +90,12 @@ class StubNativeModule(NativeModule): cmd_vel: In[Twist] +class StubIoModule(NativeModule): + config: StubNativeConfig + cmd_vel: In[Twist] + tf: IO[TFMessage] + + class StubConsumer(Module): pointcloud: In[PointCloud2] imu: In[Imu] @@ -160,6 +169,40 @@ def test_manual(dimos_cluster: ModuleCoordinator, args_file: str) -> None: } +def test_io_port_topic_reaches_the_native_process() -> None: + """An IO port is both a subscriber and a publisher, so it needs its topic.""" + module = StubIoModule(executable=_ECHO) + transports = [LCMTransport("/cmd_vel", Twist), LCMTransport("/tf", TFMessage)] + try: + module.set_transport("cmd_vel", transports[0]) + module.set_transport("tf", transports[1]) + + assert module._collect_topics() == { + "cmd_vel": "/cmd_vel#geometry_msgs.Twist", + "tf": "/tf#tf2_msgs.TFMessage", + } + finally: + module.stop() + for transport in transports: + with contextlib.suppress(Exception): + transport.stop() + + +def test_io_port_publisher_qos_reaches_the_native_process() -> None: + module = StubIoModule(executable=_ECHO) + transport = ZenohTransport(ZenohTopic("/tf", TFMessage, qos=QOS_NEVER_DROP)) + try: + module.set_transport("tf", transport) + + assert module._collect_output_qos() == { + transport.channel: {"reliability": "reliable", "congestion_control": "block"}, + } + finally: + module.stop() + with contextlib.suppress(Exception): + transport.stop() + + def test_autoconnect(args_file: str) -> None: """autoconnect passes correct topic args to the native subprocess.""" blueprint = autoconnect( diff --git a/native/rust/README.md b/native/rust/README.md index 5cc3534a15..e4e01b676e 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -26,6 +26,9 @@ struct MyModule { #[output(encode = Twist::encode)] out: Output, + #[io(decode = Twist::decode, encode = Twist::encode)] + shared: Io, + #[config] config: MyConfig, } @@ -37,6 +40,9 @@ impl MyModule { // processing function expected by cmd: Input async fn handle_cmd(&mut self, msg: Twist) { /* ... */ } + // processing function expected by shared: Io + async fn handle_shared(&mut self, msg: Twist) { /* ... */ } + // teardown / clean up logic async fn on_stop(&mut self) { /* ... */ } } @@ -57,6 +63,7 @@ Every transport is compiled into the binary. `run_with_transport` opens the one - `#[module(setup = fn, teardown = fn)]`: on the struct. Both optional. Names methods on `Self`. `setup` runs once before the input dispatch loop starts (use it to spawn background tasks or initialize resources); `teardown` runs once after the loop exits (use it for cleanup). - `#[input(decode = fn, handler = fn)]`: on a field of type `Input`. `decode` is required; `handler` defaults to `handle_`. - `#[output(encode = fn)]`: on a field of type `Output`. `encode` is required. +- `#[io(decode = fn, encode = fn, handler = fn)]`: on a field of type `Io`, a port that publishes to and subscribes on one topic. `decode` and `encode` are required; `handler` defaults to `handle_`. The transports deliver a message back to its own sender, so the handler also sees what the module publishes. Use `#[output]` instead when the module only publishes. - `#[config]`: on one field. The type must be defined with `#[native_config]` (see [Config](#config)). At most one per struct. If absent, `Config` defaults to `dimos_module::NoConfig`. - `#[tf]`: on a field of type `Tf`. Subscribes to the `tf` topic, answers transform queries, and publishes transforms (see [Transforms](#transforms)). No arguments. - Unattributed fields are initialized via `Default::default()` and treated as module state. diff --git a/native/rust/dimos-module-macros/src/lib.rs b/native/rust/dimos-module-macros/src/lib.rs index 5511356979..8679b2289c 100644 --- a/native/rust/dimos-module-macros/src/lib.rs +++ b/native/rust/dimos-module-macros/src/lib.rs @@ -17,7 +17,7 @@ use proc_macro2::TokenStream as TokenStream2; use quote::{format_ident, quote}; use syn::{parse_macro_input, Data, DeriveInput, Field, Fields, Ident, Path, Type}; -#[proc_macro_derive(Module, attributes(input, output, config, tf, module))] +#[proc_macro_derive(Module, attributes(input, output, io, config, tf, module))] pub fn derive_module(input: TokenStream) -> TokenStream { let input = parse_macro_input!(input as DeriveInput); match expand(input) { @@ -174,9 +174,22 @@ fn is_option(ty: &Type) -> bool { matches!(ty, Type::Path(p) if p.path.segments.last().is_some_and(|s| s.ident == "Option")) } +const ONE_ATTR_ONLY: &str = "field has multiple module attributes; only one of #[input], \ + #[output], #[io], #[config], #[tf] is allowed"; + enum FieldKind { - Input { decode: Path, handler: Ident }, - Output { encode: Path }, + Input { + decode: Path, + handler: Ident, + }, + Output { + encode: Path, + }, + Io { + decode: Path, + encode: Path, + handler: Ident, + }, Config, Tf, State, @@ -273,25 +286,30 @@ fn expand(input: DeriveInput) -> syn::Result { FieldKind::Output { encode } => { quote!(#name: builder.output(#name_str, #encode)) } + FieldKind::Io { decode, encode, .. } => { + quote!(#name: builder.io(#name_str, #decode, #encode)) + } FieldKind::Config => quote!(#name: config), FieldKind::Tf => quote!(#name: builder.tf()), FieldKind::State => quote!(#name: ::core::default::Default::default()), } }); - let input_fields: Vec<&ClassifiedField> = classified + // Every port that receives messages gets an arm in the select! loop. + let handled_fields: Vec<(&Ident, &Ident)> = classified .iter() - .filter(|f| matches!(f.kind, FieldKind::Input { .. })) + .filter_map(|f| match &f.kind { + FieldKind::Input { handler, .. } | FieldKind::Io { handler, .. } => { + Some((f.name, handler)) + } + _ => None, + }) .collect(); - let handle_body = if input_fields.is_empty() { + let handle_body = if handled_fields.is_empty() { quote!(::std::future::pending::<()>().await) } else { - let handle_arms = input_fields.iter().map(|f| { - let FieldKind::Input { handler, .. } = &f.kind else { - unreachable!() - }; - let name = f.name; + let handle_arms = handled_fields.iter().map(|(name, handler)| { quote!( ::core::option::Option::Some(msg) = self.#name.recv() => { self.#handler(msg).await @@ -355,10 +373,7 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { let path = attr.path(); if path.is_ident("input") { if found.is_some() { - return Err(syn::Error::new_spanned( - attr, - "field has multiple module attributes; only one of #[input], #[output], #[config], #[tf] is allowed", - )); + return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); } let mut decode: Option = None; let mut handler: Option = None; @@ -380,10 +395,7 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { found = Some(FieldKind::Input { decode, handler }); } else if path.is_ident("output") { if found.is_some() { - return Err(syn::Error::new_spanned( - attr, - "field has multiple module attributes; only one of #[input], #[output], #[config], #[tf] is allowed", - )); + return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); } let mut encode: Option = None; attr.parse_nested_meta(|meta| { @@ -400,20 +412,45 @@ fn classify_field(field: &Field, name: &Ident) -> syn::Result { syn::Error::new_spanned(attr, "#[output] requires `encode = ...`") })?; found = Some(FieldKind::Output { encode }); + } else if path.is_ident("io") { + if found.is_some() { + return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); + } + let mut decode: Option = None; + let mut encode: Option = None; + let mut handler: Option = None; + attr.parse_nested_meta(|meta| { + if meta.path.is_ident("decode") { + decode = Some(meta.value()?.parse()?); + } else if meta.path.is_ident("encode") { + encode = Some(meta.value()?.parse()?); + } else if meta.path.is_ident("handler") { + handler = Some(meta.value()?.parse()?); + } else { + return Err(meta.error( + "unrecognized #[io] argument; expected `decode = ...`, `encode = ...` or `handler = ...`", + )); + } + Ok(()) + })?; + let decode = decode + .ok_or_else(|| syn::Error::new_spanned(attr, "#[io] requires `decode = ...`"))?; + let encode = encode + .ok_or_else(|| syn::Error::new_spanned(attr, "#[io] requires `encode = ...`"))?; + let handler = handler.unwrap_or_else(|| format_ident!("handle_{}", name)); + found = Some(FieldKind::Io { + decode, + encode, + handler, + }); } else if path.is_ident("config") { if found.is_some() { - return Err(syn::Error::new_spanned( - attr, - "field has multiple module attributes; only one of #[input], #[output], #[config], #[tf] is allowed", - )); + return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); } found = Some(FieldKind::Config); } else if path.is_ident("tf") { if found.is_some() { - return Err(syn::Error::new_spanned( - attr, - "field has multiple module attributes; only one of #[input], #[output], #[config], #[tf] is allowed", - )); + return Err(syn::Error::new_spanned(attr, ONE_ATTR_ONLY)); } found = Some(FieldKind::Tf); } diff --git a/native/rust/dimos-module/src/lib.rs b/native/rust/dimos-module/src/lib.rs index b5efc805be..cd95cb182a 100644 --- a/native/rust/dimos-module/src/lib.rs +++ b/native/rust/dimos-module/src/lib.rs @@ -12,6 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. +// #[derive(Module)] emits ::dimos_module paths, so tests in this crate that +// derive Module need the crate to be nameable from inside itself. +#[cfg(test)] +extern crate self as dimos_module; + pub mod lcm; pub mod log; pub mod module; @@ -21,7 +26,7 @@ pub mod zenoh; pub use dimos_module_macros::{native_config, Module}; pub use lcm::LcmTransport; -pub use module::{run, Builder, Input, Module, ModuleConfig, NativeConfig, NoConfig, Output}; +pub use module::{run, Builder, Input, Io, Module, ModuleConfig, NativeConfig, NoConfig, Output}; pub use tf::{Tf, Transform}; pub use transport::Transport; pub use zenoh::ZenohTransport; diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index 399d8c19b9..59f29d4d0c 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -125,14 +125,41 @@ pub struct Output { impl Output { pub async fn publish(&self, msg: &T) -> io::Result<()> { - let data = (self.encode)(msg); - self.sender - .send(data) - .await - .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "background task gone")) + publish_encoded(&self.sender, (self.encode)(msg)).await } } +/// A port that publishes to and subscribes on the same topic. +/// +/// The transports deliver a message back to its own sender, so an `Io` port +/// sees the whole topic including its own publishes. +pub struct Io { + pub topic: String, + receiver: mpsc::Receiver, + encode: fn(&T) -> Vec, + sender: mpsc::Sender>, +} + +impl Io { + pub async fn recv(&mut self) -> Option { + self.receiver.recv().await + } + + pub async fn publish(&self, msg: &T) -> io::Result<()> { + publish_encoded(&self.sender, (self.encode)(msg)).await + } +} + +pub(crate) async fn publish_encoded( + sender: &mpsc::Sender>, + data: Vec, +) -> io::Result<()> { + sender + .send(data) + .await + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "background task gone")) +} + /// Extract `(topics, config)` from an already-parsed config object. `run` /// parses the line once and also reads `qos` from it, so this takes the value. fn parse_config_value( @@ -298,37 +325,68 @@ impl Builder { .unwrap_or_else(|| format!("/{port}")) } - pub fn input( + // Registers a decoding route on the topic and hands back the receiving end. + fn add_route( &mut self, - port: &str, + topic: &str, decode: fn(&[u8]) -> io::Result, - ) -> Input { - let topic = self.topic_for(port); + ) -> mpsc::Receiver { let (tx, rx) = mpsc::channel(INPUT_CHANNEL_CAPACITY); self.routes - .entry(topic.clone()) + .entry(topic.to_string()) .or_default() .push(Box::new(TypedRoute { - topic: topic.clone(), + topic: topic.to_string(), decode, sender: tx, drop_count: AtomicU64::new(0), last_log_ns: AtomicU64::new(0), })); - Input { - topic, - receiver: rx, - } + rx + } + + // Adds a publish queue for the topic and hands back the sending end. + fn add_publisher(&mut self, topic: &str) -> mpsc::Sender> { + let (tx, rx) = mpsc::channel(PUBLISH_CHANNEL_CAPACITY); + self.outputs.push((topic.to_string(), rx)); + tx + } + + pub fn input( + &mut self, + port: &str, + decode: fn(&[u8]) -> io::Result, + ) -> Input { + let topic = self.topic_for(port); + let receiver = self.add_route(&topic, decode); + Input { topic, receiver } } pub fn output(&mut self, port: &str, encode: fn(&T) -> Vec) -> Output { let topic = self.topic_for(port); - let (tx, rx) = mpsc::channel(PUBLISH_CHANNEL_CAPACITY); - self.outputs.push((topic.clone(), rx)); + let sender = self.add_publisher(&topic); Output { topic, encode, - sender: tx, + sender, + } + } + + /// A port that both subscribes and publishes on one topic. + pub fn io( + &mut self, + port: &str, + decode: fn(&[u8]) -> io::Result, + encode: fn(&T) -> Vec, + ) -> Io { + let topic = self.topic_for(port); + let receiver = self.add_route(&topic, decode); + let sender = self.add_publisher(&topic); + Io { + topic, + receiver, + encode, + sender, } } @@ -341,10 +399,9 @@ impl Builder { return tf.clone(); } let topic = self.topic_for("tf"); - let (tx, rx) = mpsc::channel(PUBLISH_CHANNEL_CAPACITY); - self.outputs.push((topic.clone(), rx)); + let sender = self.add_publisher(&topic); let (tf, route) = - crate::tf::tf_subscription(topic.clone(), crate::tf::DEFAULT_TF_WINDOW_SECS, tx); + crate::tf::tf_subscription(topic.clone(), crate::tf::DEFAULT_TF_WINDOW_SECS, sender); self.routes.entry(topic).or_default().push(route); self.tf = Some(tf.clone()); tf @@ -790,6 +847,47 @@ mod tests { assert_eq!(output.topic, "/robot/cmd_vel"); } + #[test] + fn io_uses_mapped_topic() { + let mut builder = builder_with_topics(&[("cmd", "/robot/cmd")]); + let io = builder.io("cmd", |b| Ok(b.to_vec()), |b: &Vec| b.clone()); + assert_eq!(io.topic, "/robot/cmd"); + } + + #[test] + fn io_registers_one_route_and_one_publisher_on_the_same_topic() { + let mut builder = builder_with_topics(&[("cmd", "/robot/cmd")]); + let _io = builder.io("cmd", |b| Ok(b.to_vec()), |b: &Vec| b.clone()); + assert_eq!(builder.routes.get("/robot/cmd").map(Vec::len), Some(1)); + assert_eq!(builder.outputs.len(), 1); + assert_eq!(builder.outputs[0].0, "/robot/cmd"); + } + + #[tokio::test] + async fn io_receives_on_its_route_and_publishes_to_its_queue() { + let mut builder = builder_with_topics(&[("cmd", "/robot/cmd")]); + let mut io = builder.io("cmd", |b| Ok(b.to_vec()), |b: &Vec| b.clone()); + + builder.routes["/robot/cmd"][0].try_dispatch(b"inbound"); + assert_eq!(io.recv().await.expect("inbound message"), b"inbound"); + + io.publish(&b"outbound".to_vec()).await.expect("publish"); + let (_, rx) = &mut builder.outputs[0]; + assert_eq!(rx.recv().await.expect("published bytes"), b"outbound"); + } + + #[tokio::test] + async fn io_publish_errors_when_the_publish_worker_is_gone() { + let mut builder = builder_with_topics(&[]); + let io = builder.io("cmd", |b| Ok(b.to_vec()), |b: &Vec| b.clone()); + builder.outputs.clear(); + let err = io + .publish(&b"x".to_vec()) + .await + .expect_err("publish should fail with no worker"); + assert_eq!(err.kind(), io::ErrorKind::BrokenPipe); + } + #[test] fn tf_uses_mapped_topic() { let mut builder = builder_with_topics(&[("tf", "/robot/tf")]); @@ -1025,4 +1123,52 @@ mod tests { assert_eq!(route.drop_count.load(Ordering::Relaxed), 1); assert!(logs_contain("handler was full")); } + + // Exercises the code #[derive(Module)] generates for an #[io] field. + mod derive_io { + use super::*; + use crate::Io; + + struct Msg(Vec); + + fn decode(bytes: &[u8]) -> io::Result { + Ok(Msg(bytes.to_vec())) + } + + fn encode(msg: &Msg) -> Vec { + msg.0.clone() + } + + #[derive(crate::Module)] + struct Echo { + #[io(decode = decode, encode = encode)] + cmd: Io, + } + + impl Echo { + async fn handle_cmd(&mut self, msg: Msg) { + if msg.0 == b"ping" { + self.cmd + .publish(&Msg(b"pong".to_vec())) + .await + .expect("publish"); + } + } + } + + #[tokio::test] + async fn io_field_is_wired_to_its_handler_and_can_publish() { + let mut builder = Builder::new(topics(&[("cmd", "/robot/cmd")])); + let mut echo = Echo::build(&mut builder, NoConfig); + + builder.routes["/robot/cmd"][0].try_dispatch(b"ping"); + // Dropping the routes closes the sender, so handle() drains and returns. + builder.routes.clear(); + echo.handle().await; + + let (topic, rx) = &mut builder.outputs[0]; + assert_eq!(topic, "/robot/cmd"); + assert_eq!(rx.recv().await.expect("handler reply"), b"pong"); + } + } } diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index 64d7d4e8f5..0ee535b77a 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -331,10 +331,7 @@ impl Tf { let msg = lcm_msgs::tf2_msgs::TFMessage { transforms: transforms.iter().map(to_stamped).collect(), }; - self.sender - .send(msg.encode()) - .await - .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "background task gone")) + crate::module::publish_encoded(&self.sender, msg.encode()).await } } From 15e52565df17a28137a993584e2affce4a2692d5 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Fri, 31 Jul 2026 11:17:50 -0700 Subject: [PATCH 18/26] Rewrite the semantics a bit --- native/rust/README.md | 10 +++- native/rust/dimos-module/src/lib.rs | 2 +- native/rust/dimos-module/src/tf.rs | 88 +++++++++++++++++++++-------- 3 files changed, 72 insertions(+), 28 deletions(-) diff --git a/native/rust/README.md b/native/rust/README.md index e4e01b676e..5f8d98e105 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -131,9 +131,15 @@ impl VoxelMap { } ``` -`Tf` is a cheap-to-clone handle; the graph fills in the background as `tf` messages arrive. `get(parent, child, time, tolerance)` selects the sample nearest `time` (latest when `None`) and returns `None` when no path connects the frames or no sample falls within `tolerance` seconds. `get_latest` is the no-time shorthand. The result exposes its `nalgebra` parts via `translation()` (a `Vector3`) and `rotation()` (a `UnitQuaternion`). Lookups are nearest-in-time, not interpolated. +`Tf` is a cheap-to-clone handle; the graph fills in the background as `tf` messages arrive. `get_latest(parent, child)` is the common case. For a query against a particular stamp, `lookup(parent, child)` starts one that `.at(time)` points at the sample nearest that stamp and `.tolerance(secs)` bounds how far that sample may sit from it, finished with `.get()`: -`publish` sends transforms onto the same `tf` topic, the counterpart to Python's `tf.publish()`. Published transforms also feed the module's own graph, so a `get` right after the publish sees them. Build the isometry from `dimos_module::nalgebra`, re-exported so the version matches the SDK's types: +```rust +let at_scan = self.tf.lookup("map", "base_link").at(scan_ts).tolerance(0.1).get(); +``` + +Either way the result is `None` when no path connects the frames or no sample falls within the tolerance. It exposes its `nalgebra` parts via `translation()` (a `Vector3`) and `rotation()` (a `UnitQuaternion`). Lookups are nearest-in-time, not interpolated. + +`publish` sends transforms onto the same `tf` topic, the counterpart to Python's `tf.publish()`. Published transforms also feed the module's own graph, so a lookup right after the publish sees them. Build the isometry from `dimos_module::nalgebra`, re-exported so the version matches the SDK's types: ```rust use dimos_module::nalgebra::Isometry3; diff --git a/native/rust/dimos-module/src/lib.rs b/native/rust/dimos-module/src/lib.rs index cd95cb182a..8a97eddfe2 100644 --- a/native/rust/dimos-module/src/lib.rs +++ b/native/rust/dimos-module/src/lib.rs @@ -27,7 +27,7 @@ pub mod zenoh; pub use dimos_module_macros::{native_config, Module}; pub use lcm::LcmTransport; pub use module::{run, Builder, Input, Io, Module, ModuleConfig, NativeConfig, NoConfig, Output}; -pub use tf::{Tf, Transform}; +pub use tf::{Lookup, Tf, Transform}; pub use transport::Transport; pub use zenoh::ZenohTransport; diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index 0ee535b77a..7b40290034 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -14,7 +14,7 @@ //! Transform client for native modules. //! -//! Each `/tf` edge is buffered per `(parent, child)`, and [`Tf::get`] composes +//! Each `/tf` edge is buffered per `(parent, child)`, and [`Tf::lookup`] composes //! the shortest path through the frame graph. Lookups are nearest-in-time within //! a tolerance, not interpolated. [`Tf::publish`] sends transforms onto the same //! topic and feeds the local graph. @@ -294,32 +294,34 @@ pub struct Tf { } impl Tf { - /// The transform from `parent` to `child`. + /// Start a lookup of the transform from `parent` to `child`. /// - /// `time` selects the sample nearest that stamp (latest sample when `None`), - /// and `tolerance` bounds how far that sample may be in seconds. Returns - /// `None` when no path connects the frames or no sample is within tolerance. - pub fn get( - &self, - parent: &str, - child: &str, - time: Option, - tolerance: Option, - ) -> Option { - self.buffer - .read() - .expect("tf buffer lock poisoned") - .get(parent, child, time, tolerance) + /// Refine it with [`Lookup::at`] and [`Lookup::tolerance`], then finish with + /// [`Lookup::get`]. Use [`Tf::get_latest`] when no refinement is needed. + /// + /// ```ignore + /// let at_scan = tf.lookup("map", "base_link").at(scan_ts).tolerance(0.1).get(); + /// ``` + pub fn lookup<'a>(&'a self, parent: &'a str, child: &'a str) -> Lookup<'a> { + Lookup { + tf: self, + parent, + child, + time: None, + tolerance: None, + } } - /// The latest available transform from `parent` to `child`. + /// The latest transform from `parent` to `child`. + /// + /// Shorthand for `lookup(parent, child).get()`. pub fn get_latest(&self, parent: &str, child: &str) -> Option { - self.get(parent, child, None, None) + self.lookup(parent, child).get() } /// Publish transforms on the `tf` topic. /// - /// The transforms also feed the local graph, so a `get` right after sees + /// The transforms also feed the local graph, so a lookup right after sees /// them without waiting for the transport round trip. pub async fn publish(&self, transforms: &[Transform]) -> io::Result<()> { { @@ -335,6 +337,42 @@ impl Tf { } } +/// A transform lookup being built. Created by [`Tf::lookup`]. +pub struct Lookup<'a> { + tf: &'a Tf, + parent: &'a str, + child: &'a str, + time: Option, + tolerance: Option, +} + +impl Lookup<'_> { + /// Take the sample nearest `time` rather than the latest one. + pub fn at(mut self, time: f64) -> Self { + self.time = Some(time); + self + } + + /// Bound how far, in seconds, the chosen sample may sit from [`Lookup::at`]. + pub fn tolerance(mut self, tolerance: f64) -> Self { + self.tolerance = Some(tolerance); + self + } + + /// Resolve the lookup against the transforms buffered so far. + /// + /// `None` when no path connects the frames, or when the nearest sample is + /// outside the tolerance. + pub fn get(self) -> Option { + self.tf.buffer.read().expect("tf buffer lock poisoned").get( + self.parent, + self.child, + self.time, + self.tolerance, + ) + } +} + fn to_stamped(t: &Transform) -> lcm_msgs::geometry_msgs::TransformStamped { let mut sec = t.ts.floor(); let mut nsec = ((t.ts - sec) * 1e9).round(); @@ -534,7 +572,7 @@ mod tests { t.ts ); // Explicit query time is echoed back. - let at = tf.get("base_link", "base_link", Some(42.0), None).unwrap(); + let at = tf.lookup("base_link", "base_link").at(42.0).get().unwrap(); assert!((at.ts - 42.0).abs() < 1e-9); } @@ -543,9 +581,9 @@ mod tests { let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); h.add("a", "b", 10.0, (1.0, 0.0, 0.0), 0.0); h.add("a", "b", 20.0, (2.0, 0.0, 0.0), 0.0); - let near_10 = tf.get("a", "b", Some(11.0), None).unwrap(); + let near_10 = tf.lookup("a", "b").at(11.0).get().unwrap(); assert!((near_10.translation().x - 1.0).abs() < 1e-9); - let near_20 = tf.get("a", "b", Some(18.0), None).unwrap(); + let near_20 = tf.lookup("a", "b").at(18.0).get().unwrap(); assert!((near_20.translation().x - 2.0).abs() < 1e-9); } @@ -553,8 +591,8 @@ mod tests { fn time_query_outside_tolerance_returns_none() { let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); h.add("a", "b", 10.0, (1.0, 0.0, 0.0), 0.0); - assert!(tf.get("a", "b", Some(50.0), Some(1.0)).is_none()); - assert!(tf.get("a", "b", Some(10.5), Some(1.0)).is_some()); + assert!(tf.lookup("a", "b").at(50.0).tolerance(1.0).get().is_none()); + assert!(tf.lookup("a", "b").at(10.5).tolerance(1.0).get().is_some()); } #[test] @@ -686,7 +724,7 @@ mod tests { let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); h.add("a", "b", 10.0, (1.0, 0.0, 0.0), 0.0); h.add("a", "b", 12.0, (2.0, 0.0, 0.0), 0.0); - let t = tf.get("a", "b", Some(11.0), None).unwrap(); + let t = tf.lookup("a", "b").at(11.0).get().unwrap(); assert!((t.translation().x - 2.0).abs() < 1e-9); } From 09f998d623fdd7decf1a7ab8099644ec000150cb Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Fri, 31 Jul 2026 11:40:18 -0700 Subject: [PATCH 19/26] Explicitly subscribe to tf io --- dimos/core/native_module.py | 24 +++++++++++++----------- dimos/core/test_native_module.py | 9 +++++++++ examples/native-modules/rust_tf.py | 2 ++ 3 files changed, 24 insertions(+), 11 deletions(-) diff --git a/dimos/core/native_module.py b/dimos/core/native_module.py index 5aca07a363..10bf4395a4 100644 --- a/dimos/core/native_module.py +++ b/dimos/core/native_module.py @@ -12,12 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""NativeModule: blueprint-integrated wrapper for native (C/C++) executables. +"""NativeModule: blueprint-integrated wrapper for native executables. -A NativeModule is a thin Python Module subclass that declares In/Out ports +A NativeModule is a thin Python Module subclass that declares In/Out/IO ports for blueprint wiring but delegates all real work to a managed subprocess. -The native process receives its LCM topic names via CLI args and does -pub/sub directly on the LCM multicast bus. +The native process receives its topic names via CLI args, or as a JSON line on +stdin when ``stdin_config`` is set, and does pub/sub on them directly. Example usage:: @@ -60,7 +60,6 @@ class MyCppModule(NativeModule): from dimos.core.core import rpc from dimos.core.global_config import global_config from dimos.core.module import Module, ModuleConfig -from dimos.core.transport_factory import tf_channel from dimos.utils.logging_config import setup_logger if sys.platform.startswith("linux"): @@ -179,15 +178,19 @@ class NativeModule(Module): """ Module that wraps a native executable as a managed subprocess. - Subclass this, declare In/Out ports, and annotate ``config`` with a + Subclass this, declare In/Out/IO ports, and annotate ``config`` with a :class:`NativeModuleConfig` subclass pointing at the executable. On ``start()``, the binary is launched with CLI args:: - -- ... + -- ... -- ... - The native process should parse these args and pub/sub on the given - LCM topics directly. On ``stop()``, the process receives SIGTERM. + Each topic is the wire channel for that port on the transport named by the + ``DIMOS_TRANSPORT`` env var. With ``stdin_config``, those same topics plus + the config and any publisher QoS also arrive as one JSON line on stdin. + + The native process should parse whichever it uses and pub/sub on the given + topics directly. On ``stop()``, the process receives SIGTERM. """ config: NativeModuleConfig @@ -241,8 +244,7 @@ def start(self) -> None: stdin_blob: bytes | None = None if self.config.stdin_config: config_dict = self.config.to_config_dict() - stdin_topics = {**topics, "tf": tf_channel()} - blob: dict[str, Any] = {"topics": stdin_topics, "config": config_dict or None} + blob: dict[str, Any] = {"topics": topics, "config": config_dict or None} qos = self._collect_output_qos() if qos: blob["qos"] = qos diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index 49ca8e8729..96e6c2cad1 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -188,6 +188,15 @@ def test_io_port_topic_reaches_the_native_process() -> None: transport.stop() +def test_tf_topic_comes_from_the_declared_port_only() -> None: + """No tf port declared means no tf topic, rather than a silently injected one.""" + module = StubNativeModule(executable=_ECHO) + try: + assert "tf" not in module._collect_topics() + finally: + module.stop() + + def test_io_port_publisher_qos_reaches_the_native_process() -> None: module = StubIoModule(executable=_ECHO) transport = ZenohTransport(ZenohTopic("/tf", TFMessage, qos=QOS_NEVER_DROP)) diff --git a/examples/native-modules/rust_tf.py b/examples/native-modules/rust_tf.py index 0a4c2d5179..362f59d9c5 100644 --- a/examples/native-modules/rust_tf.py +++ b/examples/native-modules/rust_tf.py @@ -97,6 +97,7 @@ class TfListenerModule(NativeModule): """ config: TfListenerConfig + tf: IO[TFMessage] class TfBroadcasterConfig(NativeModuleConfig): @@ -110,6 +111,7 @@ class TfBroadcasterModule(NativeModule): """Rust module that publishes the c -> d transform.""" config: TfBroadcasterConfig + tf: IO[TFMessage] if __name__ == "__main__": From 7d1f11cd528e8e368b5aa712819a9aa6666da36a Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Fri, 31 Jul 2026 11:55:45 -0700 Subject: [PATCH 20/26] Enforce one to one topic naming --- native/rust/dimos-module/src/module.rs | 75 ++++++++++++++++++++++++-- 1 file changed, 72 insertions(+), 3 deletions(-) diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index 59f29d4d0c..59fb20e58c 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -302,6 +302,8 @@ pub trait Module: Sized + Send + 'static { pub struct Builder { topics: HashMap, + // Every port the module asked for a topic, matched against `topics` after build. + requested: BTreeSet, routes: HashMap>>, // One publish queue per output channel, drained by its own worker. outputs: Vec<(String, mpsc::Receiver>)>, @@ -312,19 +314,41 @@ impl Builder { pub(crate) fn new(topics: HashMap) -> Self { Self { topics, + requested: BTreeSet::new(), routes: HashMap::new(), outputs: Vec::new(), tf: None, } } - fn topic_for(&self, port: &str) -> String { + fn topic_for(&mut self, port: &str) -> String { + self.requested.insert(port.to_string()); self.topics .get(port) .cloned() .unwrap_or_else(|| format!("/{port}")) } + // The coordinator sends one topic per declared Python port, so the ports + // this module claimed must be exactly that set. A port the module never + // asks for is dead wiring, and one the coordinator never sent would fall + // back to a bare `/{port}` that nothing else on the bus is using. + pub(crate) fn enforce_topics_match_ports(&self) -> io::Result<()> { + let provided: BTreeSet<&String> = self.topics.keys().collect(); + let requested: BTreeSet<&String> = self.requested.iter().collect(); + if provided == requested { + return Ok(()); + } + let missing: Vec<&&String> = requested.difference(&provided).collect(); + let unexpected: Vec<&&String> = provided.difference(&requested).collect(); + Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "topics do not match module ports: missing {missing:?}, unexpected {unexpected:?}" + ), + )) + } + // Registers a decoding route on the topic and hands back the receiving end. fn add_route( &mut self, @@ -500,6 +524,7 @@ where let mut builder = Builder::new(topics); let mut module = M::build(&mut builder, config); + builder.enforce_topics_match_ports()?; subscribe_routes(&transport, builder.routes).await?; // Kept alive until teardown so the subscriptions stay live. @@ -816,13 +841,13 @@ mod tests { #[test] fn unmapped_port_falls_back_to_slash_port() { - let builder = builder_with_topics(&[]); + let mut builder = builder_with_topics(&[]); assert_eq!(builder.topic_for("cmd_vel"), "/cmd_vel"); } #[test] fn mapped_port_uses_given_topic() { - let builder = builder_with_topics(&[("cmd_vel", "/robot/cmd_vel")]); + let mut builder = builder_with_topics(&[("cmd_vel", "/robot/cmd_vel")]); assert_eq!(builder.topic_for("cmd_vel"), "/robot/cmd_vel"); } @@ -847,6 +872,50 @@ mod tests { assert_eq!(output.topic, "/robot/cmd_vel"); } + #[test] + fn topics_matching_ports_exactly_pass() { + let mut builder = builder_with_topics(&[("cmd", "/robot/cmd"), ("odom", "/robot/odom")]); + builder.input("cmd", |b| Ok(b.to_vec())); + builder.output("odom", |b: &Vec| b.clone()); + builder.enforce_topics_match_ports().expect("exact match"); + } + + #[test] + fn a_port_the_coordinator_never_sent_is_rejected() { + let mut builder = builder_with_topics(&[("cmd", "/robot/cmd")]); + builder.input("cmd", |b| Ok(b.to_vec())); + builder.output("odom", |b: &Vec| b.clone()); + let err = builder + .enforce_topics_match_ports() + .expect_err("odom has no topic"); + assert!(err.to_string().contains("missing [\"odom\"]"), "{err}"); + } + + #[test] + fn a_topic_no_port_claimed_is_rejected() { + let mut builder = builder_with_topics(&[("cmd", "/robot/cmd"), ("stale", "/robot/stale")]); + builder.input("cmd", |b| Ok(b.to_vec())); + let err = builder + .enforce_topics_match_ports() + .expect_err("stale is unclaimed"); + assert!(err.to_string().contains("unexpected [\"stale\"]"), "{err}"); + } + + #[test] + fn a_tf_field_claims_the_tf_topic() { + let mut builder = builder_with_topics(&[("tf", "/tf#tf2_msgs.TFMessage")]); + builder.tf(); + builder.enforce_topics_match_ports().expect("tf claimed"); + } + + #[test] + fn a_module_with_no_ports_and_no_topics_passes() { + let builder = builder_with_topics(&[]); + builder + .enforce_topics_match_ports() + .expect("nothing to match"); + } + #[test] fn io_uses_mapped_topic() { let mut builder = builder_with_topics(&[("cmd", "/robot/cmd")]); From 5ed06483ecfca8114f80843229620696fa61d037 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Fri, 31 Jul 2026 12:01:34 -0700 Subject: [PATCH 21/26] Bound look up time by buffer size --- native/rust/dimos-module/src/tf.rs | 39 +++++++++++++++++++++++++++++- 1 file changed, 38 insertions(+), 1 deletion(-) diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index 7b40290034..3a51f12a25 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -162,7 +162,7 @@ impl TBuffer { ) -> Option { let s = match time { None => self.last()?, - Some(t) => self.find_closest(t, tolerance)?, + Some(t) => self.find_closest(t, Some(tolerance.unwrap_or(self.window_secs)))?, }; Some(Transform { parent: parent.to_string(), @@ -595,6 +595,43 @@ mod tests { assert!(tf.lookup("a", "b").at(10.5).tolerance(1.0).get().is_some()); } + #[test] + fn time_query_beyond_the_window_returns_none_without_a_tolerance() { + let (tf, h) = tf_with(10.0); + h.add("a", "b", 100.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.lookup("a", "b").at(50.0).get().is_none()); + } + + #[test] + fn time_query_inside_the_window_resolves_without_a_tolerance() { + let (tf, h) = tf_with(10.0); + h.add("a", "b", 100.0, (1.0, 0.0, 0.0), 0.0); + let t = tf + .lookup("a", "b") + .at(95.0) + .get() + .expect("within the window"); + assert!((t.translation().x - 1.0).abs() < 1e-9); + } + + // An explicit tolerance is the caller opting into staleness, so it widens + // past the window rather than being clamped by it. + #[test] + fn an_explicit_tolerance_reaches_past_the_window() { + let (tf, h) = tf_with(10.0); + h.add("a", "b", 100.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.lookup("a", "b").at(50.0).tolerance(60.0).get().is_some()); + } + + // The window bounds queries against a stamp, not the latest sample: with no + // `at`, the newest edge is returned however old it is. + #[test] + fn latest_is_not_bounded_by_the_window() { + let (tf, h) = tf_with(10.0); + h.add("a", "b", 100.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.get_latest("a", "b").is_some()); + } + #[test] fn prunes_samples_outside_window() { let mut buf = TBuffer::new(5.0); From 8947d701a64f89b4c202324966b11753d1306d92 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Fri, 31 Jul 2026 12:48:42 -0700 Subject: [PATCH 22/26] Support within, at, tolerance --- native/rust/README.md | 14 ++ native/rust/dimos-module/src/tf.rs | 266 +++++++++++++++++++++++++---- 2 files changed, 249 insertions(+), 31 deletions(-) diff --git a/native/rust/README.md b/native/rust/README.md index 5f8d98e105..df161fe78e 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -137,6 +137,20 @@ impl VoxelMap { let at_scan = self.tf.lookup("map", "base_link").at(scan_ts).tolerance(0.1).get(); ``` +A message and the transform it needs arrive on separate topics, so the transform for a given stamp is often merely late. `.within(duration)` replaces `.get()` to wait for one, returning as soon as the lookup succeeds or `None` at the deadline: + +```rust +let at_scan = self.tf.lookup("odom", &cloud.header.frame_id) + .at(scan_ts) + .tolerance(0.02) + .within(Duration::from_millis(200)) + .await; +``` + +`.tolerance()` and `.within()` are different clocks. Tolerance bounds how far the chosen sample may sit from `.at()` in message stamps — accuracy. `.within()` bounds how long to wait in wall time — patience. Always set a tolerance when waiting, or the lookup is satisfied by anything inside the buffer window and returns a stale transform immediately. + +`.within()` suspends the caller, and awaiting it inside a `handle_*` method parks that module's whole dispatch loop, so every other topic it subscribes to stops being served until it returns. Prefer `.get()` there; move waiting onto its own task when the wait may be long. + Either way the result is `None` when no path connects the frames or no sample falls within the tolerance. It exposes its `nalgebra` parts via `translation()` (a `Vector3`) and `rotation()` (a `UnitQuaternion`). Lookups are nearest-in-time, not interpolated. `publish` sends transforms onto the same `tf` topic, the counterpart to Python's `tf.publish()`. Published transforms also feed the module's own graph, so a lookup right after the publish sees them. Build the isometry from `dimos_module::nalgebra`, re-exported so the version matches the SDK's types: diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index 3a51f12a25..e0ac261339 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -25,7 +25,7 @@ use std::sync::{Arc, RwLock}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use nalgebra::{Isometry3, Quaternion, Translation3, UnitQuaternion, Vector3}; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, Notify}; use crate::module::Route; @@ -283,13 +283,47 @@ impl MultiTBuffer { } } +// The graph plus the signal that it changed. Writes go through `update` so a +// writer cannot leave a waiter asleep on a transform that has already landed. +struct Graph { + buffer: RwLock, + changed: Notify, +} + +impl Graph { + fn new(window_secs: f64) -> Self { + Self { + buffer: RwLock::new(MultiTBuffer::new(window_secs)), + changed: Notify::new(), + } + } + + fn update(&self, edits: impl FnOnce(&mut MultiTBuffer)) { + edits(&mut self.buffer.write().expect("tf buffer lock poisoned")); + self.changed.notify_waiters(); + } + + fn get( + &self, + parent: &str, + child: &str, + time: Option, + tolerance: Option, + ) -> Option { + self.buffer + .read() + .expect("tf buffer lock poisoned") + .get(parent, child, time, tolerance) + } +} + /// A cheap-to-clone handle for querying and publishing transforms. /// /// Obtain one from `Builder::tf` (or a `#[tf]` field on a `#[derive(Module)]` /// struct). The graph is filled in the background as `/tf` messages arrive. #[derive(Clone)] pub struct Tf { - buffer: Arc>, + graph: Arc, sender: mpsc::Sender>, } @@ -324,12 +358,11 @@ impl Tf { /// The transforms also feed the local graph, so a lookup right after sees /// them without waiting for the transport round trip. pub async fn publish(&self, transforms: &[Transform]) -> io::Result<()> { - { - let mut buffer = self.buffer.write().expect("tf buffer lock poisoned"); + self.graph.update(|buffer| { for t in transforms { buffer.receive(&t.parent, &t.child, t.ts, t.iso); } - } + }); let msg = lcm_msgs::tf2_msgs::TFMessage { transforms: transforms.iter().map(to_stamped).collect(), }; @@ -359,17 +392,57 @@ impl Lookup<'_> { self } + fn resolve(&self) -> Option { + self.tf + .graph + .get(self.parent, self.child, self.time, self.tolerance) + } + /// Resolve the lookup against the transforms buffered so far. /// /// `None` when no path connects the frames, or when the nearest sample is /// outside the tolerance. pub fn get(self) -> Option { - self.tf.buffer.read().expect("tf buffer lock poisoned").get( - self.parent, - self.child, - self.time, - self.tolerance, - ) + self.resolve() + } + + /// Resolve the lookup, waiting up to `timeout` for it to become possible. + /// + /// A message and the transform it needs arrive on separate topics, so the + /// transform for a given stamp is often merely late rather than absent. + /// Returns as soon as the lookup succeeds, or `None` at the deadline. + /// + /// [`Lookup::tolerance`] and this are different clocks: tolerance bounds how + /// far the chosen sample may sit from [`Lookup::at`] in message stamps, + /// `timeout` bounds how long to wait in wall time. Set a tolerance whenever + /// waiting, or the lookup is satisfied by anything inside the buffer window + /// and returns a stale transform immediately. + /// + /// This suspends the caller. Awaiting it inside a `handle_*` method parks + /// that module's whole dispatch loop, so every other topic it subscribes to + /// stops being served for the duration. Prefer [`Lookup::get`] there, and + /// move waiting onto a task of its own when the wait may be long. + pub async fn within(self, timeout: Duration) -> Option { + let deadline = tokio::time::Instant::now() + timeout; + loop { + // Registered before the resolve below, so a transform landing between + // the two still wakes this waiter instead of it sleeping out the + // whole timeout. + let changed = self.tf.graph.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + + if let Some(transform) = self.resolve() { + return Some(transform); + } + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return None; + } + if tokio::time::timeout(remaining, changed).await.is_err() { + return None; + } + } } } @@ -412,7 +485,7 @@ fn to_stamped(t: &Transform) -> lcm_msgs::geometry_msgs::TransformStamped { // module's existing recv loop dispatches tf traffic to it. struct TfRoute { topic: String, - buffer: Arc>, + graph: Arc, } impl Route for TfRoute { @@ -429,17 +502,18 @@ impl Route for TfRoute { return; } }; - let mut buffer = self.buffer.write().expect("tf buffer lock poisoned"); - for st in &msg.transforms { - let t = &st.transform.translation; - let q = &st.transform.rotation; - let iso = Isometry3::from_parts( - Translation3::new(t.x, t.y, t.z), - UnitQuaternion::from_quaternion(Quaternion::new(q.w, q.x, q.y, q.z)), - ); - let ts = st.header.stamp.sec as f64 + st.header.stamp.nsec as f64 * 1e-9; - buffer.receive(&st.header.frame_id, &st.child_frame_id, ts, iso); - } + self.graph.update(|buffer| { + for st in &msg.transforms { + let t = &st.transform.translation; + let q = &st.transform.rotation; + let iso = Isometry3::from_parts( + Translation3::new(t.x, t.y, t.z), + UnitQuaternion::from_quaternion(Quaternion::new(q.w, q.x, q.y, q.z)), + ); + let ts = st.header.stamp.sec as f64 + st.header.stamp.nsec as f64 * 1e-9; + buffer.receive(&st.header.frame_id, &st.child_frame_id, ts, iso); + } + }); } } @@ -450,12 +524,12 @@ pub(crate) fn tf_subscription( buffer_size: f64, sender: mpsc::Sender>, ) -> (Tf, Box) { - let buffer = Arc::new(RwLock::new(MultiTBuffer::new(buffer_size))); + let graph = Arc::new(Graph::new(buffer_size)); let tf = Tf { - buffer: Arc::clone(&buffer), + graph: Arc::clone(&graph), sender, }; - let route = Box::new(TfRoute { topic, buffer }); + let route = Box::new(TfRoute { topic, graph }); (tf, route) } @@ -470,21 +544,21 @@ mod tests { } fn tf_with_publish(buffer_size: f64) -> (Tf, mpsc::Receiver>, MultiHandle) { - let buffer = Arc::new(RwLock::new(MultiTBuffer::new(buffer_size))); + let graph = Arc::new(Graph::new(buffer_size)); let (tx, rx) = mpsc::channel(8); ( Tf { - buffer: Arc::clone(&buffer), + graph: Arc::clone(&graph), sender: tx, }, rx, - MultiHandle { buffer }, + MultiHandle { graph }, ) } // Test-only writer that bypasses LCM and pushes edges straight into the graph. struct MultiHandle { - buffer: Arc>, + graph: Arc, } impl MultiHandle { @@ -493,7 +567,8 @@ mod tests { Translation3::new(xyz.0, xyz.1, xyz.2), UnitQuaternion::from_euler_angles(0.0, 0.0, yaw), ); - self.buffer.write().unwrap().receive(parent, child, ts, iso); + self.graph + .update(|buffer| buffer.receive(parent, child, ts, iso)); } } @@ -632,6 +707,135 @@ mod tests { assert!(tf.get_latest("a", "b").is_some()); } + #[tokio::test] + async fn within_returns_without_waiting_when_already_buffered() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("a", "b", 5.0, (1.0, 0.0, 0.0), 0.0); + let t = tf + .lookup("a", "b") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_secs(30)) + .await + .expect("already available"); + assert!((t.translation().x - 1.0).abs() < 1e-9); + } + + // Returns as soon as the transform lands, not at the deadline. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn within_resolves_when_the_transform_arrives_late() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(30)).await; + h.add("a", "b", 5.0, (1.0, 0.0, 0.0), 0.0); + }); + let started = tokio::time::Instant::now(); + let t = tf + .lookup("a", "b") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_secs(30)) + .await + .expect("arrived inside the budget"); + assert!( + started.elapsed() < Duration::from_secs(5), + "waited {:?}, should have returned on arrival", + started.elapsed() + ); + assert!((t.translation().x - 1.0).abs() < 1e-9); + } + + #[tokio::test] + async fn within_times_out_when_nothing_arrives() { + let (tf, _h) = tf_with(DEFAULT_TF_WINDOW_SECS); + let started = tokio::time::Instant::now(); + let t = tf + .lookup("a", "b") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_millis(50)) + .await; + assert!(t.is_none()); + assert!( + started.elapsed() >= Duration::from_millis(50), + "returned early" + ); + } + + // The waiter is on a -> c, but what lands is b -> c. Waking only on the + // queried edge would sleep through the composition becoming possible. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn within_wakes_when_a_later_edge_completes_the_chain() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("a", "b", 5.0, (1.0, 0.0, 0.0), 0.0); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(30)).await; + h.add("b", "c", 5.0, (2.0, 0.0, 0.0), 0.0); + }); + let t = tf + .lookup("a", "c") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_secs(30)) + .await + .expect("chain completed inside the budget"); + assert!( + (t.translation().x - 3.0).abs() < 1e-9, + "{}", + t.translation().x + ); + } + + // A waiting handler is woken by the transport's dispatch task, which its own + // stall cannot block. Guards against the wait deadlocking against tf intake. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn within_is_woken_by_a_transform_dispatched_through_the_route() { + let (tx, _rx) = mpsc::channel(8); + let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_WINDOW_SECS, tx); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(30)).await; + route.try_dispatch(&stamped_message("a", "b", 5.0, 1.0)); + }); + let t = tf + .lookup("a", "b") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_secs(30)) + .await + .expect("route dispatch woke the waiter"); + assert!((t.translation().x - 1.0).abs() < 1e-9); + } + + fn stamped_message(parent: &str, child: &str, ts: f64, x: f64) -> Vec { + use lcm_msgs::geometry_msgs::{ + Quaternion as LQuat, Transform as LTransform, TransformStamped, Vector3 as LVec3, + }; + use lcm_msgs::std_msgs::{Header, Time}; + lcm_msgs::tf2_msgs::TFMessage { + transforms: vec![TransformStamped { + header: Header { + seq: 0, + stamp: Time { + sec: ts as i32, + nsec: 0, + }, + frame_id: parent.to_string(), + }, + child_frame_id: child.to_string(), + transform: LTransform { + translation: LVec3 { x, y: 0.0, z: 0.0 }, + rotation: LQuat { + x: 0.0, + y: 0.0, + z: 0.0, + w: 1.0, + }, + }, + }], + } + .encode() + } + #[test] fn prunes_samples_outside_window() { let mut buf = TBuffer::new(5.0); From 120ff33b2351ccdb1e01c9da876beee99fac9758 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Fri, 31 Jul 2026 13:02:09 -0700 Subject: [PATCH 23/26] Warn on unresolved transforms --- native/rust/dimos-module/src/tf.rs | 70 +++++++++++++++++++++++++++++- 1 file changed, 69 insertions(+), 1 deletion(-) diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index e0ac261339..02aaa94d3f 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -26,6 +26,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use nalgebra::{Isometry3, Quaternion, Translation3, UnitQuaternion, Vector3}; use tokio::sync::{mpsc, Notify}; +use tracing::warn; use crate::module::Route; @@ -398,12 +399,28 @@ impl Lookup<'_> { .get(self.parent, self.child, self.time, self.tolerance) } + // A lookup that resolves to nothing is otherwise invisible: the caller sees + // None and the buffer says nothing about which frames or stamp missed. + fn warn_unresolved(&self) { + warn!( + parent = %self.parent, + child = %self.child, + at = self.time.unwrap_or_else(now_secs), + tolerance = self.tolerance.unwrap_or(f64::NAN), + "No transform found between frames", + ); + } + /// Resolve the lookup against the transforms buffered so far. /// /// `None` when no path connects the frames, or when the nearest sample is /// outside the tolerance. pub fn get(self) -> Option { - self.resolve() + let found = self.resolve(); + if found.is_none() { + self.warn_unresolved(); + } + found } /// Resolve the lookup, waiting up to `timeout` for it to become possible. @@ -435,11 +452,15 @@ impl Lookup<'_> { if let Some(transform) = self.resolve() { return Some(transform); } + // Only the deadline warns. An intermediate miss is the normal state + // of a wait, not a failure. let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); if remaining.is_zero() { + self.warn_unresolved(); return None; } if tokio::time::timeout(remaining, changed).await.is_err() { + self.warn_unresolved(); return None; } } @@ -836,6 +857,53 @@ mod tests { .encode() } + #[test] + #[tracing_test::traced_test] + fn a_lookup_that_finds_nothing_warns_with_the_frames() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("world", "robot", 1.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.get_latest("world", "gripper").is_none()); + assert!(logs_contain("No transform found between frames")); + assert!(logs_contain("gripper")); + } + + #[test] + #[tracing_test::traced_test] + fn a_resolved_lookup_stays_quiet() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("world", "robot", 1.0, (1.0, 0.0, 0.0), 0.0); + assert!(tf.get_latest("world", "robot").is_some()); + assert!(!logs_contain("No transform found between frames")); + } + + // A wait in progress is not a failure, so only the deadline warns. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + #[tracing_test::traced_test] + async fn within_warns_only_once_it_gives_up() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(20)).await; + h.add("a", "b", 5.0, (1.0, 0.0, 0.0), 0.0); + }); + assert!(tf + .lookup("a", "b") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_secs(30)) + .await + .is_some()); + assert!(!logs_contain("No transform found between frames")); + + assert!(tf + .lookup("a", "missing") + .at(5.0) + .tolerance(0.1) + .within(Duration::from_millis(20)) + .await + .is_none()); + assert!(logs_contain("No transform found between frames")); + } + #[test] fn prunes_samples_outside_window() { let mut buf = TBuffer::new(5.0); From ae98f4bb9e052ffc986f5f17416d827e93194eaf Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Fri, 31 Jul 2026 14:22:22 -0700 Subject: [PATCH 24/26] Clean up --- dimos/core/test_native_module.py | 7 +- dimos/core/test_transport_factory.py | 9 -- dimos/core/transport_factory.py | 10 -- .../nav_3d/mls_planner/goal_relay.py | 3 - .../nav_3d/mls_planner/mls_planner_native.py | 5 +- .../mls_planner/rust/src/mls_planner.rs | 5 +- dimos/navigation/tf_pose.py | 13 +- .../go2/go2_mid360_static_transforms.py | 6 +- native/rust/dimos-module/src/module.rs | 10 +- native/rust/dimos-module/src/tf.rs | 129 +++++++++++++----- 10 files changed, 112 insertions(+), 85 deletions(-) diff --git a/dimos/core/test_native_module.py b/dimos/core/test_native_module.py index 96e6c2cad1..3425d8d7e8 100644 --- a/dimos/core/test_native_module.py +++ b/dimos/core/test_native_module.py @@ -191,10 +191,15 @@ def test_io_port_topic_reaches_the_native_process() -> None: def test_tf_topic_comes_from_the_declared_port_only() -> None: """No tf port declared means no tf topic, rather than a silently injected one.""" module = StubNativeModule(executable=_ECHO) + transport = LCMTransport("/cmd_vel", Twist) try: - assert "tf" not in module._collect_topics() + module.set_transport("cmd_vel", transport) + + assert module._collect_topics() == {"cmd_vel": "/cmd_vel#geometry_msgs.Twist"} finally: module.stop() + with contextlib.suppress(Exception): + transport.stop() def test_io_port_publisher_qos_reaches_the_native_process() -> None: diff --git a/dimos/core/test_transport_factory.py b/dimos/core/test_transport_factory.py index 2a925d6801..419deb44f1 100644 --- a/dimos/core/test_transport_factory.py +++ b/dimos/core/test_transport_factory.py @@ -28,7 +28,6 @@ default_zenoh_qos, make_transport, rpc_backend, - tf_channel, transport_topic, ) from dimos.msgs.geometry_msgs.Twist import Twist @@ -54,14 +53,6 @@ def test_transport_topic_zenoh() -> None: assert transport_topic("/coordinator/joint_state", ZENOH) == "dimos/coordinator/joint_state" -def test_tf_channel_lcm() -> None: - assert tf_channel(LCM) == "/tf#tf2_msgs.TFMessage" - - -def test_tf_channel_zenoh() -> None: - assert tf_channel(ZENOH) == "dimos/tf/tf2_msgs.TFMessage" - - def test_make_transport_lcm_typed() -> None: t = make_transport("/camera/color", Image, g=LCM) assert type(t) is LCMTransport diff --git a/dimos/core/transport_factory.py b/dimos/core/transport_factory.py index a7382e85f1..2a69c270ec 100644 --- a/dimos/core/transport_factory.py +++ b/dimos/core/transport_factory.py @@ -27,8 +27,6 @@ pLCMTransport, pZenohTransport, ) -from dimos.msgs.tf2_msgs.TFMessage import TFMessage -from dimos.protocol.pubsub.impl.lcmpubsub import Topic as LCMTopic from dimos.protocol.pubsub.impl.zenohpubsub import ( QOS_LATEST_WINS, QOS_NEVER_DROP, @@ -55,14 +53,6 @@ def transport_topic(name: str, g: GlobalConfig = global_config) -> str: return name if name.startswith("/") else "/" + name -def tf_channel(g: GlobalConfig = global_config) -> str: - """Wire channel for the tf stream on the active backend.""" - name = transport_topic("/tf", g) - if g.transport == "zenoh": - return ZenohTopic(name, TFMessage).key_expr - return str(LCMTopic(name, TFMessage)) - - # High-rate sensor streams: drop stale frames under congestion, never stall the # publisher. Matched by message type since that is what makes them high-rate. _LATEST_WINS_TYPES = ("sensor_msgs.Image", "sensor_msgs.PointCloud2") diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py index e803537d6f..2f990e6284 100644 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -26,9 +26,6 @@ from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.navigation.tf_pose import OdomBasePose, base_height_above_ground -from dimos.utils.logging_config import setup_logger - -logger = setup_logger() class GoalRelayConfig(ModuleConfig): 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 2af72f9532..a6a79431d2 100644 --- a/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py +++ b/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py @@ -33,9 +33,8 @@ class MLSPlannerNativeConfig(NativeModuleConfig): world_frame: str = "map" voxel_size: float = 0.08 robot_height: float = 0.3 - # Subtracted from the start pose z before snapping to a surface. For start - # poses stamped at the sensor instead of the ground; leave 0 when the - # publisher already ground-projects (GoalRelay does). + # Subtracted from the start pose z before snapping to a surface. Leave 0 + # when the publisher already ground-projects. start_z_offset_m: float = 0.0 max_overhead_m: float = 2.0 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 c5d02ae5ce..314c0adc84 100644 --- a/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs +++ b/dimos/navigation/nav_3d/mls_planner/rust/src/mls_planner.rs @@ -37,9 +37,8 @@ pub struct Config { pub voxel_size: f32, #[validate(range(exclusive_min = 0.0))] pub robot_height: f32, - /// Subtracted from the start pose z before snapping to a surface. For start - /// poses stamped at the sensor instead of the ground; 0 when the publisher - /// already ground-projects. + /// Subtracted from the start pose z before snapping to a surface. 0 when + /// the publisher already ground-projects. #[validate(range(min = 0.0))] pub start_z_offset_m: f32, /// Ignore surface more than this far above the sensor. diff --git a/dimos/navigation/tf_pose.py b/dimos/navigation/tf_pose.py index 05549ab712..a74c5cbaaf 100644 --- a/dimos/navigation/tf_pose.py +++ b/dimos/navigation/tf_pose.py @@ -31,21 +31,12 @@ def base_height_above_ground(lidar_height: float, base_to_sensor: Transform) -> float: - """Height of the base frame origin above the ground while standing. - - The lidar sits ``lidar_height`` above the ground and - ``base_to_sensor.translation.z`` above the base origin. - """ + """Height of the base frame origin above the ground while standing.""" return lidar_height - base_to_sensor.translation.z class OdomBasePose: - """Turn odometry messages into the base-frame pose they imply. - - The odometry message itself carries the live parent -> child edge, so only the - static leg from the message's child frame to the base frame is looked up on tf, - once, then reused for every following message. - """ + """Turn odometry messages into the base-frame pose they imply.""" # While the leg is missing, retry the lookup at most this often. The buffer # warns on every miss, so per-message retries would flood the log. diff --git a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py index 0d9cb26923..f9615bf6a0 100644 --- a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py +++ b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py @@ -25,10 +25,8 @@ - front_camera -> camera_optical: the standard ROS optical rotation (x-right, y-down, z-forward). -The published tree is rooted at ``mid360_link``: Point-LIO owns the live -``odom -> mid360_link`` edge, and the rerun bridge keys tf entities by child frame, so -publishing the two edges above the lidar inverted keeps the static tree off the entity -the live edge writes. The tf buffer composes either direction. +The published tree is rooted at mid360_link so the static edges stay off the entities +the live odom -> mid360_link edge writes. The tf buffer composes either direction. """ from __future__ import annotations diff --git a/native/rust/dimos-module/src/module.rs b/native/rust/dimos-module/src/module.rs index 59fb20e58c..dc953bd24b 100644 --- a/native/rust/dimos-module/src/module.rs +++ b/native/rust/dimos-module/src/module.rs @@ -302,7 +302,7 @@ pub trait Module: Sized + Send + 'static { pub struct Builder { topics: HashMap, - // Every port the module asked for a topic, matched against `topics` after build. + // Every port the module asked for a topic, matched against topics after build. requested: BTreeSet, routes: HashMap>>, // One publish queue per output channel, drained by its own worker. @@ -329,10 +329,8 @@ impl Builder { .unwrap_or_else(|| format!("/{port}")) } - // The coordinator sends one topic per declared Python port, so the ports - // this module claimed must be exactly that set. A port the module never - // asks for is dead wiring, and one the coordinator never sent would fall - // back to a bare `/{port}` that nothing else on the bus is using. + // A mismatch is dead wiring: an unclaimed topic reaches no port, and an + // unsent one leaves the port on a fallback name nothing else publishes to. pub(crate) fn enforce_topics_match_ports(&self) -> io::Result<()> { let provided: BTreeSet<&String> = self.topics.keys().collect(); let requested: BTreeSet<&String> = self.requested.iter().collect(); @@ -349,7 +347,6 @@ impl Builder { )) } - // Registers a decoding route on the topic and hands back the receiving end. fn add_route( &mut self, topic: &str, @@ -369,7 +366,6 @@ impl Builder { rx } - // Adds a publish queue for the topic and hands back the sending end. fn add_publisher(&mut self, topic: &str) -> mpsc::Sender> { let (tx, rx) = mpsc::channel(PUBLISH_CHANNEL_CAPACITY); self.outputs.push((topic.to_string(), rx)); diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index 02aaa94d3f..36207f7432 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -21,7 +21,8 @@ use std::collections::{HashMap, HashSet, VecDeque}; use std::io; -use std::sync::{Arc, RwLock}; +use std::sync::atomic::AtomicU64; +use std::sync::{Arc, Mutex, RwLock}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use nalgebra::{Isometry3, Quaternion, Translation3, UnitQuaternion, Vector3}; @@ -33,6 +34,8 @@ use crate::module::Route; /// How many seconds of history each edge keeps. pub(crate) const DEFAULT_TF_WINDOW_SECS: f64 = 10.0; +const WARN_INTERVAL: Duration = Duration::from_secs(1); + fn now_secs() -> f64 { SystemTime::now() .duration_since(UNIX_EPOCH) @@ -153,7 +156,7 @@ impl TBuffer { } } - // One transform for this edge: the latest sample, or the one nearest `time`. + // One transform for this edge: the latest sample, or the one nearest time. fn sample( &self, parent: &str, @@ -284,11 +287,14 @@ impl MultiTBuffer { } } -// The graph plus the signal that it changed. Writes go through `update` so a -// writer cannot leave a waiter asleep on a transform that has already landed. +// The graph plus the signal that it changed. Every write notifies, so a writer +// cannot leave a waiter asleep on a transform that has already landed. struct Graph { buffer: RwLock, changed: Notify, + // warn_throttled! keys its throttle on the call site and there is one site + // for every lookup, so one missing frame pair would mute all the others. + warned: Mutex>, } impl Graph { @@ -296,9 +302,18 @@ impl Graph { Self { buffer: RwLock::new(MultiTBuffer::new(window_secs)), changed: Notify::new(), + warned: Mutex::new(HashMap::new()), } } + fn should_warn(&self, parent: &str, child: &str) -> bool { + let mut warned = self.warned.lock().expect("tf warn map lock poisoned"); + let last = warned + .entry((parent.to_string(), child.to_string())) + .or_default(); + crate::log::check_and_record(last, WARN_INTERVAL.as_nanos() as u64) + } + fn update(&self, edits: impl FnOnce(&mut MultiTBuffer)) { edits(&mut self.buffer.write().expect("tf buffer lock poisoned")); self.changed.notify_waiters(); @@ -402,6 +417,9 @@ impl Lookup<'_> { // A lookup that resolves to nothing is otherwise invisible: the caller sees // None and the buffer says nothing about which frames or stamp missed. fn warn_unresolved(&self) { + if !self.tf.graph.should_warn(self.parent, self.child) { + return; + } warn!( parent = %self.parent, child = %self.child, @@ -423,22 +441,12 @@ impl Lookup<'_> { found } - /// Resolve the lookup, waiting up to `timeout` for it to become possible. + /// Resolve the lookup, waiting up to `timeout` for a late transform. /// - /// A message and the transform it needs arrive on separate topics, so the - /// transform for a given stamp is often merely late rather than absent. /// Returns as soon as the lookup succeeds, or `None` at the deadline. - /// - /// [`Lookup::tolerance`] and this are different clocks: tolerance bounds how - /// far the chosen sample may sit from [`Lookup::at`] in message stamps, - /// `timeout` bounds how long to wait in wall time. Set a tolerance whenever - /// waiting, or the lookup is satisfied by anything inside the buffer window - /// and returns a stale transform immediately. - /// - /// This suspends the caller. Awaiting it inside a `handle_*` method parks - /// that module's whole dispatch loop, so every other topic it subscribes to - /// stops being served for the duration. Prefer [`Lookup::get`] there, and - /// move waiting onto a task of its own when the wait may be long. + /// Awaiting this inside a `handle_*` method parks the module's whole + /// dispatch loop, so prefer [`Lookup::get`] there and move a long wait onto + /// a task of its own. pub async fn within(self, timeout: Duration) -> Option { let deadline = tokio::time::Instant::now() + timeout; loop { @@ -527,10 +535,21 @@ impl Route for TfRoute { for st in &msg.transforms { let t = &st.transform.translation; let q = &st.transform.rotation; - let iso = Isometry3::from_parts( - Translation3::new(t.x, t.y, t.z), - UnitQuaternion::from_quaternion(Quaternion::new(q.w, q.x, q.y, q.z)), - ); + // Normalizing a zero-norm quaternion yields NaN, which would + // then resolve every lookup composing through this edge. + let Some(rotation) = + UnitQuaternion::try_new(Quaternion::new(q.w, q.x, q.y, q.z), 1e-9) + else { + crate::error_throttled!( + Duration::from_secs(1), + topic = %self.topic, + parent = %st.header.frame_id, + child = %st.child_frame_id, + "tf rotation is not a valid quaternion" + ); + continue; + }; + let iso = Isometry3::from_parts(Translation3::new(t.x, t.y, t.z), rotation); let ts = st.header.stamp.sec as f64 + st.header.stamp.nsec as f64 * 1e-9; buffer.receive(&st.header.frame_id, &st.child_frame_id, ts, iso); } @@ -542,10 +561,10 @@ impl Route for TfRoute { // sender carries published messages to the tf topic's publish worker. pub(crate) fn tf_subscription( topic: String, - buffer_size: f64, + window_secs: f64, sender: mpsc::Sender>, ) -> (Tf, Box) { - let graph = Arc::new(Graph::new(buffer_size)); + let graph = Arc::new(Graph::new(window_secs)); let tf = Tf { graph: Arc::clone(&graph), sender, @@ -559,13 +578,13 @@ mod tests { use super::*; use std::f64::consts::PI; - fn tf_with(buffer_size: f64) -> (Tf, MultiHandle) { - let (tf, _rx, handle) = tf_with_publish(buffer_size); + fn tf_with(window_secs: f64) -> (Tf, MultiHandle) { + let (tf, _rx, handle) = tf_with_publish(window_secs); (tf, handle) } - fn tf_with_publish(buffer_size: f64) -> (Tf, mpsc::Receiver>, MultiHandle) { - let graph = Arc::new(Graph::new(buffer_size)); + fn tf_with_publish(window_secs: f64) -> (Tf, mpsc::Receiver>, MultiHandle) { + let graph = Arc::new(Graph::new(window_secs)); let (tx, rx) = mpsc::channel(8); ( Tf { @@ -719,8 +738,8 @@ mod tests { assert!(tf.lookup("a", "b").at(50.0).tolerance(60.0).get().is_some()); } - // The window bounds queries against a stamp, not the latest sample: with no - // `at`, the newest edge is returned however old it is. + // The window bounds queries against a stamp, not the latest sample. With no + // requested time, the newest edge is returned however old it is. #[test] fn latest_is_not_bounded_by_the_window() { let (tf, h) = tf_with(10.0); @@ -828,10 +847,21 @@ mod tests { } fn stamped_message(parent: &str, child: &str, ts: f64, x: f64) -> Vec { + rotated_message(parent, child, ts, x, (0.0, 0.0, 0.0, 1.0)) + } + + fn rotated_message( + parent: &str, + child: &str, + ts: f64, + x: f64, + quat: (f64, f64, f64, f64), + ) -> Vec { use lcm_msgs::geometry_msgs::{ Quaternion as LQuat, Transform as LTransform, TransformStamped, Vector3 as LVec3, }; use lcm_msgs::std_msgs::{Header, Time}; + let (x_q, y_q, z_q, w_q) = quat; lcm_msgs::tf2_msgs::TFMessage { transforms: vec![TransformStamped { header: Header { @@ -846,10 +876,10 @@ mod tests { transform: LTransform { translation: LVec3 { x, y: 0.0, z: 0.0 }, rotation: LQuat { - x: 0.0, - y: 0.0, - z: 0.0, - w: 1.0, + x: x_q, + y: y_q, + z: z_q, + w: w_q, }, }, }], @@ -857,6 +887,18 @@ mod tests { .encode() } + #[test] + fn a_zero_rotation_on_the_wire_is_dropped_rather_than_stored_as_nan() { + let (tx, _rx) = mpsc::channel(4); + let (tf, route) = tf_subscription("/tf".to_string(), DEFAULT_TF_WINDOW_SECS, tx); + route.try_dispatch(&rotated_message("a", "b", 5.0, 1.0, (0.0, 0.0, 0.0, 0.0))); + assert!(tf.get_latest("a", "b").is_none()); + + route.try_dispatch(&stamped_message("a", "b", 6.0, 1.0)); + let t = tf.get_latest("a", "b").expect("valid rotation is accepted"); + assert!(t.rotation().coords.iter().all(|c| c.is_finite())); + } + #[test] #[tracing_test::traced_test] fn a_lookup_that_finds_nothing_warns_with_the_frames() { @@ -867,6 +909,25 @@ mod tests { assert!(logs_contain("gripper")); } + #[test] + #[tracing_test::traced_test] + fn repeated_misses_warn_once_per_frame_pair() { + let (tf, _h) = tf_with(DEFAULT_TF_WINDOW_SECS); + for _ in 0..5 { + assert!(tf.get_latest("world", "gripper").is_none()); + } + logs_assert(|lines: &[&str]| { + match lines.iter().filter(|l| l.contains("gripper")).count() { + 1 => Ok(()), + n => Err(format!("expected 1 warning for the repeated pair, got {n}")), + } + }); + + // A pair that is throttled must not mute an unrelated one. + assert!(tf.get_latest("world", "camera").is_none()); + assert!(logs_contain("camera")); + } + #[test] #[tracing_test::traced_test] fn a_resolved_lookup_stays_quiet() { From 72b8ebcc208f083b837f2a30dfe038d7dd0632d6 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Fri, 31 Jul 2026 14:42:51 -0700 Subject: [PATCH 25/26] Toggle base link from go2 --- .../nav_3d/mls_planner/goal_relay.py | 15 ++++++ .../nav_3d/mls_planner/test_goal_relay.py | 13 +++++ .../basic/unitree_go2_mid360_record.py | 2 +- .../navigation/unitree_go2_nav_3d.py | 32 +++++------- dimos/robot/unitree/go2/connection.py | 8 ++- .../go2/go2_mid360_static_transforms.py | 13 +++-- dimos/robot/unitree/go2/test_connection.py | 52 +++++++++++++++++++ native/rust/dimos-module/src/tf.rs | 50 ++++++++++++++++-- 8 files changed, 154 insertions(+), 31 deletions(-) diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py index 2f990e6284..49f7f9b909 100644 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -26,6 +26,9 @@ from dimos.msgs.nav_msgs.Odometry import Odometry from dimos.msgs.tf2_msgs.TFMessage import TFMessage from dimos.navigation.tf_pose import OdomBasePose, base_height_above_ground +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() class GoalRelayConfig(ModuleConfig): @@ -51,6 +54,7 @@ def __init__(self, **kwargs: Any) -> None: super().__init__(**kwargs) self._base_pose: OdomBasePose | None = None self._base_height: float | None = None + self._warned_base_frame = False @rpc def start(self) -> None: @@ -72,6 +76,17 @@ def _on_odometry(self, msg: Odometry) -> None: self.start_pose.publish(start) def _resolve_base_height(self, sensor_frame: str, lidar_height: float) -> float | None: + # The base height comes from subtracting the mount leg, which + # base-frame odometry does not have. + if sensor_frame == self.config.base_frame: + if not self._warned_base_frame: + self._warned_base_frame = True + logger.warning( + "Odometry is stamped at %s, so lidar_height cannot ground-project it. " + "Dropping frames until odometry arrives stamped at a sensor.", + sensor_frame, + ) + return None if self._base_height is None: assert self._base_pose is not None leg = self._base_pose.sensor_to_base(sensor_frame) diff --git a/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py b/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py index abcf5ea7d7..97eb203314 100644 --- a/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py @@ -88,6 +88,19 @@ def test_drops_frames_without_the_mount_tf(): module.stop() +def test_base_frame_odometry_is_dropped_rather_than_over_projected(): + tf = FakeTF() + tf.receive_transform(_mount()) + module, captured = _relay(tf, lidar_height=0.45) + try: + odom = _odom() + odom.child_frame_id = "base_link" + module._on_odometry(odom) + assert captured == [] + finally: + module.stop() + + def test_no_lidar_height_skips_the_ground_correction(): tf = FakeTF() tf.receive_transform(_mount()) diff --git a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py index e77cdb2d1d..92b84714e6 100644 --- a/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py +++ b/dimos/robot/unitree/go2/blueprints/basic/unitree_go2_mid360_record.py @@ -69,7 +69,7 @@ def _default_recording_dir() -> Path: unitree_go2_mid360_record = autoconnect( MovementManager.blueprint(), - GO2Connection.blueprint().remappings( + GO2Connection.blueprint(publish_tf=False).remappings( [ (GO2Connection, "lidar", "go2_lidar"), (GO2Connection, "odom", "go2_odom"), diff --git a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py index e054d927c5..2c03b51142 100644 --- a/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py +++ b/dimos/robot/unitree/go2/blueprints/navigation/unitree_go2_nav_3d.py @@ -16,7 +16,6 @@ """3d navigation on Go2 with ray tracing and MLS planning""" from datetime import datetime -import math import os from pathlib import Path from typing import Any @@ -40,10 +39,7 @@ from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_basic import rerun_config from dimos.robot.unitree.go2.connection import GO2Connection from dimos.robot.unitree.go2.constants import ROBOT_HEIGHT, ROBOT_LENGTH, ROBOT_WIDTH -from dimos.robot.unitree.go2.go2_mid360_static_transforms import ( - MID360_PITCH_DOWN, - Go2Mid360StaticTf, -) +from dimos.robot.unitree.go2.go2_mid360_static_transforms import Go2Mid360StaticTf from dimos.visualization.vis_module import vis_module voxel_size = 0.08 @@ -98,16 +94,13 @@ def _render_path(msg: Any) -> Any: def _static_robot_body(rr: Any) -> list[Any]: - """Go2-shaped box on pointlio's sensor frame, counter-rotated for the lidar pitch.""" + """Go2-shaped box on the body frame.""" return [ rr.Boxes3D( half_sizes=[ROBOT_LENGTH / 2, ROBOT_WIDTH / 2, ROBOT_HEIGHT / 2], colors=[(0, 255, 127)], ), - rr.Transform3D( - parent_frame="tf#/mid360_link", - rotation=rr.RotationAxisAngle(axis=(0, 1, 0), degrees=-math.degrees(MID360_PITCH_DOWN)), - ), + rr.Transform3D(parent_frame="tf#/base_link"), ] @@ -126,7 +119,7 @@ def _axis_triad(rr: Any) -> Any: def _static_body_axes(rr: Any) -> Any: - """XYZ triad on the leveled robot body (child of the counter-rotated box).""" + """XYZ triad on the robot body (child of the box).""" return _axis_triad(rr) @@ -145,10 +138,8 @@ def _static_sensor_axes(rr: Any) -> list[Any]: }, # Ring buffer replayed to a connecting viewer. Small so connect catches up fast. "memory_limit": "64MB", - # base_link tf comes from the go2 internal odometry, which is not the map - # frame. Anchor the robot box to pointlio's mid360_link frame instead and hide - # the camera frustum that rides base_link. The box lives on its own entity: - # a static transform on world/tf/mid360_link itself would override the live tf. + # The robot box hangs off base_link. It lives on its own entity: a static + # transform on world/tf/base_link would override the live tf. "static": { "world/robot_body": _static_robot_body, "world/robot_body/axes": _static_body_axes, @@ -169,7 +160,11 @@ def _static_sensor_axes(rr: Any) -> list[Any]: vis_module(viewer_backend=global_config.viewer, rerun_config=_nav_rerun_config), # "mcf" for stair traversal GO2Connection.blueprint( - lidar=False, camera=False, motion_mode="mcf", odom_frame_id="go2_odom" + lidar=False, + camera=False, + motion_mode="mcf", + odom_frame_id="go2_odom", + publish_tf=False, ).remappings( [ (GO2Connection, "lidar", "lidar_l1"), @@ -205,9 +200,8 @@ def _static_sensor_axes(rr: Any) -> list[Any]: MovementManager.blueprint(), ).global_config(n_workers=10, robot_model="unitree_go2", obstacle_avoidance=False) -# The nav blueprint leaves PointLio on its default lidar / odometry topics, so -# remap the recorder's ports onto them. Streams are recorded under the port -# names pointlio_lidar / pointlio_odometry regardless of the topic. +# PointLio keeps its default topics here, so point the recorder's ports at them. +# Streams are recorded under the port names regardless of the topic. if _RECORD: unitree_go2_nav_3d = autoconnect( unitree_go2_nav_3d, diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index ee9224bbdb..b285273f23 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -75,6 +75,9 @@ class ConnectionConfig(ModuleConfig): # TF parent frame of the internal odometry (odom_frame_id -> base_link). # Rename (e.g. "go2_odom") when another odom source owns the tree root odom_frame_id: str = "world" + # Turn off where another module owns the base_link edge. The odom port + # keeps publishing either way. + publish_tf: bool = True class Go2ConnectionProtocol(Protocol): @@ -397,8 +400,9 @@ def _odom_to_tf(cls, odom: PoseStamped, prefix: str = "") -> list[Transform]: def _publish_tf(self, msg: PoseStamped) -> None: msg.frame_id = self.config.odom_frame_id - transforms = self._odom_to_tf(msg, prefix=self.config.frame_id_prefix or "") - self.tf.publish(TFMessage(*transforms)) + if self.config.publish_tf: + transforms = self._odom_to_tf(msg, prefix=self.config.frame_id_prefix or "") + self.tf.publish(TFMessage(*transforms)) if self.odom.transport: self.odom.publish(msg) diff --git a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py index f9615bf6a0..cf42173e45 100644 --- a/dimos/robot/unitree/go2/go2_mid360_static_transforms.py +++ b/dimos/robot/unitree/go2/go2_mid360_static_transforms.py @@ -21,7 +21,7 @@ Mount geometry (measured on the physical rig) --------------------------------------------- - base_link -> front_camera: 32.7cm forward, ~4.3cm up (URDF front_camera mount). -- front_camera -> mid360_link: lidar is 3.2cm back, 12cm up, pitched 44 deg down. +- front_camera -> mid360_link: lidar is 3.2cm back, 12cm up, pitched 60 deg down. - front_camera -> camera_optical: the standard ROS optical rotation (x-right, y-down, z-forward). @@ -40,7 +40,7 @@ frames_to_edge_transforms, ) -MID360_PITCH_DOWN = math.radians(44.0) +MID360_PITCH_DOWN = math.radians(60.0) # rpy that maps a sensor frame to its optical frame (z-forward, x-right, y-down) OPTICAL_RPY = (-math.pi / 2, 0.0, -math.pi / 2) @@ -53,9 +53,14 @@ ] +def mount_transforms() -> list[Transform]: + """The mount tree as published: rooted at mid360_link.""" + edges = {t.child_frame_id: t for t in frames_to_edge_transforms(FRAMES)} + return [-edges["mid360_link"], -edges["front_camera"], edges["camera_optical"]] + + class Go2Mid360StaticTf(StaticTfPublisher): """Publishes the Go2/Mid-360 mount tree onto tf on a fixed interval.""" def transforms(self) -> list[Transform]: - edges = {t.child_frame_id: t for t in frames_to_edge_transforms(FRAMES)} - return [-edges["mid360_link"], -edges["front_camera"], edges["camera_optical"]] + return mount_transforms() diff --git a/dimos/robot/unitree/go2/test_connection.py b/dimos/robot/unitree/go2/test_connection.py index 533254fe2e..ca41f1580a 100644 --- a/dimos/robot/unitree/go2/test_connection.py +++ b/dimos/robot/unitree/go2/test_connection.py @@ -23,10 +23,18 @@ import pytest +from dimos.core.coordination.blueprints import Blueprint from dimos.core.global_config import GlobalConfig from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.robot.unitree.go2 import connection as go2_conn +from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_mid360_record import ( + unitree_go2_mid360_record, +) +from dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_nav_3d import ( + unitree_go2_nav_3d, +) from dimos.robot.unitree.go2.connection import ConnectionConfig, GO2Connection +from dimos.robot.unitree.go2.go2_mid360_static_transforms import mount_transforms @pytest.fixture @@ -66,6 +74,50 @@ def test_odom_to_tf_unprefixed_by_default() -> None: ) +def _connection(publish_tf: bool) -> GO2Connection: + conn = object.__new__(GO2Connection) + conn.config = ConnectionConfig( + g=GlobalConfig(robot_ip="127.0.0.1"), publish_tf=publish_tf, odom_frame_id="go2_odom" + ) + conn.tf = MagicMock() + conn.odom = MagicMock() + return conn + + +def test_publish_tf_off_keeps_odometry_on_its_port() -> None: + """Turning tf off hands the base_link edge to another publisher, not the odom port.""" + conn = _connection(publish_tf=False) + conn._publish_tf(PoseStamped(ts=1.0, frame_id="ignored")) + assert conn.tf.publish.call_count == 0 + assert conn.odom.publish.call_count == 1 + + +def test_publish_tf_on_by_default() -> None: + conn = _connection(publish_tf=True) + conn._publish_tf(PoseStamped(ts=1.0, frame_id="ignored")) + assert conn.tf.publish.call_count == 1 + assert conn.odom.publish.call_count == 1 + + +def _go2_connection_publishes_tf(blueprint: Blueprint) -> bool | None: + for atom in blueprint.blueprints: + if atom.module is GO2Connection: + return bool(atom.kwargs.get("publish_tf", True)) + return None + + +def test_static_tree_and_connection_never_share_a_child_frame() -> None: + """One publisher per edge: rerun keys tf entities by child, so a frame written + by two sources flaps between them.""" + odom = PoseStamped(ts=1.0, frame_id="go2_odom") + connection_children = {t.child_frame_id for t in GO2Connection._odom_to_tf(odom)} + static_children = {t.child_frame_id for t in mount_transforms()} + assert connection_children & static_children == {"base_link", "camera_optical"} + + for blueprint in (unitree_go2_nav_3d, unitree_go2_mid360_record): + assert _go2_connection_publishes_tf(blueprint) is False + + def test_odom_to_tf_prefixed() -> None: """.namespace() sets frame_id_prefix: robot-local frames get prefixed, the odom parent frame stays global so all robots hang off one tree root.""" diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index 36207f7432..aa15582034 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -116,9 +116,19 @@ impl TBuffer { } fn add(&mut self, ts: f64, iso: Isometry3) { + // A stamp a whole window behind the newest is a clock reset, not jitter. + if self + .samples + .back() + .is_some_and(|s| ts < s.ts - self.window_secs) + { + self.samples.clear(); + } let pos = self.samples.partition_point(|s| s.ts <= ts); self.samples.insert(pos, Sample { ts, iso }); - self.prune(ts - self.window_secs); + // Anchored to the newest sample so a late message cannot widen the window. + let newest = self.samples.back().map_or(ts, |s| s.ts); + self.prune(newest - self.window_secs); } fn prune(&mut self, min_ts: f64) { @@ -292,8 +302,8 @@ impl MultiTBuffer { struct Graph { buffer: RwLock, changed: Notify, - // warn_throttled! keys its throttle on the call site and there is one site - // for every lookup, so one missing frame pair would mute all the others. + // Keyed per pair: warn_throttled! throttles per call site, and one site + // serves every lookup, so a missing pair would mute all the others. warned: Mutex>, } @@ -535,8 +545,7 @@ impl Route for TfRoute { for st in &msg.transforms { let t = &st.transform.translation; let q = &st.transform.rotation; - // Normalizing a zero-norm quaternion yields NaN, which would - // then resolve every lookup composing through this edge. + // Normalizing a zero-norm quaternion yields a NaN rotation. let Some(rotation) = UnitQuaternion::try_new(Quaternion::new(q.w, q.x, q.y, q.z), 1e-9) else { @@ -976,6 +985,37 @@ mod tests { assert!((buf.last().unwrap().ts - 10.0).abs() < 1e-9); } + #[test] + fn a_late_sample_does_not_spare_ones_the_window_has_aged_out() { + let mut buf = TBuffer::new(5.0); + buf.add(10.0, Isometry3::identity()); + buf.add(11.0, Isometry3::identity()); + // Late, but still inside the window. + buf.add(7.0, Isometry3::identity()); + assert_eq!(buf.samples.len(), 3); + + buf.add(20.0, Isometry3::identity()); + assert_eq!(buf.samples.len(), 1); + assert!((buf.last().unwrap().ts - 20.0).abs() < 1e-9); + } + + #[test] + fn a_clock_reset_drops_the_pre_jump_samples() { + let mut buf = TBuffer::new(5.0); + for i in 0..20 { + buf.add(1000.0 + i as f64, Isometry3::identity()); + } + for i in 0..20 { + buf.add(100.0 + i as f64, Isometry3::identity()); + } + assert!( + buf.samples.len() <= 6, + "buffer grew to {}", + buf.samples.len() + ); + assert!((buf.last().unwrap().ts - 119.0).abs() < 1e-9); + } + #[test] fn tf_route_decodes_into_graph() { use lcm_msgs::geometry_msgs::{ From a35ae8e84ad2f08ab474f334aa2220a5edaecd31 Mon Sep 17 00:00:00 2001 From: Andrew Lauer Date: Fri, 31 Jul 2026 16:06:37 -0700 Subject: [PATCH 26/26] Remove bad test --- .../basic_path_follower/test_module.py | 16 ++-- .../nav_3d/mls_planner/goal_relay.py | 3 +- .../nav_3d/mls_planner/mls_planner_native.py | 3 +- .../nav_3d/mls_planner/test_goal_relay.py | 48 ++++++++---- dimos/navigation/test_tf_pose.py | 42 +++++++---- .../go2/blueprints/test_tf_topology.py | 73 +++++++++++++++++++ dimos/robot/unitree/go2/test_connection.py | 69 +++++++----------- .../go2/test_go2_mid360_static_transforms.py | 64 ++++++++++++++++ native/rust/README.md | 2 + native/rust/dimos-module/src/tf.rs | 23 ++++-- 10 files changed, 257 insertions(+), 86 deletions(-) create mode 100644 dimos/robot/unitree/go2/blueprints/test_tf_topology.py create mode 100644 dimos/robot/unitree/go2/test_go2_mid360_static_transforms.py diff --git a/dimos/navigation/basic_path_follower/test_module.py b/dimos/navigation/basic_path_follower/test_module.py index 5cbe2223e9..c9be69766f 100644 --- a/dimos/navigation/basic_path_follower/test_module.py +++ b/dimos/navigation/basic_path_follower/test_module.py @@ -20,6 +20,8 @@ from dimos.navigation.basic_path_follower.module import BasicPathFollower, lookahead_distance from dimos.protocol.tf.tf import MultiTBuffer +MOUNT_Z = 0.163 + class FakeTF(MultiTBuffer): def dispose(self) -> None: @@ -35,11 +37,11 @@ def _odom() -> Odometry: ) -def test_on_odometry_steers_from_the_base_pose(): +def test_on_odometry_steers_from_the_base_pose() -> None: tf = FakeTF() tf.receive_transform( Transform( - translation=Vector3(0.0, 0.0, 0.163), + translation=Vector3(0.0, 0.0, MOUNT_Z), rotation=Quaternion(0.0, 0.0, 0.0, 1.0), frame_id="base_link", child_frame_id="mid360_link", @@ -51,12 +53,12 @@ def test_on_odometry_steers_from_the_base_pose(): try: module._on_odometry(_odom()) assert module._current_pose is not None - assert abs(module._current_pose.position.z - (3.0 - 0.163)) < 1e-9 + assert abs(module._current_pose.position.z - (3.0 - MOUNT_Z)) < 1e-9 finally: module.stop() -def test_on_odometry_drops_frames_without_the_mount_tf(): +def test_on_odometry_drops_frames_without_the_mount_tf() -> None: module = BasicPathFollower() module._tf = FakeTF() try: @@ -66,13 +68,13 @@ def test_on_odometry_drops_frames_without_the_mount_tf(): module.stop() -def test_lookahead_floor_at_low_speed(): +def test_lookahead_floor_at_low_speed() -> None: assert lookahead_distance(0.1, 1.5, 0.4, 1.5) == 0.4 -def test_lookahead_scales_in_linear_region(): +def test_lookahead_scales_in_linear_region() -> None: assert lookahead_distance(0.5, 1.5, 0.4, 1.5) == 0.75 -def test_lookahead_clamped_at_ceiling(): +def test_lookahead_clamped_at_ceiling() -> None: assert lookahead_distance(2.0, 1.5, 0.4, 1.5) == 1.5 diff --git a/dimos/navigation/nav_3d/mls_planner/goal_relay.py b/dimos/navigation/nav_3d/mls_planner/goal_relay.py index 49f7f9b909..778597cec8 100644 --- a/dimos/navigation/nav_3d/mls_planner/goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/goal_relay.py @@ -34,7 +34,8 @@ class GoalRelayConfig(ModuleConfig): base_frame: str = "base_link" # Lidar height above the ground while standing. None skips the ground - # correction. + # correction, leaving it to the planner's start_z_offset_m. Set exactly one + # of the two: both set drops the start pose to the ground twice. lidar_height: float | None = None 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 a6a79431d2..c89c9accbf 100644 --- a/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py +++ b/dimos/navigation/nav_3d/mls_planner/mls_planner_native.py @@ -34,7 +34,8 @@ class MLSPlannerNativeConfig(NativeModuleConfig): voxel_size: float = 0.08 robot_height: float = 0.3 # Subtracted from the start pose z before snapping to a surface. Leave 0 - # when the publisher already ground-projects. + # when the publisher already ground-projects via GoalRelay lidar_height. + # Set exactly one of the two: both set drops the start pose twice. start_z_offset_m: float = 0.0 max_overhead_m: float = 2.0 diff --git a/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py b/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py index 97eb203314..67f4b6ff73 100644 --- a/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py +++ b/dimos/navigation/nav_3d/mls_planner/test_goal_relay.py @@ -12,7 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any + from dimos.msgs.geometry_msgs.Pose import Pose +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion from dimos.msgs.geometry_msgs.Transform import Transform from dimos.msgs.geometry_msgs.Vector3 import Vector3 @@ -21,6 +24,7 @@ from dimos.protocol.tf.tf import MultiTBuffer MOUNT_Z = 0.163 +LIDAR_HEIGHT = 0.45 class FakeTF(MultiTBuffer): @@ -30,9 +34,23 @@ def __init__(self) -> None: super().__init__() self.gets = 0 - def get(self, *args, **kwargs): # type: ignore[no-untyped-def] + def get( + self, + parent_frame: str, + child_frame: str, + time_point: float | None = None, + time_tolerance: float | None = None, + *, + forward_tolerance: float = 0.0, + ) -> Transform | None: self.gets += 1 - return super().get(*args, **kwargs) + return super().get( + parent_frame, + child_frame, + time_point, + time_tolerance, + forward_tolerance=forward_tolerance, + ) def dispose(self) -> None: pass @@ -57,30 +75,30 @@ def _odom(z: float = 3.0) -> Odometry: ) -def _relay(tf: FakeTF, **config) -> tuple[GoalRelay, list]: # type: ignore[no-untyped-def] +def _relay(tf: FakeTF, **config: Any) -> tuple[GoalRelay, list[PoseStamped]]: module = GoalRelay(**config) module._tf = tf - captured: list = [] + captured: list[PoseStamped] = [] module.start_pose.subscribe(captured.append) return module, captured -def test_start_pose_is_ground_projected(): +def test_start_pose_is_ground_projected() -> None: tf = FakeTF() tf.receive_transform(_mount()) - module, captured = _relay(tf, lidar_height=0.45) + module, captured = _relay(tf, lidar_height=LIDAR_HEIGHT) try: module._on_odometry(_odom()) # Base sits MOUNT_Z below the sensor, then drops by the base's height # above ground (0.45 - MOUNT_Z): together exactly the lidar height. assert len(captured) == 1 - assert abs(captured[0].position.z - (3.0 - 0.45)) < 1e-9 + assert abs(captured[0].position.z - (3.0 - LIDAR_HEIGHT)) < 1e-9 finally: module.stop() -def test_drops_frames_without_the_mount_tf(): - module, captured = _relay(FakeTF(), lidar_height=0.45) +def test_drops_frames_without_the_mount_tf() -> None: + module, captured = _relay(FakeTF(), lidar_height=LIDAR_HEIGHT) try: module._on_odometry(_odom()) assert captured == [] @@ -88,10 +106,10 @@ def test_drops_frames_without_the_mount_tf(): module.stop() -def test_base_frame_odometry_is_dropped_rather_than_over_projected(): +def test_base_frame_odometry_is_dropped_rather_than_over_projected() -> None: tf = FakeTF() tf.receive_transform(_mount()) - module, captured = _relay(tf, lidar_height=0.45) + module, captured = _relay(tf, lidar_height=LIDAR_HEIGHT) try: odom = _odom() odom.child_frame_id = "base_link" @@ -101,7 +119,7 @@ def test_base_frame_odometry_is_dropped_rather_than_over_projected(): module.stop() -def test_no_lidar_height_skips_the_ground_correction(): +def test_no_lidar_height_skips_the_ground_correction() -> None: tf = FakeTF() tf.receive_transform(_mount()) module, captured = _relay(tf) @@ -113,15 +131,15 @@ def test_no_lidar_height_skips_the_ground_correction(): module.stop() -def test_mount_is_looked_up_once(): +def test_mount_is_looked_up_once() -> None: tf = FakeTF() tf.receive_transform(_mount()) - module, captured = _relay(tf, lidar_height=0.45) + module, captured = _relay(tf, lidar_height=LIDAR_HEIGHT) try: module._on_odometry(_odom()) module._on_odometry(_odom(z=4.0)) assert len(captured) == 2 - assert abs(captured[1].position.z - (4.0 - 0.45)) < 1e-9 + assert abs(captured[1].position.z - (4.0 - LIDAR_HEIGHT)) < 1e-9 assert tf.gets == 1 finally: module.stop() diff --git a/dimos/navigation/test_tf_pose.py b/dimos/navigation/test_tf_pose.py index 0f1f487447..1ded4b6980 100644 --- a/dimos/navigation/test_tf_pose.py +++ b/dimos/navigation/test_tf_pose.py @@ -21,6 +21,8 @@ from dimos.protocol.tf.tf import MultiTBuffer IDENTITY = Quaternion(0.0, 0.0, 0.0, 1.0) +MOUNT_Z = 0.163 +LIDAR_HEIGHT = 0.45 class CountingTF(MultiTBuffer): @@ -28,12 +30,26 @@ def __init__(self) -> None: super().__init__() self.gets = 0 - def get(self, *args, **kwargs): # type: ignore[no-untyped-def] + def get( + self, + parent_frame: str, + child_frame: str, + time_point: float | None = None, + time_tolerance: float | None = None, + *, + forward_tolerance: float = 0.0, + ) -> Transform | None: self.gets += 1 - return super().get(*args, **kwargs) + return super().get( + parent_frame, + child_frame, + time_point, + time_tolerance, + forward_tolerance=forward_tolerance, + ) -def _mount(z: float = 0.163, pitch: float = 0.0) -> Transform: +def _mount(z: float = MOUNT_Z, pitch: float = 0.0) -> Transform: return Transform( translation=Vector3(0.0, 0.0, z), rotation=Quaternion.from_euler(Vector3(0.0, pitch, 0.0)), @@ -52,7 +68,7 @@ def _odom(orientation: Quaternion = IDENTITY) -> Odometry: ) -def test_translates_to_base_frame(): +def test_translates_to_base_frame() -> None: tf = MultiTBuffer() tf.receive_transform(_mount()) pose = OdomBasePose(tf, "base_link").resolve(_odom()) @@ -61,10 +77,10 @@ def test_translates_to_base_frame(): assert pose.ts == 1.0 assert abs(pose.position.x - 1.0) < 1e-9 assert abs(pose.position.y - 2.0) < 1e-9 - assert abs(pose.position.z - (3.0 - 0.163)) < 1e-9 + assert abs(pose.position.z - (3.0 - MOUNT_Z)) < 1e-9 -def test_composes_out_the_mount_pitch(): +def test_composes_out_the_mount_pitch() -> None: # A level body reads its own mount tilt as the sensor's world orientation, so # composing the mount out returns identity. mount = _mount(pitch=0.3) @@ -75,7 +91,7 @@ def test_composes_out_the_mount_pitch(): assert pose.orientation.angle_to(IDENTITY) < 1e-5 -def test_preserves_body_yaw_under_mount_tilt(): +def test_preserves_body_yaw_under_mount_tilt() -> None: mount = _mount(pitch=0.3) body = Quaternion.from_euler(Vector3(0.0, 0.0, 0.7)) tf = MultiTBuffer() @@ -85,7 +101,7 @@ def test_preserves_body_yaw_under_mount_tilt(): assert pose.orientation.angle_to(body) < 1e-5 -def test_drops_frames_until_the_mount_leg_arrives(): +def test_drops_frames_until_the_mount_leg_arrives() -> None: tf = MultiTBuffer() resolver = OdomBasePose(tf, "base_link") assert resolver.resolve(_odom()) is None @@ -94,7 +110,7 @@ def test_drops_frames_until_the_mount_leg_arrives(): assert resolver.resolve(_odom()) is not None -def test_missing_leg_lookups_are_throttled(): +def test_missing_leg_lookups_are_throttled() -> None: tf = CountingTF() resolver = OdomBasePose(tf, "base_link") assert resolver.resolve(_odom()) is None @@ -102,7 +118,7 @@ def test_missing_leg_lookups_are_throttled(): assert tf.gets == 1 -def test_mount_leg_is_looked_up_once(): +def test_mount_leg_is_looked_up_once() -> None: tf = CountingTF() tf.receive_transform(_mount()) resolver = OdomBasePose(tf, "base_link") @@ -111,7 +127,7 @@ def test_mount_leg_is_looked_up_once(): assert tf.gets == 1 -def test_base_frame_odometry_passes_through(): +def test_base_frame_odometry_passes_through() -> None: resolver = OdomBasePose(MultiTBuffer(), "base_link") msg = Odometry(ts=1.0, frame_id="odom", child_frame_id="base_link") pose = resolver.resolve(msg) @@ -119,5 +135,5 @@ def test_base_frame_odometry_passes_through(): assert pose.frame_id == "odom" -def test_base_height_above_ground(): - assert abs(base_height_above_ground(0.45, _mount()) - (0.45 - 0.163)) < 1e-9 +def test_base_height_above_ground() -> None: + assert abs(base_height_above_ground(LIDAR_HEIGHT, _mount()) - (LIDAR_HEIGHT - MOUNT_Z)) < 1e-9 diff --git a/dimos/robot/unitree/go2/blueprints/test_tf_topology.py b/dimos/robot/unitree/go2/blueprints/test_tf_topology.py new file mode 100644 index 0000000000..eef20aa954 --- /dev/null +++ b/dimos/robot/unitree/go2/blueprints/test_tf_topology.py @@ -0,0 +1,73 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Every Go2 blueprint publishes a tf tree, not a tf graph. + +tf gives each frame one parent. Two publishers writing the same child frame make +the buffer a graph instead, and a lookup then resolves by hop count rather than by +which source is authoritative, so the wrong odometry can win silently. Each +blueprint keeps the invariant a different way: nav_3d turns GO2Connection's tf off +so the static mount tree owns base_link, and the static tree is rooted at +mid360_link so it never writes the frame PointLio owns. +""" + +import pytest + +from dimos.core.coordination.blueprints import Blueprint +from dimos.hardware.sensors.lidar.pointlio.module import PointLio +from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped +from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_mid360_record import ( + unitree_go2_mid360_record, +) +from dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_nav_3d import ( + unitree_go2_nav_3d, +) +from dimos.robot.unitree.go2.connection import GO2Connection +from dimos.robot.unitree.go2.go2_mid360_static_transforms import ( + Go2Mid360StaticTf, + mount_transforms, +) + +BLUEPRINTS = [unitree_go2_nav_3d, unitree_go2_mid360_record] + + +def _tf_children_by_publisher(blueprint: Blueprint) -> dict[str, set[str]]: + """Child frames each tf publisher the blueprint actually enables will write.""" + odom = PoseStamped(ts=1.0, frame_id="go2_odom") + children: dict[str, set[str]] = {} + for atom in blueprint.blueprints: + if atom.module is GO2Connection and atom.kwargs.get("publish_tf", True): + children["GO2Connection"] = {t.child_frame_id for t in GO2Connection._odom_to_tf(odom)} + if atom.module is Go2Mid360StaticTf: + children["Go2Mid360StaticTf"] = {t.child_frame_id for t in mount_transforms()} + if atom.module is PointLio: + sensor_frame = atom.kwargs.get("sensor_frame_id", "mid360_link") + children["PointLio"] = {sensor_frame} + return children + + +@pytest.mark.parametrize("blueprint", BLUEPRINTS) +def test_no_frame_has_two_tf_parents(blueprint: Blueprint) -> None: + by_publisher = _tf_children_by_publisher(blueprint) + assert by_publisher, "blueprint publishes no tf, so this asserts nothing" + claimed: set[str] = set() + for publisher, frames in by_publisher.items(): + clash = claimed & frames + assert not clash, f"{publisher} also writes {sorted(clash)}" + claimed |= frames + + +def test_static_tree_does_not_write_the_pointlio_frame() -> None: + """Rooting the mount tree at mid360_link is what keeps it off PointLio's edge.""" + assert "mid360_link" not in {t.child_frame_id for t in mount_transforms()} diff --git a/dimos/robot/unitree/go2/test_connection.py b/dimos/robot/unitree/go2/test_connection.py index ca41f1580a..79e99fd402 100644 --- a/dimos/robot/unitree/go2/test_connection.py +++ b/dimos/robot/unitree/go2/test_connection.py @@ -18,23 +18,16 @@ dimos/robot/unitree/test_connection.py; this pins the go2-local routing. """ +from collections.abc import Callable, Iterator from types import SimpleNamespace from unittest.mock import MagicMock import pytest -from dimos.core.coordination.blueprints import Blueprint from dimos.core.global_config import GlobalConfig from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.robot.unitree.go2 import connection as go2_conn -from dimos.robot.unitree.go2.blueprints.basic.unitree_go2_mid360_record import ( - unitree_go2_mid360_record, -) -from dimos.robot.unitree.go2.blueprints.navigation.unitree_go2_nav_3d import ( - unitree_go2_nav_3d, -) from dimos.robot.unitree.go2.connection import ConnectionConfig, GO2Connection -from dimos.robot.unitree.go2.go2_mid360_static_transforms import mount_transforms @pytest.fixture @@ -74,50 +67,44 @@ def test_odom_to_tf_unprefixed_by_default() -> None: ) -def _connection(publish_tf: bool) -> GO2Connection: - conn = object.__new__(GO2Connection) - conn.config = ConnectionConfig( - g=GlobalConfig(robot_ip="127.0.0.1"), publish_tf=publish_tf, odom_frame_id="go2_odom" - ) - conn.tf = MagicMock() - conn.odom = MagicMock() - return conn - - -def test_publish_tf_off_keeps_odometry_on_its_port() -> None: +@pytest.fixture +def connection(stub_webrtc: MagicMock) -> Iterator[Callable[[bool], GO2Connection]]: + """Build GO2Connections with the tf and odom ports stubbed, and stop them after.""" + built: list[GO2Connection] = [] + + def build(publish_tf: bool) -> GO2Connection: + conn = GO2Connection( + g=GlobalConfig(robot_ip="127.0.0.1"), + publish_tf=publish_tf, + odom_frame_id="go2_odom", + ) + conn.tf = MagicMock() + conn.odom = MagicMock() + built.append(conn) + return conn + + yield build + for conn in built: + conn.stop() + + +def test_publish_tf_off_keeps_odometry_on_its_port( + connection: Callable[[bool], GO2Connection], +) -> None: """Turning tf off hands the base_link edge to another publisher, not the odom port.""" - conn = _connection(publish_tf=False) + conn = connection(publish_tf=False) conn._publish_tf(PoseStamped(ts=1.0, frame_id="ignored")) assert conn.tf.publish.call_count == 0 assert conn.odom.publish.call_count == 1 -def test_publish_tf_on_by_default() -> None: - conn = _connection(publish_tf=True) +def test_publish_tf_on_by_default(connection: Callable[[bool], GO2Connection]) -> None: + conn = connection(publish_tf=True) conn._publish_tf(PoseStamped(ts=1.0, frame_id="ignored")) assert conn.tf.publish.call_count == 1 assert conn.odom.publish.call_count == 1 -def _go2_connection_publishes_tf(blueprint: Blueprint) -> bool | None: - for atom in blueprint.blueprints: - if atom.module is GO2Connection: - return bool(atom.kwargs.get("publish_tf", True)) - return None - - -def test_static_tree_and_connection_never_share_a_child_frame() -> None: - """One publisher per edge: rerun keys tf entities by child, so a frame written - by two sources flaps between them.""" - odom = PoseStamped(ts=1.0, frame_id="go2_odom") - connection_children = {t.child_frame_id for t in GO2Connection._odom_to_tf(odom)} - static_children = {t.child_frame_id for t in mount_transforms()} - assert connection_children & static_children == {"base_link", "camera_optical"} - - for blueprint in (unitree_go2_nav_3d, unitree_go2_mid360_record): - assert _go2_connection_publishes_tf(blueprint) is False - - def test_odom_to_tf_prefixed() -> None: """.namespace() sets frame_id_prefix: robot-local frames get prefixed, the odom parent frame stays global so all robots hang off one tree root.""" diff --git a/dimos/robot/unitree/go2/test_go2_mid360_static_transforms.py b/dimos/robot/unitree/go2/test_go2_mid360_static_transforms.py new file mode 100644 index 0000000000..7cf999cd0f --- /dev/null +++ b/dimos/robot/unitree/go2/test_go2_mid360_static_transforms.py @@ -0,0 +1,64 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The published mount tree composes back to the measured rig geometry. + +mount_transforms() inverts two of the four FRAMES edges to root the tree at +mid360_link, so the geometry a consumer reads back is not the geometry written +in FRAMES. These pin the composed result, which is what nav actually uses. +""" + +import math + +from dimos.protocol.tf.tf import MultiTBuffer +from dimos.robot.unitree.go2.go2_mid360_static_transforms import ( + MID360_PITCH_DOWN, + mount_transforms, +) + +# base_link -> mid360_link, summed down the FRAMES chain. +MOUNT_X = 0.32715 - 0.032 +MOUNT_Z = 0.04297 + 0.12 + + +def _buffer() -> MultiTBuffer: + buffer = MultiTBuffer() + buffer.receive_transform(*mount_transforms()) + return buffer + + +def test_mount_height_survives_the_inversion() -> None: + """The lidar sits MOUNT_Z above base_link, the offset every ground projection uses.""" + leg = _buffer().get("mid360_link", "base_link") + assert leg is not None + base_to_sensor = -leg + assert abs(base_to_sensor.translation.z - MOUNT_Z) < 1e-6 + assert abs(base_to_sensor.translation.x - MOUNT_X) < 1e-6 + + +def test_mount_pitch_survives_the_inversion() -> None: + """A sign flip here steers the follower off-heading rather than failing loudly.""" + leg = _buffer().get("mid360_link", "base_link") + assert leg is not None + pitch = (-leg).rotation.euler.y + assert abs(pitch - MID360_PITCH_DOWN) < 1e-6 + assert abs(math.degrees(pitch) - 60.0) < 1e-6 + + +def test_camera_optical_hangs_off_base_link() -> None: + """The tree is rooted at mid360_link, so the camera edge is reachable by composition.""" + optical = _buffer().get("base_link", "camera_optical") + assert optical is not None + assert abs(optical.translation.x - 0.32715) < 1e-6 + assert abs(optical.translation.z - 0.04297) < 1e-6 diff --git a/native/rust/README.md b/native/rust/README.md index df161fe78e..4c31fc1d81 100644 --- a/native/rust/README.md +++ b/native/rust/README.md @@ -153,6 +153,8 @@ let at_scan = self.tf.lookup("odom", &cloud.header.frame_id) Either way the result is `None` when no path connects the frames or no sample falls within the tolerance. It exposes its `nalgebra` parts via `translation()` (a `Vector3`) and `rotation()` (a `UnitQuaternion`). Lookups are nearest-in-time, not interpolated. +A result composed over several edges carries the stamp of the stalest edge on the path, so `ts` reads as the age of the whole answer rather than of one hop. A chain mixing a live edge with a static one is only as fresh as the live edge, in either direction. + `publish` sends transforms onto the same `tf` topic, the counterpart to Python's `tf.publish()`. Published transforms also feed the module's own graph, so a lookup right after the publish sees them. Build the isometry from `dimos_module::nalgebra`, re-exported so the version matches the SDK's types: ```rust diff --git a/native/rust/dimos-module/src/tf.rs b/native/rust/dimos-module/src/tf.rs index aa15582034..15ffcd5a26 100644 --- a/native/rust/dimos-module/src/tf.rs +++ b/native/rust/dimos-module/src/tf.rs @@ -166,7 +166,6 @@ impl TBuffer { } } - // One transform for this edge: the latest sample, or the one nearest time. fn sample( &self, parent: &str, @@ -222,7 +221,6 @@ impl MultiTBuffer { out } - // A single forward or reverse edge. Reverse returns the inverse. fn edge( &self, parent: &str, @@ -260,12 +258,15 @@ impl MultiTBuffer { return Some(direct); } let path = self.bfs(parent, child, time, tolerance)?; + // A composition is only as fresh as its stalest edge. + let oldest = path.iter().map(|t| t.ts).fold(f64::INFINITY, f64::min); let mut steps = path.into_iter(); let first = steps.next()?; - Some(steps.fold(first, |acc, step| acc.compose(&step))) + let mut composed = steps.fold(first, |acc, step| acc.compose(&step)); + composed.ts = oldest; + Some(composed) } - // Shortest path of edges from parent to child. fn bfs( &self, parent: &str, @@ -644,7 +645,6 @@ mod tests { assert_eq!(inv.child, "base_link"); } - // A 30-degree yaw then a pure translation. #[test] fn composes_ros_example_chain() { let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); @@ -665,7 +665,6 @@ mod tests { assert_eq!(t.child, "end_effector"); } - // world->robot->sensor multi-hop composition. #[test] fn composes_multi_hop_chain() { let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); @@ -677,6 +676,16 @@ mod tests { assert!((t.translation().z - 3.2).abs() < 1e-3); } + #[test] + fn composed_stamp_is_the_stalest_edge() { + let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); + h.add("world", "robot", 700.0, (1.0, 0.0, 0.0), 0.0); + h.add("robot", "sensor", 1000.0, (0.5, 0.0, 0.0), 0.0); + assert_eq!(tf.get_latest("world", "sensor").unwrap().ts, 700.0); + // Both directions, so the answer does not depend on which end is queried. + assert_eq!(tf.get_latest("sensor", "world").unwrap().ts, 700.0); + } + #[test] fn missing_path_returns_none() { let (tf, h) = tf_with(DEFAULT_TF_WINDOW_SECS); @@ -980,7 +989,6 @@ mod tests { buf.add(1.0, Isometry3::identity()); buf.add(2.0, Isometry3::identity()); buf.add(10.0, Isometry3::identity()); - // The window is [5.0, 10.0]. The 1.0 and 2.0 samples are dropped. assert_eq!(buf.samples.len(), 1); assert!((buf.last().unwrap().ts - 10.0).abs() < 1e-9); } @@ -1079,7 +1087,6 @@ mod tests { assert!((t.ts - 7.0).abs() < 1e-9); } - // Publish on one handle, dispatch the wire bytes into another graph. #[tokio::test] async fn publish_round_trips_through_route() { let (tf_out, mut rx, _h) = tf_with_publish(DEFAULT_TF_WINDOW_SECS);