feat: ray tracer and 3d planner mem2 transforms #2368
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,65 @@ | ||
| # Copyright 2026 Dimensional Inc. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import numpy as np | ||
| from numpy.typing import NDArray | ||
| import pytest | ||
|
|
||
| pytest.importorskip("dimos_voxel_ray_tracing") | ||
|
Contributor
There was a problem hiding this comment. Noooooo, don't skip the tests 😭
Collaborator
There was a problem hiding this comment. Think this'll be necessary if it's an optional dependency (also see the other thread about merging into the regular test job).
Contributor
There was a problem hiding this comment. But if the test is never executed then why have it at all?
Collaborator
Author
There was a problem hiding this comment. The tests do get ran in the
Collaborator
Author
There was a problem hiding this comment. They just get skipped in all the other jobs that don't build the binaries
Contributor
There was a problem hiding this comment. But they do get skipped locally by default, right? That is, if I run this locally, do they get executed:
Collaborator
There was a problem hiding this comment. Because you need to build the dependency, yeah.
Collaborator
Author
There was a problem hiding this comment. Yeah they'll get skipped by default. But part of the dev process is building the binaries, so I think that's alright at least for now? Once we do what Sam mentioned, then I think we should be able to remove this |
||
|
|
||
| from dimos.mapping.ray_tracing.transformer import RayTraceMap | ||
| from dimos.memory2.type.observation import Observation | ||
| from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 | ||
|
|
||
|
|
||
| def _obs( | ||
| points: NDArray[np.float32], ts: float, pose: tuple[float, float, float] | ||
| ) -> Observation[PointCloud2]: | ||
| return Observation( | ||
| id=0, | ||
| ts=ts, | ||
| pose=pose, | ||
| _data=PointCloud2.from_numpy(points), | ||
| ) | ||
|
|
||
|
|
||
| def _cube(n: int = 100) -> NDArray[np.float32]: | ||
| rng = np.random.default_rng(0) | ||
| return rng.random((n, 3)).astype(np.float32) | ||
|
|
||
|
|
||
| def test_emit_every_n_yields_on_cadence_and_flushes_remainder() -> None: | ||
| points = _cube() | ||
| obs = [_obs(points, ts=float(i), pose=(0.0, 0.0, 0.0)) for i in range(7)] | ||
|
|
||
| results = list(RayTraceMap(emit_every=3)(iter(obs))) | ||
|
|
||
| assert [r.tags["frame_count"] for r in results] == [3, 6, 7] | ||
|
|
||
|
|
||
| def test_poseless_obs_are_skipped() -> None: | ||
| points = _cube() | ||
| poseless = Observation(id=1, ts=0.0, pose=None, _data=PointCloud2.from_numpy(points)) | ||
| posed = _obs(points, ts=1.0, pose=(0.0, 0.0, 0.0)) | ||
|
|
||
| results = list(RayTraceMap()(iter([poseless, posed]))) | ||
|
|
||
| assert [r.tags["frame_count"] for r in results] == [1] | ||
|
|
||
|
|
||
| def test_negative_emit_every_is_rejected() -> None: | ||
| with pytest.raises(ValueError): | ||
| RayTraceMap(emit_every=-1) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,87 @@ | ||
| # 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 TYPE_CHECKING, Any | ||
|
|
||
| import open3d as o3d # type: ignore[import-untyped] | ||
| import open3d.core as o3c # type: ignore[import-untyped] | ||
|
|
||
| from dimos.mapping.ray_tracing.voxel_map import VoxelRayMapper | ||
| from dimos.memory2.transform import Transformer | ||
| from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 | ||
|
|
||
| if TYPE_CHECKING: | ||
| from collections.abc import Iterator | ||
|
|
||
| from dimos.memory2.type.observation import Observation | ||
|
|
||
|
|
||
| class RayTraceMap(Transformer[PointCloud2, PointCloud2]): | ||
| """Accumulate world-frame lidar into a voxel map with raycast clearing.""" | ||
|
|
||
| def __init__( | ||
| self, | ||
| *, | ||
| voxel_size: float = 0.1, | ||
| max_range: float = 30.0, | ||
| emit_every: int = 1, | ||
| **mapper_kwargs: Any, | ||
|
aclauer marked this conversation as resolved.
|
||
| ) -> None: | ||
| if emit_every < 0: | ||
| raise ValueError(f"emit_every must be >= 0, got {emit_every}") | ||
| self.voxel_size = voxel_size | ||
| self.max_range = max_range | ||
| self.emit_every = emit_every | ||
| self._mapper_kwargs = mapper_kwargs | ||
|
|
||
| def _make_obs( | ||
| self, | ||
| mapper: VoxelRayMapper, | ||
| last_obs: Observation[PointCloud2], | ||
| count: int, | ||
| ) -> Observation[PointCloud2]: | ||
| positions = mapper.global_map() | ||
| pcd = o3d.t.geometry.PointCloud() | ||
| pcd.point["positions"] = o3c.Tensor.from_numpy(positions) | ||
| cloud = PointCloud2(pointcloud=pcd, frame_id="world", ts=last_obs.ts) | ||
| return last_obs.derive( | ||
| data=cloud, | ||
| tags={**last_obs.tags, "frame_count": count}, | ||
| ) | ||
|
|
||
| def __call__( | ||
| self, | ||
| upstream: Iterator[Observation[PointCloud2]], | ||
| ) -> Iterator[Observation[PointCloud2]]: | ||
| mapper = VoxelRayMapper( | ||
| voxel_size=self.voxel_size, max_range=self.max_range, **self._mapper_kwargs | ||
| ) | ||
| last_obs: Observation[PointCloud2] | None = None | ||
| count = 0 | ||
|
|
||
| for obs in upstream: | ||
| if obs.pose_tuple is None: | ||
| continue | ||
| x, y, z, *_ = obs.pose_tuple | ||
| mapper.add_frame(obs.data.points_f32(), (x, y, z)) | ||
| last_obs = obs | ||
| count += 1 | ||
|
|
||
| if self.emit_every > 0 and count % self.emit_every == 0: | ||
| yield self._make_obs(mapper, last_obs, count) | ||
|
|
||
| if last_obs is not None and (self.emit_every == 0 or count % self.emit_every != 0): | ||
| yield self._make_obs(mapper, last_obs, count) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| # Copyright 2026 Dimensional Inc. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import numpy as np | ||
| from numpy.typing import NDArray | ||
| import pytest | ||
|
|
||
| pytest.importorskip("dimos_mls_planner") | ||
|
|
||
| from dimos.memory2.type.observation import Observation | ||
| from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 | ||
| from dimos.navigation.nav_3d.mls_planner.transformer import MLSPlan | ||
|
|
||
|
|
||
| def _obs(points: NDArray[np.float32], pose: tuple[float, float, float]) -> Observation[PointCloud2]: | ||
| return Observation(id=0, ts=0.0, pose=pose, _data=PointCloud2.from_numpy(points)) | ||
|
|
||
|
|
||
| def _flat_floor(half_extent: float = 3.0, spacing: float = 0.1) -> NDArray[np.float32]: | ||
| coords = np.arange(-half_extent, half_extent, spacing, dtype=np.float32) | ||
| xs, ys = np.meshgrid(coords, coords) | ||
| zs = np.zeros_like(xs) | ||
| return np.stack([xs.ravel(), ys.ravel(), zs.ravel()], axis=1) | ||
|
|
||
|
|
||
| def test_flat_floor_yields_populated_path_and_planned_true() -> None: | ||
| obs = _obs(_flat_floor(), pose=(-2.0, -2.0, 1.0)) | ||
|
|
||
| [out] = list(MLSPlan(goal=(2.0, 2.0, 0.0), voxel_size=0.2, robot_height=1.0)(iter([obs]))) | ||
|
|
||
| assert out.tags["planned"] is True | ||
| assert len(out.data.poses) >= 2 | ||
| assert out.tags["voxel_map"] is obs.data | ||
| assert out.tags["nodes"].shape[1] == 3 | ||
| assert out.tags["surface_map"].shape[1] == 3 | ||
|
|
||
|
|
||
| def test_no_route_yields_empty_path_with_planned_false() -> None: | ||
| rng = np.random.default_rng(0) | ||
| obs = _obs(rng.random((50, 3)).astype(np.float32), pose=(0.0, 0.0, 0.0)) | ||
|
|
||
| [out] = list(MLSPlan(goal=(100.0, 100.0, 100.0))(iter([obs]))) | ||
|
|
||
| assert out.tags["planned"] is False | ||
| assert out.data.poses == [] | ||
|
|
||
|
|
||
| def test_poseless_obs_is_skipped_and_following_posed_obs_plans() -> None: | ||
| poseless = Observation(id=1, ts=0.0, pose=None, _data=PointCloud2.from_numpy(_flat_floor())) | ||
| posed = _obs(_flat_floor(), pose=(-2.0, -2.0, 1.0)) | ||
|
|
||
| results = list( | ||
| MLSPlan(goal=(2.0, 2.0, 0.0), voxel_size=0.2, robot_height=1.0)(iter([poseless, posed])) | ||
| ) | ||
|
|
||
| assert len(results) == 1 | ||
| assert results[0].tags["planned"] is True |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,95 @@ | ||
| # 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 TYPE_CHECKING, Any | ||
|
|
||
| from dimos.memory2.transform import Transformer | ||
| from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped | ||
| from dimos.msgs.nav_msgs.Path import Path | ||
| from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 | ||
| from dimos.navigation.nav_3d.mls_planner.mls_planner import MLSPlanner | ||
|
|
||
| if TYPE_CHECKING: | ||
| from collections.abc import Iterator | ||
|
|
||
| import numpy as np | ||
| from numpy.typing import NDArray | ||
|
|
||
| from dimos.memory2.type.observation import Observation | ||
|
|
||
|
|
||
| class MLSPlan(Transformer[PointCloud2, Path]): | ||
| """Plan paths from current pose to a fixed goal over an accumulating voxel map.""" | ||
|
|
||
| def __init__( | ||
| self, | ||
| *, | ||
| goal: tuple[float, float, float], | ||
| voxel_size: float = 0.1, | ||
| robot_height: float = 1.5, | ||
| **planner_kwargs: Any, | ||
| ) -> None: | ||
| self.goal = goal | ||
| self.voxel_size = voxel_size | ||
| self.robot_height = robot_height | ||
| self._planner_kwargs = planner_kwargs | ||
|
|
||
| def _path_from_waypoints(self, waypoints: NDArray[np.float32] | None, ts: float) -> Path: | ||
| poses: list[PoseStamped] = [] | ||
| if waypoints is not None: | ||
| for x, y, z in waypoints: | ||
| poses.append( | ||
| PoseStamped( | ||
| ts=ts, | ||
| frame_id="world", | ||
| position=(float(x), float(y), float(z)), | ||
| orientation=(0.0, 0.0, 0.0, 1.0), | ||
| ) | ||
| ) | ||
| return Path(ts=ts, frame_id="world", poses=poses) | ||
|
|
||
| def __call__( | ||
| self, | ||
| upstream: Iterator[Observation[PointCloud2]], | ||
| ) -> Iterator[Observation[Path]]: | ||
| planner = MLSPlanner( | ||
| voxel_size=self.voxel_size, | ||
| robot_height=self.robot_height, | ||
| **self._planner_kwargs, | ||
| ) | ||
| for obs in upstream: | ||
| if obs.pose_tuple is None: | ||
| continue | ||
| x, y, z, *_ = obs.pose_tuple | ||
| start = (float(x), float(y), float(z) - self.robot_height) | ||
|
|
||
| voxel_map = obs.data | ||
| planner.update_global_map(voxel_map.points_f32()) | ||
| waypoints = planner.plan(start, self.goal) | ||
| path = self._path_from_waypoints(waypoints, obs.ts) | ||
|
|
||
| yield obs.derive( | ||
| data=path, | ||
| tags={ | ||
| **obs.tags, | ||
| "voxel_map": voxel_map, | ||
| "surface_map": planner.surface_map(), | ||
| "nodes": planner.nodes(), | ||
| "node_edges": planner.node_edges(), | ||
| "start": start, | ||
| "planned": waypoints is not None, | ||
| }, | ||
| ) |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.