diff --git a/src/lerobot/navigation/README.md b/src/lerobot/navigation/README.md index e772ed97f..623cdba7f 100644 --- a/src/lerobot/navigation/README.md +++ b/src/lerobot/navigation/README.md @@ -70,6 +70,13 @@ python -m lerobot.navigation.dog_cli --dry-run --command "go to the couch" # On a real Unitree Go2 (DDS + LingBot-Map + SigLIP2 on the GPU host): python -m lerobot.navigation.dog_cli --live --network-interface enp2s0 --device cuda + +# Add --viz to stream the map into a Rerun viewer as it builds/updates +# (pip install 'lerobot[viz]'). --color-mode recency shows observation age; +# carved voxels (moved/removed objects) flash red then vanish. +python -m lerobot.navigation.dog_cli --dry-run --viz +python -m lerobot.navigation.dog_cli --map-only --viz --color-mode recency \ + --network-interface enp2s0 --device cuda ``` Idle (no prompt) ⇒ autonomous exploration; a typed object name ⇒ navigate diff --git a/src/lerobot/navigation/dog_cli.py b/src/lerobot/navigation/dog_cli.py index 3cda75e0e..67bcdf9d7 100644 --- a/src/lerobot/navigation/dog_cli.py +++ b/src/lerobot/navigation/dog_cli.py @@ -65,10 +65,23 @@ class DogController: skills: SpatialSkills, agent: DeterministicAgent | None = None, parser: HardcodedTaskParser | None = None, + viz=None, ) -> None: self.skills = skills self.agent = agent or DeterministicAgent(skills) self.parser = parser or HardcodedTaskParser() + self.viz = viz # optional MapVisualizer + + def refresh_viz(self, target_xyz=None, path_xyz=None) -> None: + """Log the current map, occupancy, robot pose (+ optional target/path).""" + if self.viz is None: + return + self.viz.log_map(self.skills.voxel_map.snapshot()) + self.viz.log_occupancy(self.skills.occupancy()) + self.viz.log_robot(self.skills.base.pose()) + self.viz.log_target(target_xyz) + if path_xyz is not None: + self.viz.log_path(path_xyz) def handle_prompt(self, text: str) -> AgentResult: """Query the map and navigate to the target (exploring if needed).""" @@ -79,6 +92,8 @@ class DogController: LOG.info(" reached %r at %s (conf %.3f)", tr.target, tr.final_xyz, tr.confidence) else: LOG.info(" did not reach %r: %s (conf %.3f)", tr.target, tr.reason, tr.confidence) + last = result.target_results[-1] if result.target_results else None + self.refresh_viz(target_xyz=last.final_xyz if last and last.reached else None) return result def report_location(self, text: str): @@ -92,6 +107,8 @@ class DogController: LOG.info(" %r is at %s (conf %.3f, %d voxels)", text, loc.xyz, loc.confidence, loc.n_voxels) else: LOG.info(" %r not found yet (conf %.3f) — map more of the area", text, loc.confidence) + if self.viz is not None: + self.viz.log_target(loc.xyz if loc.found else None) return loc def idle_tick(self) -> ExploreResult: @@ -102,13 +119,14 @@ class DogController: self.skills.goto(ex.target_xyz) else: LOG.debug("idle: no frontier to explore (%s)", ex.reason) + self.refresh_viz() return ex def stop(self) -> None: self.skills.base.stop() -def _build_dry_run() -> DogController: +def _build_dry_run(viz=None) -> DogController: """Wire the controller against the synthetic kitchen scene.""" from lerobot.navigation.base_controller import StubBaseController from lerobot.navigation.sim import kitchen_scene @@ -131,7 +149,9 @@ def _build_dry_run() -> DogController: agent = DeterministicAgent(skills, AgentConfig(max_explore_iters=4)) objs = ", ".join(o.name for o in scene.objects) LOG.info("dry-run kitchen scene ready — try one of: %s", objs) - return DogController(skills, agent) + controller = DogController(skills, agent, viz=viz) + controller.refresh_viz() # show the prebuilt map immediately + return controller class LiveMapper: @@ -151,7 +171,7 @@ class LiveMapper: until the first tick. """ - def __init__(self, robot, base, geometry, siglip, voxel_map, pcfg=None) -> None: + def __init__(self, robot, base, geometry, siglip, voxel_map, pcfg=None, viz=None) -> None: self.robot = robot self.base = base # RobotBaseController (unwrapped) for feed_observation/pose self.safe = None # optional SafeBaseController for the watchdog @@ -159,6 +179,7 @@ class LiveMapper: self.siglip = siglip self.voxel_map = voxel_map self.pcfg = pcfg + self.viz = viz # optional MapVisualizer self._frame = 0 def tick(self, t_sec: float) -> None: @@ -201,10 +222,15 @@ class LiveMapper: pose=pose, feat_map=feat_map, ) - integrate_keyframe(self.voxel_map, ctx, self.pcfg or PipelineConfig()) + carve, _ = integrate_keyframe(self.voxel_map, ctx, self.pcfg or PipelineConfig()) self._frame += 1 if self.safe is not None: self.safe.feed_watchdog() + if self.viz is not None: + self.viz.set_time(t_sec) + self.viz.log_map(self.voxel_map.snapshot(), now=t_sec) + self.viz.log_removed(carve.removed_xyz) # dynamic: carved voxels flashed red + self.viz.log_robot(pose) def _build_live( @@ -213,6 +239,7 @@ def _build_live( camera_hfov_deg: float = 90.0, max_lin_speed: float = 0.4, max_yaw_rate: float = 0.8, + viz=None, ) -> tuple[DogController, LiveMapper]: """Wire the controller + live mapper against a real Unitree Go2. @@ -253,8 +280,8 @@ def _build_live( pcfg = PipelineConfig(focal_px=focal_px) skills = SpatialSkills(voxel_map, safe, siglip, SkillsConfig(cell_size=0.05)) - controller = DogController(skills, DeterministicAgent(skills, AgentConfig())) - mapper = LiveMapper(robot, inner, geometry, siglip, voxel_map, pcfg=pcfg) + controller = DogController(skills, DeterministicAgent(skills, AgentConfig()), viz=viz) + mapper = LiveMapper(robot, inner, geometry, siglip, voxel_map, pcfg=pcfg, viz=viz) mapper.safe = safe LOG.info( "live stack wired (iface=%s, device=%s, focal=%.1fpx, vmax=%.2f m/s) — connect the dog and run", @@ -382,6 +409,18 @@ def main(argv: list[str] | None = None) -> int: ) ap.add_argument("--max-lin-speed", type=float, default=0.4, help="Body linear speed cap (m/s).") ap.add_argument("--max-yaw-rate", type=float, default=0.8, help="Yaw-rate cap (rad/s).") + ap.add_argument( + "--viz", + action="store_true", + help="Open a Rerun viewer and stream the map live as it builds/updates " + "(needs `pip install 'lerobot[viz]'`).", + ) + ap.add_argument( + "--color-mode", + default="rgb", + choices=["rgb", "recency"], + help="Voxel coloring in the viewer: rgb, or recency (recent=cyan, old=red).", + ) ap.add_argument("--command", default=None, help="Run a single command non-interactively, then exit.") ap.add_argument("--log-level", default="INFO", choices=["DEBUG", "INFO", "WARNING"]) args = ap.parse_args(argv) @@ -390,6 +429,12 @@ def main(argv: list[str] | None = None) -> int: level=getattr(logging, args.log_level), format="%(levelname)-7s %(name)s: %(message)s" ) + viz = None + if args.viz: + from lerobot.navigation.viz import MapVisualizer + + viz = MapVisualizer(color_mode=args.color_mode) + if args.live or args.map_only: controller, mapper = _build_live( args.network_interface, @@ -397,13 +442,14 @@ def main(argv: list[str] | None = None) -> int: camera_hfov_deg=args.camera_hfov_deg, max_lin_speed=args.max_lin_speed, max_yaw_rate=args.max_yaw_rate, + viz=viz, ) return run_live_repl(controller, mapper, map_only=args.map_only) if not args.dry_run: raise SystemExit("Choose a mode: --dry-run (synthetic scene) or --live (real Unitree Go2).") - controller = _build_dry_run() + controller = _build_dry_run(viz=viz) if args.command is not None: result = controller.handle_prompt(args.command) controller.stop() diff --git a/src/lerobot/navigation/viz.py b/src/lerobot/navigation/viz.py new file mode 100644 index 000000000..5599bd42a --- /dev/null +++ b/src/lerobot/navigation/viz.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# 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. + +"""Live Rerun visualization of the spatial-memory map. + +Shows the voxel map as it is built and updated: the point cloud (colored +by RGB or by observation recency), the robot pose, the top-down occupancy +grid, the planned path, query hits, and — the dynamic part — voxels that +were carved out this keyframe (moved/removed objects), flashed in red. + +Because the full current voxel snapshot is re-logged under one entity path +each keyframe, carved voxels simply disappear from the cloud on the next +frame, so DynaMem-style dynamic updates are visible in real time. Rerun +(`rerun-sdk`) is imported lazily — ``pip install 'lerobot[viz]'`` — so the +rest of the stack never depends on it. + +Requires ``rerun-sdk``; install with ``pip install 'lerobot[viz]'``. +""" + +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np + +LOG = logging.getLogger(__name__) + +_TIMELINE = "t" + + +def _recency_colors(last_time: np.ndarray, now: float, horizon_s: float = 30.0) -> np.ndarray: + """Map per-voxel age to an (M, 3) uint8 color: recent = cyan, old = red.""" + age = np.clip((now - last_time.astype(np.float64)) / max(horizon_s, 1e-6), 0.0, 1.0) + r = (60 + 195 * age).astype(np.uint8) + g = (200 * (1.0 - age)).astype(np.uint8) + b = (200 * (1.0 - age) + 40).astype(np.uint8) + return np.stack([r, g, b], axis=-1) + + +class MapVisualizer: + """Rerun visualizer for the navigation map. Lazily starts the viewer.""" + + def __init__( + self, + app_id: str = "dog-nav", + spawn: bool = True, + color_mode: str = "rgb", + voxel_radius: float = 0.03, + ) -> None: + self.app_id = app_id + self.spawn = spawn + self.color_mode = color_mode # "rgb" | "recency" + self.voxel_radius = float(voxel_radius) + self._rr: Any | None = None + + def _ensure_started(self): + if self._rr is not None: + return self._rr + import rerun as rr + + rr.init(self.app_id, spawn=self.spawn) + # OpenCV world convention: X right, Y down, Z forward (RDF). + rr.log("world", rr.ViewCoordinates.RDF, static=True) + self._rr = rr + return rr + + def set_time(self, t_sec: float) -> None: + rr = self._ensure_started() + rr.set_time(_TIMELINE, timestamp=float(t_sec)) + + # ----- map + dynamics -------------------------------------------------- + + def log_map(self, snapshot, now: float | None = None) -> None: + """Log the current voxel cloud. Re-logging replaces the previous + frame, so carved voxels vanish — that's the dynamic update.""" + rr = self._ensure_started() + xyz = snapshot.xyz + if xyz.size == 0: + rr.log("world/map", rr.Clear(recursive=False)) + return + if self.color_mode == "recency" and now is not None: + colors = _recency_colors(snapshot.last_time, now) + else: + colors = snapshot.rgb + rr.log( + "world/map", + rr.Points3D(xyz.astype(np.float32), colors=colors, radii=self.voxel_radius), + ) + + def log_removed(self, xyz: np.ndarray, radius: float | None = None) -> None: + """Flash this keyframe's carved (removed) voxels in red — the + moved/vanished objects DynaMem carves out.""" + rr = self._ensure_started() + r = radius if radius is not None else self.voxel_radius * 1.6 + if xyz is None or len(xyz) == 0: + rr.log("world/carved", rr.Clear(recursive=False)) + return + red = np.tile(np.array([[230, 40, 40]], dtype=np.uint8), (len(xyz), 1)) + rr.log("world/carved", rr.Points3D(xyz.astype(np.float32), colors=red, radii=r)) + + # ----- robot + planning ------------------------------------------------ + + def log_robot(self, pose: np.ndarray, body_radius: float = 0.15) -> None: + rr = self._ensure_started() + rr.log( + "world/robot", + rr.Transform3D( + translation=pose[:3, 3].astype(np.float32), mat3x3=pose[:3, :3].astype(np.float32) + ), + ) + rr.log( + "world/robot/body", + rr.Points3D( + np.zeros((1, 3), dtype=np.float32), + colors=np.array([[60, 140, 255]], dtype=np.uint8), + radii=body_radius, + ), + ) + + def log_occupancy(self, grid) -> None: + rr = self._ensure_started() + from lerobot.navigation.occupancy import occupancy_to_rgb + + rr.log("plan/occupancy", rr.Image(occupancy_to_rgb(grid))) + + def log_path(self, path_xyz: list[tuple[float, float, float]], radius: float = 0.02) -> None: + rr = self._ensure_started() + if not path_xyz: + rr.log("world/path", rr.Clear(recursive=False)) + return + pts = np.asarray(path_xyz, dtype=np.float32) + rr.log("world/path", rr.LineStrips3D([pts], radii=radius, colors=[[255, 210, 60]])) + + def log_target(self, xyz: tuple[float, float, float] | None) -> None: + """Highlight the located target (green) or clear it when not found.""" + rr = self._ensure_started() + if xyz is None: + rr.log("world/target", rr.Clear(recursive=False)) + return + rr.log( + "world/target", + rr.Points3D( + np.asarray([xyz], dtype=np.float32), + colors=np.array([[40, 230, 90]], dtype=np.uint8), + radii=0.12, + ), + ) diff --git a/tests/navigation/test_viz.py b/tests/navigation/test_viz.py new file mode 100644 index 000000000..3b720db95 --- /dev/null +++ b/tests/navigation/test_viz.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# 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. + +"""Headless tests for the Rerun map visualizer. + +``spawn=False`` buffers to an in-memory recording, so these run without a +display and skip cleanly when rerun-sdk isn't installed. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("rerun", reason="rerun-sdk not installed (pip install 'lerobot[viz]')") + +from lerobot.navigation.dog_cli import _build_dry_run # noqa: E402 +from lerobot.navigation.viz import MapVisualizer, _recency_colors # noqa: E402 +from lerobot.navigation.voxel_map import VoxelMap # noqa: E402 + + +def test_recency_colors_recent_vs_old(): + last = np.array([10.0, 0.0]) # one recent, one old + colors = _recency_colors(last, now=10.0, horizon_s=10.0) + assert colors.shape == (2, 3) + assert colors.dtype == np.uint8 + # Recent voxel (age 0) is more green/cyan; old (age 1) is more red. + assert colors[0, 1] > colors[1, 1] # green channel higher for recent + assert colors[1, 0] > colors[0, 0] # red channel higher for old + + +def _viz() -> MapVisualizer: + return MapVisualizer(app_id="test-dog-nav", spawn=False) + + +def test_log_map_and_dynamics_do_not_raise(): + viz = _viz() + vm = VoxelMap(voxel_size=0.1) + pts = np.array([[[0.0, 0.0, 1.0]], [[0.1, 0.0, 1.0]]], dtype=np.float64) + rgb = np.full((2, 1, 3), 200, dtype=np.uint8) + conf = np.ones((2, 1), dtype=np.float32) + vm.add(pts, rgb, conf, frame=0, t=0.0) + + viz.set_time(0.0) + viz.log_map(vm.snapshot(), now=0.0) + viz.log_removed(np.array([[0.5, 0.0, 1.0]], dtype=np.float32)) # a carved voxel + viz.log_robot(np.eye(4)) + viz.log_path([(0.0, 0.0, 0.0), (0.5, 0.0, 0.5)]) + viz.log_target((0.1, 0.0, 1.0)) + + +def test_log_empty_map_clears_cleanly(): + viz = _viz() + viz.set_time(1.0) + viz.log_map(VoxelMap().snapshot()) # empty + viz.log_removed(np.zeros((0, 3), dtype=np.float32)) + viz.log_target(None) + + +def test_recency_color_mode(): + viz = MapVisualizer(app_id="test-recency", spawn=False, color_mode="recency") + vm = VoxelMap(voxel_size=0.1) + pts = np.array([[[0.0, 0.0, 1.0]]], dtype=np.float64) + vm.add(pts, np.full((1, 1, 3), 100, np.uint8), np.ones((1, 1), np.float32), frame=0, t=5.0) + viz.set_time(5.0) + viz.log_map(vm.snapshot(), now=5.0) + + +def test_dry_run_controller_with_viz_navigates(): + """End-to-end: the dry-run stack with a headless visualizer still reaches + the couch, and every viz call along the way is exercised.""" + viz = _viz() + controller = _build_dry_run(viz=viz) + result = controller.handle_prompt("couch") + assert result.fully_successful