navigation: live Rerun visualization of the map (+ dynamic carving)

MapVisualizer (lazy rerun-sdk, lerobot[viz]) streams the map as it builds:
voxel cloud (RGB or recency colormap), robot pose, top-down occupancy,
planned path, and the located target. The dynamic part is visible — each
keyframe re-logs the full snapshot so carved voxels vanish, and this
frame's removed voxels flash red under world/carved.

Wired into LiveMapper.tick (per-keyframe map/removed/robot) and
DogController (map/occupancy/robot/target/path after each action), behind
a --viz flag (+ --color-mode rgb|recency) on dog-nav; works in --dry-run,
--map-only, and --live. Headless tests (spawn=False) cover the log paths
and a full dry-run navigate-with-viz. 140 tests in tests/navigation/.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

pyproject.toml

@@ -187,6 +187,11 @@ unitree_g1 = [
     "lerobot[matplotlib-dep]",
     "lerobot[pygame-dep]",
 ]
+# Go2 talks plain DDS from the host — no bridge server, no extra deps beyond
+# the SDK itself (cyclonedds-based, hence Linux-only).
+unitree_go2 = [
+    "unitree_sdk2py>=1.0.1; sys_platform == 'linux'",
+]
 # reachy2-sdk caps grpcio<=1.73.1 and protobuf<=6.32.0; quarantined here so downstream users aren't held back. reachy2-sdk is unlikely to release new versions.
 reachy2 = [
     "reachy2_sdk>=1.0.15,<1.1.0",
@@ -357,6 +362,7 @@ lerobot-edit-dataset="lerobot.scripts.lerobot_edit_dataset:main"
 lerobot-setup-can="lerobot.scripts.lerobot_setup_can:main"
 lerobot-annotate="lerobot.scripts.lerobot_annotate:main"
 lerobot-rollout="lerobot.scripts.lerobot_rollout:main"
+dog-nav="lerobot.navigation.dog_cli:main"
 
 # ---------------- Tool Configurations ----------------
 
@@ -374,11 +380,7 @@ torch = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
 torchvision = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
 
 [tool.setuptools.package-data]
-lerobot = [
-    "envs/*.json",
-    "annotations/steerable_pipeline/prompts/*.txt",
-    "teleoperators/pico_headset/assets/*.npz",
-]
+lerobot = ["envs/*.json", "annotations/steerable_pipeline/prompts/*.txt"]
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/lerobot/navigation/BRINGUP.md

@@ -0,0 +1,117 @@
+# dog-nav on a real Unitree Go2 — bring-up guide
+
+The synthetic scene (`--dry-run`) exists only to test the logic without a
+robot. To run for real you need the dog, the GPU host, and the steps
+below. Bring it up **in stages** — never start with autonomous motion on
+untested hardware.
+
+## 0. Prerequisites
+
+- Unitree Go2 **EDU** (SDK access; the consumer Go2 can't be commanded).
+- A GPU host (your 5090) on the **same network as the dog**. Over
+  Ethernet the dog is on `192.168.123.x`; find your interface with
+  `ip link` (e.g. `enp2s0`).
+- A remote/controller in hand for a hardware e-stop at all times.
+
+## 1. Get the branch onto the 5090
+
+The branch `feat/unitree-go2` is local (not pushed to upstream). Either:
+
+**Option A — your fork:**
+```bash
+# on the mac, one time:
+git remote add fork git@github.com:<you>/lerobot.git
+git push fork feat/unitree-go2
+# on the 5090:
+git clone git@github.com:<you>/lerobot.git && cd lerobot
+git checkout feat/unitree-go2
+```
+
+**Option B — git bundle (no remote needed):**
+```bash
+# on the mac:
+git bundle create go2-nav.bundle origin/main..feat/unitree-go2
+# copy go2-nav.bundle to the 5090, then:
+git clone https://github.com/huggingface/lerobot.git && cd lerobot
+git fetch ../go2-nav.bundle feat/unitree-go2:feat/unitree-go2
+git checkout feat/unitree-go2
+```
+
+## 2. Environment on the 5090
+
+```bash
+uv venv --python 3.12 .venv
+uv pip install -e .                         # lerobot core (torch, etc.)
+uv pip install transformers                 # SigLIP2
+uv pip install unitree_sdk2py               # DDS to the dog (Linux only)
+# LingBot-Map (geometry) — source install:
+pip install -e 'git+https://github.com/robbyant/lingbot-map#egg=lingbot-map'
+```
+
+Smoke-test the code path with no dog:
+```bash
+.venv/bin/python -m lerobot.navigation.dog_cli --dry-run --command "go to the couch"
+```
+
+## 3. Stage 1 — verify DDS + sensors (NO motion)
+
+Confirm the host talks to the dog and reads odometry + camera before
+anything moves:
+```python
+from lerobot.robots.unitree_go2 import UnitreeGo2, UnitreeGo2Config
+r = UnitreeGo2(UnitreeGo2Config(network_interface="enp2s0", stand_on_connect=False))
+r.connect()
+obs = r.get_observation()
+print({k: (v.shape if hasattr(v, "shape") else v) for k, v in obs.items()})
+r.disconnect()
+```
+You want a real `front` image `(720, 1280, 3)` and non-garbage
+`x.pos/y.pos/theta.pos`. If `theta.pos` doesn't change sign the way you
+expect when you turn the dog by hand, tell me — the odometry sign
+conventions may need a tweak for your firmware.
+
+## 4. Stage 2 — teleop (low speed, hand on e-stop)
+
+```bash
+lerobot-teleoperate --robot.type=unitree_go2 \
+    --robot.network_interface=enp2s0 --teleop.type=gamepad
+```
+Confirm forward/left/turn go the right way. This validates
+`send_action`/`SportClient.Move` before the nav loop drives.
+
+## 5. Stage 3 — MAP-ONLY (still no autonomous motion)
+
+Build the map by teleoperating the dog around while the models run.
+Query where things are; the dog never drives itself:
+```bash
+.venv/bin/python -m lerobot.navigation.dog_cli --map-only \
+    --network-interface enp2s0 --device cuda --camera-hfov-deg 90
+# teleop the dog around the room, then type object names:
+#   couch      -> "couch is at (x, y, z) ..."  or  "not found yet"
+```
+Tune `--camera-hfov-deg` to your Go2 front camera so free-space carving
+is correct (a wrong value only hurts dynamic removal, not the map).
+
+## 6. Stage 4 — autonomous nav (open space, low speed, e-stop ready)
+
+Only after 1–3 look right. Start in a clear area:
+```bash
+.venv/bin/python -m lerobot.navigation.dog_cli --live \
+    --network-interface enp2s0 --device cuda \
+    --max-lin-speed 0.3 --max-yaw-rate 0.6
+# empty line -> one exploration step; type an object -> navigate to it.
+```
+`SafeBaseController` clamps speed, refuses moves into obstacle cells, and
+latches an e-stop if keyframes go stale (>2 s). Ctrl-C stops the base.
+
+## Known things to expect / tune on first hardware contact
+
+- **Odometry sign conventions** (`position[0/1]`, `imu_state.rpy[2]`):
+  verified in sim, not yet against live firmware — check in Stage 1.
+- **Camera FOV / focal**: set `--camera-hfov-deg` from your camera.
+- **Gait bob**: pose is planarized (yaw only); pitch/roll wobble is
+  ignored for now. Fine at low speed; a full-SE(3) camera pose is the
+  refinement if the map smears vertically.
+- **Keyframe rate**: SAM2 isn't in this path; the per-tick cost is
+  LingBot-Map + SigLIP2 on the 5090 (~tens of ms each). If ticks lag,
+  drop camera resolution.

src/lerobot/navigation/README.md

@@ -0,0 +1,96 @@
+# `lerobot.navigation` — spatial-memory navigation
+
+Online spatio-semantic mapping (DynaMem-style), A* planning, obstacle
+avoidance and open-vocabulary goto/explore for LeRobot mobile bases.
+Ported from the dyna360 research stack; the physical robot layer lives in
+`lerobot.robots` (e.g. [`unitree_go2`](../robots/unitree_go2)).
+
+## Idea
+
+Drive any LeRobot `Robot` on the standard REP-103 mobile-base contract —
+body-velocity actions `x.vel`/`y.vel`/`theta.vel` and planar odometry
+`x.pos`/`y.pos`/`theta.pos` — from a spatial memory that is built and
+updated online from the robot's camera. With no prompt the base explores
+autonomously; given a text prompt it queries the map and navigates to the
+matching object, or explores to find it if it isn't there (or has moved).
+
+## Architecture
+
+The navigation layer talks to hardware only through LeRobot's own `Robot`
+interface, so it is robot-agnostic and carries no SDK dependency.
+
+```
+BaseController (protocol)         world-frame move()/pose() seam
+├── StubBaseController            kinematic integrator (sim, tests)
+├── RobotBaseController           wraps any Robot; world<->body +
+│                                 odometry<->world frame math
+└── SafeBaseController            velocity clamp, occupancy gate,
+                                  keyframe watchdog, e-stop latch
+```
+
+World frame is OpenCV (x right, y down, z forward); the base moves in the
+XZ plane. `RobotBaseController.feed_observation(obs)` updates pose from
+the observation the navigation loop already fetches (closed-loop
+odometry), avoiding an extra camera read; absent odometry it integrates
+open-loop so sim matches hardware.
+
+## Status (branch `feat/unitree-go2`)
+
+Implemented:
+- `base_controller.py` — the controller seam (protocol, stub, safety
+  wrapper, robot-backed controller + frame math).
+- `voxel_map.py` — 5 cm sparse-hash `VoxelMap`: count-weighted geometry,
+  free-space `carve` (dynamic updates), per-voxel feature + `query`. No
+  point-cloud retention.
+- `occupancy.py` — 3-class top-down grid + A* (no corner-cutting) +
+  obstacle inflation + frontier extraction.
+- `value_map.py` — DynaMem §3.4 recency (V_T) + similarity (V_S)
+  exploration scoring.
+- `features.py` — `SiglipFeatureExtractor` (lazy transformers) +
+  `FeatureExtractor` protocol + `BasisVectorFeatureExtractor` stand-in.
+- `geometry.py` — `GeometryRunner` protocol + `LingBotMapRunner` (lazy) +
+  `FakeGeometryRunner`; `align_trajectory_to_odometry` (Umeyama) anchors
+  the monocular scale to sport-mode odometry.
+- `pipeline.py` — viz-free `integrate_keyframe` (carve → add) +
+  feature upsampling.
+- `skills.py` / `agent.py` — `SpatialSkills` (locate/goto/explore) +
+  `DeterministicAgent` + regex parser.
+- `sim.py` — self-contained synthetic scenes for model-free dry-runs.
+- `dog_cli.py` — the `dog-nav` REPL (the deliverable).
+
+Everything is model/hardware-free-testable (191 tests across the branch).
+The one thing that needs the real dog + GPU models is `--live`.
+
+## Running
+
+```bash
+# Synthetic scene, no robot/camera/models:
+python -m lerobot.navigation.dog_cli --dry-run
+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
+to it, exploring to find it if it isn't mapped yet.
+
+## Target platform
+
+Unitree Go2 EDU, no companion computer: the workstation (single RTX 5090)
+talks DDS straight to the dog; geometry is monocular LingBot-Map from the
+built-in front camera, scale-anchored to sport-mode odometry; the map is
+5 cm voxels. See [`robots/unitree_go2`](../robots/unitree_go2).
+
+## Not yet ported (optional enhancement)
+
+`SegmentVoxelMap` (object-centric per-segment features via SAM 2) is a
+storage/precision optimization over the plain per-voxel features used
+here; the locate/goto/explore stack is fully functional without it.

src/lerobot/navigation/__init__.py

@@ -0,0 +1,118 @@
+#!/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.
+
+"""Spatial-memory navigation for LeRobot mobile bases.
+
+Online spatio-semantic mapping (DynaMem-style), A* planning, obstacle
+avoidance and open-vocabulary goto/explore, driving any LeRobot ``Robot``
+that exposes body-velocity actions and planar odometry. Ported from the
+dyna360 research stack; the physical robot layer lives in
+``lerobot.robots`` (e.g. ``unitree_go2``).
+"""
+
+from .agent import (
+    AgentConfig,
+    AgentResult,
+    DeterministicAgent,
+    HardcodedTaskParser,
+    Task,
+    TaskParser,
+)
+from .base_controller import (
+    BaseController,
+    RobotBaseController,
+    SafeBaseController,
+    StubBaseController,
+    odometry_to_world_pose,
+    world_velocity_to_body,
+)
+from .features import (
+    BasisVectorFeatureExtractor,
+    FeatureExtractor,
+    SiglipFeatureExtractor,
+)
+from .geometry import (
+    FakeGeometryRunner,
+    GeometryOutput,
+    GeometryRunner,
+    LingBotMapRunner,
+    align_trajectory_to_odometry,
+)
+from .occupancy import (
+    NAVIGABLE,
+    OBSTACLE,
+    UNOBSERVED,
+    OccupancyGrid,
+    astar,
+    find_frontier_cells,
+    project_voxel_map_to_grid,
+)
+from .pipeline import KeyframeContext, PipelineConfig, integrate_keyframe
+from .skills import (
+    ExploreResult,
+    GotoResult,
+    LocateResult,
+    SkillsConfig,
+    SpatialSkills,
+)
+from .value_map import ValueMapConfig, ValueMaps, compute_value_maps, pick_best_frontier_cell
+from .voxel_map import CarveResult, QueryResult, VoxelMap, VoxelSnapshot
+
+__all__ = [
+    "NAVIGABLE",
+    "OBSTACLE",
+    "UNOBSERVED",
+    "AgentConfig",
+    "AgentResult",
+    "BaseController",
+    "BasisVectorFeatureExtractor",
+    "CarveResult",
+    "DeterministicAgent",
+    "ExploreResult",
+    "FakeGeometryRunner",
+    "FeatureExtractor",
+    "GeometryOutput",
+    "GeometryRunner",
+    "GotoResult",
+    "HardcodedTaskParser",
+    "KeyframeContext",
+    "LingBotMapRunner",
+    "LocateResult",
+    "OccupancyGrid",
+    "PipelineConfig",
+    "QueryResult",
+    "RobotBaseController",
+    "SafeBaseController",
+    "SiglipFeatureExtractor",
+    "SkillsConfig",
+    "SpatialSkills",
+    "StubBaseController",
+    "Task",
+    "TaskParser",
+    "ValueMapConfig",
+    "ValueMaps",
+    "VoxelMap",
+    "VoxelSnapshot",
+    "align_trajectory_to_odometry",
+    "astar",
+    "compute_value_maps",
+    "integrate_keyframe",
+    "find_frontier_cells",
+    "odometry_to_world_pose",
+    "pick_best_frontier_cell",
+    "project_voxel_map_to_grid",
+    "world_velocity_to_body",
+]

src/lerobot/navigation/agent.py

@@ -0,0 +1,262 @@
+#!/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.
+
+"""Deterministic agent wrapper + language-parser interface.
+
+Ported from the dyna360 research stack. The high-level agent is a thin
+deterministic wrapper, not LLM-driven: explore-vs-go control lives here
+in plain Python. A language model (when wired up) only parses a
+natural-language command into a typed :class:`Task`; the deterministic
+wrapper then executes it. Swapping the parser (regex vs a real LLM) must
+not change the spatial behaviour.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from lerobot.navigation.skills import SpatialSkills
+
+LOG = logging.getLogger(__name__)
+
+
+# ============== task data structures ====================================== #
+
+
+@dataclass(frozen=True)
+class Task:
+    """Parsed command, ready for the deterministic wrapper to execute.
+
+    ``go to X`` yields ``Task(targets=['X'])``; ``go to X then Y`` yields
+    ``Task(targets=['X', 'Y'])``, executed sequentially.
+    """
+
+    targets: list[str]
+    raw: str = ""
+
+
+@dataclass(frozen=True)
+class TargetResult:
+    """Outcome of executing the policy for a single target."""
+
+    target: str
+    reached: bool
+    final_xyz: tuple[float, float, float] | None
+    n_explore_iters: int
+    confidence: float
+    reason: str
+    """'ok' | 'no_path' | 'budget_exhausted' | 'no_frontier' | 'parse_empty'."""
+
+
+@dataclass(frozen=True)
+class AgentResult:
+    """Outcome of executing a full Task (one or more sequential targets)."""
+
+    task: Task
+    target_results: list[TargetResult] = field(default_factory=list)
+
+    @property
+    def fully_successful(self) -> bool:
+        return bool(self.target_results) and all(r.reached for r in self.target_results)
+
+
+# ============== language parser ========================================== #
+
+
+@runtime_checkable
+class TaskParser(Protocol):
+    """Anything that turns a free-text command into a :class:`Task`."""
+
+    def parse(self, command: str) -> Task: ...
+
+
+class HardcodedTaskParser:
+    """Regex-only parser — fast, dependency-free, good enough to validate
+    the deterministic policy without loading a language model.
+
+    Handles ``go to (the) X`` / ``find (the) X`` → single target, ``go to
+    X then Y`` → multi-step, and falls back to "the whole command is the
+    target" if no pattern matches.
+    """
+
+    _SINGLE_PATTERNS = (
+        re.compile(
+            r"^\s*(?:go to|navigate to|find|locate|look for)\s+(?:the\s+)?(.+?)\s*$",
+            re.IGNORECASE,
+        ),
+    )
+    _SPLIT_PATTERN = re.compile(r"\s+(?:then|and then)\s+|\s*,\s*", re.IGNORECASE)
+
+    def parse(self, command: str) -> Task:
+        raw = command.strip()
+        if not raw:
+            return Task(targets=[], raw=raw)
+
+        parts = self._SPLIT_PATTERN.split(raw)
+        targets: list[str] = []
+        for part in parts:
+            t = self._extract_target(part)
+            if t:
+                targets.append(t)
+        return Task(targets=targets, raw=raw)
+
+    def _extract_target(self, text: str) -> str:
+        text = text.strip().rstrip(".?!")
+        for p in self._SINGLE_PATTERNS:
+            m = p.match(text)
+            if m:
+                return m.group(1).strip()
+        prefix = re.match(r"^\s*(?:the\s+)?(.+)$", text, re.IGNORECASE)
+        if prefix:
+            return prefix.group(1).strip()
+        return text
+
+
+# ============== deterministic agent ====================================== #
+
+
+@dataclass(frozen=True)
+class AgentConfig:
+    """Agent policy knobs."""
+
+    max_explore_iters: int = 5
+    """How many ``explore → relocate`` loops before giving up on a target."""
+
+    explore_step_uses_goto: bool = True
+    """Drive to the explore frontier via closed-loop ``goto``. False
+    teleports instead (fast offline eval)."""
+
+
+class DeterministicAgent:
+    """Executes a :class:`Task` via a fixed policy.
+
+    For each target: locate; if found, goto and done; else explore(query),
+    goto the frontier, and relocate — up to ``max_explore_iters``, then give
+    up. The control flow is plain Python; no LLM in the loop.
+    """
+
+    def __init__(self, skills: SpatialSkills, cfg: AgentConfig | None = None) -> None:
+        self.skills = skills
+        self.cfg = cfg or AgentConfig()
+
+    def execute(self, task: Task) -> AgentResult:
+        out: list[TargetResult] = []
+        for target in task.targets:
+            out.append(self._execute_target(target))
+            if not out[-1].reached:
+                # Don't auto-skip after a failed multi-step leg; bail so the
+                # caller sees the failure clearly.
+                break
+        return AgentResult(task=task, target_results=out)
+
+    def execute_command(self, command: str, parser: TaskParser) -> AgentResult:
+        """Parse a free-text command, then execute."""
+        task = parser.parse(command)
+        if not task.targets:
+            return AgentResult(
+                task=task,
+                target_results=[
+                    TargetResult(
+                        target="",
+                        reached=False,
+                        final_xyz=None,
+                        n_explore_iters=0,
+                        confidence=-1.0,
+                        reason="parse_empty",
+                    )
+                ],
+            )
+        return self.execute(task)
+
+    # ----- single-target inner loop ----------------------------------------
+
+    def _execute_target(self, target: str) -> TargetResult:
+        last_conf = -1.0
+        for it in range(self.cfg.max_explore_iters + 1):
+            loc = self.skills.locate(target)
+            last_conf = loc.confidence
+            if loc.found and loc.xyz is not None:
+                LOG.info(
+                    "agent: locate(%r) found at %s (conf %.3f); goto",
+                    target,
+                    loc.xyz,
+                    loc.confidence,
+                )
+                gr = self.skills.goto(loc.xyz)
+                return TargetResult(
+                    target=target,
+                    reached=gr.reached,
+                    final_xyz=gr.final_xyz,
+                    n_explore_iters=it,
+                    confidence=loc.confidence,
+                    reason="ok" if gr.reached else gr.reason,
+                )
+
+            if it >= self.cfg.max_explore_iters:
+                LOG.info(
+                    "agent: locate(%r) NOT_FOUND (conf %.3f) and explore budget exhausted",
+                    target,
+                    loc.confidence,
+                )
+                return TargetResult(
+                    target=target,
+                    reached=False,
+                    final_xyz=None,
+                    n_explore_iters=it,
+                    confidence=loc.confidence,
+                    reason="budget_exhausted",
+                )
+
+            # NOT_FOUND → explore once, then loop and re-locate.
+            LOG.info(
+                "agent: locate(%r) NOT_FOUND (conf %.3f) → explore iter %d",
+                target,
+                loc.confidence,
+                it + 1,
+            )
+            ex = self.skills.explore(query=target)
+            if not ex.found_frontier or ex.target_xyz is None:
+                return TargetResult(
+                    target=target,
+                    reached=False,
+                    final_xyz=None,
+                    n_explore_iters=it,
+                    confidence=loc.confidence,
+                    reason="no_frontier",
+                )
+            if self.cfg.explore_step_uses_goto:
+                self.skills.goto(ex.target_xyz)
+            else:
+                # Teleport for offline-eval speed.
+                self.skills.base.move(0.0, 0.0, dt=0.0)
+                pose = self.skills.base.pose()
+                pose[0, 3] = ex.target_xyz[0]
+                pose[2, 3] = ex.target_xyz[2]
+                if hasattr(self.skills.base, "_pose"):
+                    self.skills.base._pose = pose  # noqa: SLF001
+
+        return TargetResult(
+            target=target,
+            reached=False,
+            final_xyz=None,
+            n_explore_iters=self.cfg.max_explore_iters,
+            confidence=last_conf,
+            reason="budget_exhausted",
+        )

src/lerobot/navigation/base_controller.py

@@ -0,0 +1,389 @@
+#!/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.
+
+"""Base controller for spatial-memory navigation.
+
+The navigation/skills layer commands motion in a single **world frame**
+(OpenCV convention: x right, y down, z forward — the base lives in the XZ
+plane, y is gravity) and reads back an SE(3) pose. :class:`BaseController`
+is that seam. Three implementations:
+
+  - :class:`StubBaseController` — kinematic integrator, no hardware; sim +
+    unit tests.
+  - :class:`RobotBaseController` — drives any LeRobot :class:`Robot` whose
+    action space is body-frame velocities ``x.vel`` (forward, m/s),
+    ``y.vel`` (left, m/s), ``theta.vel`` (CCW yaw, rad/s) and whose
+    observation carries planar odometry ``x.pos``/``y.pos``/``theta.pos``
+    (REP-103: x forward, y left, yaw CCW). The Unitree Go2 satisfies this
+    out of the box; so would a LeKiwi base.
+  - :class:`SafeBaseController` — wraps any of the above with velocity
+    clamping, an optional occupancy gate, a keyframe watchdog and an
+    e-stop latch.
+
+All frame conversions between the world frame and a robot's body/odometry
+frame live in :func:`world_velocity_to_body` and
+:func:`odometry_to_world_pose`; nothing else needs to know the mapping.
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+import time
+from abc import abstractmethod
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from lerobot.robots import Robot
+
+logger = logging.getLogger(__name__)
+
+
+# --------------------------------------------------------------------- #
+# BaseController protocol
+# --------------------------------------------------------------------- #
+
+
+@runtime_checkable
+class BaseController(Protocol):
+    """Mobile-base interface used by the navigation/skills layer.
+
+    Velocities are in **world** frame XZ (m/s); ``yaw_rate`` is rad/s
+    about the world's −Y axis (turning around the up vector). ``pose``
+    is 4×4 SE(3) camera-to-world (OpenCV).
+    """
+
+    @abstractmethod
+    def move(self, vx: float, vz: float, yaw_rate: float = 0.0, dt: float = 0.05) -> None: ...
+
+    @abstractmethod
+    def stop(self) -> None: ...
+
+    @abstractmethod
+    def pose(self) -> np.ndarray: ...
+
+    @abstractmethod
+    def position(self) -> tuple[float, float, float]: ...
+
+
+# --------------------------------------------------------------------- #
+# Frame math (pure functions)
+# --------------------------------------------------------------------- #
+
+
+def world_velocity_to_body(
+    vx_world: float,
+    vz_world: float,
+    yaw_rate_rad_s: float,
+    heading_rad: float,
+) -> tuple[float, float, float]:
+    """World-frame velocity → body-frame ``(x.vel, y.vel, theta.vel)``.
+
+    Returns ``(vx_forward, vy_left, vyaw)`` in m/s, m/s, rad/s — the
+    action a REP-103 base expects. At heading ``h`` the body axes in the
+    world XZ plane are forward = (sin h, cos h), left = (−cos h, sin h)
+    (left = up × forward, up = −y). The navigation world's positive yaw
+    is clockwise about the up vector; a REP-103 base's ``theta.vel`` is
+    counter-clockwise, hence the sign flip.
+    """
+    s, c = math.sin(heading_rad), math.cos(heading_rad)
+    vx_fwd = vx_world * s + vz_world * c
+    vy_left = -vx_world * c + vz_world * s
+    return vx_fwd, vy_left, -yaw_rate_rad_s
+
+
+def odometry_to_world_pose(
+    x_fwd: float,
+    y_left: float,
+    yaw: float,
+    origin: tuple[float, float, float],
+) -> tuple[np.ndarray, float]:
+    """Planar odometry ``(x_fwd, y_left, yaw)`` → world pose + heading.
+
+    ``origin`` is the ``(x_fwd, y_left, yaw)`` sample captured when the
+    controller first saw odometry, so the run starts at identity
+    regardless of where the robot's odometry origin sits. The result is
+    the OpenCV world convention, planarized: height Y is 0 and only yaw
+    survives of the orientation — pitch/roll gait wobble is the camera's
+    concern, not the base's.
+
+    Odometry frame is REP-103 (x forward, y left, yaw CCW about z-up).
+    Mapping to OpenCV world: ``x_world = −y_odom``, ``z_world = x_odom``,
+    ``heading = −yaw``.
+    """
+    ox, oy, oyaw = origin
+    dx, dy = x_fwd - ox, y_left - oy
+    c0, s0 = math.cos(-oyaw), math.sin(-oyaw)
+    x_rel = c0 * dx - s0 * dy
+    y_rel = s0 * dx + c0 * dy
+    yaw_rel = yaw - oyaw
+
+    x_world, z_world = -y_rel, x_rel
+    heading = -yaw_rel
+
+    ch, sh = math.cos(heading), math.sin(heading)
+    pose = np.eye(4, dtype=np.float64)
+    pose[0, 0], pose[0, 2] = ch, sh
+    pose[2, 0], pose[2, 2] = -sh, ch
+    pose[0, 3], pose[2, 3] = x_world, z_world
+    return pose, heading
+
+
+def _heading_pose(x: float, z: float, heading: float) -> np.ndarray:
+    """Build a planar world pose from position + heading."""
+    c, s = math.cos(heading), math.sin(heading)
+    pose = np.eye(4, dtype=np.float64)
+    pose[0, 0], pose[0, 2] = c, s
+    pose[2, 0], pose[2, 2] = -s, c
+    pose[0, 3], pose[2, 3] = x, z
+    return pose
+
+
+# --------------------------------------------------------------------- #
+# Stub controller (kinematic, no hardware)
+# --------------------------------------------------------------------- #
+
+
+@dataclass
+class StubBaseController:
+    """Kinematic stub: integrates each ``move()`` into pose exactly.
+
+    No latency, slip or dynamics — for sim and skill-layer unit tests.
+    """
+
+    initial_pose: np.ndarray | None = None
+    max_lin_speed: float = 1.0
+    max_yaw_rate: float = 1.0
+
+    def __post_init__(self) -> None:
+        self._pose = (
+            np.asarray(self.initial_pose, dtype=np.float64).copy()
+            if self.initial_pose is not None
+            else np.eye(4, dtype=np.float64)
+        )
+        if self._pose.shape != (4, 4):
+            raise ValueError(f"initial_pose must be (4, 4); got {self._pose.shape}")
+        self._heading = 0.0
+        self._stopped = False
+
+    def move(self, vx: float, vz: float, yaw_rate: float = 0.0, dt: float = 0.05) -> None:
+        vx = float(np.clip(vx, -self.max_lin_speed, self.max_lin_speed))
+        vz = float(np.clip(vz, -self.max_lin_speed, self.max_lin_speed))
+        yaw_rate = float(np.clip(yaw_rate, -self.max_yaw_rate, self.max_yaw_rate))
+        if dt <= 0:
+            return
+        self._pose[0, 3] += vx * dt
+        self._pose[2, 3] += vz * dt
+        if yaw_rate != 0.0:
+            self._heading += yaw_rate * dt
+            self._pose = _heading_pose(self._pose[0, 3], self._pose[2, 3], self._heading)
+        self._stopped = False
+
+    def stop(self) -> None:
+        self._stopped = True
+
+    def pose(self) -> np.ndarray:
+        return self._pose.copy()
+
+    def position(self) -> tuple[float, float, float]:
+        p = self._pose[:3, 3]
+        return float(p[0]), float(p[1]), float(p[2])
+
+    @property
+    def is_stopped(self) -> bool:
+        return self._stopped
+
+
+# --------------------------------------------------------------------- #
+# Robot-backed controller
+# --------------------------------------------------------------------- #
+
+
+@dataclass(frozen=True)
+class RobotBaseControllerConfig:
+    """Behaviour knobs for :class:`RobotBaseController`."""
+
+    max_lin_speed: float = 0.6
+    """Hard cap on per-axis world linear velocity (m/s)."""
+
+    max_yaw_rate: float = 1.2
+    """Hard cap on yaw rate (rad/s)."""
+
+    pose_from_odometry: bool = True
+    """Report pose from the robot's odometry (closed-loop). When False,
+    integrate pose open-loop from commanded velocities."""
+
+
+class RobotBaseController(BaseController):
+    """:class:`BaseController` over any LeRobot :class:`Robot`.
+
+    The robot must accept body-velocity actions ``x.vel`` (forward),
+    ``y.vel`` (left), ``theta.vel`` (CCW yaw) and — for closed-loop pose
+    — report odometry ``x.pos``/``y.pos``/``theta.pos`` in its
+    observation. This is the standard REP-103 mobile-base contract, which
+    ``UnitreeGo2`` implements.
+
+    Pose is refreshed from observations the navigation loop already
+    fetches: call :meth:`feed_observation` each keyframe rather than
+    having the controller poll the robot (which would trigger an extra
+    camera read). Absent any fed observation, pose falls back to
+    open-loop integration so sim/dry-run behaves like the stub.
+    """
+
+    def __init__(self, robot: Robot, cfg: RobotBaseControllerConfig | None = None) -> None:
+        self.robot = robot
+        self.cfg = cfg or RobotBaseControllerConfig()
+        self._pose = np.eye(4, dtype=np.float64)
+        self._heading = 0.0
+        self._stopped = False
+        self._odom_origin: tuple[float, float, float] | None = None
+        self._have_odom = False
+
+    # ----- odometry feed --------------------------------------------------
+
+    def feed_observation(self, obs: dict) -> None:
+        """Update pose from an observation the nav loop already fetched."""
+        if not self.cfg.pose_from_odometry:
+            return
+        if not {"x.pos", "y.pos", "theta.pos"} <= obs.keys():
+            return
+        sample = (float(obs["x.pos"]), float(obs["y.pos"]), float(obs["theta.pos"]))
+        if self._odom_origin is None:
+            self._odom_origin = sample
+        self._pose, self._heading = odometry_to_world_pose(*sample, self._odom_origin)
+        self._have_odom = True
+
+    # ----- BaseController API --------------------------------------------
+
+    def move(self, vx: float, vz: float, yaw_rate: float = 0.0, dt: float = 0.05) -> None:
+        vx = float(np.clip(vx, -self.cfg.max_lin_speed, self.cfg.max_lin_speed))
+        vz = float(np.clip(vz, -self.cfg.max_lin_speed, self.cfg.max_lin_speed))
+        yaw_rate = float(np.clip(yaw_rate, -self.cfg.max_yaw_rate, self.cfg.max_yaw_rate))
+        if dt <= 0:
+            return
+
+        vx_fwd, vy_left, vyaw = world_velocity_to_body(vx, vz, yaw_rate, self._heading)
+        self.robot.send_action({"x.vel": vx_fwd, "y.vel": vy_left, "theta.vel": vyaw})
+
+        # Open-loop pose only when we have no odometry to trust.
+        if not (self.cfg.pose_from_odometry and self._have_odom):
+            self._pose[0, 3] += vx * dt
+            self._pose[2, 3] += vz * dt
+            if yaw_rate != 0.0:
+                self._heading += yaw_rate * dt
+            self._pose = _heading_pose(self._pose[0, 3], self._pose[2, 3], self._heading)
+        self._stopped = False
+
+    def stop(self) -> None:
+        self._stopped = True
+        try:
+            self.robot.send_action({"x.vel": 0.0, "y.vel": 0.0, "theta.vel": 0.0})
+        except Exception:
+            logger.exception("stop(): failed to send zero-velocity action")
+
+    def pose(self) -> np.ndarray:
+        return self._pose.copy()
+
+    def position(self) -> tuple[float, float, float]:
+        p = self._pose[:3, 3]
+        return float(p[0]), float(p[1]), float(p[2])
+
+    @property
+    def is_stopped(self) -> bool:
+        return self._stopped
+
+
+# --------------------------------------------------------------------- #
+# Safety wrapper
+# --------------------------------------------------------------------- #
+
+
+@dataclass
+class SafeBaseController(BaseController):
+    """Wrap any :class:`BaseController` with safety layers:
+
+    - **velocity clamp** on every ``move()``;
+    - **occupancy gate**: when ``occupancy_provider`` is set, predict
+      the next position and refuse (latch e-stop) if it lands in an
+      obstacle cell. The provider returns an object exposing
+      ``world_to_cell(x, z) -> (iz, ix)`` and an ``is_obstacle(iz, ix)
+      -> bool`` predicate; ``None`` means "no map yet, allow";
+    - **watchdog**: if no keyframe has been fed in
+      ``watchdog_timeout_s`` (caller ticks :meth:`feed_watchdog` per
+      map update), ``move()`` latches stop until :meth:`reset_watchdog`.
+    """
+
+    inner: BaseController
+    max_lin_speed: float = 0.6
+    max_yaw_rate: float = 1.2
+    occupancy_provider: object = None  # callable[[], grid | None] when set
+    watchdog_timeout_s: float = 2.0
+    e_stop_latched: bool = False
+    _last_keyframe_walltime: float = field(default_factory=time.monotonic, init=False)
+
+    def feed_watchdog(self) -> None:
+        self._last_keyframe_walltime = time.monotonic()
+
+    def reset_watchdog(self) -> None:
+        self.e_stop_latched = False
+        self._last_keyframe_walltime = time.monotonic()
+
+    def latch_estop(self, reason: str = "external") -> None:
+        logger.warning("SafeBaseController e-stop latched: %s", reason)
+        self.e_stop_latched = True
+        self.inner.stop()
+
+    def move(self, vx: float, vz: float, yaw_rate: float = 0.0, dt: float = 0.05) -> None:
+        if self.e_stop_latched:
+            return
+        if (time.monotonic() - self._last_keyframe_walltime) > self.watchdog_timeout_s:
+            self.latch_estop(f"watchdog: no keyframe in last {self.watchdog_timeout_s:.2f}s")
+            return
+
+        vx = float(np.clip(vx, -self.max_lin_speed, self.max_lin_speed))
+        vz = float(np.clip(vz, -self.max_lin_speed, self.max_lin_speed))
+        yaw_rate = float(np.clip(yaw_rate, -self.max_yaw_rate, self.max_yaw_rate))
+
+        if self.occupancy_provider is not None:
+            try:
+                grid = self.occupancy_provider()
+            except Exception:
+                logger.exception("occupancy_provider raised; refusing move")
+                return
+            if grid is not None and self._would_enter_obstacle(grid, vx, vz, dt):
+                self.latch_estop("about to enter obstacle cell")
+                return
+
+        self.inner.move(vx, vz, yaw_rate, dt)
+
+    def stop(self) -> None:
+        self.inner.stop()
+
+    def pose(self) -> np.ndarray:
+        return self.inner.pose()
+
+    def position(self) -> tuple[float, float, float]:
+        return self.inner.position()
+
+    def _would_enter_obstacle(self, grid, vx: float, vz: float, dt: float) -> bool:
+        pos = self.inner.position()
+        next_x = pos[0] + vx * dt
+        next_z = pos[2] + vz * dt
+        iz, ix = grid.world_to_cell(next_x, next_z)
+        return bool(grid.is_obstacle(iz, ix))

src/lerobot/navigation/dog_cli.py

@@ -0,0 +1,461 @@
+#!/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.
+
+"""``dog-nav`` — interactive spatial-memory navigation REPL.
+
+Behaviour:
+  - **No prompt** (idle) → the base explores autonomously: value-map
+    frontier selection, A* on the live occupancy map, obstacle-gated
+    motion. The map grows/refreshes as it goes.
+  - **Typed prompt** (e.g. ``find the couch``) → query the map; if a
+    confident match exists, navigate to it; otherwise explore until it is
+    found (or the budget is exhausted), then resume idle exploring.
+
+A new prompt preempts the current goal. Ctrl-C latches an e-stop and
+exits. ``--dry-run`` runs the whole loop against a synthetic scene with no
+robot, camera, or models — the default until the live geometry pipeline
+(LingBot-Map) is wired.
+
+Run: ``python -m lerobot.navigation.dog_cli --dry-run`` and type object
+names; empty line ⇒ one exploration step; ``quit`` ⇒ exit.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import select
+import sys
+
+from lerobot.navigation.agent import (
+    AgentConfig,
+    AgentResult,
+    DeterministicAgent,
+    HardcodedTaskParser,
+)
+from lerobot.navigation.skills import ExploreResult, SkillsConfig, SpatialSkills
+
+LOG = logging.getLogger("dog-nav")
+
+
+class DogController:
+    """The behaviour loop over a :class:`SpatialSkills` toolset.
+
+    Construct with a ready ``SpatialSkills`` (real robot or synthetic
+    scene). :meth:`handle_prompt` runs a full locate/goto/explore task;
+    :meth:`idle_tick` runs one autonomous exploration step. Both are
+    plain calls, so the REPL and the tests share the same code.
+    """
+
+    def __init__(
+        self,
+        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)."""
+        LOG.info("prompt: %r", text)
+        result = self.agent.execute_command(text, self.parser)
+        for tr in result.target_results:
+            if tr.reached:
+                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):
+        """Locate a target and report where it is — no motion commanded.
+
+        The safe query for map-only bring-up: build the map by teleop, then
+        ask where an object is without the dog driving itself.
+        """
+        loc = self.skills.locate(text)
+        if loc.found:
+            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:
+        """One autonomous exploration step: pick a frontier and drive to it."""
+        ex = self.skills.explore(query=None)
+        if ex.found_frontier and ex.target_xyz is not None:
+            LOG.info("idle: exploring toward %s (value %.3f)", ex.target_xyz, ex.value)
+            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(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
+
+    scene = kitchen_scene()
+    base = StubBaseController()
+    siglip = scene.feature_extractor()
+    skills = SpatialSkills(
+        scene.voxel_map,
+        base,
+        siglip,
+        SkillsConfig(
+            cell_size=0.2,
+            obstacle_inflate_cells=0,
+            goto_threshold=1.0,
+            goto_max_steps=300,
+            locate_threshold=0.5,
+        ),
+    )
+    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)
+    controller = DogController(skills, agent, viz=viz)
+    controller.refresh_viz()  # show the prebuilt map immediately
+    return controller
+
+
+class LiveMapper:
+    """One perceive→integrate step of live mapping on the robot.
+
+    Each :meth:`tick` reads an observation (front camera + odometry),
+    updates the base pose from odometry, runs the geometry model + feature
+    extractor on the frame, and integrates the keyframe.
+
+    Frame convention (important): the **odometry frame is the one world
+    frame**. The geometry model supplies only relative camera-frame
+    geometry (``local_points``/depth); those points are projected through
+    the base's odometry pose, so the voxel map and the robot pose live in
+    the same coordinates and ``goto`` drives to the right place. The
+    model's own ``camera_poses`` (its internal monocular frame) are not
+    used as the world frame. Constructed lazily — no SDK/model touched
+    until the first tick.
+    """
+
+    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
+        self.geometry = geometry
+        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:
+        import numpy as np
+
+        from lerobot.navigation.pipeline import (
+            KeyframeContext,
+            PipelineConfig,
+            integrate_keyframe,
+            local_points_to_world,
+            upsample_features_to_view,
+        )
+
+        obs = self.robot.get_observation()
+        self.base.feed_observation(obs)  # updates the odometry world pose
+        pose = self.base.pose()  # camera-to-world in the odometry frame
+
+        frame = obs.get("front")
+        if frame is None:
+            return
+        views = np.asarray(frame)[None].astype(np.uint8)  # (1, H, W, 3)
+        geo = self.geometry(views)
+        h, w = frame.shape[:2]
+
+        feat_map = None
+        if self.siglip is not None:
+            patches = self.siglip.encode_views(views)[0]  # (Hp, Wp, D)
+            feat_map = upsample_features_to_view(patches, h, w)
+
+        # World points come from the model's camera-frame geometry projected
+        # through the odometry pose — NOT the model's own world frame.
+        points_world = local_points_to_world(geo.local_points[0], pose)
+        ctx = KeyframeContext(
+            frame_idx=self._frame,
+            t_sec=t_sec,
+            rgb_uint8=views[0],
+            points_world=points_world,
+            local_points=geo.local_points[0],
+            conf=geo.conf[0],
+            pose=pose,
+            feat_map=feat_map,
+        )
+        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(
+    network_interface: str = "eth0",
+    device: str = "cuda",
+    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.
+
+    Nothing here touches the SDK or loads a model — construction is lazy;
+    the DDS connection and model loads happen on first use.
+
+    ``camera_hfov_deg`` sets the pinhole focal length used for free-space
+    carving (``focal = W / (2·tan(HFOV/2))``). Calibrate it to the Go2
+    front camera for correct carving; a wrong value only degrades dynamic
+    removal, not the additive map. Speed caps are deliberately low for
+    first bring-up.
+    """
+    import math
+
+    from lerobot.navigation.base_controller import (
+        RobotBaseController,
+        RobotBaseControllerConfig,
+        SafeBaseController,
+    )
+    from lerobot.navigation.features import SiglipFeatureExtractor
+    from lerobot.navigation.geometry import LingBotMapRunner
+    from lerobot.navigation.pipeline import PipelineConfig
+    from lerobot.navigation.voxel_map import VoxelMap
+    from lerobot.robots.unitree_go2 import UnitreeGo2, UnitreeGo2Config
+
+    robot_cfg = UnitreeGo2Config(network_interface=network_interface)
+    robot = UnitreeGo2(robot_cfg)
+    inner = RobotBaseController(
+        robot, RobotBaseControllerConfig(max_lin_speed=max_lin_speed, max_yaw_rate=max_yaw_rate)
+    )
+    safe = SafeBaseController(inner=inner, max_lin_speed=max_lin_speed, max_yaw_rate=max_yaw_rate)
+    voxel_map = VoxelMap(voxel_size=0.05)
+    siglip = SiglipFeatureExtractor(device=device)
+    geometry = LingBotMapRunner(device=device)
+
+    w = robot_cfg.front_camera_width
+    focal_px = w / (2.0 * math.tan(math.radians(camera_hfov_deg) / 2.0))
+    pcfg = PipelineConfig(focal_px=focal_px)
+
+    skills = SpatialSkills(voxel_map, safe, siglip, SkillsConfig(cell_size=0.05))
+    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",
+        network_interface,
+        device,
+        focal_px,
+        max_lin_speed,
+    )
+    return controller, mapper
+
+
+def _stdin_line_ready(timeout_s: float) -> bool:
+    """True when a full line is available on stdin within ``timeout_s``.
+
+    Uses ``select`` so idle ticks keep running while we wait for input.
+    Falls back to blocking reads where ``select`` on stdin isn't supported
+    (e.g. some Windows terminals).
+    """
+    try:
+        ready, _, _ = select.select([sys.stdin], [], [], timeout_s)
+        return bool(ready)
+    except (OSError, ValueError):
+        return True
+
+
+def run_repl(controller: DogController, idle_period_s: float = 0.5) -> int:
+    """Interactive loop: explore while idle, run a task on each typed line."""
+    print("dog-nav ready. Type an object to find it, empty line to explore, 'quit' to exit.")
+    try:
+        while True:
+            if _stdin_line_ready(idle_period_s):
+                line = sys.stdin.readline()
+                if not line:  # EOF
+                    break
+                text = line.strip()
+                if text.lower() in {"quit", "exit"}:
+                    break
+                if text:
+                    controller.handle_prompt(text)  # a new prompt preempts idle
+                else:
+                    controller.idle_tick()
+            else:
+                controller.idle_tick()
+    except KeyboardInterrupt:
+        LOG.warning("interrupted — stopping base")
+    finally:
+        controller.stop()
+    return 0
+
+
+def run_live_repl(
+    controller: DogController,
+    mapper: LiveMapper,
+    idle_period_s: float = 0.2,
+    map_only: bool = False,
+) -> int:
+    """Live loop on the robot: map continuously, act on typed lines.
+
+    Each iteration integrates one keyframe (perceive → geometry → features →
+    voxel map). In ``map_only`` mode the dog is never commanded to move —
+    you teleop it while the map builds, and a typed object name reports
+    where it is (safe first bring-up). Otherwise a typed name runs a full
+    locate/goto task and an empty line takes one autonomous exploration
+    step. The DDS connection is opened here so ``--help`` stays model-free.
+    """
+    import time
+
+    mapper.robot.connect()
+    controller.skills.base.reset_watchdog()
+    if map_only:
+        print("dog-nav (live, MAP-ONLY — no autonomous motion). Teleop the dog; type an")
+        print("object to ask where it is; 'quit' to exit.")
+    else:
+        print("dog-nav (live). Type an object to find it, empty line to explore, 'quit' to exit.")
+    t0 = time.monotonic()
+    try:
+        while True:
+            mapper.tick(time.monotonic() - t0)
+            if _stdin_line_ready(idle_period_s):
+                line = sys.stdin.readline()
+                if not line:
+                    break
+                text = line.strip()
+                if text.lower() in {"quit", "exit"}:
+                    break
+                if text:
+                    controller.report_location(text) if map_only else controller.handle_prompt(text)
+                elif not map_only:
+                    controller.idle_tick()
+            elif not map_only:
+                controller.idle_tick()
+    except KeyboardInterrupt:
+        LOG.warning("interrupted — stopping base")
+    finally:
+        controller.stop()
+        mapper.robot.disconnect()
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    ap = argparse.ArgumentParser(prog="dog-nav", description=__doc__)
+    ap.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Run against a synthetic scene (no robot/camera/models).",
+    )
+    ap.add_argument(
+        "--live",
+        action="store_true",
+        help="Run on a real Unitree Go2 (DDS + LingBot-Map + SigLIP2 on the GPU host).",
+    )
+    ap.add_argument(
+        "--map-only",
+        action="store_true",
+        help="Live mode with NO autonomous motion: teleop the dog, build the map, "
+        "and query where objects are. Recommended for first bring-up.",
+    )
+    ap.add_argument("--network-interface", default="eth0", help="Host interface wired to the dog.")
+    ap.add_argument("--device", default="cuda", help="Torch device for the geometry/feature models.")
+    ap.add_argument(
+        "--camera-hfov-deg",
+        type=float,
+        default=90.0,
+        help="Go2 front-camera horizontal FOV, for the carve focal length. Calibrate to your camera.",
+    )
+    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)
+
+    logging.basicConfig(
+        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,
+            args.device,
+            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(viz=viz)
+    if args.command is not None:
+        result = controller.handle_prompt(args.command)
+        controller.stop()
+        return 0 if result.fully_successful else 1
+    return run_repl(controller)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

src/lerobot/navigation/features.py

@@ -0,0 +1,231 @@
+#!/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.
+
+"""SigLIP2 dense patch features (MaskCLIP-style) + text query encoding.
+
+Ported from the dyna360 research stack. Default checkpoint
+``google/siglip2-so400m-patch16-384``. For per-patch dense matching
+against text, raw ``last_hidden_state`` is the wrong space: SigLIP2's
+image-text matching lives in the MAP (Multihead Attention Pooling) head
+output. We use the MaskCLIP recipe — apply the MAP head's value
+projection + output projection + LayerNorm + MLP residual to each patch
+token, skipping the attention reduction — so each patch lands in
+(approximately) the shared text/vision space. Outputs are L2-normalized
+fp16.
+
+For dry-run and tests, :class:`BasisVectorFeatureExtractor` provides a
+deterministic name→vector stand-in with the same interface, no models
+required.
+"""
+
+from __future__ import annotations
+
+import logging
+from contextlib import nullcontext
+from typing import Any, Protocol, runtime_checkable
+
+import numpy as np
+
+LOG = logging.getLogger(__name__)
+
+DEFAULT_CHECKPOINT = "google/siglip2-so400m-patch16-384"
+
+
+@runtime_checkable
+class FeatureExtractor(Protocol):
+    """What the navigation stack needs from a vision-language encoder.
+
+    ``encode_text`` is required (used by ``locate``/``explore`` queries);
+    ``feature_dim`` reports the embedding size. Dense image encoding
+    (``encode_views``) is only needed by the live mapping pipeline.
+    """
+
+    @property
+    def feature_dim(self) -> int: ...
+
+    def encode_text(self, text: str) -> np.ndarray: ...
+
+
+def _select_autocast(device: str) -> tuple[Any, str]:
+    """Pick an autocast context + label for the given device."""
+    import torch
+
+    if device != "cuda":
+        return nullcontext(), "no-autocast"
+    if not torch.cuda.is_available():
+        raise RuntimeError("device='cuda' requested but torch.cuda.is_available() is False")
+    cap = torch.cuda.get_device_capability()[0]
+    dtype = torch.bfloat16 if cap >= 8 else torch.float16
+    return torch.amp.autocast("cuda", dtype=dtype), f"cuda/{str(dtype).split('.')[-1]}"
+
+
+class SiglipFeatureExtractor:
+    """Lazy-loaded SigLIP2 wrapper for dense patch features + text query."""
+
+    def __init__(
+        self,
+        checkpoint: str = DEFAULT_CHECKPOINT,
+        device: str = "cuda",
+        max_batch: int = 8,
+    ) -> None:
+        self.checkpoint = checkpoint
+        self.device = device
+        self.max_batch = int(max_batch)
+        self._model: Any | None = None
+        self._processor: Any | None = None
+        self._patch_grid: tuple[int, int] | None = None
+        self._feature_dim: int | None = None
+
+    @property
+    def feature_dim(self) -> int:
+        if self._feature_dim is None:
+            raise RuntimeError("SigLIP2 not loaded yet; call encode_views first")
+        return self._feature_dim
+
+    @property
+    def patch_grid(self) -> tuple[int, int]:
+        if self._patch_grid is None:
+            raise RuntimeError("SigLIP2 not loaded yet; call encode_views first")
+        return self._patch_grid
+
+    def _ensure_loaded(self) -> None:
+        if self._model is not None:
+            return
+        from transformers import AutoModel, AutoProcessor
+
+        LOG.info("loading SigLIP2 (%s) on %s ...", self.checkpoint, self.device)
+        self._processor = AutoProcessor.from_pretrained(self.checkpoint)
+        self._model = AutoModel.from_pretrained(self.checkpoint).to(self.device).eval()
+        LOG.info("SigLIP2 loaded")
+
+    def _maskclip_project(self, patches):
+        """Push raw patch tokens through the MAP head with the attention
+        reduction removed — value-projects + post-processes each patch so it
+        lives in the shared text/vision space. ``patches``: (B, P, D)."""
+        import torch
+
+        assert self._model is not None
+        head = self._model.vision_model.head
+        mha = head.attention  # nn.MultiheadAttention
+        embed_dim = patches.shape[-1]
+
+        # in_proj_weight is concatenated [Q | K | V], (3*D, D). Slice out V.
+        v_weight = mha.in_proj_weight[2 * embed_dim : 3 * embed_dim]
+        v_bias = mha.in_proj_bias[2 * embed_dim : 3 * embed_dim] if mha.in_proj_bias is not None else None
+        v = torch.nn.functional.linear(patches, v_weight, v_bias)  # (B, P, D)
+        v = mha.out_proj(v)
+
+        residual = v
+        v = head.layernorm(v)
+        v = residual + head.mlp(v)
+        return v
+
+    def encode_views(self, views_rgb_uint8: np.ndarray) -> np.ndarray:
+        """Encode ``(N, H, W, 3)`` RGB uint8 views to ``(N, Hp, Wp, D)`` fp16
+        dense patch features in the shared text/vision space, L2-normalized."""
+        import torch
+
+        if views_rgb_uint8.ndim != 4 or views_rgb_uint8.shape[-1] != 3:  # noqa: N806
+            raise ValueError(f"expected (N, H, W, 3), got {views_rgb_uint8.shape}")
+        if views_rgb_uint8.dtype != np.uint8:
+            raise ValueError(f"expected uint8, got {views_rgb_uint8.dtype}")
+        self._ensure_loaded()
+        assert self._model is not None and self._processor is not None
+
+        autocast_ctx, autocast_label = _select_autocast(self.device)
+        LOG.info(
+            "SigLIP2 forward (MaskCLIP-projected patches): N=%d (batched up to %d), %s",
+            views_rgb_uint8.shape[0],
+            self.max_batch,
+            autocast_label,
+        )
+
+        out_list: list[np.ndarray] = []
+        for s in range(0, views_rgb_uint8.shape[0], self.max_batch):
+            e = s + self.max_batch
+            chunk = [views_rgb_uint8[i] for i in range(s, min(e, views_rgb_uint8.shape[0]))]
+            inputs = self._processor(images=chunk, return_tensors="pt").to(self.device)
+            with torch.no_grad(), autocast_ctx:
+                vision = self._model.vision_model(**inputs)
+                patches = vision.last_hidden_state  # (B, P, D)
+                patches = self._maskclip_project(patches)  # (B, P, D) shared-space
+            patches = torch.nn.functional.normalize(patches.float(), dim=-1)
+            out_list.append(patches.to(torch.float16).cpu().numpy())
+
+        feats = np.concatenate(out_list, axis=0)  # (N, P, D)
+        n, p, d = feats.shape
+        side = int(round(p**0.5))
+        if side * side != p:
+            raise RuntimeError(
+                f"SigLIP2 returned a non-square patch grid (P={p}); non-square inputs aren't supported yet"
+            )
+        self._patch_grid = (side, side)
+        self._feature_dim = d
+        return feats.reshape(n, side, side, d)
+
+    def encode_text(self, text: str) -> np.ndarray:
+        """Encode a text query to a single (D,) fp16 unit vector.
+
+        SigLIP2 uses last-token ([EOS]) pooling for text. We extract it
+        explicitly because ``get_text_features`` behaves differently across
+        ``transformers`` versions.
+        """
+        import torch
+
+        self._ensure_loaded()
+        assert self._model is not None and self._processor is not None
+        autocast_ctx, _ = _select_autocast(self.device)
+        inputs = self._processor(text=[text], return_tensors="pt", padding="max_length").to(self.device)
+        with torch.no_grad(), autocast_ctx:
+            text_outputs = self._model.text_model(**inputs)
+
+        pooled = getattr(text_outputs, "pooler_output", None)
+        if pooled is not None and pooled.dim() == 2:
+            feat = pooled[0]
+        else:
+            feat = text_outputs.last_hidden_state[0, -1]
+
+        feat = feat.float()
+        feat = torch.nn.functional.normalize(feat, dim=-1)
+        return feat.to(torch.float16).cpu().numpy()
+
+
+class BasisVectorFeatureExtractor:
+    """Deterministic name→vector stand-in for :class:`SiglipFeatureExtractor`.
+
+    Maps known names to their stored feature vectors; unknown queries get a
+    deterministic per-text pseudo-random unit vector (same string → same
+    vector), so a locate threshold reliably rejects absent objects. Used by
+    the synthetic-scene dry-run and by tests — no models required.
+    """
+
+    def __init__(self, name_to_vec: dict[str, np.ndarray], feature_dim: int) -> None:
+        self.name_to_vec = name_to_vec
+        self._feature_dim = int(feature_dim)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def encode_text(self, text: str) -> np.ndarray:
+        v = self.name_to_vec.get(text)
+        if v is None:
+            seed = abs(hash(text)) % (2**32)
+            rng = np.random.default_rng(seed)
+            v = rng.normal(size=self._feature_dim).astype(np.float32)
+        v = v.astype(np.float32)
+        v = v / max(float(np.linalg.norm(v)), 1e-6)
+        return v

src/lerobot/navigation/geometry.py

@@ -0,0 +1,222 @@
+#!/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.
+
+"""Monocular geometry runners for the mapping pipeline.
+
+A :class:`GeometryRunner` turns a stack of RGB views into per-pixel world
+points, camera-frame points (depth), confidence, and camera-to-world
+poses — the four arrays the voxel-map pipeline consumes.
+:class:`LingBotMapRunner` wraps Ant Group's streaming LingBot-Map model
+(feed-forward 3D reconstruction with persistent memory); the SDK import
+is lazy so configs/tests/``--help`` don't pay the model cost.
+
+Because LingBot-Map is monocular, its world frame has an unknown metric
+scale. On a robot with wheel/leg odometry (the Unitree Go2 sport-mode
+state), :func:`align_trajectory_to_odometry` fits a similarity transform
+(scale + rotation + translation) from the model's camera trajectory to
+the odometry trajectory, so the voxel map comes out metric and A* speeds
+are real m/s. :class:`FakeGeometryRunner` produces deterministic planar
+geometry for hardware-free tests.
+"""
+
+# ruff: noqa: N806  — R, U, S, Vt, D: conventional linear-algebra / array-dimension names
+from __future__ import annotations
+
+import logging
+from contextlib import nullcontext
+from dataclasses import dataclass
+from typing import Any, Protocol, runtime_checkable
+
+import numpy as np
+
+LOG = logging.getLogger(__name__)
+
+DEFAULT_LINGBOT_CHECKPOINT = "robbyant/lingbot-map"
+
+
+@dataclass(frozen=True)
+class GeometryOutput:
+    """Per-view geometry outputs (fp32, on CPU).
+
+    The contract every :class:`GeometryRunner` emits and the voxel-map
+    pipeline consumes.
+    """
+
+    points: np.ndarray  # (N, H, W, 3) world points
+    local_points: np.ndarray  # (N, H, W, 3) camera-frame points; depth = [..., 2]
+    conf: np.ndarray  # (N, H, W) in [0, 1]
+    camera_poses: np.ndarray  # (N, 4, 4) camera-to-world, OpenCV convention
+
+
+@runtime_checkable
+class GeometryRunner(Protocol):
+    """Turns ``(N, H, W, 3)`` uint8 RGB views into a :class:`GeometryOutput`."""
+
+    def __call__(self, views_rgb_uint8: np.ndarray) -> GeometryOutput: ...
+
+
+def _select_autocast(device: str) -> tuple[Any, str]:
+    """Return (autocast context manager, label for logging)."""
+    import torch
+
+    if device != "cuda":
+        return nullcontext(), "no-autocast"
+    if not torch.cuda.is_available():
+        raise RuntimeError("device='cuda' requested but torch.cuda.is_available() is False")
+    cap = torch.cuda.get_device_capability()[0]
+    dtype = torch.bfloat16 if cap >= 8 else torch.float16
+    return torch.amp.autocast("cuda", dtype=dtype), f"cuda/{str(dtype).split('.')[-1]}"
+
+
+class LingBotMapRunner:
+    """Lazy-loaded LingBot-Map streaming reconstruction runner.
+
+    Streaming feed-forward reconstruction with a persistent KV-cache keeps
+    every view anchored to one consistent world frame — so, unlike
+    window-based models, no cross-window pose stitching is needed. The
+    model download/load is deferred to the first call.
+    """
+
+    def __init__(
+        self,
+        device: str = "cuda",
+        checkpoint: str = DEFAULT_LINGBOT_CHECKPOINT,
+    ) -> None:
+        self.device = device
+        self.checkpoint = checkpoint
+        self._model: Any | None = None
+
+    def _ensure_loaded(self) -> None:
+        if self._model is not None:
+            return
+        try:
+            from lingbot_map import LingBotMap  # type: ignore[import-not-found]
+        except ImportError as exc:
+            raise RuntimeError(
+                f"lingbot-map is not importable ({exc}). Install it from "
+                "github.com/robbyant/lingbot-map on the GPU host."
+            ) from exc
+        LOG.info("loading LingBot-Map (%s) on %s ...", self.checkpoint, self.device)
+        self._model = LingBotMap.from_pretrained(self.checkpoint).to(self.device).eval()
+        LOG.info("LingBot-Map loaded")
+
+    def __call__(self, views_rgb_uint8: np.ndarray) -> GeometryOutput:
+        import torch
+
+        if views_rgb_uint8.ndim != 4 or views_rgb_uint8.shape[-1] != 3:
+            raise ValueError(f"expected (N, H, W, 3), got {views_rgb_uint8.shape}")
+        if views_rgb_uint8.dtype != np.uint8:
+            raise ValueError(f"expected uint8, got {views_rgb_uint8.dtype}")
+        self._ensure_loaded()
+        assert self._model is not None
+
+        imgs = (
+            torch.from_numpy(views_rgb_uint8)
+            .to(self.device)
+            .float()
+            .div_(255.0)
+            .permute(0, 3, 1, 2)
+            .contiguous()
+        )  # (N, 3, H, W)
+
+        autocast_ctx, label = _select_autocast(self.device)
+        LOG.info("LingBot-Map forward: N=%d, %s", views_rgb_uint8.shape[0], label)
+        with torch.no_grad(), autocast_ctx:
+            res = self._model(imgs[None])  # (1, N, ...)
+
+        def _np(t) -> np.ndarray:
+            return t.detach().float().cpu().numpy()
+
+        points = _np(res["points"][0])
+        local_points = _np(res["local_points"][0])
+        conf = _np(res["conf"][0])
+        if conf.ndim == 4:  # (N, H, W, 1) → (N, H, W)
+            conf = conf[..., 0]
+        camera_poses = _np(res["camera_poses"][0])
+        return GeometryOutput(points, local_points, conf, camera_poses)
+
+
+class FakeGeometryRunner:
+    """Deterministic planar geometry for hardware-free tests.
+
+    Emits a flat floor at ``depth`` metres in front of the camera with a
+    pinhole model, unit confidence, and identity (or supplied) poses — no
+    model required.
+    """
+
+    def __init__(self, depth: float = 3.0, focal_px: float = 100.0) -> None:
+        self.depth = float(depth)
+        self.focal_px = float(focal_px)
+
+    def __call__(self, views_rgb_uint8: np.ndarray) -> GeometryOutput:
+        if views_rgb_uint8.ndim != 4 or views_rgb_uint8.shape[-1] != 3:
+            raise ValueError(f"expected (N, H, W, 3), got {views_rgb_uint8.shape}")
+        n, h, w, _ = views_rgb_uint8.shape
+        cx, cy = (w - 1) / 2.0, (h - 1) / 2.0
+        us, vs = np.meshgrid(np.arange(w), np.arange(h))
+        x = (us - cx) * self.depth / self.focal_px
+        y = (vs - cy) * self.depth / self.focal_px
+        z = np.full_like(x, self.depth, dtype=np.float64)
+        local = np.stack([x, y, z], axis=-1).astype(np.float32)  # (H, W, 3)
+        local_points = np.broadcast_to(local, (n, h, w, 3)).copy()
+        # Identity poses → world == camera frame.
+        points = local_points.copy()
+        conf = np.ones((n, h, w), dtype=np.float32)
+        poses = np.broadcast_to(np.eye(4, dtype=np.float32), (n, 4, 4)).copy()
+        return GeometryOutput(points, local_points, conf, poses)
+
+
+def umeyama_similarity(src: np.ndarray, dst: np.ndarray) -> tuple[float, np.ndarray, np.ndarray]:
+    """Least-squares similarity (scale s, rotation R, translation t) mapping
+    ``src`` onto ``dst`` such that ``dst ≈ s · R @ src + t``.
+
+    ``src``/``dst`` are ``(K, 3)``. Returns ``(s, R, t)``. Used to anchor a
+    monocular trajectory to metric odometry.
+    """
+    src = np.asarray(src, dtype=np.float64)
+    dst = np.asarray(dst, dtype=np.float64)
+    if src.shape != dst.shape or src.ndim != 2 or src.shape[1] != 3:
+        raise ValueError(f"src/dst must be matching (K, 3); got {src.shape}, {dst.shape}")
+    k = src.shape[0]
+    mu_src = src.mean(axis=0)
+    mu_dst = dst.mean(axis=0)
+    sc = src - mu_src
+    dc = dst - mu_dst
+    cov = (dc.T @ sc) / k
+    U, D, Vt = np.linalg.svd(cov)
+    S = np.eye(3)
+    if np.linalg.det(U) * np.linalg.det(Vt) < 0:
+        S[2, 2] = -1.0
+    R = U @ S @ Vt
+    var_src = (sc**2).sum() / k
+    s = float((D * np.diag(S)).sum() / max(var_src, 1e-12))
+    t = mu_dst - s * R @ mu_src
+    return s, R, t
+
+
+def align_trajectory_to_odometry(
+    camera_positions: np.ndarray,
+    odom_positions: np.ndarray,
+) -> tuple[float, np.ndarray, np.ndarray]:
+    """Fit the similarity transform from a monocular camera trajectory to a
+    metric odometry trajectory (both ``(K, 3)``, time-aligned).
+
+    Returns ``(scale, R, t)`` to apply to model world points/poses so the
+    voxel map is metric. Needs at least 3 non-degenerate points.
+    """
+    if camera_positions.shape[0] < 3:
+        raise ValueError("need at least 3 corresponding poses to fit a similarity")
+    return umeyama_similarity(camera_positions, odom_positions)

src/lerobot/navigation/occupancy.py

@@ -0,0 +1,371 @@
+#!/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.
+
+"""2D occupancy projection of the voxel map + A* path planning.
+
+Ported from the dyna360 research stack. Derived, not maintained: every
+call to :func:`project_voxel_map_to_grid` rebuilds the 3-class grid from
+a fresh ``VoxelMap.snapshot()``, so the projection reflects whatever the
+keyframe loop most recently carved or added — no separate obstacle
+structure to keep in sync.
+
+Coordinate convention: OpenCV (X right, Y *down*, Z forward), matching
+the navigation world frame. "Up" is the −Y direction. The top-down grid
+indexes the XZ plane; cell ``(iz, ix)`` covers world rectangle
+``[origin_x + ix·cell, origin_x + (ix+1)·cell]`` ×
+``[origin_z + iz·cell, origin_z + (iz+1)·cell]``.
+
+Classes:
+  - ``UNOBSERVED`` (0): no voxel projects here. The base must not plan
+    through it (might be an unseen obstacle), but explorers treat it as
+    the goal class.
+  - ``NAVIGABLE`` (1): observed ground / open space.
+  - ``OBSTACLE`` (2): at least one voxel in the robot-height band
+    projects here.
+"""
+
+# ruff: noqa: N806  — H, W, D are conventional array-dimension names (and appear verbatim in error strings)
+from __future__ import annotations
+
+import heapq
+import logging
+import math
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from lerobot.navigation.voxel_map import VoxelMap
+
+LOG = logging.getLogger(__name__)
+
+# Class constants — picked so a colormap can index directly.
+UNOBSERVED = np.int8(0)
+NAVIGABLE = np.int8(1)
+OBSTACLE = np.int8(2)
+
+
+@dataclass(frozen=True)
+class OccupancyGrid:
+    """3-class top-down grid plus its world↔cell mapping."""
+
+    classes: np.ndarray  # (H, W) int8 — H = z-extent, W = x-extent
+    cell_size: float  # m per cell
+    origin_x: float  # world x of the LEFT edge of column 0
+    origin_z: float  # world z of the TOP edge of row 0
+    ground_y: float  # world y of the (auto-estimated or given) ground plane
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        return self.classes.shape  # (H, W)
+
+    def world_to_cell(self, x: float, z: float) -> tuple[int, int]:
+        """Return ``(iz, ix)``, clipped to grid extents."""
+        ix = int(np.clip(math.floor((x - self.origin_x) / self.cell_size), 0, self.shape[1] - 1))
+        iz = int(np.clip(math.floor((z - self.origin_z) / self.cell_size), 0, self.shape[0] - 1))
+        return iz, ix
+
+    def cell_to_world(self, iz: int, ix: int) -> tuple[float, float]:
+        """Return ``(x, z)`` at the *centre* of cell ``(iz, ix)``."""
+        x = self.origin_x + (ix + 0.5) * self.cell_size
+        z = self.origin_z + (iz + 0.5) * self.cell_size
+        return x, z
+
+    def is_navigable(self, iz: int, ix: int) -> bool:
+        H, W = self.shape
+        return 0 <= iz < H and 0 <= ix < W and self.classes[iz, ix] == NAVIGABLE
+
+    def is_obstacle(self, iz: int, ix: int) -> bool:
+        """Whether cell ``(iz, ix)`` is a known obstacle. Out-of-bounds is
+        not an obstacle (it is simply unobservable) — used by
+        ``SafeBaseController``'s occupancy gate."""
+        H, W = self.shape
+        return 0 <= iz < H and 0 <= ix < W and self.classes[iz, ix] == OBSTACLE
+
+    def is_in_bounds(self, iz: int, ix: int) -> bool:
+        H, W = self.shape
+        return 0 <= iz < H and 0 <= ix < W
+
+    def nearest_navigable_cell(self, iz: int, ix: int, max_radius: int = 50) -> tuple[int, int] | None:
+        """BFS outward until a navigable cell is found, or give up."""
+        if self.is_navigable(iz, ix):
+            return iz, ix
+        for r in range(1, max_radius + 1):
+            for diz in range(-r, r + 1):
+                for dix in range(-r, r + 1):
+                    if max(abs(diz), abs(dix)) != r:
+                        continue  # ring only, not the interior
+                    if self.is_navigable(iz + diz, ix + dix):
+                        return iz + diz, ix + dix
+        return None
+
+
+def estimate_ground_y(xyz: np.ndarray, percentile: float = 95.0) -> float:
+    """Estimate the world-frame y of the ground plane.
+
+    Y is down (OpenCV), so the ground is at the LARGEST y values. Using a
+    high percentile (default 95) is robust to outliers below the ground.
+    """
+    if xyz.size == 0:
+        return 0.0
+    return float(np.percentile(xyz[:, 1], percentile))
+
+
+def project_voxel_map_to_grid(
+    voxel_map: VoxelMap,
+    *,
+    cell_size: float = 0.1,
+    ground_y: float | None = None,
+    obstacle_y_range: tuple[float, float] = (-2.0, -0.1),
+    bbox: tuple[float, float, float, float] | None = None,
+    bbox_pad: float = 1.0,
+    inflate_cells: int = 0,
+) -> OccupancyGrid:
+    """Snapshot the voxel map and project it into a 2D occupancy grid.
+
+    ``obstacle_y_range`` is interpreted *relative* to ``ground_y`` with the
+    Y-down convention, so the default ``(-2.0, -0.1)`` means "voxels
+    between 2.0 m and 0.1 m above the ground are obstacles". Anything above
+    the ceiling band or below ground level is silently ignored.
+
+    ``inflate_cells`` dilates the OBSTACLE class by N cells of clearance
+    (square morphology) — a body-radius safety margin for the base without
+    resampling the voxel map.
+    """
+    snap = voxel_map.snapshot()
+    xyz = snap.xyz
+
+    if xyz.size == 0:
+        # Empty map → a 1×1 grid of UNOBSERVED at world origin.
+        return OccupancyGrid(
+            classes=np.zeros((1, 1), dtype=np.int8),
+            cell_size=float(cell_size),
+            origin_x=0.0,
+            origin_z=0.0,
+            ground_y=ground_y if ground_y is not None else 0.0,
+        )
+
+    if ground_y is None:
+        ground_y = estimate_ground_y(xyz)
+
+    # Promote to float64 — VoxelMap snapshots are float32, and naive
+    # ``(float32_array <= float64_scalar)`` lets numpy downcast the scalar
+    # back to float32, which causes edge-case bugs (e.g. 1.0 <= 0.999999999
+    # becomes True because the threshold rounds up to 1.0 in float32).
+    x_arr = xyz[:, 0].astype(np.float64)
+    y_arr = xyz[:, 1].astype(np.float64)
+    z_arr = xyz[:, 2].astype(np.float64)
+
+    abs_y_top = ground_y + obstacle_y_range[0]  # most-negative y (highest above ground)
+    abs_y_bottom = ground_y + obstacle_y_range[1]  # closer to ground
+    is_obstacle = (y_arr >= abs_y_top) & (y_arr <= abs_y_bottom)
+
+    if bbox is None:
+        x_min = float(x_arr.min()) - bbox_pad
+        x_max = float(x_arr.max()) + bbox_pad
+        z_min = float(z_arr.min()) - bbox_pad
+        z_max = float(z_arr.max()) + bbox_pad
+    else:
+        x_min, z_min, x_max, z_max = bbox
+
+    W = max(1, int(math.ceil((x_max - x_min) / cell_size)))
+    H = max(1, int(math.ceil((z_max - z_min) / cell_size)))
+    classes = np.zeros((H, W), dtype=np.int8)  # default UNOBSERVED
+
+    # eps absorbs float32→float64 representation drift so points that should
+    # land exactly on a cell boundary aren't randomly bumped into the
+    # previous cell. 1e-3 of a cell width is well above float32's ~1e-7
+    # relative precision and well below the 0.5-cell misclassification
+    # threshold.
+    eps = cell_size * 1e-3
+    ix = np.clip(np.floor((x_arr - x_min) / cell_size + eps).astype(np.int32), 0, W - 1)
+    iz = np.clip(np.floor((z_arr - z_min) / cell_size + eps).astype(np.int32), 0, H - 1)
+
+    # Two-pass labelling: any voxel makes a cell observed (-> NAVIGABLE);
+    # obstacle voxels then upgrade those cells to OBSTACLE.
+    classes[iz, ix] = NAVIGABLE
+    obs_iz = iz[is_obstacle]
+    obs_ix = ix[is_obstacle]
+    classes[obs_iz, obs_ix] = OBSTACLE
+
+    if inflate_cells > 0:
+        classes = _inflate_obstacles(classes, inflate_cells)
+
+    return OccupancyGrid(
+        classes=classes,
+        cell_size=float(cell_size),
+        origin_x=x_min,
+        origin_z=z_min,
+        ground_y=float(ground_y),
+    )
+
+
+def _inflate_obstacles(classes: np.ndarray, radius: int) -> np.ndarray:
+    """Dilate OBSTACLE cells by `radius` cells (Chebyshev). Pure-numpy
+    morphological dilation — fine for our grid sizes."""
+    out = classes.copy()
+    obs = classes == OBSTACLE
+    H, W = classes.shape
+    for diz in range(-radius, radius + 1):
+        for dix in range(-radius, radius + 1):
+            if diz == 0 and dix == 0:
+                continue
+            sl_src_iz = slice(max(0, -diz), H - max(0, diz))
+            sl_src_ix = slice(max(0, -dix), W - max(0, dix))
+            sl_dst_iz = slice(max(0, diz), H - max(0, -diz))
+            sl_dst_ix = slice(max(0, dix), W - max(0, -dix))
+            inflated = obs[sl_src_iz, sl_src_ix]
+            # Only upgrade NAVIGABLE → OBSTACLE; never overwrite UNOBSERVED
+            # so the frontier (NAVIGABLE↔UNOBSERVED boundary) survives.
+            target = out[sl_dst_iz, sl_dst_ix]
+            promote = inflated & (target == NAVIGABLE)
+            target[promote] = OBSTACLE
+            out[sl_dst_iz, sl_dst_ix] = target
+    return out
+
+
+# --------------------------------------------------------------------- A*
+
+_DIAG_COST = math.sqrt(2.0)
+_NEIGHBOURS_ORTHO = ((-1, 0), (1, 0), (0, -1), (0, 1))
+_NEIGHBOURS_DIAG = ((-1, -1), (-1, 1), (1, -1), (1, 1))
+
+
+def astar(
+    grid: OccupancyGrid,
+    start_world: tuple[float, float],
+    goal_world: tuple[float, float],
+    *,
+    allow_unobserved_goal: bool = True,
+) -> list[tuple[float, float]] | None:
+    """Plan a path from ``start_world`` to ``goal_world`` in (x, z) world m.
+
+    Returns a list of world ``(x, z)`` waypoints, or ``None`` if no path
+    exists. Start/goal are snapped to the nearest navigable cell.
+    """
+    H, W = grid.shape
+    if H == 0 or W == 0:
+        return None
+
+    s_iz, s_ix = grid.world_to_cell(*start_world)
+    g_iz, g_ix = grid.world_to_cell(*goal_world)
+
+    if not grid.is_navigable(s_iz, s_ix):
+        snapped = grid.nearest_navigable_cell(s_iz, s_ix)
+        if snapped is None:
+            return None
+        s_iz, s_ix = snapped
+    if not grid.is_navigable(g_iz, g_ix):
+        if not allow_unobserved_goal:
+            return None
+        snapped = grid.nearest_navigable_cell(g_iz, g_ix)
+        if snapped is None:
+            return None
+        g_iz, g_ix = snapped
+
+    def heuristic(iz: int, ix: int) -> float:
+        d_iz = abs(iz - g_iz)
+        d_ix = abs(ix - g_ix)
+        return (max(d_iz, d_ix) - min(d_iz, d_ix)) + _DIAG_COST * min(d_iz, d_ix)
+
+    open_heap: list[tuple[float, int, tuple[int, int]]] = []
+    counter = 0  # tiebreaker so heapq doesn't compare tuples on ties
+    heapq.heappush(open_heap, (0.0, counter, (s_iz, s_ix)))
+    came_from: dict[tuple[int, int], tuple[int, int]] = {}
+    g_score: dict[tuple[int, int], float] = {(s_iz, s_ix): 0.0}
+
+    while open_heap:
+        _, _, current = heapq.heappop(open_heap)
+        if current == (g_iz, g_ix):
+            return _reconstruct_path(came_from, current, grid)
+
+        cur_iz, cur_ix = current
+        cur_g = g_score[current]
+
+        for diz, dix in _NEIGHBOURS_ORTHO:
+            n = (cur_iz + diz, cur_ix + dix)
+            if not grid.is_navigable(*n):
+                continue
+            tentative = cur_g + 1.0
+            if tentative < g_score.get(n, float("inf")):
+                came_from[n] = current
+                g_score[n] = tentative
+                counter += 1
+                heapq.heappush(open_heap, (tentative + heuristic(*n), counter, n))
+
+        for diz, dix in _NEIGHBOURS_DIAG:
+            n = (cur_iz + diz, cur_ix + dix)
+            if not grid.is_navigable(*n):
+                continue
+            # Prevent corner-cutting: both perpendicular neighbours must be
+            # navigable, or we'd squeeze through an obstacle's diagonal.
+            if not grid.is_navigable(cur_iz + diz, cur_ix):
+                continue
+            if not grid.is_navigable(cur_iz, cur_ix + dix):
+                continue
+            tentative = cur_g + _DIAG_COST
+            if tentative < g_score.get(n, float("inf")):
+                came_from[n] = current
+                g_score[n] = tentative
+                counter += 1
+                heapq.heappush(open_heap, (tentative + heuristic(*n), counter, n))
+
+    return None
+
+
+def _reconstruct_path(
+    came_from: dict[tuple[int, int], tuple[int, int]],
+    end: tuple[int, int],
+    grid: OccupancyGrid,
+) -> list[tuple[float, float]]:
+    cells = [end]
+    while cells[-1] in came_from:
+        cells.append(came_from[cells[-1]])
+    cells.reverse()
+    return [grid.cell_to_world(iz, ix) for iz, ix in cells]
+
+
+# ----------------------------------------------------------------- frontier
+
+
+def find_frontier_cells(grid: OccupancyGrid) -> np.ndarray:
+    """Return cells on the NAVIGABLE↔UNOBSERVED boundary, as ``(K, 2)`` int.
+
+    These are the cells exploration aims for: places we already know we
+    can stand at, but with unknown adjacent territory worth visiting.
+    """
+    nav = grid.classes == NAVIGABLE
+    unobs = grid.classes == UNOBSERVED
+    if not nav.any() or not unobs.any():
+        return np.zeros((0, 2), dtype=np.int32)
+
+    boundary = np.zeros_like(nav)
+    boundary[1:, :] |= nav[1:, :] & unobs[:-1, :]
+    boundary[:-1, :] |= nav[:-1, :] & unobs[1:, :]
+    boundary[:, 1:] |= nav[:, 1:] & unobs[:, :-1]
+    boundary[:, :-1] |= nav[:, :-1] & unobs[:, 1:]
+    iz, ix = np.where(boundary)
+    return np.stack([iz, ix], axis=-1).astype(np.int32)
+
+
+def occupancy_to_rgb(grid: OccupancyGrid) -> np.ndarray:
+    """Render the 3-class grid as an (H, W, 3) uint8 image."""
+    img = np.zeros((*grid.shape, 3), dtype=np.uint8)
+    img[grid.classes == UNOBSERVED] = (40, 40, 50)
+    img[grid.classes == NAVIGABLE] = (200, 200, 200)
+    img[grid.classes == OBSTACLE] = (220, 60, 60)
+    return img

src/lerobot/navigation/pipeline.py

@@ -0,0 +1,133 @@
+#!/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.
+
+"""Keyframe integration loop core.
+
+Ported from the dyna360 research stack (viz-free). One keyframe is
+carved then added into the voxel map — carve first so we never remove
+voxels we just created this frame. This is the shared step behind live
+mapping on the robot.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from lerobot.navigation.voxel_map import CarveResult, VoxelMap, VoxelMapStats
+
+LOG = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class KeyframeContext:
+    """Everything one keyframe needs to contribute to the voxel map.
+
+    ``rgb_uint8`` is RGB order (same layout fed to the geometry model and
+    the feature extractor). ``points_world`` / ``local_points`` come from
+    the geometry runner; ``feat_map`` is the bilinearly-upsampled patch
+    grid at ``(H, W, D)`` fp16, or ``None`` for a geometry-only frame.
+    """
+
+    frame_idx: int
+    t_sec: float
+    rgb_uint8: np.ndarray  # (H, W, 3) RGB uint8
+    points_world: np.ndarray  # (H, W, 3) float32
+    local_points: np.ndarray  # (H, W, 3) float32
+    conf: np.ndarray  # (H, W) in [0, 1]
+    pose: np.ndarray  # (4, 4) cam-to-world
+    feat_map: np.ndarray | None  # (H, W, D) fp16, L2-normalized per pixel
+
+
+@dataclass(frozen=True)
+class PipelineConfig:
+    """Knobs that change per-run but not per-keyframe."""
+
+    conf_thresh: float = 0.5
+    carve_margin: float = 0.05
+    focal_px: float = 100.0
+
+
+def integrate_keyframe(
+    voxel_map: VoxelMap,
+    ctx: KeyframeContext,
+    pcfg: PipelineConfig | None = None,
+) -> tuple[CarveResult, VoxelMapStats]:
+    """Carve observed free space, then add this keyframe's points.
+
+    Carve runs before add. Returns the carve result + add stats so callers
+    can surface them in their own progress UI.
+    """
+    pcfg = pcfg or PipelineConfig()
+    carve = voxel_map.carve(
+        local_points=ctx.local_points,
+        conf=ctx.conf,
+        pose=ctx.pose,
+        focal_px=pcfg.focal_px,
+        frame=ctx.frame_idx,
+        t=ctx.t_sec,
+        conf_thresh=pcfg.conf_thresh,
+        margin=pcfg.carve_margin,
+    )
+    stats = voxel_map.add(
+        points=ctx.points_world,
+        rgb=ctx.rgb_uint8,
+        conf=ctx.conf,
+        frame=ctx.frame_idx,
+        t=ctx.t_sec,
+        conf_thresh=pcfg.conf_thresh,
+        feat_map=ctx.feat_map,
+    )
+    return carve, stats
+
+
+def local_points_to_world(local_points: np.ndarray, pose: np.ndarray) -> np.ndarray:
+    """Transform camera-frame points ``(H, W, 3)`` into the world frame using
+    a 4×4 camera-to-world ``pose``.
+
+    On the robot the world frame is the odometry frame (from the base
+    controller), and the geometry model supplies only relative
+    camera-frame geometry — so projecting through the odometry pose keeps
+    the voxel map and the robot pose in ONE consistent frame. Returns
+    ``(H, W, 3)`` float32.
+    """
+    if local_points.ndim != 3 or local_points.shape[-1] != 3:
+        raise ValueError(f"expected (H, W, 3), got {local_points.shape}")
+    if pose.shape != (4, 4):
+        raise ValueError(f"pose must be (4, 4); got {pose.shape}")
+    r = pose[:3, :3].astype(np.float64)
+    t = pose[:3, 3].astype(np.float64)
+    flat = local_points.reshape(-1, 3).astype(np.float64)
+    world = flat @ r.T + t
+    return world.reshape(local_points.shape).astype(np.float32)
+
+
+def upsample_features_to_view(
+    patch_feats_one_view: np.ndarray,
+    view_h: int,
+    view_w: int,
+) -> np.ndarray:
+    """Bilinearly upsample one keyframe's ``(Hp, Wp, D)`` patch features to
+    view resolution ``(H, W, D)`` fp16."""
+    import torch
+
+    fp = torch.from_numpy(patch_feats_one_view).permute(2, 0, 1).unsqueeze(0).float()  # (1, D, Hp, Wp)
+    fp_up = torch.nn.functional.interpolate(fp, size=(view_h, view_w), mode="bilinear", align_corners=False)
+    return fp_up.squeeze(0).permute(1, 2, 0).to(torch.float16).numpy()

src/lerobot/navigation/sim.py

@@ -0,0 +1,207 @@
+#!/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.
+
+"""Synthetic scenes for hardware-free dry-runs and tests.
+
+Ported from the dyna360 eval harness. A :class:`SyntheticScene` is a
+deterministic hand-crafted :class:`~lerobot.navigation.voxel_map.VoxelMap`
+— a navigable floor plus labelled objects each carrying a unit feature
+vector — paired with a
+:class:`~lerobot.navigation.features.BasisVectorFeatureExtractor` whose
+text encodings live in the same space. This lets ``dog_cli --dry-run``
+(and the tests) exercise the full locate/goto/explore stack with no
+models, camera, or robot.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+
+import numpy as np
+
+from lerobot.navigation.features import BasisVectorFeatureExtractor
+from lerobot.navigation.voxel_map import VoxelMap
+
+LOG = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class SyntheticObject:
+    """One labelled object. ``feature_vec`` lives in the same space as the
+    text embeddings fed to ``VoxelMap.query`` (one-hot basis vectors, so a
+    query hits the right cluster cleanly)."""
+
+    name: str
+    xyz: tuple[float, float, float]
+    half_extent_m: float
+    feature_vec: np.ndarray
+
+
+@dataclass(frozen=True)
+class SyntheticScene:
+    """A ground-truth scene: voxel map + object metadata."""
+
+    voxel_map: VoxelMap
+    objects: list[SyntheticObject]
+    floor_extent_m: float
+    voxel_size: float
+    feature_dim: int
+
+    def name_to_xyz(self) -> dict[str, tuple[float, float, float]]:
+        return {o.name: o.xyz for o in self.objects}
+
+    def object(self, name: str) -> SyntheticObject | None:
+        for o in self.objects:
+            if o.name == name:
+                return o
+        return None
+
+    def feature_extractor(self) -> BasisVectorFeatureExtractor:
+        """A text encoder whose vectors match this scene's object features."""
+        table = {o.name: o.feature_vec for o in self.objects}
+        return BasisVectorFeatureExtractor(table, self.feature_dim)
+
+
+@dataclass(frozen=True)
+class SceneSpec:
+    """Declarative recipe used by :func:`build_scene`."""
+
+    objects: list[SyntheticObject]
+    floor_extent_m: float = 6.0
+    voxel_size: float = 0.1
+    feature_dim: int = 8
+    ground_y: float = 1.0
+    object_density_per_dim: int = 5
+    wall_xz_range: tuple[float, float, float, float] | None = None
+    """Optional axis-aligned wall ``(x_min, z_min, x_max, z_max)`` of
+    obstacle voxels at robot height — to test ``goto`` against a block."""
+    feature_noise: float = 0.0
+    rng_seed: int = 0
+
+
+def basis_vec(dim: int, idx: int) -> np.ndarray:
+    """A unit basis vector of length ``dim`` with a 1 at ``idx``."""
+    v = np.zeros(dim, dtype=np.float32)
+    v[idx] = 1.0
+    return v
+
+
+def build_scene(spec: SceneSpec) -> SyntheticScene:
+    """Construct a deterministic :class:`SyntheticScene` from a spec."""
+    rng = np.random.default_rng(spec.rng_seed)
+    vm = VoxelMap(voxel_size=spec.voxel_size)
+
+    # ----- floor (NAVIGABLE) -----
+    half = spec.voxel_size / 2.0
+    floor_pts: list[tuple[float, float, float]] = []
+    for x in np.arange(-spec.floor_extent_m + half, spec.floor_extent_m + half, spec.voxel_size):
+        for z in np.arange(-spec.floor_extent_m + half, spec.floor_extent_m + half, spec.voxel_size):
+            floor_pts.append((float(x), spec.ground_y, float(z)))
+    arr = np.asarray(floor_pts, dtype=np.float64).reshape(-1, 1, 3)
+    rgb = np.full((len(floor_pts), 1, 3), 180, dtype=np.uint8)
+    conf = np.ones((len(floor_pts), 1), dtype=np.float32)
+    if spec.feature_dim >= 1:
+        floor_vec = np.zeros(spec.feature_dim, dtype=np.float16)
+        floor_vec[-1] = 1.0
+        floor_feat = np.tile(floor_vec, (len(floor_pts), 1, 1))
+        vm.add(arr, rgb, conf, frame=0, t=0.0, feat_map=floor_feat)
+    else:
+        vm.add(arr, rgb, conf, frame=0, t=0.0)
+
+    # ----- objects -----
+    for i, obj in enumerate(spec.objects, start=1):
+        d = obj.half_extent_m
+        n = spec.object_density_per_dim
+        coords = np.linspace(-d + half, d - half, n)
+        pts = np.array(
+            [
+                (float(obj.xyz[0] + dx), float(obj.xyz[1] + dy), float(obj.xyz[2] + dz))
+                for dx in coords
+                for dy in coords
+                for dz in coords
+            ],
+            dtype=np.float64,
+        ).reshape(-1, 1, 3)
+        rgb_o = np.full((pts.shape[0], 1, 3), 100 + (i * 30) % 156, dtype=np.uint8)
+        conf_o = np.ones((pts.shape[0], 1), dtype=np.float32)
+
+        if obj.feature_vec.shape != (spec.feature_dim,):
+            raise ValueError(
+                f"object {obj.name!r} feature_vec has shape {obj.feature_vec.shape}, "
+                f"expected ({spec.feature_dim},) to match SceneSpec.feature_dim"
+            )
+        base = obj.feature_vec.astype(np.float32).reshape(1, 1, -1)
+        feats = np.tile(base, (pts.shape[0], 1, 1))
+        if spec.feature_noise > 0:
+            noise = rng.normal(scale=spec.feature_noise, size=feats.shape).astype(np.float32)
+            feats = feats + noise
+            norms = np.linalg.norm(feats, axis=-1, keepdims=True)
+            feats = feats / np.maximum(norms, 1e-6)
+        vm.add(pts, rgb_o, conf_o, frame=i, t=float(i), feat_map=feats.astype(np.float16))
+
+    # ----- optional wall (OBSTACLE) -----
+    if spec.wall_xz_range is not None:
+        wx0, wz0, wx1, wz1 = spec.wall_xz_range
+        wall_pts = [
+            (float(x), float(y), float(z))
+            for x in np.arange(wx0 + half, wx1, spec.voxel_size)
+            for z in np.arange(wz0 + half, wz1, spec.voxel_size)
+            for y in np.arange(spec.ground_y - 1.0, spec.ground_y - 0.1, spec.voxel_size)
+        ]
+        if wall_pts:
+            pts = np.asarray(wall_pts, dtype=np.float64).reshape(-1, 1, 3)
+            rgb_w = np.full((len(wall_pts), 1, 3), 80, dtype=np.uint8)
+            conf_w = np.ones((len(wall_pts), 1), dtype=np.float32)
+            vm.add(pts, rgb_w, conf_w, frame=99, t=99.0)
+
+    LOG.info(
+        "built scene: %d voxels, %d objects, floor extent %.1f m, D=%d",
+        len(vm),
+        len(spec.objects),
+        spec.floor_extent_m,
+        spec.feature_dim,
+    )
+    return SyntheticScene(
+        voxel_map=vm,
+        objects=list(spec.objects),
+        floor_extent_m=spec.floor_extent_m,
+        voxel_size=spec.voxel_size,
+        feature_dim=spec.feature_dim,
+    )
+
+
+_KITCHEN_DIM = 64  # Feature dim sized so the random-direction noise floor
+# (≈1/sqrt(D) ≈ 0.125) sits well below a sane locate threshold, so an absent
+# object reliably ABSTAINS instead of hitting a known basis vector.
+
+
+def kitchen_scene(wall: tuple[float, float, float, float] | None = None) -> SyntheticScene:
+    """A 6×6 m floor with four labelled objects at distinctive corners."""
+    spec = SceneSpec(
+        objects=[
+            SyntheticObject("couch", (3.0, 0.5, 2.0), 0.3, basis_vec(_KITCHEN_DIM, 0)),
+            SyntheticObject("chair", (-2.0, 0.5, -1.5), 0.2, basis_vec(_KITCHEN_DIM, 1)),
+            SyntheticObject("lamp", (2.5, 0.5, -2.0), 0.15, basis_vec(_KITCHEN_DIM, 2)),
+            SyntheticObject("plant", (-2.5, 0.5, 2.5), 0.25, basis_vec(_KITCHEN_DIM, 3)),
+        ],
+        floor_extent_m=6.0,
+        voxel_size=0.1,
+        feature_dim=_KITCHEN_DIM,
+        ground_y=1.0,
+        wall_xz_range=wall,
+    )
+    return build_scene(spec)

src/lerobot/navigation/skills.py

@@ -0,0 +1,321 @@
+#!/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.
+
+"""SpatialSkills tool layer.
+
+Ported from the dyna360 research stack. The agent calls these as a fixed
+toolset:
+
+  - :meth:`SpatialSkills.locate` — text → 3D position (or NOT_FOUND)
+  - :meth:`SpatialSkills.goto` — base navigation to a 3D target
+  - :meth:`SpatialSkills.explore` — pick a frontier to drive toward
+
+The skills compose a :class:`~lerobot.navigation.voxel_map.VoxelMap`
+(geometry + semantic features) with a
+:class:`~lerobot.navigation.base_controller.BaseController` (motion) and a
+text encoder. Stateless-per-call: each call snapshots the world, does its
+work, and hands control back. The agent decides what to call next.
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from lerobot.navigation.occupancy import (
+    OccupancyGrid,
+    astar,
+    find_frontier_cells,
+    project_voxel_map_to_grid,
+)
+from lerobot.navigation.value_map import (
+    ValueMapConfig,
+    compute_value_maps,
+    pick_best_frontier_cell,
+)
+
+if TYPE_CHECKING:
+    from lerobot.navigation.base_controller import BaseController
+    from lerobot.navigation.features import FeatureExtractor
+    from lerobot.navigation.voxel_map import VoxelMap
+
+LOG = logging.getLogger(__name__)
+
+
+# ----- typed results returned to the agent -------------------------------
+
+
+@dataclass(frozen=True)
+class LocateResult:
+    """Output of :meth:`SpatialSkills.locate`.
+
+    ``found=False`` is load-bearing — the signal the agent uses to pick
+    :meth:`explore` over :meth:`goto`. Don't fabricate an ``xyz`` when
+    abstaining.
+    """
+
+    found: bool
+    xyz: tuple[float, float, float] | None
+    confidence: float  # top cosine score; -1.0 if no features
+    n_voxels: int  # how many voxels supported the cluster
+    text: str
+
+
+@dataclass(frozen=True)
+class GotoResult:
+    """Output of :meth:`SpatialSkills.goto`."""
+
+    reached: bool
+    final_xyz: tuple[float, float, float]
+    distance_to_target: float
+    n_steps: int
+    reason: str  # "ok" | "no path" | "max steps" | "blocked"
+    path_xyz: list[tuple[float, float, float]]  # for viz / debugging
+
+
+@dataclass(frozen=True)
+class ExploreResult:
+    """Output of :meth:`SpatialSkills.explore`."""
+
+    target_xyz: tuple[float, float, float] | None
+    found_frontier: bool
+    distance_to_target: float  # 0.0 when no frontier
+    reason: str  # "ok" | "no frontier" | ...
+    value: float = 0.0
+    """Combined V_T + α·V_S value of the chosen frontier — useful for
+    debugging exploration bias and as a give-up signal for the agent."""
+
+
+# ----- configuration ------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class SkillsConfig:
+    """Knobs shared across the skills."""
+
+    # Occupancy projection
+    cell_size: float = 0.1
+    ground_y: float | None = None  # None ⇒ auto-estimate from voxels
+    obstacle_y_range: tuple[float, float] = (-2.0, -0.1)  # m above ground (y-down)
+    obstacle_inflate_cells: int = 1
+
+    # locate()
+    locate_top_k: int = 128
+    locate_threshold: float = 0.15  # min cosine for found=True
+    locate_outlier_quantile: float = 0.5
+    locate_outlier_scale: float = 2.0
+
+    # goto()
+    goto_threshold: float = 0.3
+    goto_step_size: float = 0.2  # m advanced per controller tick
+    goto_max_steps: int = 500
+    goto_replan_every: int = 5
+    goto_dt: float = 0.1
+
+    # explore()
+    explore_max_frontiers: int = 256
+    value_cfg: ValueMapConfig = field(default_factory=ValueMapConfig)
+    """DynaMem-style V_T (recency) + V_S (similarity) knobs."""
+
+
+# ----- the skills layer ---------------------------------------------------
+
+
+class SpatialSkills:
+    """Composes the voxel memory + base + text encoder into the agent toolset."""
+
+    def __init__(
+        self,
+        voxel_map: VoxelMap,
+        base: BaseController,
+        siglip: FeatureExtractor | None = None,
+        cfg: SkillsConfig | None = None,
+    ) -> None:
+        self.voxel_map = voxel_map
+        self.base = base
+        self.siglip = siglip
+        self.cfg = cfg or SkillsConfig()
+
+    # ----- shared helper ---------------------------------------------------
+
+    def occupancy(self) -> OccupancyGrid:
+        """Project the *current* voxel map into a 2D occupancy grid."""
+        return project_voxel_map_to_grid(
+            self.voxel_map,
+            cell_size=self.cfg.cell_size,
+            ground_y=self.cfg.ground_y,
+            obstacle_y_range=self.cfg.obstacle_y_range,
+            inflate_cells=self.cfg.obstacle_inflate_cells,
+        )
+
+    # ----- locate(text) ----------------------------------------------------
+
+    def locate(self, text: str) -> LocateResult:
+        text = text.strip()
+        if not text:
+            return LocateResult(False, None, -1.0, 0, text)
+        if self.siglip is None:
+            return LocateResult(False, None, -1.0, 0, text)
+        if self.voxel_map.feature_dim is None:
+            return LocateResult(False, None, -1.0, 0, text)
+
+        text_emb = self.siglip.encode_text(text)
+        qr = self.voxel_map.query(text_emb, top_k=self.cfg.locate_top_k)
+        if qr.score.size == 0:
+            return LocateResult(False, None, -1.0, 0, text)
+        top_score = float(qr.score.max())
+        if top_score < self.cfg.locate_threshold:
+            LOG.info(
+                "locate(%r): top score %.3f < threshold %.3f → NOT_FOUND",
+                text,
+                top_score,
+                self.cfg.locate_threshold,
+            )
+            return LocateResult(False, None, top_score, 0, text)
+
+        # Score-weighted centroid, then outlier rejection (anchor against the
+        # cluster median distance so a couple of stray voxels in the top-k
+        # can't drag the centroid into empty space).
+        scores = qr.score.astype(np.float64)
+        weights = scores - scores.min() + 1e-6
+        centroid = (qr.xyz * weights[:, None]).sum(axis=0) / weights.sum()
+        d = np.linalg.norm(qr.xyz - centroid, axis=1)
+        thresh = max(
+            self.cfg.cell_size * 4,
+            float(np.quantile(d, self.cfg.locate_outlier_quantile)) * self.cfg.locate_outlier_scale,
+        )
+        inliers = d <= thresh
+        if inliers.sum() >= 3:
+            inlier_xyz = qr.xyz[inliers]
+            inlier_w = weights[inliers]
+            centroid = (inlier_xyz * inlier_w[:, None]).sum(axis=0) / inlier_w.sum()
+        return LocateResult(
+            True,
+            (float(centroid[0]), float(centroid[1]), float(centroid[2])),
+            top_score,
+            int(inliers.sum()),
+            text,
+        )
+
+    # ----- goto(xyz) -------------------------------------------------------
+
+    def goto(
+        self,
+        target_xyz: tuple[float, float, float],
+        *,
+        max_steps: int | None = None,
+        threshold: float | None = None,
+    ) -> GotoResult:
+        """Closed-loop nav: A* → step a few cells → replan → repeat.
+
+        The replan cadence makes this a staleness governor — a moving
+        obstacle (or a previously-mapped one that got carved out) is picked
+        up at the next replan.
+        """
+        max_steps = max_steps if max_steps is not None else self.cfg.goto_max_steps
+        threshold = threshold if threshold is not None else self.cfg.goto_threshold
+
+        path_xyz_global: list[tuple[float, float, float]] = []
+        n_steps = 0
+        last_path: list[tuple[float, float]] = []
+
+        for step in range(max_steps):
+            pos = self.base.position()
+            d = math.hypot(pos[0] - target_xyz[0], pos[2] - target_xyz[2])
+            if d <= threshold:
+                return GotoResult(True, pos, d, n_steps, "ok", path_xyz_global)
+
+            if step % self.cfg.goto_replan_every == 0 or not last_path:
+                grid = self.occupancy()
+                last_path = (
+                    astar(
+                        grid,
+                        start_world=(pos[0], pos[2]),
+                        goal_world=(target_xyz[0], target_xyz[2]),
+                    )
+                    or []
+                )
+                if not last_path or len(last_path) < 2:
+                    return GotoResult(False, pos, d, n_steps, "no path", path_xyz_global)
+
+            # Head toward the next-but-one cell to smooth corners.
+            next_idx = min(2, len(last_path) - 1)
+            target_xz = last_path[next_idx]
+            dx = target_xz[0] - pos[0]
+            dz = target_xz[1] - pos[2]
+            n = math.hypot(dx, dz)
+            if n < 1e-6:
+                last_path.pop(0)
+                continue
+            vx = self.cfg.goto_step_size / max(self.cfg.goto_dt, 1e-6) * dx / n
+            vz = self.cfg.goto_step_size / max(self.cfg.goto_dt, 1e-6) * dz / n
+            self.base.move(vx=vx, vz=vz, dt=self.cfg.goto_dt)
+            pos = self.base.position()
+            path_xyz_global.append(pos)
+            n_steps += 1
+
+            # Pop waypoint when we've crossed it.
+            if math.hypot(target_xz[0] - pos[0], target_xz[1] - pos[2]) < self.cfg.cell_size:
+                last_path.pop(0)
+                if not last_path:
+                    last_path = []  # force replan
+
+        pos = self.base.position()
+        d = math.hypot(pos[0] - target_xyz[0], pos[2] - target_xyz[2])
+        return GotoResult(False, pos, d, n_steps, "max steps", path_xyz_global)
+
+    # ----- explore() -------------------------------------------------------
+
+    def explore(self, query: str | None = None) -> ExploreResult:
+        """Pick a frontier to drive toward via the DynaMem §3.4 value map.
+
+        With no query this is pure recency (visit oldest-observed or
+        UNOBSERVED frontiers first); with a query + features it biases
+        toward semantic matches.
+        """
+        grid = self.occupancy()
+        cells = find_frontier_cells(grid)
+        if cells.shape[0] == 0:
+            return ExploreResult(None, False, 0.0, "no frontier")
+
+        # Subsample if huge so the loop stays fast even on big maps.
+        if cells.shape[0] > self.cfg.explore_max_frontiers:
+            idx = np.random.default_rng(0).choice(
+                cells.shape[0], self.cfg.explore_max_frontiers, replace=False
+            )
+            cells = cells[idx]
+
+        text_emb = None
+        if query is not None and self.siglip is not None and self.voxel_map.feature_dim is not None:
+            text_emb = self.siglip.encode_text(query)
+
+        values = compute_value_maps(self.voxel_map, grid, text_emb=text_emb, cfg=self.cfg.value_cfg)
+
+        pos = self.base.position()
+        _, (xt, zt), dist, score = pick_best_frontier_cell(
+            grid, cells, values, robot_position_xz=(pos[0], pos[2]), cfg=self.cfg.value_cfg
+        )
+        return ExploreResult(
+            target_xyz=(xt, grid.ground_y, zt),
+            found_frontier=True,
+            distance_to_target=dist,
+            reason="ok",
+            value=score,
+        )

src/lerobot/navigation/value_map.py

@@ -0,0 +1,221 @@
+#!/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.
+
+"""DynaMem-style value maps for exploration.
+
+Ported from the dyna360 research stack. Two scalar fields over the same
+occupancy grid as :mod:`occupancy`:
+
+  - **V_T (time-recency)** — sigmoid of "how long ago was this cell last
+    observed?" Cells not seen in a while (or never) score high; freshly
+    observed cells score low. This biases exploration away from
+    just-covered territory.
+  - **V_S (query-similarity)** — sigmoid of the cosine between the cell's
+    aggregated feature and a text query. Only defined when a query is
+    given AND the voxel map carries features.
+
+Linear combination ``V = (1 − α)·V_T + α·V_S`` gates exploration. With no
+query it is a pure recency-driven frontier walk; with a query it biases
+toward regions semantically consistent with the target (DynaMem §3.4).
+Maps are derived per-call from ``VoxelMap.snapshot`` so they inherit
+carving for free.
+"""
+
+# ruff: noqa: N806  — H, W, D are conventional array-dimension names
+from __future__ import annotations
+
+import logging
+import math
+from dataclasses import dataclass
+
+import numpy as np
+
+from lerobot.navigation.occupancy import OccupancyGrid
+
+LOG = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class ValueMapConfig:
+    """Knobs shared between recency and similarity value maps."""
+
+    recency_mid_s: float = 10.0
+    """Age (s) at which V_T crosses 0.5 — older = more interesting."""
+
+    recency_scale_s: float = 8.0
+    """How sharply V_T transitions around the mid age. Smaller = sharper."""
+
+    similarity_mid: float = 0.15
+    """Cosine score at which V_S crosses 0.5."""
+
+    similarity_scale: float = 0.05
+    """How sharply V_S transitions around the mid cosine."""
+
+    alpha_similarity: float = 0.6
+    """Weight of V_S in the combined value when a query is given.
+    0.0 = pure recency, 1.0 = pure similarity."""
+
+    unknown_value: float = 1.0
+    """V_T for UNOBSERVED cells — they are maximally interesting."""
+
+    distance_discount_per_meter: float = 0.05
+    """Multiplicative discount on far frontiers so the base does not
+    ping-pong across the map. 0 disables."""
+
+
+@dataclass(frozen=True)
+class ValueMaps:
+    """The scalar fields, all shaped ``(H, W)`` like the occupancy grid."""
+
+    last_time: np.ndarray  # float64 — −inf where UNOBSERVED
+    recency: np.ndarray  # float32 V_T in [0, 1]
+    similarity: np.ndarray | None  # float32 V_S in [0, 1], None when no query
+    combined: np.ndarray  # float32 V — what explore() optimizes
+
+
+# --------------------------------------------------------------------- #
+
+
+def _eps_for_cell(cell_size: float) -> float:
+    """Same float32-drift epsilon as :mod:`occupancy` so the two
+    projections agree on which voxels land in which cells."""
+    return cell_size * 1e-3
+
+
+def _project_voxels_to_cells(voxel_map, grid: OccupancyGrid, want_features: bool):
+    """Project every voxel into its XZ cell.
+
+    Returns ``(last_time_per_cell, feat_per_cell)`` where last_time is
+    (H, W) float64 (−inf for empty cells) and feat_per_cell is
+    (H, W, D) float32 or None.
+    """
+    snap = voxel_map.snapshot(include_features=want_features)
+    H, W = grid.shape
+    last_time = np.full((H, W), -math.inf, dtype=np.float64)
+    if snap.xyz.size == 0:
+        return last_time, None
+
+    x = snap.xyz[:, 0].astype(np.float64)
+    z = snap.xyz[:, 2].astype(np.float64)
+    eps = _eps_for_cell(grid.cell_size)
+    ix = np.clip(np.floor((x - grid.origin_x) / grid.cell_size + eps).astype(np.int32), 0, W - 1)
+    iz = np.clip(np.floor((z - grid.origin_z) / grid.cell_size + eps).astype(np.int32), 0, H - 1)
+
+    # Per-cell max last_time. `np.maximum.at` is the unbuffered ufunc version,
+    # which correctly handles duplicate (iz, ix) targets.
+    np.maximum.at(last_time, (iz, ix), snap.last_time.astype(np.float64))
+
+    feat_per_cell: np.ndarray | None = None
+    if want_features and snap.feat is not None and snap.feat.size > 0:
+        D = snap.feat.shape[1]
+        feat_sum = np.zeros((H, W, D), dtype=np.float32)
+        np.add.at(feat_sum, (iz, ix), snap.feat.astype(np.float32))
+        counts = np.zeros((H, W), dtype=np.int32)
+        np.add.at(counts, (iz, ix), 1)
+        # Normalize per-cell — count is the number of CONTRIBUTING voxels.
+        denom = np.maximum(counts, 1).astype(np.float32)[..., None]
+        feat_per_cell = feat_sum / denom
+
+    return last_time, feat_per_cell
+
+
+def _recency_value(last_time_per_cell: np.ndarray, now_t: float, cfg: ValueMapConfig) -> np.ndarray:
+    """V_T per cell. Unobserved cells get ``cfg.unknown_value``."""
+    out = np.full(last_time_per_cell.shape, cfg.unknown_value, dtype=np.float32)
+    observed = last_time_per_cell > -math.inf
+    if not observed.any():
+        return out
+    age = (now_t - last_time_per_cell[observed]).astype(np.float32)
+    out[observed] = 1.0 / (1.0 + np.exp(-(age - cfg.recency_mid_s) / cfg.recency_scale_s))
+    return out
+
+
+def _similarity_value(
+    feat_per_cell: np.ndarray | None,
+    text_emb: np.ndarray | None,
+    cfg: ValueMapConfig,
+) -> np.ndarray | None:
+    """V_S per cell. ``None`` when there are no features or no query."""
+    if feat_per_cell is None or text_emb is None:
+        return None
+    text = text_emb.astype(np.float32)
+    text = text / max(float(np.linalg.norm(text)), 1e-6)
+    # Per-cell mean feat may not be unit-norm — renormalize so the dot product
+    # behaves like a cosine. Empty cells stay a 0 vector, so renorm clamps to 0.
+    norms = np.linalg.norm(feat_per_cell, axis=-1, keepdims=True)
+    feat_normed = feat_per_cell / np.maximum(norms, 1e-6)
+    with np.errstate(invalid="ignore", over="ignore", divide="ignore"):
+        cosine = np.nan_to_num((feat_normed @ text).astype(np.float32))
+    sim = 1.0 / (1.0 + np.exp(-(cosine - cfg.similarity_mid) / cfg.similarity_scale))
+    sim = np.where(norms.squeeze(-1) > 1e-6, sim, 0.0).astype(np.float32)
+    return sim
+
+
+def compute_value_maps(
+    voxel_map,
+    grid: OccupancyGrid,
+    *,
+    text_emb: np.ndarray | None = None,
+    now_t: float | None = None,
+    cfg: ValueMapConfig | None = None,
+) -> ValueMaps:
+    """Build the full value-map bundle for one ``explore`` call."""
+    cfg = cfg or ValueMapConfig()
+    last_time, feat_per_cell = _project_voxels_to_cells(voxel_map, grid, want_features=(text_emb is not None))
+    if now_t is None:
+        observed_mask = last_time > -math.inf
+        now_t = float(last_time[observed_mask].max()) if observed_mask.any() else 0.0
+
+    v_t = _recency_value(last_time, now_t, cfg)
+    v_s = _similarity_value(feat_per_cell, text_emb, cfg)
+
+    if v_s is not None:
+        combined = ((1.0 - cfg.alpha_similarity) * v_t + cfg.alpha_similarity * v_s).astype(np.float32)
+    else:
+        combined = v_t
+
+    return ValueMaps(last_time=last_time, recency=v_t, similarity=v_s, combined=combined)
+
+
+def pick_best_frontier_cell(
+    grid: OccupancyGrid,
+    frontier_cells: np.ndarray,
+    values: ValueMaps,
+    robot_position_xz: tuple[float, float],
+    cfg: ValueMapConfig | None = None,
+) -> tuple[int, tuple[float, float], float, float]:
+    """Score every frontier cell by ``values.combined`` (with a distance
+    discount) and return the winner.
+
+    Returns ``(index_into_frontier_cells, (x, z), distance_m, score)``.
+    """
+    if frontier_cells.shape[0] == 0:
+        raise ValueError("frontier_cells is empty")
+    cfg = cfg or ValueMapConfig()
+
+    iz_f = frontier_cells[:, 0]
+    ix_f = frontier_cells[:, 1]
+    raw = values.combined[iz_f, ix_f]
+
+    xs = grid.origin_x + (ix_f.astype(np.float64) + 0.5) * grid.cell_size
+    zs = grid.origin_z + (iz_f.astype(np.float64) + 0.5) * grid.cell_size
+    rx, rz = robot_position_xz
+    d = np.hypot(xs - rx, zs - rz)
+    discount = 1.0 / (1.0 + cfg.distance_discount_per_meter * d)
+    scored = raw * discount
+
+    best = int(np.argmax(scored))
+    return best, (float(xs[best]), float(zs[best])), float(d[best]), float(scored[best])

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,
+            ),
+        )

src/lerobot/navigation/voxel_map.py

@@ -0,0 +1,511 @@
+#!/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.
+
+"""Sparse-hash voxel memory with free-space carving + semantic features.
+
+Ported from the dyna360 research stack. Per occupied voxel: voxel index,
+running-mean xyz (count-weighted), running-mean rgb (count-weighted),
+count, last_frame, last_time, and — once vision-language features have
+been fed in — a conf-weighted running-mean feature in fp16 plus the
+weight sum.
+
+Storage is hybrid: a Python dict maps voxel index ``(ix, iy, iz)`` to a
+row in column-stored numpy arrays so lookup is O(1) and bulk arithmetic
+stays vectorized. ``carve`` removes voxels that fall inside a view's
+observed free space (DynaMem-style dynamic updates); ``query`` returns
+the top-k cosine matches against a text embedding.
+
+Default voxel size is 5 cm. The map is geometry-only until
+``add(..., feat_map=...)`` supplies per-pixel features; occupancy /
+planning use only the geometry, so they work without any features.
+"""
+
+# ruff: noqa: N806  — H, W, D are conventional array-dimension names (and appear verbatim in error strings)
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+
+import numpy as np
+
+LOG = logging.getLogger(__name__)
+
+
+@dataclass
+class VoxelMapStats:
+    """Per-keyframe deltas, surfaced to scalar logs."""
+
+    n_voxels: int
+    n_added: int
+    n_updated: int
+    n_removed: int = 0
+
+
+@dataclass(frozen=True)
+class VoxelSnapshot:
+    """Current voxel map state, materialized for visualization / export."""
+
+    xyz: np.ndarray  # (M, 3) float32 — count-weighted mean position
+    rgb: np.ndarray  # (M, 3) uint8 — count-weighted mean color (RGB)
+    count: np.ndarray  # (M,)  int64
+    last_frame: np.ndarray  # (M,)  int64
+    last_time: np.ndarray  # (M,)  float64
+    feat: np.ndarray | None = None  # (M, D) fp16 — L2-normalized per-voxel mean
+
+
+@dataclass(frozen=True)
+class CarveResult:
+    """Output of one ``carve`` pass."""
+
+    n_removed: int
+    removed_xyz: np.ndarray  # (K, 3) float32 — centres of removed voxels, for viz
+
+
+@dataclass(frozen=True)
+class QueryResult:
+    """Top-k cosine matches against a text embedding."""
+
+    xyz: np.ndarray  # (k, 3) float32
+    score: np.ndarray  # (k,) float32 — cosine similarity in [-1, 1]
+    voxel_indices: np.ndarray  # (k,) int64 — row indices into the map
+
+
+_MAX_ABS_VOXEL_INDEX = 1 << 20
+_FEAT_CHUNK_PIXELS = 16384  # bound peak per-keyframe feature contribution memory
+
+
+class VoxelMap:
+    """Sparse-hash voxel grid with count-weighted means and semantic features."""
+
+    def __init__(self, voxel_size: float = 0.05) -> None:
+        if voxel_size <= 0:
+            raise ValueError("voxel_size must be > 0")
+        self.voxel_size = float(voxel_size)
+
+        self._lookup: dict[tuple[int, int, int], int] = {}
+        self._idx = np.zeros((0, 3), dtype=np.int64)
+        self._count = np.zeros(0, dtype=np.int64)
+        self._xyz_sum = np.zeros((0, 3), dtype=np.float64)
+        self._rgb_sum = np.zeros((0, 3), dtype=np.float64)
+        self._last_frame = np.zeros(0, dtype=np.int64)
+        self._last_time = np.zeros(0, dtype=np.float64)
+
+        # Lazily allocated on first add() with feat_map.
+        self._feature_dim: int | None = None
+        self._feat_sum: np.ndarray | None = None  # (M, D) fp16
+        self._feat_weight: np.ndarray | None = None  # (M,) fp32
+
+    def __len__(self) -> int:
+        return int(self._count.shape[0])
+
+    @property
+    def feature_dim(self) -> int | None:
+        return self._feature_dim
+
+    # ------------------------------------------------------------------ add
+
+    def add(
+        self,
+        points: np.ndarray,
+        rgb: np.ndarray,
+        conf: np.ndarray,
+        frame: int,
+        t: float,
+        conf_thresh: float = 0.5,
+        feat_map: np.ndarray | None = None,
+    ) -> VoxelMapStats:
+        """Insert / update voxels from a per-pixel observation.
+
+        ``points``:   ``(..., 3)`` world xyz, fp32.
+        ``rgb``:      ``(..., 3)`` uint8 (RGB order).
+        ``conf``:     ``(...,)`` in [0, 1].
+        ``feat_map``: optional ``(..., D)`` fp16 per-pixel feature, already
+                      bilinearly upsampled to the points/conf grid. First
+                      call with features locks the feature dimension;
+                      subsequent calls must match.
+        """
+        pts = np.asarray(points).reshape(-1, 3)
+        cols = np.asarray(rgb).reshape(-1, 3)
+        cnf = np.asarray(conf).reshape(-1)
+        if not (len(pts) == len(cols) == len(cnf)):
+            raise ValueError(f"length mismatch: points={len(pts)}, rgb={len(cols)}, conf={len(cnf)}")
+
+        features: np.ndarray | None = None
+        if feat_map is not None:
+            features = np.asarray(feat_map).reshape(-1, feat_map.shape[-1])
+            if len(features) != len(pts):
+                raise ValueError(f"feat_map length {len(features)} != points length {len(pts)}")
+            D = features.shape[-1]
+            if self._feature_dim is None:
+                self._feature_dim = int(D)
+                # Pad pre-existing voxels (added before features arrived) with zeros.
+                self._feat_sum = np.zeros((len(self), D), dtype=np.float16)
+                self._feat_weight = np.zeros(len(self), dtype=np.float32)
+                LOG.info("VoxelMap features enabled: D=%d (fp16 storage)", D)
+            elif self._feature_dim != D:
+                raise ValueError(f"feature dim mismatch: existing={self._feature_dim}, got={D}")
+
+        mask = (cnf >= conf_thresh) & np.isfinite(pts).all(axis=1)
+        pts = pts[mask]
+        cols = cols[mask]
+        cnf_kept = cnf[mask]
+        if features is not None:
+            features = features[mask]
+        if pts.size == 0:
+            return VoxelMapStats(n_voxels=len(self), n_added=0, n_updated=0)
+
+        idx = np.floor(pts / self.voxel_size).astype(np.int64)
+        sane = (np.abs(idx) < _MAX_ABS_VOXEL_INDEX).all(axis=1)
+        if not sane.all():
+            n_drop = int((~sane).sum())
+            LOG.debug("dropping %d points with extreme voxel index", n_drop)
+            idx = idx[sane]
+            pts = pts[sane]
+            cols = cols[sane]
+            cnf_kept = cnf_kept[sane]
+            if features is not None:
+                features = features[sane]
+        if idx.size == 0:
+            return VoxelMapStats(n_voxels=len(self), n_added=0, n_updated=0)
+
+        unique_idx, inverse = np.unique(idx, axis=0, return_inverse=True)
+        inverse = inverse.reshape(-1)
+        n_unique = unique_idx.shape[0]
+        kf_count = np.bincount(inverse, minlength=n_unique).astype(np.int64)
+        kf_xyz_sum = np.zeros((n_unique, 3), dtype=np.float64)
+        kf_rgb_sum = np.zeros((n_unique, 3), dtype=np.float64)
+        np.add.at(kf_xyz_sum, inverse, pts.astype(np.float64))
+        np.add.at(kf_rgb_sum, inverse, cols.astype(np.float64))
+
+        kf_feat_sum: np.ndarray | None = None
+        kf_feat_weight: np.ndarray | None = None
+        if features is not None:
+            kf_feat_sum = np.zeros((n_unique, self._feature_dim), dtype=np.float32)
+            kf_feat_weight = np.zeros(n_unique, dtype=np.float32)
+            cnf_f = cnf_kept.astype(np.float32)
+            # Chunked accumulation — keeps the (chunk, D) intermediate small.
+            for s in range(0, features.shape[0], _FEAT_CHUNK_PIXELS):
+                e = s + _FEAT_CHUNK_PIXELS
+                w = cnf_f[s:e]
+                contrib = w[:, None] * features[s:e].astype(np.float32)
+                np.add.at(kf_feat_sum, inverse[s:e], contrib)
+                np.add.at(kf_feat_weight, inverse[s:e], w)
+
+        existing_rows: list[int] = []
+        existing_local: list[int] = []
+        new_local: list[int] = []
+        new_keys: list[tuple[int, int, int]] = []
+        for i in range(n_unique):
+            key = (int(unique_idx[i, 0]), int(unique_idx[i, 1]), int(unique_idx[i, 2]))
+            row = self._lookup.get(key)
+            if row is None:
+                new_local.append(i)
+                new_keys.append(key)
+            else:
+                existing_rows.append(row)
+                existing_local.append(i)
+
+        if existing_rows:
+            rows = np.asarray(existing_rows, dtype=np.int64)
+            local = np.asarray(existing_local, dtype=np.int64)
+            self._count[rows] += kf_count[local]
+            self._xyz_sum[rows] += kf_xyz_sum[local]
+            self._rgb_sum[rows] += kf_rgb_sum[local]
+            self._last_frame[rows] = frame
+            self._last_time[rows] = t
+            if kf_feat_sum is not None:
+                assert self._feat_sum is not None and self._feat_weight is not None
+                # fp32 accumulator -> fp16 storage; cast on store to match storage dtype.
+                self._feat_sum[rows] = (self._feat_sum[rows].astype(np.float32) + kf_feat_sum[local]).astype(
+                    np.float16
+                )
+                self._feat_weight[rows] += kf_feat_weight[local]
+
+        if new_local:
+            base = len(self)
+            local = np.asarray(new_local, dtype=np.int64)
+            self._idx = np.concatenate([self._idx, unique_idx[local]], axis=0)
+            self._count = np.concatenate([self._count, kf_count[local]])
+            self._xyz_sum = np.concatenate([self._xyz_sum, kf_xyz_sum[local]], axis=0)
+            self._rgb_sum = np.concatenate([self._rgb_sum, kf_rgb_sum[local]], axis=0)
+            self._last_frame = np.concatenate(
+                [self._last_frame, np.full(len(new_local), frame, dtype=np.int64)]
+            )
+            self._last_time = np.concatenate([self._last_time, np.full(len(new_local), t, dtype=np.float64)])
+            if self._feature_dim is not None:
+                assert self._feat_sum is not None and self._feat_weight is not None
+                if kf_feat_sum is not None:
+                    new_feats = kf_feat_sum[local].astype(np.float16)
+                    new_weights = kf_feat_weight[local]
+                else:
+                    # Features enabled, but this call didn't bring any — pad zeros
+                    # so array sizes stay aligned with _count.
+                    new_feats = np.zeros((len(new_local), self._feature_dim), dtype=np.float16)
+                    new_weights = np.zeros(len(new_local), dtype=np.float32)
+                self._feat_sum = np.concatenate([self._feat_sum, new_feats], axis=0)
+                self._feat_weight = np.concatenate([self._feat_weight, new_weights])
+            for offset, key in enumerate(new_keys):
+                self._lookup[key] = base + offset
+
+        return VoxelMapStats(
+            n_voxels=len(self),
+            n_added=len(new_local),
+            n_updated=len(existing_rows),
+        )
+
+    # ------------------------------------------------------- hard-delete
+    def remove_voxels_in_box(
+        self,
+        xyz_min: tuple[float, float, float],
+        xyz_max: tuple[float, float, float],
+    ) -> int:
+        """Surgical hard-delete of every voxel whose mean position lies inside
+        the axis-aligned bounding box.
+
+        Different from :meth:`carve` (DynaMem-style free-space removal from a
+        camera frustum + depth). This one is for simulated scene mutation:
+        "the couch moved away" is removing the box around the old couch then
+        ``add()``-ing one at the new position.
+        """
+        if len(self) == 0:
+            return 0
+        cnt = self._count.astype(np.float64).reshape(-1, 1)
+        means = self._xyz_sum / cnt
+        in_box = (
+            (means[:, 0] >= xyz_min[0])
+            & (means[:, 0] <= xyz_max[0])
+            & (means[:, 1] >= xyz_min[1])
+            & (means[:, 1] <= xyz_max[1])
+            & (means[:, 2] >= xyz_min[2])
+            & (means[:, 2] <= xyz_max[2])
+        )
+        if not in_box.any():
+            return 0
+        keep = ~in_box
+        n_removed = int(in_box.sum())
+        for k in self._idx[in_box]:
+            del self._lookup[(int(k[0]), int(k[1]), int(k[2]))]
+        self._idx = self._idx[keep]
+        self._count = self._count[keep]
+        self._xyz_sum = self._xyz_sum[keep]
+        self._rgb_sum = self._rgb_sum[keep]
+        self._last_frame = self._last_frame[keep]
+        self._last_time = self._last_time[keep]
+        if self._feat_sum is not None and self._feat_weight is not None:
+            self._feat_sum = self._feat_sum[keep]
+            self._feat_weight = self._feat_weight[keep]
+        # Row indices shifted — rebuild the lookup.
+        self._lookup = {
+            (int(self._idx[i, 0]), int(self._idx[i, 1]), int(self._idx[i, 2])): i
+            for i in range(len(self._idx))
+        }
+        return n_removed
+
+    # ---------------------------------------------------------------- carve
+
+    def carve(
+        self,
+        local_points: np.ndarray,
+        conf: np.ndarray,
+        pose: np.ndarray,
+        focal_px: float,
+        frame: int,
+        t: float,
+        conf_thresh: float = 0.5,
+        margin: float = 0.05,
+    ) -> CarveResult:
+        """Remove voxels inside this view's observed free space.
+
+        A voxel is carved when it projects into the image, sits in front of
+        the camera, and lies closer than the observed depth (minus a margin)
+        at that pixel — i.e. we can see through where it claims to be. Carve
+        runs before ``add`` each keyframe so moved/removed objects disappear.
+        """
+        if len(self) == 0:
+            return CarveResult(0, np.zeros((0, 3), dtype=np.float32))
+
+        if local_points.ndim != 3 or local_points.shape[-1] != 3:
+            raise ValueError(f"expected (H, W, 3), got {local_points.shape}")
+        if conf.shape != local_points.shape[:2]:
+            raise ValueError(f"conf shape {conf.shape} != local_points (H, W) {local_points.shape[:2]}")
+        if pose.shape != (4, 4):
+            raise ValueError(f"pose must be (4, 4); got {pose.shape}")
+
+        H, W = local_points.shape[:2]
+        cx = (W - 1) / 2.0
+        cy = (H - 1) / 2.0
+        depth_map = local_points[..., 2]
+
+        cnt = self._count.astype(np.float64).reshape(-1, 1)
+        xyz_world = self._xyz_sum / cnt
+
+        R = pose[:3, :3].astype(np.float64)
+        t_vec = pose[:3, 3].astype(np.float64)
+        xyz_cam = (xyz_world - t_vec[None, :]) @ R
+
+        d_voxel = xyz_cam[:, 2]
+        front = d_voxel > 1e-3
+
+        d_safe = np.where(front, d_voxel, 1.0)
+        u = focal_px * xyz_cam[:, 0] / d_safe + cx
+        v = focal_px * xyz_cam[:, 1] / d_safe + cy
+        in_bounds = (u >= 0.0) & (u < W) & (v >= 0.0) & (v < H)
+        valid = front & in_bounds
+
+        u_i = np.clip(np.floor(u).astype(np.int64), 0, W - 1)
+        v_i = np.clip(np.floor(v).astype(np.int64), 0, H - 1)
+        D_at = depth_map[v_i, u_i]
+        C_at = conf[v_i, u_i]
+
+        finite_D = np.isfinite(D_at) & (D_at > 0.0)
+        free_space = valid & finite_D & (C_at >= conf_thresh) & (d_voxel < (D_at - margin))
+
+        n_removed = int(free_space.sum())
+        if n_removed == 0:
+            return CarveResult(0, np.zeros((0, 3), dtype=np.float32))
+
+        removed_xyz = xyz_world[free_space].astype(np.float32)
+        removed_keys = self._idx[free_space]
+        for k in removed_keys:
+            del self._lookup[(int(k[0]), int(k[1]), int(k[2]))]
+
+        keep = ~free_space
+        self._idx = self._idx[keep]
+        self._count = self._count[keep]
+        self._xyz_sum = self._xyz_sum[keep]
+        self._rgb_sum = self._rgb_sum[keep]
+        self._last_frame = self._last_frame[keep]
+        self._last_time = self._last_time[keep]
+        if self._feat_sum is not None and self._feat_weight is not None:
+            self._feat_sum = self._feat_sum[keep]
+            self._feat_weight = self._feat_weight[keep]
+
+        self._lookup = {
+            (int(self._idx[i, 0]), int(self._idx[i, 1]), int(self._idx[i, 2])): i
+            for i in range(len(self._idx))
+        }
+        LOG.debug("carve frame=%d t=%.3fs removed=%d", frame, t, n_removed)
+        return CarveResult(n_removed=n_removed, removed_xyz=removed_xyz)
+
+    # ------------------------------------------------------------- snapshot
+
+    def snapshot(self, include_features: bool = False) -> VoxelSnapshot:
+        """Materialize the current map.
+
+        ``include_features``: pay the cost of normalizing the per-voxel
+        feature mean. Off by default — visualization doesn't need features.
+        """
+        if len(self) == 0:
+            return VoxelSnapshot(
+                xyz=np.zeros((0, 3), dtype=np.float32),
+                rgb=np.zeros((0, 3), dtype=np.uint8),
+                count=np.zeros(0, dtype=np.int64),
+                last_frame=np.zeros(0, dtype=np.int64),
+                last_time=np.zeros(0, dtype=np.float64),
+                feat=None,
+            )
+        cnt = self._count.astype(np.float64).reshape(-1, 1)
+        xyz = (self._xyz_sum / cnt).astype(np.float32)
+        rgb = np.clip(self._rgb_sum / cnt, 0, 255).astype(np.uint8)
+
+        feat = None
+        if include_features and self._feat_sum is not None and self._feat_weight is not None:
+            feat = self._normalized_features()
+
+        return VoxelSnapshot(
+            xyz=xyz,
+            rgb=rgb,
+            count=self._count.copy(),
+            last_frame=self._last_frame.copy(),
+            last_time=self._last_time.copy(),
+            feat=feat,
+        )
+
+    def _normalized_features(self) -> np.ndarray:
+        """Per-voxel L2-normalized feature mean. (M, D) fp16."""
+        assert self._feat_sum is not None and self._feat_weight is not None
+        w = np.maximum(self._feat_weight, 1e-6).reshape(-1, 1)
+        mean = self._feat_sum.astype(np.float32) / w
+        norms = np.linalg.norm(mean, axis=1, keepdims=True)
+        mean = mean / np.maximum(norms, 1e-6)
+        return mean.astype(np.float16)
+
+    # ----------------------------------------------------------------- query
+
+    def query(self, text_embedding: np.ndarray, top_k: int = 32) -> QueryResult:
+        """Top-k cosine matches against ``text_embedding``.
+
+        ``text_embedding``: ``(D,)`` array — does NOT need to be unit norm;
+        we re-normalize.
+        """
+        if self._feat_sum is None or self._feature_dim is None:
+            raise RuntimeError("VoxelMap has no semantic features yet — call add(..., feat_map=...) first")
+        if len(self) == 0:
+            return QueryResult(
+                xyz=np.zeros((0, 3), dtype=np.float32),
+                score=np.zeros(0, dtype=np.float32),
+                voxel_indices=np.zeros(0, dtype=np.int64),
+            )
+        if text_embedding.shape != (self._feature_dim,):
+            raise ValueError(f"text_embedding shape {text_embedding.shape} != ({self._feature_dim},)")
+
+        voxel_feat = self._normalized_features().astype(np.float32)
+        text_unit = text_embedding.astype(np.float32)
+        text_unit = text_unit / max(float(np.linalg.norm(text_unit)), 1e-6)
+
+        # fp16 feature storage can carry the odd inf/nan from a saturated
+        # running sum; the cosine stays well-defined, so don't warn on it.
+        with np.errstate(invalid="ignore", over="ignore", divide="ignore"):
+            scores = np.nan_to_num(voxel_feat @ text_unit)  # (M,)
+        k = min(int(top_k), len(scores))
+        # Partition-and-sort for the top-k.
+        top_idx = np.argpartition(scores, -k)[-k:]
+        order = np.argsort(-scores[top_idx])
+        top_idx = top_idx[order]
+
+        snap_xyz = (self._xyz_sum[top_idx] / self._count[top_idx].astype(np.float64).reshape(-1, 1)).astype(
+            np.float32
+        )
+        return QueryResult(
+            xyz=snap_xyz,
+            score=scores[top_idx].astype(np.float32),
+            voxel_indices=top_idx.astype(np.int64),
+        )
+
+    # --------------------------------------------------------- introspection
+
+    def memory_bytes(self) -> dict[str, int]:
+        """Return per-array memory footprint."""
+        out = {
+            "idx": self._idx.nbytes,
+            "count": self._count.nbytes,
+            "xyz_sum": self._xyz_sum.nbytes,
+            "rgb_sum": self._rgb_sum.nbytes,
+            "last_frame": self._last_frame.nbytes,
+            "last_time": self._last_time.nbytes,
+            "lookup_dict": _approx_dict_bytes(self._lookup),
+        }
+        if self._feat_sum is not None:
+            out["feat_sum"] = self._feat_sum.nbytes
+            assert self._feat_weight is not None
+            out["feat_weight"] = self._feat_weight.nbytes
+        out["total"] = sum(v for k, v in out.items() if k != "total")
+        return out
+
+
+def _approx_dict_bytes(d: dict) -> int:
+    """Rough lower-bound estimate; ~100 bytes/entry is a fine ballpark."""
+    return 100 * len(d)

src/lerobot/policies/eo1/modeling_eo1.py

@@ -18,7 +18,6 @@ from __future__ import annotations
 
 import contextlib
 import logging
-import math
 from collections import deque
 from typing import TYPE_CHECKING, Any
 
@@ -31,6 +30,8 @@ from torch import Tensor
 from lerobot.utils.constants import ACTION, OBS_STATE
 from lerobot.utils.import_utils import _transformers_available, require_package
 
+from ..common.flow_matching import euler_integrate, sample_noise, sample_time_beta
+from ..common.vla_utils import create_sinusoidal_pos_embedding, pad_vector
 from ..pretrained import PreTrainedPolicy
 from .configuration_eo1 import EO1Config
 
@@ -46,17 +47,6 @@ else:
 logger = logging.getLogger(__name__)
 
 
-def pad_vector(vector, new_dim):
-    """Pad the last dimension of a vector to new_dim with zeros.
-
-    Can be (batch_size x sequence_length x features_dimension)
-    or (batch_size x features_dimension)
-    """
-    if vector.shape[-1] >= new_dim:
-        return vector
-    return F.pad(vector, (0, new_dim - vector.shape[-1]))
-
-
 class EO1Policy(PreTrainedPolicy):
     """EO1 policy wrapper for LeRobot robot-only training/evaluation."""
 
@@ -136,47 +126,6 @@ class EO1Policy(PreTrainedPolicy):
         return self.parameters()
 
 
-def get_safe_dtype(target_dtype, device_type):
-    """Get a safe dtype for the given device type."""
-    if device_type == "mps" and target_dtype == torch.float64:
-        return torch.float32
-    if device_type == "cpu":
-        # CPU doesn't support bfloat16, use float32 instead
-        if target_dtype == torch.bfloat16:
-            return torch.float32
-        if target_dtype == torch.float64:
-            return torch.float64
-    return target_dtype
-
-
-def create_sinusoidal_pos_embedding(  # see openpi `create_sinusoidal_pos_embedding` (exact copy)
-    time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
-) -> Tensor:
-    """Computes sine-cosine positional embedding vectors for scalar positions."""
-    if dimension % 2 != 0:
-        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
-
-    if time.ndim != 1:
-        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
-
-    dtype = get_safe_dtype(torch.float64, device.type)
-    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
-    period = min_period * (max_period / min_period) ** fraction
-
-    # Compute the outer product
-    scaling_factor = 1.0 / period * 2 * math.pi
-    sin_input = scaling_factor[None, :] * time[:, None]
-    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
-
-
-def sample_beta(alpha, beta, bsize, device):  # see openpi `sample_beta` (exact copy)
-    # Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
-    alpha_t = torch.tensor(alpha, dtype=torch.float32)
-    beta_t = torch.tensor(beta, dtype=torch.float32)
-    dist = torch.distributions.Beta(alpha_t, beta_t)
-    return dist.sample((bsize,)).to(device)
-
-
 class EO1VisionActionProjector(torch.nn.Sequential):
     """This block implements the multi-layer perceptron (MLP) module."""
 
@@ -267,21 +216,17 @@ class EO1VisionFlowMatchingModel(nn.Module):
         return func(*args, **kwargs)
 
     def sample_noise(self, shape, device):
-        noise = torch.normal(
-            mean=0.0,
-            std=1.0,
-            size=shape,
-            dtype=torch.float32,
-            device=device,
-        )
-        return noise
+        return sample_noise(shape, device)
 
     def sample_time(self, bsize, device):
-        time_beta = sample_beta(
-            self.config.time_sampling_beta_alpha, self.config.time_sampling_beta_beta, bsize, device
+        return sample_time_beta(
+            bsize,
+            device,
+            alpha=self.config.time_sampling_beta_alpha,
+            beta=self.config.time_sampling_beta_beta,
+            scale=self.config.time_sampling_scale,
+            offset=self.config.time_sampling_offset,
         )
-        time = time_beta * self.config.time_sampling_scale + self.config.time_sampling_offset
-        return time.to(dtype=torch.float32, device=device)
 
     def get_placeholder_mask(
         self,
@@ -587,18 +532,11 @@ class EO1VisionFlowMatchingModel(nn.Module):
             (batch_size, chunk_size, self.config.max_action_dim),
             device,
         ).to(dtype=self.action_in_proj.weight.dtype)
-        dt = -1.0 / self.config.num_denoise_steps
         past_key_values = outputs.past_key_values
 
         # 3. Denoise only the action chunk while keeping the prefix cache invariant.
-        for step in range(self.config.num_denoise_steps):
-            time = torch.full(
-                (batch_size,),
-                1.0 + step * dt,
-                device=device,
-                dtype=torch.float32,
-            )
-            action_time_embs = self.embed_suffix(time, x_t)
+        def denoise_fn(input_x_t, current_timestep):
+            action_time_embs = self.embed_suffix(current_timestep, input_x_t)
             inputs_embeds[:, act_slice] = action_time_embs.to(inputs_embeds.dtype)
 
             # Keep the prefix KV cache invariant across denoising steps.
@@ -615,7 +553,7 @@ class EO1VisionFlowMatchingModel(nn.Module):
                 hidden_states = outputs.last_hidden_state[:, :chunk_size]
                 hidden_states = hidden_states.to(dtype=self.action_out_proj.dtype)
                 v_t = self.action_out_proj(hidden_states)
+            return v_t.reshape(input_x_t.shape).to(input_x_t.dtype)
 
-            x_t += dt * v_t.reshape(x_t.shape)
-
+        x_t = euler_integrate(denoise_fn, x_t, self.config.num_denoise_steps)
         return x_t

src/lerobot/policies/pi0/modeling_pi0.py

@@ -16,7 +16,6 @@
 
 import builtins
 import logging
-import math
 from collections import deque
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal, TypedDict, Unpack
@@ -29,7 +28,6 @@ from lerobot.utils.import_utils import _transformers_available, require_package
 
 # Conditional import for type checking and lazy loading
 if TYPE_CHECKING or _transformers_available:
-    from transformers.cache_utils import DynamicCache
     from transformers.models.auto import CONFIG_MAPPING
     from transformers.models.gemma import modeling_gemma
 
@@ -41,7 +39,6 @@ if TYPE_CHECKING or _transformers_available:
     )
 else:
     CONFIG_MAPPING = None
-    DynamicCache = None
     modeling_gemma = None
     PiGemmaForCausalLM = None
     _gated_residual = None
@@ -55,9 +52,17 @@ from lerobot.utils.constants import (
     OBS_LANGUAGE_ATTENTION_MASK,
     OBS_LANGUAGE_TOKENS,
     OBS_STATE,
-    OPENPI_ATTENTION_MASK_VALUE,
 )
 
+from ..common.flow_matching import euler_integrate, sample_noise, sample_time_beta
+from ..common.vla_utils import (
+    clone_past_key_values,
+    create_sinusoidal_pos_embedding,
+    make_att_2d_masks,
+    pad_vector,
+    prepare_attention_masks_4d,
+    resize_with_pad_torch,
+)
 from ..pretrained import PreTrainedPolicy, T
 from ..rtc.modeling_rtc import RTCProcessor
 from .configuration_pi0 import DEFAULT_IMAGE_SIZE, PI0Config
@@ -69,173 +74,6 @@ class ActionSelectKwargs(TypedDict, total=False):
     execution_horizon: int | None
 
 
-def get_safe_dtype(target_dtype, device_type):
-    """Get a safe dtype for the given device type."""
-    if device_type == "mps" and target_dtype == torch.float64:
-        return torch.float32
-    if device_type == "cpu":
-        # CPU doesn't support bfloat16, use float32 instead
-        if target_dtype == torch.bfloat16:
-            return torch.float32
-        if target_dtype == torch.float64:
-            return torch.float64
-    return target_dtype
-
-
-def create_sinusoidal_pos_embedding(  # see openpi `create_sinusoidal_pos_embedding` (exact copy)
-    time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
-) -> Tensor:
-    """Computes sine-cosine positional embedding vectors for scalar positions."""
-    if dimension % 2 != 0:
-        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
-
-    if time.ndim != 1:
-        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
-
-    dtype = get_safe_dtype(torch.float64, device.type)
-    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
-    period = min_period * (max_period / min_period) ** fraction
-
-    # Compute the outer product
-    scaling_factor = 1.0 / period * 2 * math.pi
-    sin_input = scaling_factor[None, :] * time[:, None]
-    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
-
-
-def sample_beta(alpha, beta, bsize, device):  # see openpi `sample_beta` (exact copy)
-    # Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
-    alpha_t = torch.tensor(alpha, dtype=torch.float32)
-    beta_t = torch.tensor(beta, dtype=torch.float32)
-    dist = torch.distributions.Beta(alpha_t, beta_t)
-    return dist.sample((bsize,)).to(device)
-
-
-def make_att_2d_masks(pad_masks, att_masks):  # see openpi `make_att_2d_masks` (exact copy)
-    """Copied from big_vision.
-
-    Tokens can attend to valid inputs tokens which have a cumulative mask_ar
-    smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to
-    setup several types of attention, for example:
-
-      [[1 1 1 1 1 1]]: pure causal attention.
-
-      [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
-          themselves and the last 3 tokens have a causal attention. The first
-          entry could also be a 1 without changing behaviour.
-
-      [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
-          block can attend all previous blocks and all tokens on the same block.
-
-    Args:
-      input_mask: bool[B, N] true if its part of the input, false if padding.
-      mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on
-        it and 0 where it shares the same attention mask as the previous token.
-    """
-    if att_masks.ndim != 2:
-        raise ValueError(att_masks.ndim)
-    if pad_masks.ndim != 2:
-        raise ValueError(pad_masks.ndim)
-
-    cumsum = torch.cumsum(att_masks, dim=1)
-    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
-    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
-    return att_2d_masks & pad_2d_masks
-
-
-def clone_past_key_values(past_key_values):
-    """Clone the DynamicCache returned by prefix prefill for compiled denoising."""
-    return DynamicCache(
-        tuple(
-            (keys.clone(), values.clone(), sliding_window) for keys, values, sliding_window in past_key_values
-        )
-    )
-
-
-def pad_vector(vector, new_dim):
-    """Pad the last dimension of a vector to new_dim with zeros.
-
-    Can be (batch_size x sequence_length x features_dimension)
-    or (batch_size x features_dimension)
-    """
-    if vector.shape[-1] >= new_dim:
-        return vector
-    return F.pad(vector, (0, new_dim - vector.shape[-1]))
-
-
-def resize_with_pad_torch(  # see openpi `resize_with_pad_torch` (exact copy)
-    images: torch.Tensor,
-    height: int,
-    width: int,
-    mode: str = "bilinear",
-) -> torch.Tensor:
-    """PyTorch version of resize_with_pad. Resizes an image to a target height and width without distortion
-    by padding with black. If the image is float32, it must be in the range [-1, 1].
-
-    Args:
-        images: Tensor of shape [*b, h, w, c] or [*b, c, h, w]
-        height: Target height
-        width: Target width
-        mode: Interpolation mode ('bilinear', 'nearest', etc.)
-
-    Returns:
-        Resized and padded tensor with same shape format as input
-    """
-    # Check if input is in channels-last format [*b, h, w, c] or channels-first [*b, c, h, w]
-    if images.shape[-1] <= 4:  # Assume channels-last format
-        channels_last = True
-        if images.dim() == 3:
-            images = images.unsqueeze(0)  # Add batch dimension
-        images = images.permute(0, 3, 1, 2)  # [b, h, w, c] -> [b, c, h, w]
-    else:
-        channels_last = False
-        if images.dim() == 3:
-            images = images.unsqueeze(0)  # Add batch dimension
-
-    batch_size, channels, cur_height, cur_width = images.shape
-
-    # Calculate resize ratio
-    ratio = max(cur_width / width, cur_height / height)
-    resized_height = int(cur_height / ratio)
-    resized_width = int(cur_width / ratio)
-
-    # Resize
-    resized_images = F.interpolate(
-        images,
-        size=(resized_height, resized_width),
-        mode=mode,
-        align_corners=False if mode == "bilinear" else None,
-    )
-
-    # Handle dtype-specific clipping
-    if images.dtype == torch.uint8:
-        resized_images = torch.round(resized_images).clamp(0, 255).to(torch.uint8)
-    elif images.dtype == torch.float32:
-        resized_images = resized_images.clamp(0.0, 1.0)
-    else:
-        raise ValueError(f"Unsupported image dtype: {images.dtype}")
-
-    # Calculate padding
-    pad_h0, remainder_h = divmod(height - resized_height, 2)
-    pad_h1 = pad_h0 + remainder_h
-    pad_w0, remainder_w = divmod(width - resized_width, 2)
-    pad_w1 = pad_w0 + remainder_w
-
-    # Pad
-    constant_value = 0 if images.dtype == torch.uint8 else 0.0
-    padded_images = F.pad(
-        resized_images,
-        (pad_w0, pad_w1, pad_h0, pad_h1),  # left, right, top, bottom
-        mode="constant",
-        value=constant_value,
-    )
-
-    # Convert back to original format if needed
-    if channels_last:
-        padded_images = padded_images.permute(0, 2, 3, 1)  # [b, c, h, w] -> [b, h, w, c]
-
-    return padded_images
-
-
 # Define the complete layer computation function for gradient checkpointing
 def compute_layer_complete(inputs_embeds, attention_mask, position_ids, adarms_cond, layers, rotary_emb):
     query_states = []
@@ -633,26 +471,18 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             )
         return func(*args, **kwargs)
 
-    def _prepare_attention_masks_4d(self, att_2d_masks):
-        """Helper method to prepare 4D attention masks for transformer."""
-        att_2d_masks_4d = att_2d_masks[:, None, :, :]
-        return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
-
     def sample_noise(self, shape, device):
-        return torch.normal(
-            mean=0.0,
-            std=1.0,
-            size=shape,
-            dtype=torch.float32,
-            device=device,
-        )
+        return sample_noise(shape, device)
 
     def sample_time(self, bsize, device):
-        time_beta = sample_beta(
-            self.config.time_sampling_beta_alpha, self.config.time_sampling_beta_beta, bsize, device
+        return sample_time_beta(
+            bsize,
+            device,
+            alpha=self.config.time_sampling_beta_alpha,
+            beta=self.config.time_sampling_beta_beta,
+            scale=self.config.time_sampling_scale,
+            offset=self.config.time_sampling_offset,
         )
-        time = time_beta * self.config.time_sampling_scale + self.config.time_sampling_offset
-        return time.to(dtype=torch.float32, device=device)
 
     def embed_prefix(
         self, images, img_masks, lang_tokens, lang_masks
@@ -783,7 +613,7 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
         position_ids = torch.cumsum(pad_masks, dim=1) - 1
 
-        att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks)
+        att_2d_masks_4d = prepare_attention_masks_4d(att_2d_masks)
 
         def forward_func(prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond):
             (_, suffix_out), _ = self.paligemma_with_expert.forward(
@@ -844,7 +674,7 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
         prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
 
-        prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks)
+        prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d_masks)
         self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager"  # noqa: SLF001
 
         _, past_key_values = self.paligemma_with_expert.forward(
@@ -855,44 +685,22 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             use_cache=True,
         )
 
-        dt = -1.0 / num_steps
-
-        x_t = noise
-        for step in range(num_steps):
-            time = 1.0 + step * dt
-            time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
-
-            def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
-                return self.denoise_step(
-                    state=state,
-                    prefix_pad_masks=prefix_pad_masks,
-                    past_key_values=past_key_values,
-                    x_t=input_x_t,
-                    timestep=current_timestep,
-                )
-
-            if self._rtc_enabled():
-                inference_delay = kwargs.get("inference_delay")
-                prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
-                execution_horizon = kwargs.get("execution_horizon")
-
-                v_t = self.rtc_processor.denoise_step(
-                    x_t=x_t,
-                    prev_chunk_left_over=prev_chunk_left_over,
-                    inference_delay=inference_delay,
-                    time=time,
-                    original_denoise_step_partial=denoise_step_partial_call,
-                    execution_horizon=execution_horizon,
-                )
-            else:
-                v_t = denoise_step_partial_call(x_t)
-
-            x_t = x_t + dt * v_t
-
-            if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled():
-                self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
-
-        return x_t
+        return euler_integrate(
+            lambda input_x_t, current_timestep: self.denoise_step(
+                state=state,
+                prefix_pad_masks=prefix_pad_masks,
+                past_key_values=past_key_values,
+                x_t=input_x_t,
+                timestep=current_timestep,
+            ),
+            noise,
+            num_steps,
+            rtc_processor=self.rtc_processor,
+            rtc_enabled=self._rtc_enabled(),
+            inference_delay=kwargs.get("inference_delay"),
+            prev_chunk_left_over=kwargs.get("prev_chunk_left_over"),
+            execution_horizon=kwargs.get("execution_horizon"),
+        )
 
     def denoise_step(
         self,
@@ -916,7 +724,7 @@ class PI0Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
         position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
 
-        full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks)
+        full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks)
         self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager"  # noqa: SLF001
 
         past_key_values = clone_past_key_values(past_key_values)

src/lerobot/policies/pi05/modeling_pi05.py

@@ -16,7 +16,6 @@
 
 import builtins
 import logging
-import math
 from collections import deque
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal, TypedDict, Unpack
@@ -29,7 +28,6 @@ from lerobot.utils.import_utils import _transformers_available, require_package
 
 # Conditional import for type checking and lazy loading
 if TYPE_CHECKING or _transformers_available:
-    from transformers.cache_utils import DynamicCache
     from transformers.models.auto import CONFIG_MAPPING
     from transformers.models.gemma import modeling_gemma
 
@@ -41,7 +39,6 @@ if TYPE_CHECKING or _transformers_available:
     )
 else:
     CONFIG_MAPPING = None
-    DynamicCache = None
     modeling_gemma = None
     PiGemmaForCausalLM = None
     _gated_residual = None
@@ -52,9 +49,17 @@ from lerobot.utils.constants import (
     ACTION,
     OBS_LANGUAGE_ATTENTION_MASK,
     OBS_LANGUAGE_TOKENS,
-    OPENPI_ATTENTION_MASK_VALUE,
 )
 
+from ..common.flow_matching import euler_integrate, sample_noise, sample_time_beta
+from ..common.vla_utils import (
+    clone_past_key_values,
+    create_sinusoidal_pos_embedding,
+    make_att_2d_masks,
+    pad_vector,
+    prepare_attention_masks_4d,
+    resize_with_pad_torch,
+)
 from ..pretrained import PreTrainedPolicy, T
 from ..rtc.modeling_rtc import RTCProcessor
 from .configuration_pi05 import DEFAULT_IMAGE_SIZE, PI05Config
@@ -66,173 +71,6 @@ class ActionSelectKwargs(TypedDict, total=False):
     execution_horizon: int | None
 
 
-def get_safe_dtype(target_dtype, device_type):
-    """Get a safe dtype for the given device type."""
-    if device_type == "mps" and target_dtype == torch.float64:
-        return torch.float32
-    if device_type == "cpu":
-        # CPU doesn't support bfloat16, use float32 instead
-        if target_dtype == torch.bfloat16:
-            return torch.float32
-        if target_dtype == torch.float64:
-            return torch.float64
-    return target_dtype
-
-
-def create_sinusoidal_pos_embedding(  # see openpi `create_sinusoidal_pos_embedding` (exact copy)
-    time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
-) -> Tensor:
-    """Computes sine-cosine positional embedding vectors for scalar positions."""
-    if dimension % 2 != 0:
-        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
-
-    if time.ndim != 1:
-        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
-
-    dtype = get_safe_dtype(torch.float64, device.type)
-    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
-    period = min_period * (max_period / min_period) ** fraction
-
-    # Compute the outer product
-    scaling_factor = 1.0 / period * 2 * math.pi
-    sin_input = scaling_factor[None, :] * time[:, None]
-    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
-
-
-def sample_beta(alpha, beta, bsize, device):  # see openpi `sample_beta` (exact copy)
-    # Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
-    alpha_t = torch.tensor(alpha, dtype=torch.float32)
-    beta_t = torch.tensor(beta, dtype=torch.float32)
-    dist = torch.distributions.Beta(alpha_t, beta_t)
-    return dist.sample((bsize,)).to(device)
-
-
-def make_att_2d_masks(pad_masks, att_masks):  # see openpi `make_att_2d_masks` (exact copy)
-    """Copied from big_vision.
-
-    Tokens can attend to valid inputs tokens which have a cumulative mask_ar
-    smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to
-    setup several types of attention, for example:
-
-      [[1 1 1 1 1 1]]: pure causal attention.
-
-      [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
-          themselves and the last 3 tokens have a causal attention. The first
-          entry could also be a 1 without changing behaviour.
-
-      [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
-          block can attend all previous blocks and all tokens on the same block.
-
-    Args:
-      input_mask: bool[B, N] true if its part of the input, false if padding.
-      mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on
-        it and 0 where it shares the same attention mask as the previous token.
-    """
-    if att_masks.ndim != 2:
-        raise ValueError(att_masks.ndim)
-    if pad_masks.ndim != 2:
-        raise ValueError(pad_masks.ndim)
-
-    cumsum = torch.cumsum(att_masks, dim=1)
-    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
-    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
-    return att_2d_masks & pad_2d_masks
-
-
-def clone_past_key_values(past_key_values):
-    """Clone the DynamicCache returned by prefix prefill for compiled denoising."""
-    return DynamicCache(
-        tuple(
-            (keys.clone(), values.clone(), sliding_window) for keys, values, sliding_window in past_key_values
-        )
-    )
-
-
-def pad_vector(vector, new_dim):
-    """Pad the last dimension of a vector to new_dim with zeros.
-
-    Can be (batch_size x sequence_length x features_dimension)
-    or (batch_size x features_dimension)
-    """
-    if vector.shape[-1] >= new_dim:
-        return vector
-    return F.pad(vector, (0, new_dim - vector.shape[-1]))
-
-
-def resize_with_pad_torch(  # see openpi `resize_with_pad_torch` (exact copy)
-    images: torch.Tensor,
-    height: int,
-    width: int,
-    mode: str = "bilinear",
-) -> torch.Tensor:
-    """PyTorch version of resize_with_pad. Resizes an image to a target height and width without distortion
-    by padding with black. If the image is float32, it must be in the range [-1, 1].
-
-    Args:
-        images: Tensor of shape [*b, h, w, c] or [*b, c, h, w]
-        height: Target height
-        width: Target width
-        mode: Interpolation mode ('bilinear', 'nearest', etc.)
-
-    Returns:
-        Resized and padded tensor with same shape format as input
-    """
-    # Check if input is in channels-last format [*b, h, w, c] or channels-first [*b, c, h, w]
-    if images.shape[-1] <= 4:  # Assume channels-last format
-        channels_last = True
-        if images.dim() == 3:
-            images = images.unsqueeze(0)  # Add batch dimension
-        images = images.permute(0, 3, 1, 2)  # [b, h, w, c] -> [b, c, h, w]
-    else:
-        channels_last = False
-        if images.dim() == 3:
-            images = images.unsqueeze(0)  # Add batch dimension
-
-    batch_size, channels, cur_height, cur_width = images.shape
-
-    # Calculate resize ratio
-    ratio = max(cur_width / width, cur_height / height)
-    resized_height = int(cur_height / ratio)
-    resized_width = int(cur_width / ratio)
-
-    # Resize
-    resized_images = F.interpolate(
-        images,
-        size=(resized_height, resized_width),
-        mode=mode,
-        align_corners=False if mode == "bilinear" else None,
-    )
-
-    # Handle dtype-specific clipping
-    if images.dtype == torch.uint8:
-        resized_images = torch.round(resized_images).clamp(0, 255).to(torch.uint8)
-    elif images.dtype == torch.float32:
-        resized_images = resized_images.clamp(0.0, 1.0)
-    else:
-        raise ValueError(f"Unsupported image dtype: {images.dtype}")
-
-    # Calculate padding
-    pad_h0, remainder_h = divmod(height - resized_height, 2)
-    pad_h1 = pad_h0 + remainder_h
-    pad_w0, remainder_w = divmod(width - resized_width, 2)
-    pad_w1 = pad_w0 + remainder_w
-
-    # Pad
-    constant_value = 0 if images.dtype == torch.uint8 else 0.0
-    padded_images = F.pad(
-        resized_images,
-        (pad_w0, pad_w1, pad_h0, pad_h1),  # left, right, top, bottom
-        mode="constant",
-        value=constant_value,
-    )
-
-    # Convert back to original format if needed
-    if channels_last:
-        padded_images = padded_images.permute(0, 2, 3, 1)  # [b, c, h, w] -> [b, h, w, c]
-
-    return padded_images
-
-
 # Define the complete layer computation function for gradient checkpointing
 def compute_layer_complete(inputs_embeds, attention_mask, position_ids, adarms_cond, layers, rotary_emb):
     query_states = []
@@ -629,26 +467,18 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             )
         return func(*args, **kwargs)
 
-    def _prepare_attention_masks_4d(self, att_2d_masks):
-        """Helper method to prepare 4D attention masks for transformer."""
-        att_2d_masks_4d = att_2d_masks[:, None, :, :]
-        return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
-
     def sample_noise(self, shape, device):
-        return torch.normal(
-            mean=0.0,
-            std=1.0,
-            size=shape,
-            dtype=torch.float32,
-            device=device,
-        )
+        return sample_noise(shape, device)
 
     def sample_time(self, bsize, device):
-        time_beta = sample_beta(
-            self.config.time_sampling_beta_alpha, self.config.time_sampling_beta_beta, bsize, device
+        return sample_time_beta(
+            bsize,
+            device,
+            alpha=self.config.time_sampling_beta_alpha,
+            beta=self.config.time_sampling_beta_beta,
+            scale=self.config.time_sampling_scale,
+            offset=self.config.time_sampling_offset,
         )
-        time = time_beta * self.config.time_sampling_scale + self.config.time_sampling_offset
-        return time.to(dtype=torch.float32, device=device)
 
     def embed_prefix(
         self, images, img_masks, tokens, masks
@@ -761,7 +591,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
         position_ids = torch.cumsum(pad_masks, dim=1) - 1
 
-        att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks)
+        att_2d_masks_4d = prepare_attention_masks_4d(att_2d_masks)
 
         def forward_func(prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond):
             (_, suffix_out), _ = self.paligemma_with_expert.forward(
@@ -819,7 +649,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
         prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
 
-        prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks)
+        prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d_masks)
         self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager"  # noqa: SLF001
 
         _, past_key_values = self.paligemma_with_expert.forward(
@@ -830,43 +660,21 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             use_cache=True,
         )
 
-        dt = -1.0 / num_steps
-
-        x_t = noise
-        for step in range(num_steps):
-            time = 1.0 + step * dt
-            time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
-
-            def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
-                return self.denoise_step(
-                    prefix_pad_masks=prefix_pad_masks,
-                    past_key_values=past_key_values,
-                    x_t=input_x_t,
-                    timestep=current_timestep,
-                )
-
-            if self._rtc_enabled():
-                inference_delay = kwargs.get("inference_delay")
-                prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
-                execution_horizon = kwargs.get("execution_horizon")
-
-                v_t = self.rtc_processor.denoise_step(
-                    x_t=x_t,
-                    prev_chunk_left_over=prev_chunk_left_over,
-                    inference_delay=inference_delay,
-                    time=time,
-                    original_denoise_step_partial=denoise_step_partial_call,
-                    execution_horizon=execution_horizon,
-                )
-            else:
-                v_t = denoise_step_partial_call(x_t)
-
-            x_t = x_t + dt * v_t
-
-            if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled():
-                self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
-
-        return x_t
+        return euler_integrate(
+            lambda input_x_t, current_timestep: self.denoise_step(
+                prefix_pad_masks=prefix_pad_masks,
+                past_key_values=past_key_values,
+                x_t=input_x_t,
+                timestep=current_timestep,
+            ),
+            noise,
+            num_steps,
+            rtc_processor=self.rtc_processor,
+            rtc_enabled=self._rtc_enabled(),
+            inference_delay=kwargs.get("inference_delay"),
+            prev_chunk_left_over=kwargs.get("prev_chunk_left_over"),
+            execution_horizon=kwargs.get("execution_horizon"),
+        )
 
     def denoise_step(
         self,
@@ -889,7 +697,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
         position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
 
-        full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks)
+        full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks)
         self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager"  # noqa: SLF001
 
         past_key_values = clone_past_key_values(past_key_values)

src/lerobot/policies/pi0_fast/modeling_pi0_fast.py

@@ -22,7 +22,6 @@ from typing import TYPE_CHECKING, Literal, TypedDict, Unpack
 
 import numpy as np
 import torch
-import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
 
 from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
@@ -55,9 +54,9 @@ from lerobot.utils.constants import (
     ACTION_TOKENS,
     OBS_LANGUAGE_ATTENTION_MASK,
     OBS_LANGUAGE_TOKENS,
-    OPENPI_ATTENTION_MASK_VALUE,
 )
 
+from ..common.vla_utils import pad_vector, prepare_attention_masks_4d, resize_with_pad_torch
 from ..pretrained import PreTrainedPolicy, T
 from ..rtc.modeling_rtc import RTCProcessor
 from .configuration_pi0_fast import PI0FastConfig
@@ -67,91 +66,6 @@ class ActionSelectKwargs(TypedDict, total=False):
     temperature: float | None
 
 
-def pad_vector(vector, new_dim):
-    """Pad the last dimension of a vector to new_dim with zeros.
-
-    Can be (batch_size x sequence_length x features_dimension)
-    or (batch_size x features_dimension)
-    """
-    if vector.shape[-1] >= new_dim:
-        return vector
-    return F.pad(vector, (0, new_dim - vector.shape[-1]))
-
-
-def resize_with_pad_torch(  # see openpi `resize_with_pad_torch` (exact copy)
-    images: torch.Tensor,
-    height: int,
-    width: int,
-    mode: str = "bilinear",
-) -> torch.Tensor:
-    """PyTorch version of resize_with_pad. Resizes an image to a target height and width without distortion
-    by padding with black. If the image is float32, it must be in the range [-1, 1].
-
-    Args:
-        images: Tensor of shape [*b, h, w, c] or [*b, c, h, w]
-        height: Target height
-        width: Target width
-        mode: Interpolation mode ('bilinear', 'nearest', etc.)
-
-    Returns:
-        Resized and padded tensor with same shape format as input
-    """
-    # Check if input is in channels-last format [*b, h, w, c] or channels-first [*b, c, h, w]
-    if images.shape[-1] <= 4:  # Assume channels-last format
-        channels_last = True
-        if images.dim() == 3:
-            images = images.unsqueeze(0)  # Add batch dimension
-        images = images.permute(0, 3, 1, 2)  # [b, h, w, c] -> [b, c, h, w]
-    else:
-        channels_last = False
-        if images.dim() == 3:
-            images = images.unsqueeze(0)  # Add batch dimension
-
-    batch_size, channels, cur_height, cur_width = images.shape
-
-    # Calculate resize ratio
-    ratio = max(cur_width / width, cur_height / height)
-    resized_height = int(cur_height / ratio)
-    resized_width = int(cur_width / ratio)
-
-    # Resize
-    resized_images = F.interpolate(
-        images,
-        size=(resized_height, resized_width),
-        mode=mode,
-        align_corners=False if mode == "bilinear" else None,
-    )
-
-    # Handle dtype-specific clipping
-    if images.dtype == torch.uint8:
-        resized_images = torch.round(resized_images).clamp(0, 255).to(torch.uint8)
-    elif images.dtype == torch.float32:
-        resized_images = resized_images.clamp(0.0, 1.0)
-    else:
-        raise ValueError(f"Unsupported image dtype: {images.dtype}")
-
-    # Calculate padding
-    pad_h0, remainder_h = divmod(height - resized_height, 2)
-    pad_h1 = pad_h0 + remainder_h
-    pad_w0, remainder_w = divmod(width - resized_width, 2)
-    pad_w1 = pad_w0 + remainder_w
-
-    # Pad
-    constant_value = 0 if images.dtype == torch.uint8 else 0.0
-    padded_images = F.pad(
-        resized_images,
-        (pad_w0, pad_w1, pad_h0, pad_h1),  # left, right, top, bottom
-        mode="constant",
-        value=constant_value,
-    )
-
-    # Convert back to original format if needed
-    if channels_last:
-        padded_images = padded_images.permute(0, 2, 3, 1)  # [b, c, h, w] -> [b, h, w, c]
-
-    return padded_images
-
-
 class GemmaConfig:  # see openpi `gemma.py: Config`
     """Configuration for Gemma model variants."""
 
@@ -357,14 +271,6 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
             )
         return func(*args, **kwargs)
 
-    def _prepare_attention_masks_4d(self, att_2d_masks, dtype=None):
-        """Helper method to prepare 4D attention masks for transformer."""
-        att_2d_masks_4d = att_2d_masks[:, None, :, :]
-        result = torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
-        if dtype is not None:
-            result = result.to(dtype=dtype)
-        return result
-
     def embed_prefix_fast(
         self,
         images,
@@ -545,7 +451,7 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         input_att_masks = prefix_att_masks
 
         position_ids = torch.cumsum(input_pad_masks, dim=1) - 1
-        att_2d_4d = self._prepare_attention_masks_4d(input_att_masks, dtype=input_embs.dtype)
+        att_2d_4d = prepare_attention_masks_4d(input_att_masks, dtype=input_embs.dtype)
 
         # forward pass through paligemma (language model)
         (prefix_out, _), _ = self.paligemma_with_expert.forward(
@@ -638,7 +544,7 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         for t in range(max_decoding_steps):
             # always re-calculate position IDs from the current pad mask
             position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
-            att_4d = self._prepare_attention_masks_4d(prefix_att_masks, dtype=prefix_embs.dtype)
+            att_4d = prepare_attention_masks_4d(prefix_att_masks, dtype=prefix_embs.dtype)
 
             # full forward pass (no kv cache)
             (prefix_out, _), _ = self.paligemma_with_expert.forward(
@@ -733,7 +639,7 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
 
         # Create 4D mask for the prefix
-        att_4d = self._prepare_attention_masks_4d(prefix_att_masks, dtype=prefix_embs.dtype)
+        att_4d = prepare_attention_masks_4d(prefix_att_masks, dtype=prefix_embs.dtype)
 
         # Forward pass (Prefill) with use_cache=True
         # We only pass [prefix_embs, None] because we aren't using the suffix (expert) model yet
@@ -782,7 +688,7 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
             # Create Attention Mask for the single new step
             # The new token attends to all valid tokens in history (captured by current_pad_mask).
             # Shape becomes (B, 1, 1, Total_Len) which works with HF's cache logic.
-            step_att_mask = self._prepare_attention_masks_4d(
+            step_att_mask = prepare_attention_masks_4d(
                 current_pad_mask.unsqueeze(1), dtype=next_token_emb.dtype
             )
 

src/lerobot/robots/unitree_g1/README.md

@@ -1,89 +0,0 @@
-# Unitree G1 — SONIC encoder/decoder whole-body control
-
-This package runs NVIDIA's **SONIC** encoder/decoder on the Unitree G1, in MuJoCo
-simulation or on real hardware, driven by a dense **34-D whole-body command** (the
-OpenHLM / pi0.5 action layout). It is a pure-Python/ONNX reimplementation of the
-reference-tracking half of the SONIC deploy stack (no `gear_sonic`/torch dependency, and
-no motion planner): the encoder compresses a reference motion window into a latent token
-and the decoder maps that token + proprioception history into 50 Hz joint-position
-targets for the robot's PD controller.
-
-## Controllers
-
-Selected with `--robot.controller=<ClassName>`:
-
-| Controller                     | Purpose                                                      |
-| ------------------------------ | ------------------------------------------------------------ |
-| `SonicWholeBodyController`     | SONIC encoder/decoder driven by a 34-D OpenHLM/pi0.5 command |
-| `GrootLocomotionController`    | GR00T locomotion policy                                      |
-| `HolosomaLocomotionController` | Holosoma locomotion policy                                   |
-
-The rest of this document covers the SONIC whole-body path.
-
-Each tick the `SonicWholeBodyController` takes a 34-D command (`wb.0.pos … wb.33.pos`) in the OpenHLM
-layout:
-
-```
-[L-arm(7), L-grip(1), R-arm(7), R-grip(1), L-leg(6), R-leg(6), waist(3),
- root roll/pitch + yaw-rate(3)]
-```
-
-The 29 joint targets become the SONIC encode-mode-0 reference (accumulated into a rolling
-50-frame trajectory with finite-difference velocities so the encoder's lookahead sees a
-real motion sequence), the root roll/pitch set the anchor orientation, and the two grip
-scalars can drive the Dex3 hands (see below). On startup the controller **interpolates**
-from the robot's measured pose into the policy's commanded target over ~3 s (no snap).
-
-## Requirements
-
-- `onnxruntime` (CPU) **or** `onnxruntime-gpu` (recommended). Verify with:
-  ```bash
-  python -c "import onnxruntime as ort; print(ort.get_available_providers())"
-  ```
-- `mujoco` for simulation (`is_simulation=True`).
-- The SONIC encoder/decoder ONNX models download automatically from the
-  `nvidia/GEAR-SONIC` Hub repo.
-
-## Running a rollout
-
-Drive the G1 with a 34-D VLA policy (OpenHLM / pi0.5) via `lerobot-rollout`:
-
-```bash
-lerobot-rollout \
-  --strategy.type=base \
-  --policy.path=<pi05_openhlm_dir> \
-  --robot.type=unitree_g1 \
-  --robot.controller=SonicWholeBodyController \
-  --robot.is_simulation=true \
-  --robot.publish_hands=true \
-  --task="<language instruction>" \
-  --duration=45 --device=cuda
-```
-
-### Cameras
-
-Image-conditioned policies need camera frames. Two options are available without live
-cameras:
-
-- **Black frames**: `--robot.empty_cameras='[base, left_wrist, right_wrist]'`.
-- **Replay a recorded episode** as the camera feed:
-  ```bash
-  --robot.replay_camera_parquet=<episode.parquet> \
-  --robot.replay_camera_map='{base: head_image_left, left_wrist: left_wrist_image, right_wrist: right_wrist_image}'
-  ```
-
-### Hands (Dex3)
-
-`--robot.publish_hands=true` publishes `rt/dex3/{left,right}/cmd` from the two grip
-scalars (`wb.7.pos` left, `wb.15.pos` right). The scalar is interpolated between
-`hand_open_grip_value` (default 1.0 = open) and `hand_closed_grip_value` (default 0.0 =
-closed) and scaled onto `hand_closed_pose` (7 joints:
-`thumb_0, thumb_1, thumb_2, middle_0, middle_1, index_0, index_1`). Flip the signs in
-`hand_closed_pose` if the fingers curl the wrong way, or raise `hand_kp` for a firmer
-grip.
-
-## Observation state
-
-When the whole-body controller is active the robot exposes a 34-D proprio state
-(`wb_state.0.pos … wb_state.33.pos`) in the same OpenHLM layout as the action, which the
-rollout aggregates into `observation.state` for the policy.

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -62,76 +62,12 @@ class UnitreeG1Config(RobotConfig):
     # Socket config for ZMQ bridge
     robot_ip: str = "192.168.123.164"  # default G1 IP
 
-    # Run the locomotion / whole-body controller ONBOARD the robot (policy on the G1
-    # itself, against local DDS at full rate) instead of on the laptop over the ZMQ
-    # socket bridge. In this mode the robot object uses the real Unitree SDK channels
-    # and expects high-level actions (arm targets + joystick axes, or 64-D SONIC
-    # tokens) fed via send_action -- e.g. by run_g1_onboard.py, which receives them
-    # from the laptop over ZMQ. Mutually exclusive with is_simulation.
-    onboard: bool = False
-    # DDS network interface for onboard mode (None = SDK default, matching
-    # run_g1_server.py's ChannelFactoryInitialize(0)).
-    dds_interface: str | None = None
-    # Onboard sub-flags. On a real G1 both are True: the built-in motion services
-    # must be released before we can write lowcmd, and locomotion axes are read from
-    # the physical wireless remote. Against a DDS sim neither applies (no
-    # MotionSwitcher, no physical remote), so set both False so the controller takes
-    # its locomotion axes purely from send_action (ZMQ) input.
-    release_motion_control: bool = True
-    physical_remote: bool = True
-
     # Cameras (ZMQ-based remote cameras)
     cameras: dict[str, CameraConfig] = field(default_factory=dict)
 
-    # Synthetic zero-image cameras exposed as ``observation.images.{name}`` (H×W×3
-    # black frames). Lets image-conditioned policies (e.g. pi0.5 / OpenHLM) run in
-    # sim before real cameras are wired. Empty = disabled.
-    empty_cameras: list[str] = field(default_factory=list)
-    empty_camera_hw: tuple[int, int] = (224, 224)
-
-    # Publish Dex3 hand commands (``rt/dex3/{left,right}/cmd``) driven by the OpenHLM
-    # gripper scalars (``wb.7.pos`` left, ``wb.15.pos`` right). Lets the 43-DoF sim
-    # (or a real Dex3-equipped G1) show grasping. The scalar in [0, 1] is remapped to
-    # a curl amount (``hand_open_grip_value`` -> open) and scaled onto
-    # ``hand_closed_pose`` (7 joints: thumb_0/1/2, middle_0/1, index_0/1). Flip signs
-    # in ``hand_closed_pose`` if fingers curl the wrong way.
-    publish_hands: bool = False
-    # When False, connect() does not start the background controller thread, so a
-    # caller can drive the controller synchronously (one decode per fed action),
-    # reproducing the deploy's single 50Hz control clock for faithful replay.
-    run_controller_thread: bool = True
-    hand_open_grip_value: float = 1.0
-    hand_closed_grip_value: float = 0.0
-    hand_closed_pose: list[float] = field(default_factory=lambda: [1.0, 0.9, 0.9, 1.3, 1.3, 1.3, 1.3])
-    hand_kp: float = 1.5
-    hand_kd: float = 0.1
-
-    # Replay recorded camera frames from a LeRobot parquet episode as the camera
-    # feed (e.g. OpenHLM-data episode). Maps a robot camera name to a parquet image
-    # column; frames advance one per observation and loop. Lets a VLA see the real
-    # task video in sim without live cameras. Empty map = disabled.
-    replay_camera_parquet: str | None = None
-    replay_camera_map: dict[str, str] = field(default_factory=dict)
-    replay_camera_loop: bool = True
-
-    # Token-output VLA interface for the SONIC decoder. When True (and the controller
-    # is ``SonicWholeBodyController``), the robot advertises a 64-D latent-token action
-    # space (``motion_token.{i}.pos``) instead of the 34-D whole-body command, and
-    # exposes the last commanded token as a 64-D ``observation.state``
-    # (``motion_token_state.{i}.pos``). This lets ``lerobot-rollout`` drive a policy
-    # that was trained with 64-D SONIC motion tokens as both state and action
-    # (e.g. nepyope/sonic_walk): the decoder consumes the token directly, encoder
-    # bypassed. Ignored unless a SONIC whole-body controller is active.
-    sonic_token_action: bool = False
-
     # Compensates for gravity on the unitree's arms using the arm ik solver
     gravity_compensation: bool = False
 
-    # Locomotion controller class name, e.g. "GrootLocomotionController",
-    # "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
+    # Lower-body controller class name, e.g. "GrootLocomotionController" or
+    # "HolosomaLocomotionController". None disables it.
     controller: str | None = None
-
-    # On disconnect (e.g. Ctrl-C), seconds to hold the current pose while ramping joint
-    # stiffness (kp) to zero — a soft, damped settle instead of an instant limp /
-    # free-fall. 0 disables it (immediate zero-torque). Real robot only.
-    graceful_stop_s: float = 1.5

src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py

@@ -1,670 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 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.
-
-"""SONIC encoder/decoder pipeline for the Unitree G1 whole-body controller.
-
-Pure-Python/ONNX re-implementation of the reference-tracking half of NVIDIA's SONIC
-deploy stack (mirrors ``g1_deploy_onnx_ref.cpp``). Given a reference motion buffer
-(joint targets + body orientation per frame) it produces 50 Hz joint-position targets
-for the robot's PD controller. The upstream *motion planner* is intentionally absent:
-here the reference is supplied directly by the caller (e.g. a 34-D OpenHLM / pi0.5 VLA
-command per tick, in ``sonic_whole_body.py``).
-
-Two cooperating ONNX models:
-  * **encoder** – compresses the reference window into a 64-D latent ``token``
-    (refreshed every ``ENCODER_UPDATE_EVERY`` ticks).
-  * **decoder** – every tick, maps the token + recent proprioception history to a
-    residual action that is scaled and added to ``DEFAULT_ANGLES``.
-
-Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order)
-and **MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` convert
-between them. Quaternions are scalar-first ``(w, x, y, z)``.
-"""
-
-from __future__ import annotations
-
-import logging
-import threading
-from typing import TYPE_CHECKING
-
-import numpy as np
-
-from lerobot.utils.import_utils import _onnxruntime_available
-
-from ..g1_utils import (
-    ISAACLAB_TO_MUJOCO,
-    MUJOCO_TO_ISAACLAB,
-    G1_29_JointIndex,
-    get_gravity_orientation,
-)
-
-if TYPE_CHECKING or _onnxruntime_available:
-    import onnxruntime as ort
-else:
-    ort = None
-
-logger = logging.getLogger(__name__)
-
-# ── Constants ────────────────────────────────────────────────────────────────
-# Robot/motor physical constants and the joint-order permutation tables. All
-# 29-vectors are in IsaacLab joint order unless the name says ``_MUJOCO``.
-
-# Nominal standing pose (rad), 29 joints in IsaacLab order. Actions are residuals
-# added on top of this; also used as the planner/encoder standing reference.
-DEFAULT_ANGLES = np.array(
-    [
-        -0.312,
-        0.0,
-        0.0,
-        0.669,
-        -0.363,
-        0.0,
-        -0.312,
-        0.0,
-        0.0,
-        0.669,
-        -0.363,
-        0.0,
-        0.0,
-        0.0,
-        0.0,
-        0.2,
-        0.2,
-        0.0,
-        0.6,
-        0.0,
-        0.0,
-        0.0,
-        0.2,
-        -0.2,
-        0.0,
-        0.6,
-        0.0,
-        0.0,
-        0.0,
-    ],
-    dtype=np.float32,
-)
-
-# Per-motor-type parameters used to derive action scaling and PD gains. Keys are
-# Unitree motor model names; ARMATURE = rotor inertia, EFFORT = torque limit (N·m).
-NATURAL_FREQ = 10.0 * 2.0 * np.pi  # target closed-loop stiffness bandwidth (rad/s)
-ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425}
-EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
-
-
-def _action_scale(k):
-    """Per-motor residual-action scale (maps policy output to joint-angle delta)."""
-    return 0.25 * EFFORT[k] / (ARMATURE[k] * NATURAL_FREQ**2)
-
-
-# Per-joint motor model (IsaacLab order): legs, waist, then arms. Single source of
-# truth for both ACTION_SCALE and compute_kp_kd().
-MOTOR_MODELS = (
-    ["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
-    + ["7520_14", "5020", "5020"]
-    + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
-)
-ACTION_SCALE = np.array([_action_scale(k) for k in MOTOR_MODELS], dtype=np.float32)  # (29,) IsaacLab order
-
-CONTROL_DT = 0.02  # 50 Hz control period (s)
-DEFAULT_HEIGHT = 0.788740  # nominal pelvis height (m)
-TOKEN_DIM = 64  # encoder latent size
-ENCODER_UPDATE_EVERY = 5  # refresh the encoder token every N ticks (decoder runs every tick)
-DEBUG_PRINT_EVERY = 100  # ticks between debug prints
-
-
-def _to_mujoco(a):
-    """Apply the ``MUJOCO_TO_ISAACLAB`` gather to a 29-vector (deploy-order reorder).
-
-    NOTE: this returns ``a[MUJOCO_TO_ISAACLAB]``. The ``_mj`` suffixes and the exact
-    permutation direction throughout this module are a fixed convention validated
-    against the deployed SONIC ONNX policy (the encoder/decoder consume vectors in
-    this order). Do not "correct" the table or rename toward the opposite direction
-    without re-validating on hardware — the labels are historical, the ordering is
-    load-bearing.
-    """
-    return a[MUJOCO_TO_ISAACLAB]
-
-
-DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES)
-ENCODER_STANDING_REF = DEFAULT_ANGLES.copy()
-
-# Joint-index subsets (IsaacLab order) used to slice encoder observations.
-LOWER_BODY_IL = np.array([0, 3, 6, 9, 13, 17, 1, 4, 7, 10, 14, 18], dtype=np.int32)  # 12 leg joints
-WRIST_IL = np.array([23, 24, 25, 26, 27, 28], dtype=np.int32)  # 6 wrist joints
-VR_TARGET_DEF = np.zeros(9, dtype=np.float32)  # 3-point VR position targets (mode 1)
-VR_ORN_DEF = np.array([1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0], dtype=np.float32)  # VR orn targets (mode 1)
-SMPL_DEF = np.zeros(720, dtype=np.float32)  # SMPL whole-body window default (mode 2)
-
-# ── PD gains ─────────────────────────────────────────────────────────────────
-
-
-def compute_kp_kd():
-    """Derive per-joint PD gains (kp, kd) from motor armature and target bandwidth.
-
-    Ankle and waist joints get a x2 factor for extra stiffness. Returns two
-    (29,) float32 arrays in IsaacLab joint order.
-    """
-
-    def s(k):
-        return ARMATURE[k] * NATURAL_FREQ**2
-
-    def d(k):
-        return 2.0 * 2.0 * ARMATURE[k] * NATURAL_FREQ
-
-    _double = {4, 5, 10, 11, 13, 14}  # ankle + waist indices with factor 2
-    kp = np.array([2 * s(k) if i in _double else s(k) for i, k in enumerate(MOTOR_MODELS)], dtype=np.float32)
-    kd = np.array([2 * d(k) if i in _double else d(k) for i, k in enumerate(MOTOR_MODELS)], dtype=np.float32)
-    return kp, kd
-
-
-_kp_kd = compute_kp_kd  # backward-compatible alias
-
-
-# ── Quaternion helpers ────────────────────────────────────────────────────────
-# All quaternions are scalar-first (w, x, y, z). "heading" = yaw-only quaternion.
-
-
-def quat_conj(q):
-    """Quaternion conjugate (inverse for unit quaternions)."""
-    return np.array([q[0], -q[1], -q[2], -q[3]], dtype=np.float32)
-
-
-def quat_mul(q1, q2):
-    """Hamilton product ``q1 ⊗ q2``."""
-    w1, x1, y1, z1 = q1
-    w2, x2, y2, z2 = q2
-    return np.array(
-        [
-            w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
-            w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
-            w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
-            w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
-        ],
-        dtype=np.float32,
-    )
-
-
-def quat_to_6d(q):
-    """Quaternion → 6-D rotation representation (first two rotated basis rows)."""
-    w, x, y, z = q
-    return np.array(
-        [
-            1 - 2 * (y * y + z * z),
-            2 * (x * y - z * w),
-            2 * (x * y + z * w),
-            1 - 2 * (x * x + z * z),
-            2 * (x * z - y * w),
-            2 * (y * z + x * w),
-        ],
-        dtype=np.float32,
-    )
-
-
-def calc_heading(q):
-    """Extract the yaw (heading) angle in radians from a quaternion."""
-    w, x, y, z = q
-    return float(np.arctan2(2 * (x * y + w * z), 1 - 2 * (y * y + z * z)))
-
-
-def heading_quat(q, sign=1.0):
-    """Yaw-only quaternion for ``q``'s heading (``sign=-1`` gives its inverse)."""
-    a = sign * calc_heading(q) / 2.0
-    return np.array([np.cos(a), 0, 0, np.sin(a)], dtype=np.float64)
-
-
-def heading_quat_inv(q):
-    """Inverse yaw-only quaternion for ``q``'s heading."""
-    return heading_quat(q, -1.0)
-
-
-def quat_slerp(q0, q1, t):
-    """Spherical linear interpolation between two quaternions (scalar ``t``)."""
-    q0 = q0 / (np.linalg.norm(q0) + 1e-12)
-    q1 = q1 / (np.linalg.norm(q1) + 1e-12)
-    dot = float(np.dot(q0, q1))
-    if dot < 0:
-        q1, dot = -q1, -dot
-    dot = min(dot, 1.0)
-    if dot > 0.9995:
-        r = q0 + t * (q1 - q0)
-        return r / (np.linalg.norm(r) + 1e-12)
-    th = np.arccos(dot)
-    st = np.sin(th)
-    return (np.sin((1 - t) * th) / st) * q0 + (np.sin(t * th) / st) * q1
-
-
-def quat_slerp_batch(q0, q1, t):
-    """Vectorized slerp over arrays of quaternions with a per-row parameter ``t``."""
-    q0 = q0 / (np.linalg.norm(q0, axis=1, keepdims=True) + 1e-12)
-    q1 = q1 / (np.linalg.norm(q1, axis=1, keepdims=True) + 1e-12)
-    dot = np.sum(q0 * q1, axis=1)
-    neg = dot < 0
-    q1 = q1.copy()
-    q1[neg] = -q1[neg]
-    dot[neg] = -dot[neg]
-    dot = np.clip(dot, -1, 1)
-    lin = dot > 0.9995
-    th = np.arccos(dot)
-    st = np.where(np.sin(th) == 0, 1, np.sin(th))
-    c0 = np.sin((1 - t) * th) / st
-    c1 = np.sin(t * th) / st
-    c0[lin] = 1 - t[lin]
-    c1[lin] = t[lin]
-    r = c0[:, None] * q0 + c1[:, None] * q1
-    return r / (np.linalg.norm(r, axis=1, keepdims=True) + 1e-12)
-
-
-def ort_providers(force_cpu: bool = False) -> list[str]:
-    """Prefer CUDA for enc/dec/planner (matches deploy when onnxruntime-gpu is installed)."""
-    avail = ort.get_available_providers()
-    if not force_cpu and "CUDAExecutionProvider" in avail:
-        return ["CUDAExecutionProvider", "CPUExecutionProvider"]
-    return ["CPUExecutionProvider"]
-
-
-def make_ort_session_options(intra_op_num_threads: int | None = None,
-                             inter_op_num_threads: int | None = None):
-    """Build ONNX Runtime SessionOptions (quiet logging).
-
-    Pass thread counts to cap ORT's CPU pool. These tiny MLP policies are latency-
-    bound, not throughput-bound, so letting ORT grab every core just starves the
-    real-time control loop / torch policy / IK solver and causes stutter. 1 intra +
-    1 inter thread is plenty and lowest-latency for a per-step MLP inference.
-    """
-    so = ort.SessionOptions()
-    so.log_severity_level = 3
-    if intra_op_num_threads is not None:
-        so.intra_op_num_threads = intra_op_num_threads
-    if inter_op_num_threads is not None:
-        so.inter_op_num_threads = inter_op_num_threads
-    return so
-
-
-# ── Encoder / Decoder ─────────────────────────────────────────────────────────
-
-
-class StandingEncoderDecoder:
-    """Runs the encoder + decoder ONNX models and owns the proprioception history.
-
-    Each tick it appends the latest robot state to 10-frame history buffers, builds
-    the encoder observation (1762-D, layout depends on ``encode_mode``) to refresh
-    the 64-D ``token``, then builds the decoder observation (994-D) and maps
-    ``token + history`` to a residual action added onto ``DEFAULT_ANGLES``.
-
-    ``PlannerController`` subclasses this to source the reference from a live,
-    planner-generated motion buffer instead of a fixed standing pose.
-    """
-
-    def __init__(self, encoder, decoder):
-        self.encoder, self.decoder = encoder, decoder
-        self.encoder_input = encoder.get_inputs()[0].name
-        self.decoder_input = decoder.get_inputs()[0].name
-        enc_dim = int(encoder.get_inputs()[0].shape[1])
-        dec_dim = int(decoder.get_inputs()[0].shape[1])
-        if enc_dim != 1762 or dec_dim != 994:
-            raise RuntimeError(f"Unexpected dims encoder={enc_dim}, decoder={dec_dim}")
-        self.token = np.zeros(TOKEN_DIM, np.float32)
-        self.last_action_mj = np.zeros(29, np.float32)
-        self.h_q_mj = [np.zeros(29, np.float32)] * 10
-        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
-        self.h_ang = [np.zeros(3, np.float32)] * 10
-        self.h_act_mj = [np.zeros(29, np.float32)] * 10
-        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
-        self.init_base_quat = np.array([1, 0, 0, 0], np.float32)
-        self.init_ref_quat = np.array([1, 0, 0, 0], np.float32)
-        self._heading_init = False
-        self.encode_mode = 0
-        self.vr_3point_local_target = VR_TARGET_DEF.copy()
-        self.vr_3point_local_orn_target = VR_ORN_DEF.copy()
-        self.smpl_joints_10frame_step1 = SMPL_DEF.copy()
-        # Optional per-frame SMPL root orientation (wxyz) for the mode-2 anchor.
-        # When None, the anchor falls back to the planner reference body quat.
-        self.smpl_root_quat = None
-        self.set_zero_reference()
-
-    def reset(self):
-        """Clear the token, 10-frame proprioception history and heading init.
-
-        ``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new
-        episode are not contaminated by the previous episode's state.
-        """
-        self.token = np.zeros(TOKEN_DIM, np.float32)
-        self.last_action_mj = np.zeros(29, np.float32)
-        self.h_q_mj = [np.zeros(29, np.float32)] * 10
-        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
-        self.h_ang = [np.zeros(3, np.float32)] * 10
-        self.h_act_mj = [np.zeros(29, np.float32)] * 10
-        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
-        self.init_base_quat = np.array([1, 0, 0, 0], np.float32)
-        self.init_ref_quat = np.array([1, 0, 0, 0], np.float32)
-        self._heading_init = False
-
-    def update_history(self, q, dq, ang, quat):
-        """Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers."""
-        quat = quat / (np.linalg.norm(quat) + 1e-8)
-        q_mj = _to_mujoco(q)
-        dq_mj = _to_mujoco(dq)
-        self.h_q_mj = [q_mj - DEFAULT_ANGLES_MUJOCO] + self.h_q_mj[:-1]
-        self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
-        self.h_ang = [ang.copy()] + self.h_ang[:-1]
-        self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
-        self.h_quat = [quat.copy()] + self.h_quat[:-1]
-        if not self._heading_init:
-            self.init_base_quat = quat.copy()
-            self._heading_init = True
-
-    def _heading_quat(self, q):
-        h = calc_heading(q) / 2.0
-        return np.array([np.cos(h), 0, 0, np.sin(h)], np.float32)
-
-    def _heading_quat_inv(self, q):
-        h = calc_heading(q) / 2.0
-        return np.array([np.cos(-h), 0, 0, np.sin(-h)], np.float32)
-
-    def _anchor_6d(self, base_quat, ref_quat=None):
-        """6-D orientation error between the robot base and the (heading-aligned) reference."""
-        if ref_quat is None:
-            ref_quat = self.init_ref_quat
-        delta = quat_mul(self._heading_quat(self.init_base_quat), self._heading_quat_inv(self.init_ref_quat))
-        new_ref = quat_mul(delta, ref_quat)
-        return quat_to_6d(quat_mul(quat_conj(base_quat), new_ref))
-
-    def set_zero_reference(self):
-        """Initialize the reference to a single standing frame (used before a plan exists)."""
-        self.motion_joint_positions = [ENCODER_STANDING_REF.copy()]
-        self.motion_joint_velocities = [np.zeros(29, np.float32)]
-        self.motion_body_quats = [np.array([1, 0, 0, 0], np.float32)]
-        self.motion_body_z = [DEFAULT_HEIGHT]
-        self.motion_timesteps = 1
-        self.freeze_ref_frame = 0
-        self.init_ref_quat = self.motion_body_quats[0].copy()
-
-    def build_encoder_obs(self):
-        """Assemble the 1762-D encoder input; slot layout depends on ``encode_mode``.
-
-        mode 0 = locomotion (ref joint pos + anchor), 1 = 3-point VR teleop
-        (lower-body ref + VR targets), 2 = SMPL whole-body window + anchor/wrist.
-        """
-        obs = np.zeros(1762, np.float32)
-        obs[0] = float(self.encode_mode)
-        rf = min(self.freeze_ref_frame, self.motion_timesteps - 1)
-        ref_pos, ref_quat = self.motion_joint_positions[rf], self.motion_body_quats[rf]
-        if self.encode_mode == 0:
-            for f in range(10):
-                obs[4 + 29 * f : 4 + 29 * (f + 1)] = ref_pos
-                obs[601 + 6 * f : 601 + 6 * (f + 1)] = self._anchor_6d(self.h_quat[0], ref_quat)
-        elif self.encode_mode == 1:
-            ref_lower = ref_pos[LOWER_BODY_IL]
-            for f in range(10):
-                obs[661 + 12 * f : 661 + 12 * (f + 1)] = ref_lower
-            obs[901:910] = self.vr_3point_local_target
-            obs[910:922] = self.vr_3point_local_orn_target
-            obs[595:601] = self._anchor_6d(self.h_quat[0], ref_quat)
-        elif self.encode_mode == 2:
-            # Prefer the SMPL clip/stream root orientation for the anchor; fall
-            # back to the planner reference body quat when no root is provided.
-            anchor_ref = self.smpl_root_quat if self.smpl_root_quat is not None else ref_quat
-            obs[922:1642] = self.smpl_joints_10frame_step1
-            for f in range(10):
-                obs[1642 + 6 * f : 1642 + 6 * (f + 1)] = self._anchor_6d(self.h_quat[0], anchor_ref)
-                obs[1702 + 6 * f : 1702 + 6 * (f + 1)] = ref_pos[WRIST_IL]
-        else:
-            raise RuntimeError(f"Unsupported encoder mode: {self.encode_mode}")
-        return obs
-
-    def build_decoder_obs(self):
-        """Assemble the 994-D decoder input: token + 10-frame proprioception history + gravity."""
-        obs = np.zeros(994, np.float32)
-        off = 0
-        obs[off : off + 64] = self.token
-        off += 64
-        for h, sz in [
-            (list(reversed(self.h_ang)), 3),
-            (list(reversed(self.h_q_mj)), 29),
-            (list(reversed(self.h_dq_mj)), 29),
-            (list(reversed(self.h_act_mj)), 29),
-        ]:
-            for f in range(10):
-                obs[off : off + sz] = h[f]
-                off += sz
-        for q in reversed(self.h_quat):
-            obs[off : off + 3] = get_gravity_orientation(q)
-            off += 3
-        assert off == 994, f"Decoder obs mismatch: {off}"
-        return obs
-
-    def run_encoder(self):
-        """Run the encoder ONNX model and return the fresh 64-D token."""
-        return (
-            self.encoder.run(None, {self.encoder_input: self.build_encoder_obs().reshape(1, -1)})[0]
-            .squeeze()
-            .astype(np.float32)
-        )
-
-    def step(self, robot_obs, update_encoder, debug=False):
-        """One control tick: read robot obs, (optionally) re-encode, decode → joint targets.
-
-        Args:
-            robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
-            update_encoder: refresh the token this tick (else reuse the cached one).
-            debug: print action/delta norms.
-
-        Returns:
-            dict of ``<joint>.q`` target positions (rad) in IsaacLab joint order.
-        """
-        jnames = [m.name for m in G1_29_JointIndex]
-        q = np.array(
-            [
-                robot_obs.get(f"{n}.q", DEFAULT_ANGLES[m.value])
-                for m, n in zip(G1_29_JointIndex, jnames, strict=False)
-            ],
-            np.float32,
-        )
-        dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
-        quat = np.array(
-            [
-                robot_obs.get("imu.quat.w", 1),
-                robot_obs.get("imu.quat.x", 0),
-                robot_obs.get("imu.quat.y", 0),
-                robot_obs.get("imu.quat.z", 0),
-            ],
-            np.float32,
-        )
-        ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32)
-        self.update_history(q, dq, ang, quat)
-        if update_encoder:
-            self.token = self.run_encoder()
-        action_mj = (
-            self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0]
-            .squeeze()
-            .astype(np.float32)
-        )
-        self.last_action_mj = action_mj.copy()
-        target = DEFAULT_ANGLES + action_mj[ISAACLAB_TO_MUJOCO] * ACTION_SCALE
-        if debug:
-            delta = target - q
-            logger.debug(
-                "token_norm=%.4f action_norm=%.4f delta_max=%.4f delta_rms=%.4f",
-                np.linalg.norm(self.token),
-                np.linalg.norm(action_mj),
-                np.max(np.abs(delta)),
-                np.sqrt(np.mean(delta**2)),
-            )
-        return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
-
-
-class PlannerController(StandingEncoderDecoder):
-    """Encoder/decoder driven by a caller-supplied, rolling motion buffer.
-
-    Extends ``StandingEncoderDecoder`` so the reference comes from a motion buffer
-    (a lookahead window with per-frame velocities) instead of a single fixed pose,
-    and handles heading re-initialization on the first frame / after a reset.
-    ``motion_lock`` guards the buffer, which the whole-body controller rewrites each
-    tick from the incoming command. The class name is retained for continuity with
-    the SONIC reference; no motion planner is involved.
-    """
-
-    def __init__(self, encoder, decoder):
-        super().__init__(encoder, decoder)
-        self.ref_cursor = 0
-        self.motion_timesteps = 0
-        self.motion_joint_positions = np.zeros((1500, 29), np.float64)
-        self.motion_joint_velocities = np.zeros((1500, 29), np.float64)
-        self.motion_body_quats = np.zeros((1500, 4), np.float64)
-        self.motion_body_quats[:, 0] = 1.0
-        self.motion_body_pos = np.zeros((1500, 3), np.float64)
-        self.init_ref_quat = np.array([1, 0, 0, 0], np.float64)
-        self.heading_init_base_quat = np.array([1, 0, 0, 0], np.float64)
-        self.delta_heading = 0.0
-        self.reinit_heading = False
-        self.playing = self.first_motion = False
-        self.motion_lock = threading.Lock()
-
-    def reset(self):
-        """Full reset: clear enc/dec state (super) plus the motion buffer and heading.
-
-        Forces a heading re-init on the next ``step`` so the reference frame is
-        re-latched to the post-reset robot orientation.
-        """
-        super().reset()
-        with self.motion_lock:
-            self.ref_cursor = 0
-            self.motion_timesteps = 0
-            self.motion_joint_positions[:] = 0.0
-            self.motion_joint_velocities[:] = 0.0
-            self.motion_body_quats[:] = 0.0
-            self.motion_body_quats[:, 0] = 1.0
-            self.motion_body_pos[:] = 0.0
-        self.init_ref_quat = np.array([1, 0, 0, 0], np.float64)
-        self.heading_init_base_quat = np.array([1, 0, 0, 0], np.float64)
-        self.delta_heading = 0.0
-        self.first_motion = False
-        self.playing = False
-        self.reinit_heading = True
-
-    def _heading_apply_delta(self):
-        """Heading correction quaternion (init base-vs-ref heading + operator ``delta_heading``)."""
-        delta = quat_mul(
-            heading_quat(self.heading_init_base_quat).astype(np.float32),
-            heading_quat_inv(self.init_ref_quat).astype(np.float32),
-        )
-        if self.delta_heading:
-            h = self.delta_heading / 2.0
-            delta = quat_mul(np.array([np.cos(h), 0, 0, np.sin(h)], np.float32), delta)
-        return delta
-
-    def _anchor_6d(self, base_quat, ref_quat=None):
-        """6-D base-vs-reference orientation error, including the operator heading delta."""
-        if ref_quat is None:
-            ref_quat = self.init_ref_quat
-        new_ref = quat_mul(self._heading_apply_delta(), ref_quat.astype(np.float32))
-        return quat_to_6d(quat_mul(quat_conj(base_quat.astype(np.float32)), new_ref))
-
-    def build_encoder_obs(self):
-        """Encoder input sourced from the live motion buffer (mode 0/2), lock-protected."""
-        obs = np.zeros(1762, np.float32)
-        obs[0] = float(self.encode_mode)
-        with self.motion_lock:
-            if self.encode_mode == 2:
-                # SMPL whole-body imitation: the 720-dim SMPL window carries the
-                # target pose; the planner reference frame supplies anchor + wrist.
-                rf = min(self.ref_cursor, self.motion_timesteps - 1)
-                ref_pos = self.motion_joint_positions[rf].astype(np.float32)
-                ref_quat = self.motion_body_quats[rf].astype(np.float32)
-                # Prefer the SMPL clip/stream root orientation (if provided) so the
-                # anchor tracks the operator's/clip's heading; else planner ref.
-                if self.smpl_root_quat is not None:
-                    ref_quat = np.asarray(self.smpl_root_quat, np.float32)
-                anchor = self._anchor_6d(self.h_quat[0], ref_quat)
-                wrist = ref_pos[WRIST_IL]
-                obs[922:1642] = self.smpl_joints_10frame_step1
-                for f in range(10):
-                    obs[1642 + 6 * f : 1642 + 6 * (f + 1)] = anchor
-                    obs[1702 + 6 * f : 1702 + 6 * (f + 1)] = wrist
-                return obs
-            if self.encode_mode == 1:
-                # 3-point VR teleop: the upper body tracks the VR wrist/neck targets
-                # while the planner reference supplies the lower body + anchor. Lower
-                # body is per-frame (step 5) like mode 0; the VR targets are current.
-                rf = min(self.ref_cursor, self.motion_timesteps - 1)
-                obs[595:601] = self._anchor_6d(self.h_quat[0], self.motion_body_quats[rf].astype(np.float32))
-                for f in range(10):
-                    tf = min(
-                        self.ref_cursor + f * 5 if self.playing else self.ref_cursor,
-                        self.motion_timesteps - 1,
-                    )
-                    ref_lower = self.motion_joint_positions[tf].astype(np.float32)[LOWER_BODY_IL]
-                    obs[661 + 12 * f : 661 + 12 * (f + 1)] = ref_lower
-                obs[901:910] = self.vr_3point_local_target
-                obs[910:922] = self.vr_3point_local_orn_target
-                return obs
-            for f in range(10):
-                tf = min(
-                    self.ref_cursor + f * 5 if self.playing else self.ref_cursor, self.motion_timesteps - 1
-                )
-                obs[4 + 29 * f : 4 + 29 * (f + 1)] = self.motion_joint_positions[tf].astype(np.float32)
-                if self.playing:
-                    obs[294 + 29 * f : 294 + 29 * (f + 1)] = self.motion_joint_velocities[tf].astype(
-                        np.float32
-                    )
-                obs[601 + 6 * f : 601 + 6 * (f + 1)] = self._anchor_6d(
-                    self.h_quat[0], self.motion_body_quats[tf].astype(np.float32)
-                )
-        return obs
-
-    def step(self, robot_obs, update_encoder, debug=False):
-        """Re-init the heading reference on first frame / after a reset, then run the base step."""
-        if robot_obs and (self.first_motion or self.reinit_heading):
-            q = None
-            if "imu.quat.w" in robot_obs:
-                q = np.array(
-                    [
-                        robot_obs["imu.quat.w"],
-                        robot_obs["imu.quat.x"],
-                        robot_obs["imu.quat.y"],
-                        robot_obs["imu.quat.z"],
-                    ],
-                    np.float64,
-                )
-            else:
-                q = robot_obs.get("imu.quaternion")
-                if q is not None:
-                    q = np.array(q, np.float64)
-            if q is not None:
-                self.heading_init_base_quat = np.array(q, np.float64)
-                with self.motion_lock:
-                    rf = min(self.ref_cursor, self.motion_timesteps - 1)
-                    if self.encode_mode == 2 and self.smpl_root_quat is not None:
-                        # Anchor the heading delta to the SMPL root at init so the
-                        # robot turns *relative* to the clip/operator start heading.
-                        self.init_ref_quat = np.asarray(self.smpl_root_quat, np.float64)
-                    else:
-                        self.init_ref_quat = self.motion_body_quats[rf].copy()
-                self.delta_heading = 0.0
-                self.first_motion = False
-                self.reinit_heading = False
-                logger.debug("[Heading] init quat: %s", self.heading_init_base_quat)
-        return super().step(robot_obs, update_encoder=update_encoder, debug=debug)
-
-    def advance_cursor(self):
-        """Advance the reference cursor one frame per 50 Hz tick (no wall-clock catch-up)."""
-        if not self.playing:
-            return
-        with self.motion_lock:
-            if self.motion_timesteps > 0:
-                self.ref_cursor = min(self.ref_cursor + 1, self.motion_timesteps - 1)

src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py

@@ -1,415 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 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.
-
-"""SONIC full-body controller for Unitree G1."""
-
-from __future__ import annotations
-
-from collections import deque
-import logging
-from typing import TYPE_CHECKING
-
-from huggingface_hub import hf_hub_download
-import numpy as np
-
-from lerobot.utils.import_utils import _onnxruntime_available, require_package
-
-from ..g1_utils import (
-    MUJOCO_TO_ISAACLAB,
-    WB_ACTION_DIM,
-    G1_29_JointIndex,
-    lowstate_to_obs,
-    wb_action_key,
-)
-from .sonic_pipeline import (
-    CONTROL_DT,
-    DEFAULT_ANGLES,
-    ENCODER_UPDATE_EVERY,
-    TOKEN_DIM,
-    PlannerController,
-    compute_kp_kd,
-    make_ort_session_options,
-    ort_providers,
-)
-
-# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
-TOKEN_ACTION_PREFIX = "motion_token"
-# Proprio-state prefix for the token interface: the robot echoes the last commanded
-# token here so ``lerobot-rollout`` aggregates it into a 64-D ``observation.state``.
-TOKEN_STATE_PREFIX = "motion_token_state"
-
-
-def token_action_key(i: int) -> str:
-    """Action-dict key for the i-th component of the 64-D SONIC latent token.
-
-    The ``.pos`` suffix is required so the value flows through ``lerobot-rollout``,
-    which only routes ``.pos`` scalar features onto the policy action vector.
-    """
-    return f"{TOKEN_ACTION_PREFIX}.{i}.pos"
-
-
-def token_state_key(i: int) -> str:
-    """Observation key for the i-th component of the 64-D SONIC latent token state."""
-    return f"{TOKEN_STATE_PREFIX}.{i}.pos"
-
-if TYPE_CHECKING or _onnxruntime_available:
-    import onnxruntime as ort
-else:
-    ort = None
-
-logger = logging.getLogger(__name__)
-
-# Startup blend duration: over the first control ticks, linearly interpolate every joint
-# from the robot's initial measured pose into the policy's commanded target, so control
-# eases in without a snap on the first command.
-INIT_RAMP_S = 3.0
-
-# Neutral ("zero pose") SONIC token, held by token_mode until the first real token
-# arrives. Captured from the encoder's own output while the robot stood idle in sim
-# (capture_neutral_token.py): the encoder is an FSQ bottleneck (~5 bit/dim, 15.5 half-
-# width, Div(16)), so its tokens live on the 1/16 grid. We store the integer FSQ codes
-# and rescale by the same 1/16 step, giving an exact on-grid token -- unlike the literal
-# all-zero token, which is off the encoder's learned manifold and decodes to a slightly
-# goofy stance. This one decodes to a stable, natural standing pose.
-_NEUTRAL_TOKEN_CODES = np.array(
-    [-1, 3, 1, -1, 1, -3, 6, 1, 1, 1, -2, -4, -2, 0, -3, -1,
-     2, -1, -3, -5, 3, 1, 1, -4, -1, -1, 1, -7, 0, 1, 2, -2,
-     5, -2, -2, -4, 0, -1, 3, -1, 0, -5, -1, 0, -4, 0, 0, -1,
-     -1, 2, -2, 1, 3, 3, 1, 0, 0, 6, 0, -7, 3, 0, 2, -2],
-    dtype=np.float32,
-)
-NEUTRAL_TOKEN = _NEUTRAL_TOKEN_CODES / 16.0  # FSQ Div(16): integer codes -> on-grid token
-
-
-def _extract_wb34_from_action(action: dict | None) -> np.ndarray | None:
-    """Reassemble a dense (34,) whole-body command from ``wb.{i}.pos`` keys, or None.
-
-    This is the OpenHLM / pi0.5 joint-based interface: one 34-D vector per tick
-    (sentinel: presence of ``wb.0.pos``) carrying absolute joint targets in real
-    units. The ``.pos`` suffix lets these flow through ``lerobot-rollout`` as normal
-    joint-position action features.
-    """
-    if not action:
-        return None
-    keys = [wb_action_key(i) for i in range(WB_ACTION_DIM)]
-    # Require the full dense command: a partial action (e.g. only ``wb.0.pos``)
-    # must not be silently zero-filled, which would drive most joints toward 0.
-    if any(key not in action for key in keys):
-        return None
-    return np.fromiter(
-        (float(action[key]) for key in keys),
-        dtype=np.float32,
-        count=WB_ACTION_DIM,
-    )
-
-
-def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
-    """Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None.
-
-    This is the token-only replay interface: instead of a joint reference driving the
-    encoder, the caller supplies the 64-D encoder latent directly (e.g. a recorded
-    ``action.motion_token`` column), which the decoder consumes with the encoder
-    bypassed. Requires the full dense token; a partial one is ignored (returns None).
-    """
-    if not action:
-        return None
-    keys = [token_action_key(i) for i in range(TOKEN_DIM)]
-    if any(key not in action for key in keys):
-        return None
-    return np.fromiter(
-        (float(action[key]) for key in keys),
-        dtype=np.float32,
-        count=TOKEN_DIM,
-    )
-
-
-def _wb34_to_reference(wb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-    """Map a 34-D OpenHLM whole-body command to a SONIC mode-0 reference.
-
-    Returns ``(ref29, anchor_quat)`` where ``ref29`` is the 29 joint targets in
-    IsaacLab order (what SONIC's ``motion_joint_positions`` expects) and
-    ``anchor_quat`` (wxyz) encodes the root roll/pitch (yaw=0).
-
-    OpenHLM layout : [L-arm 0:7, L-grip 7, R-arm 8:15, R-grip 15,
-                      L-leg 16:22, R-leg 22:28, waist 28:31, root rp+yaw 31:34]
-    The 29 joints are first assembled in MuJoCo / Unitree-SDK order
-    ([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29] — the
-    ``G1_29_JointIndex`` grouping OpenHLM uses), then permuted to IsaacLab order via
-    ``MUJOCO_TO_ISAACLAB``. Grippers (7, 15) are not part of the 29-DoF SONIC
-    reference, and yaw-rate (33) is integrated into the heading by the caller (it
-    cannot be represented in this static per-tick anchor).
-    """
-    ref_mj = np.zeros(29, np.float32)  # MuJoCo / Unitree-SDK grouped order
-    ref_mj[0:6] = wb[16:22]  # left leg
-    ref_mj[6:12] = wb[22:28]  # right leg
-    ref_mj[12:15] = wb[28:31]  # waist
-    ref_mj[15:22] = wb[0:7]  # left arm
-    ref_mj[22:29] = wb[8:15]  # right arm
-    ref = ref_mj[MUJOCO_TO_ISAACLAB].astype(np.float32)  # -> IsaacLab order for SONIC
-    roll, pitch = float(wb[31]), float(wb[32])
-    cr, sr, cp, sp = np.cos(roll / 2), np.sin(roll / 2), np.cos(pitch / 2), np.sin(pitch / 2)
-    anchor = np.array([cr * cp, sr * cp, cr * sp, sr * sp], np.float32)  # Rx(roll)·Ry(pitch)
-    return ref, anchor
-
-
-class SonicRuntime:
-    """Loads the SONIC encoder/decoder ONNX models and owns the controller.
-
-    No motion planner: the reference motion buffer is written directly each tick by
-    :class:`SonicWholeBodyController` from the incoming 34-D whole-body command.
-    """
-
-    def __init__(self, force_cpu: bool = False):
-        require_package("onnxruntime", extra="unitree_g1")
-        encoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_encoder.onnx")
-        decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx")
-
-        providers = ort_providers(force_cpu=force_cpu)
-        so = make_ort_session_options()
-
-        encoder_sess = ort.InferenceSession(encoder_path, sess_options=so, providers=providers)
-        decoder_sess = ort.InferenceSession(decoder_path, sess_options=so, providers=providers)
-
-        # Report the provider actually bound, not the one requested: ORT silently falls
-        # back to CPU if CUDA can't load (e.g. libcudnn not on LD_LIBRARY_PATH), and a
-        # CPU decoder drifts the closed-loop heading. Warn loudly so it can't hide.
-        self.use_gpu = decoder_sess.get_providers()[0] == "CUDAExecutionProvider"
-        if not force_cpu and not self.use_gpu:
-            print(
-                "[SONIC] WARNING: decoder bound to CPUExecutionProvider (CUDA unavailable). "
-                "Closed-loop replay/control will drift. Ensure libcudnn is on LD_LIBRARY_PATH "
-                "(site-packages/nvidia/*/lib).",
-                flush=True,
-            )
-
-        self.kp, self.kd = compute_kp_kd()
-        self.controller = PlannerController(encoder_sess, decoder_sess)
-
-    @property
-    def pipeline(self):
-        return self.controller
-
-    def reset(self):
-        # Full pipeline reset: clears the encoder token, proprioception history and
-        # heading, and rewinds the motion buffer. reinit_heading is set so the next
-        # step re-latches the reference frame to the current robot orientation.
-        self.controller.reset()
-
-    def shutdown(self):
-        pass
-
-
-class SonicWholeBodyController:
-    """Full-body SONIC controller for UnitreeG1's background controller thread."""
-
-    control_dt = CONTROL_DT
-    full_body = True
-    # Advertise a dense 34-D whole-body action space (OpenHLM / pi0.5) so the robot
-    # exposes ``wb.{i}.pos`` action features and ``lerobot-rollout`` can drive it
-    # directly with a 34-D VLA policy.
-    wb_action = True
-
-    def __init__(self, force_cpu: bool = False):
-        logger.info("Loading SONIC whole-body controller...")
-        self._runtime = SonicRuntime(force_cpu=force_cpu)
-        self.kp = self._runtime.kp
-        self.kd = self._runtime.kd
-        self.controller = self._runtime.controller
-
-        # Startup blend: ease from the robot's initial pose into the first commanded
-        # policy targets over INIT_RAMP_S (captured on the first control tick).
-        self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
-        self._init_step = 0
-        self._start_pose: dict[str, float] = {}
-
-        # Tick counter for the dense whole-body (OpenHLM, mode-0) path's encoder cadence.
-        self._wb_step = 0
-        # Rolling 50-frame reference trajectory (ref29 + anchor quat) built from the
-        # stream of per-tick whole-body commands, fed to the encoder as a batch.
-        self._wb_traj: deque[np.ndarray] = deque(maxlen=50)
-        self._wb_quat_traj: deque[np.ndarray] = deque(maxlen=50)
-        # Integrated heading (rad) from the whole-body command's yaw-rate (index 33),
-        # forwarded to the pipeline as ``delta_heading`` so turn commands take effect.
-        self._heading = 0.0
-
-        # Token-interface state. ``token_mode`` is set True by the robot when the deploy
-        # is token-driven (``UnitreeG1Config.sonic_token_action``): the controller then
-        # holds a stable *neutral* (all-zero) token until the first real token arrives,
-        # and afterwards holds the *last* token received between ticks (the async
-        # controller runs ~50 Hz while a token VLA streams ~30 Hz). This lives here (not
-        # in the entry-point script) so it applies uniformly to run_g1_onboard,
-        # lerobot-rollout and the sim replays. ``token_mode`` stays False for the dense
-        # 34-D whole-body / OpenHLM path, which keeps its own "hold last target" idle.
-        self.token_mode = False
-        self._last_token: np.ndarray | None = None
-
-        logger.info("SONIC ready (encoder/decoder, 34-D whole-body command path)")
-
-    def _run_wholebody34(self, obs: dict, wb: np.ndarray) -> dict:
-        """Feed a dense 34-D OpenHLM whole-body command as the mode-0 encoder reference.
-
-        The 29 joint targets are held across the encoder lookahead window (zero
-        velocity) and the root roll/pitch set the anchor orientation, then the
-        encoder/decoder run directly (planner bypassed). One command per tick, so the
-        VLA's commanded pose is what SONIC tracks.
-        """
-        ref, anchor = _wb34_to_reference(wb)
-        c = self.controller
-        if c.encode_mode != 0:
-            c.encode_mode = 0
-            c.reinit_heading = True
-        # Index 33 is a yaw-rate (rad/s): integrate it into a heading offset and hand
-        # it to the pipeline as ``delta_heading`` so commanded turns are tracked rather
-        # than silently dropped (the anchor from _wb34_to_reference only carries r/p).
-        self._heading += float(wb[33]) * CONTROL_DT
-        c.delta_heading = self._heading
-        # Capture the heading/anchor reference on the first whole-body tick. The
-        # controller only latches ``init_ref_quat`` (and the base heading) inside
-        # ``step()`` when ``first_motion or reinit_heading`` — but it already boots in
-        # mode 0, so the mode-switch guard above misses the very first command and the
-        # anchor would stay identity. This mirrors the GEAR reference, which seeds
-        # ``init_ref_quat`` from the first anchor. Must run before the buffers below so
-        # ``step()`` latches ``motion_body_quats[0]`` = this tick's anchor.
-        if self._wb_step == 0:
-            c.reinit_heading = True
-
-        # Accumulate the per-tick commands into a rolling 50-frame reference
-        # trajectory so the encoder's 10-frame, step-5 lookahead sees an actual
-        # motion sequence (with velocities) instead of one repeated pose. 50 frames
-        # == chunk horizon == 10 lookahead frames × step 5.
-        self._wb_traj.append(ref)
-        self._wb_quat_traj.append(anchor)
-        traj = np.asarray(self._wb_traj, np.float32)  # (L, 29), oldest -> newest
-        quats = np.asarray(self._wb_quat_traj, np.float32)  # (L, 4)
-        n = len(traj)
-        # Per-frame velocities from finite differences (rad/s at the control rate).
-        vel = np.zeros_like(traj)
-        if n > 1:
-            vel[1:] = (traj[1:] - traj[:-1]) / CONTROL_DT
-            vel[0] = vel[1]
-        with c.motion_lock:
-            c.motion_joint_positions[:n] = traj
-            c.motion_joint_velocities[:n] = vel
-            c.motion_body_quats[:n] = quats
-            c.motion_body_pos[:n] = 0.0
-            c.motion_timesteps = n
-            c.ref_cursor = 0
-        c.playing = True
-        do_enc = self._wb_step % ENCODER_UPDATE_EVERY == 0
-        out = c.step(obs, update_encoder=do_enc, debug=False)
-        if self._wb_step % 25 == 0:
-            tgt = np.array([out[f"{m.name}.q"] for m in G1_29_JointIndex], np.float32)
-            logger.info(
-                "[WB34] step=%d |ref|mean=%.3f |target|mean=%.3f target_std=%.3f init_ref_quat=%s",
-                self._wb_step,
-                float(np.abs(ref).mean()),
-                float(np.abs(tgt).mean()),
-                float(tgt.std()),
-                np.round(c.init_ref_quat, 3).tolist(),
-            )
-        self._wb_step += 1
-        return out
-
-    def _run_token(self, obs: dict, token: np.ndarray) -> dict:
-        """Decode a supplied 64-D latent token directly (encoder bypassed).
-
-        Token-only replay: set the pipeline's cached token to the supplied one and run
-        a decode-only step (``update_encoder=False``). The decoder still closes the loop
-        on live proprioception (history is refreshed inside ``step`` from ``obs``); only
-        the encoder — which would recompute the token from a motion reference — is
-        skipped. Returns the ``<joint>.q`` target dict.
-        """
-        c = self.controller
-        c.token = np.asarray(token, np.float32)
-        self._wb_step += 1
-        return c.step(obs, update_encoder=False, debug=False)
-
-    def _startup_blend(self, obs: dict, out: dict) -> dict:
-        """Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
-        interpolate between the robot's pose captured on the first tick and the policy's
-        live commanded target, so the handoff has no snap.
-
-        ``out`` is the policy's ``<joint>.q`` target dict for this tick; the blend ratio
-        climbs 0->1 over the ramp, after which the raw policy target passes through.
-        """
-        if self._init_step >= self._init_ramp_steps or not out:
-            return out
-        if self._init_step == 0:
-            # Capture the robot's actual pose as the interpolation start point.
-            self._start_pose = {
-                f"{m.name}.q": float(obs.get(f"{m.name}.q", DEFAULT_ANGLES[m.value]))
-                for m in G1_29_JointIndex
-            }
-        self._init_step += 1
-        ratio = min(1.0, self._init_step / self._init_ramp_steps)
-        blended = {
-            k: self._start_pose.get(k, float(tgt)) * (1.0 - ratio) + float(tgt) * ratio
-            for k, tgt in out.items()
-        }
-        if self._init_step >= self._init_ramp_steps:
-            logger.info("SONIC startup blend complete -> full policy control")
-        return blended
-
-    def run_step(self, action: dict, lowstate) -> dict:
-        if lowstate is None:
-            return {}
-        obs = lowstate_to_obs(lowstate)
-
-        # Token-only interface (latent replay / token-output VLA): a dense 64-D
-        # ``motion_token.{i}`` command is decoded directly, bypassing the encoder.
-        # Checked before the joint path so a token action takes precedence.
-        token = _extract_token_from_action(action)
-        if token is not None:
-            self._last_token = token
-        elif self._last_token is None and self.token_mode:
-            # Token-driven deploy, but no token has arrived yet: hold the captured
-            # neutral token (NEUTRAL_TOKEN), which the decoder maps to a stable, natural
-            # standing pose (the encoder's own idle output; see NEUTRAL_TOKEN).
-            self._last_token = NEUTRAL_TOKEN.copy()
-        if self._last_token is not None:
-            # Either a fresh token this tick or the last one received (held between the
-            # ~30 Hz token stream and the ~50 Hz control loop).
-            return self._startup_blend(obs, self._run_token(obs, self._last_token))
-
-        # Dense 34-D whole-body command (OpenHLM / pi0.5 joint interface): a single
-        # vector per tick drives the mode-0 encoder reference directly. Until the
-        # policy produces one, hold (no command) so the robot keeps its last target.
-        wb = _extract_wb34_from_action(action)
-        if wb is None:
-            self._wb_miss = getattr(self, "_wb_miss", 0) + 1
-            if self._wb_miss % 50 == 1:
-                akeys = [k for k in action if isinstance(k, str)]
-                logger.info(
-                    "[WB34] no wb.*.pos in action this tick (miss=%d). action keys sample: %s",
-                    self._wb_miss,
-                    akeys[:8],
-                )
-            return {}
-        return self._startup_blend(obs, self._run_wholebody34(obs, wb))
-
-    def reset(self):
-        self._runtime.reset()
-        self._init_step = 0  # re-run the startup blend after a reset
-        self._start_pose = {}
-        self._wb_step = 0
-        self._wb_traj.clear()
-        self._wb_quat_traj.clear()
-        self._heading = 0.0
-        # Drop the held token so token_mode re-seeds the neutral token after a reset.
-        self._last_token = None
-
-    def shutdown(self):
-        self._runtime.shutdown()

src/lerobot/robots/unitree_g1/g1_utils.py

@@ -23,102 +23,10 @@ import numpy as np
 
 NUM_MOTORS = 29
 
-# Joint-order permutations between the two 29-DoF layouts used across the G1 stack:
-# IsaacLab (policy/training order) and MuJoCo (deploy order). ``a[ISAACLAB_TO_MUJOCO]``
-# reorders an IsaacLab-ordered vector into MuJoCo order, and vice-versa.
-ISAACLAB_TO_MUJOCO = np.array(
-    [
-        0,
-        3,
-        6,
-        9,
-        13,
-        17,
-        1,
-        4,
-        7,
-        10,
-        14,
-        18,
-        2,
-        5,
-        8,
-        11,
-        15,
-        19,
-        21,
-        23,
-        25,
-        27,
-        12,
-        16,
-        20,
-        22,
-        24,
-        26,
-        28,
-    ],
-    dtype=np.int32,
-)
-MUJOCO_TO_ISAACLAB = np.array(
-    [
-        0,
-        6,
-        12,
-        1,
-        7,
-        13,
-        2,
-        8,
-        14,
-        3,
-        9,
-        15,
-        22,
-        4,
-        10,
-        16,
-        23,
-        5,
-        11,
-        17,
-        24,
-        18,
-        25,
-        19,
-        26,
-        20,
-        27,
-        21,
-        28,
-    ],
-    dtype=np.int32,
-)
-
 REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
 REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
 REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
 
-# Reserved action-dict field used to forward the set of currently-pressed keyboard
-# keys from a KeyboardTeleop through the standard action pipeline to the SONIC
-# whole-body controller (see SonicWholeBodyController._process_keyboard).
-KEYBOARD_KEYS_FIELD = "keyboard.keys"
-
-# ── Dense whole-body joint reference (SONIC encode_mode 0, OpenHLM / pi0.5) ──────
-# A single 34-D whole-body command per tick, in the OpenHLM action layout:
-#   [L-arm(7), L-grip(1), R-arm(7), R-grip(1), L-leg(6), R-leg(6), waist(3),
-#    root roll/pitch + yaw-rate(3)]
-# Fed as flat scalars ``wb.0.pos .. wb.33.pos``. The ``.pos`` suffix makes these
-# behave like ordinary joint-position action features so ``lerobot-rollout`` routes
-# them straight from a 34-D VLA (OpenHLM / pi0.5) onto the robot.
-WB_ACTION_PREFIX = "wb."
-WB_ACTION_DIM = 34
-
-
-def wb_action_key(i: int) -> str:
-    """Action-dict key for the ``i``-th whole-body command scalar (``wb.{i}.pos``)."""
-    return f"{WB_ACTION_PREFIX}{i}.pos"
-
 
 def default_remote_input() -> dict[str, float]:
     """Return a zeroed-out remote input dict (axes + buttons)."""
@@ -155,92 +63,13 @@ class G1_29_JointArmIndex(IntEnum):
     kRightWristYaw = 28
 
 
-def lowstate_to_obs(lowstate) -> dict:
-    """Build a robot observation dict from a Unitree lowstate.
-
-    Shared by ``UnitreeG1.get_observation`` and the SONIC pipeline so the
-    lowstate -> obs mapping lives in exactly one place. Keys match the
-    ``<joint>.q``/``imu.*`` schema consumed across the controllers.
-    """
-    obs: dict = {}
-
-    for motor in G1_29_JointIndex:
-        idx = motor.value
-        obs[f"{motor.name}.q"] = lowstate.motor_state[idx].q
-        obs[f"{motor.name}.dq"] = lowstate.motor_state[idx].dq
-        obs[f"{motor.name}.tau"] = lowstate.motor_state[idx].tau_est
-
-    imu = lowstate.imu_state
-    if imu.gyroscope:
-        obs["imu.gyro.x"] = imu.gyroscope[0]
-        obs["imu.gyro.y"] = imu.gyroscope[1]
-        obs["imu.gyro.z"] = imu.gyroscope[2]
-    if imu.accelerometer:
-        obs["imu.accel.x"] = imu.accelerometer[0]
-        obs["imu.accel.y"] = imu.accelerometer[1]
-        obs["imu.accel.z"] = imu.accelerometer[2]
-    if imu.quaternion:
-        obs["imu.quat.w"] = imu.quaternion[0]
-        obs["imu.quat.x"] = imu.quaternion[1]
-        obs["imu.quat.y"] = imu.quaternion[2]
-        obs["imu.quat.z"] = imu.quaternion[3]
-    if imu.rpy:
-        obs["imu.rpy.roll"] = imu.rpy[0]
-        obs["imu.rpy.pitch"] = imu.rpy[1]
-        obs["imu.rpy.yaw"] = imu.rpy[2]
-
-    wr = getattr(lowstate, "wireless_remote", None)
-    if wr:
-        obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
-
-    return obs
-
-
-def obs_to_wb34_state(obs: dict) -> np.ndarray:
-    """Build the 34-D OpenHLM / pi0.5 proprio state from a G1 observation dict.
-
-    Mirrors the whole-body *action* layout so the policy sees state and action in
-    the same coordinates::
-
-        [L-arm(7), L-grip(1), R-arm(7), R-grip(1),
-         L-leg(6), R-leg(6), waist(3), root roll/pitch + yaw-rate(3)]
-
-    Joint positions come from the ``<joint>.q`` obs keys, which are already in
-    MuJoCo / Unitree-SDK order — the same body-part grouping OpenHLM uses
-    ([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29]) — so they are
-    regrouped directly (no IsaacLab permutation). The G1 has no grippers in its
-    29-DoF body, so both gripper slots are 0. Root roll/pitch are the IMU RPY and
-    the last slot is the IMU yaw rate (gyro z).
-    """
-    q_mj = np.array(
-        [float(obs.get(f"{m.name}.q", 0.0)) for m in G1_29_JointIndex],
-        dtype=np.float32,
-    )
-    lleg, rleg, waist = q_mj[0:6], q_mj[6:12], q_mj[12:15]
-    larm, rarm = q_mj[15:22], q_mj[22:29]
-
-    state = np.zeros(34, dtype=np.float32)
-    state[0:7] = larm
-    # state[7] left gripper — none on 29-DoF G1
-    state[8:15] = rarm
-    # state[15] right gripper — none on 29-DoF G1
-    state[16:22] = lleg
-    state[22:28] = rleg
-    state[28:31] = waist
-    state[31] = float(obs.get("imu.rpy.roll", 0.0))
-    state[32] = float(obs.get("imu.rpy.pitch", 0.0))
-    state[33] = float(obs.get("imu.gyro.z", 0.0))
-    return state
-
-
 def make_locomotion_controller(name: str | None):
     """Instantiate a locomotion controller by class name. Returns None if name is None."""
     if name is None:
         return None
     controllers = {
-        "GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
-        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
-        "SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
+        "GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
+        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
     }
     module_path = controllers.get(name)
     if module_path is None:

src/lerobot/robots/unitree_g1/gr00t_locomotion.py

@@ -14,29 +14,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import logging
 from collections import deque
-from typing import TYPE_CHECKING
 
 import numpy as np
+import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from lerobot.utils.import_utils import _onnxruntime_available, require_package
-
-from ..g1_utils import (
+from .g1_utils import (
     REMOTE_AXES,
     REMOTE_BUTTONS,
     G1_29_JointIndex,
     get_gravity_orientation,
 )
 
-if TYPE_CHECKING or _onnxruntime_available:
-    import onnxruntime as ort
-else:
-    ort = None
-
 logger = logging.getLogger(__name__)
 
 
@@ -77,15 +68,9 @@ def load_groot_policies(
         filename="GR00T-WholeBodyControl-Walk.onnx",
     )
 
-    # Load ONNX policies with a capped thread pool. GR00T runs at 50 Hz in a
-    # background thread alongside the (torch) upper-body policy, IK and sim; letting
-    # ORT grab every core starves those and makes the whole rollout stutter. These
-    # are small MLPs, so 1 thread is both enough and lowest-latency.
-    from .sonic_pipeline import make_ort_session_options
-
-    so = make_ort_session_options(intra_op_num_threads=1, inter_op_num_threads=1)
-    policy_balance = ort.InferenceSession(balance_path, sess_options=so)
-    policy_walk = ort.InferenceSession(walk_path, sess_options=so)
+    # Load ONNX policies
+    policy_balance = ort.InferenceSession(balance_path)
+    policy_walk = ort.InferenceSession(walk_path)
 
     logger.info("GR00T policies loaded successfully")
 
@@ -98,7 +83,6 @@ class GrootLocomotionController:
     control_dt = CONTROL_DT  # Expose for unitree_g1.py
 
     def __init__(self):
-        require_package("onnxruntime", extra="unitree_g1")
         # Load policies
         self.policy_balance, self.policy_walk = load_groot_policies()
 
@@ -212,16 +196,6 @@ class GrootLocomotionController:
         # Transform action back to target joint positions
         target_dof_pos_15 = GROOT_DEFAULT_ANGLES[:15] + self.groot_action * ACTION_SCALE
 
-        # Waist override: an external upper-body IK can command the 3 waist joints
-        # (indices 12/13/14) via ``kWaist{Yaw,Roll,Pitch}.q`` in the action dict. When
-        # present, we substitute the balance policy's waist target so the torso tracks
-        # the IK while the policy keeps only the legs balanced. Single-publisher stays
-        # intact (this thread still owns joints 0-14).
-        for idx in (G1_29_JointIndex.kWaistYaw, G1_29_JointIndex.kWaistRoll, G1_29_JointIndex.kWaistPitch):
-            key = f"{idx.name}.q"
-            if key in action and action[key] is not None:
-                target_dof_pos_15[idx.value] = float(action[key])
-
         # Build action dict
         action_dict = {}
         for i in range(15):

src/lerobot/robots/unitree_g1/holosoma_locomotion.py

@@ -14,34 +14,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from __future__ import annotations
-
 import json
 import logging
-from typing import TYPE_CHECKING
 
 import numpy as np
+import onnx
+import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from lerobot.utils.import_utils import _onnx_available, _onnxruntime_available, require_package
-
-from ..g1_utils import (
+from .g1_utils import (
     REMOTE_AXES,
     G1_29_JointArmIndex,
     G1_29_JointIndex,
     get_gravity_orientation,
 )
 
-if TYPE_CHECKING or _onnxruntime_available:
-    import onnxruntime as ort
-else:
-    ort = None
-
-if TYPE_CHECKING or _onnx_available:
-    import onnx
-else:
-    onnx = None
-
 logger = logging.getLogger(__name__)
 
 DEFAULT_ANGLES = np.zeros(29, dtype=np.float32)
@@ -114,8 +101,6 @@ class HolosomaLocomotionController:
     control_dt = CONTROL_DT  # Expose for unitree_g1.py
 
     def __init__(self):
-        require_package("onnxruntime", extra="unitree_g1")
-        require_package("onnx", extra="unitree_g1")
         # Load policy and gains
         self.policy, self.kp, self.kd = load_policy()
 

src/lerobot/robots/unitree_g1/infer_sonic_g1_onboard.py

@@ -1,192 +0,0 @@
-#!/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.
-
-"""Laptop-side sender for the SONIC whole-body walk policy, onboard deployment.
-
-This is the counterpart to ``run_g1_onboard.py`` (which runs the SONIC decoder on the
-robot). The heavy VLA (``nepyope/sonic_walk``, a pi0.5 token policy) runs here on the
-laptop GPU; only the resulting 64-D latent token is shipped to the robot over ZMQ:
-
-    laptop:  camera frame (ZMQ from robot :5555) + previous token
-             -> pi0.5 -> next 64-D token
-             -> PUSH JSON {motion_token.i.pos: ...} to robot :6004
-    robot:   run_g1_onboard receives the token, SonicWholeBodyController decodes it
-             into whole-body joint commands against local DDS at full rate.
-
-The policy's ``observation.state`` is the token currently being executed, so we close
-the loop by feeding back the *last token we sent* (the decoder holds it until a new one
-arrives). This mirrors what ``lerobot-rollout`` does via the robot's token echo, but
-without a controller / DDS on the laptop.
-
-The policy is pi0.5 with chunk_size=50, so a full diffusion inference runs only about
-once every 50 ticks; ``select_action`` pops one queued token per tick in between.
-
-Run ``run_g1_onboard.py --controller SonicWholeBodyController --sonic-token-action
---cameras ...`` on the robot first, then this on the laptop:
-
-    python -m lerobot.robots.unitree_g1.infer_sonic_g1_onboard \
-        --policy-path nepyope/sonic_walk --robot-ip 192.168.123.164 \
-        --task "walk back and forth"
-"""
-
-import argparse
-import contextlib
-import json
-import logging
-import signal
-import time
-
-import numpy as np
-import torch
-
-from lerobot.cameras.zmq import ZMQCamera, ZMQCameraConfig
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.policies.factory import get_policy_class, make_pre_post_processors
-from lerobot.policies.utils import prepare_observation_for_inference
-from lerobot.robots.unitree_g1.controllers.sonic_whole_body import (
-    NEUTRAL_TOKEN,
-    TOKEN_DIM,
-    token_action_key,
-)
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", force=True)
-logger = logging.getLogger("sonic_sender")
-
-ACTION_PORT = 6004  # matches run_g1_onboard.py --action-port
-IMAGE_KEY = "observation.images.ego_view"  # pi05 sonic_walk VISUAL input
-STATE_KEY = "observation.state"
-
-
-def main() -> None:
-    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    p.add_argument("--policy-path", default="nepyope/sonic_walk", help="Policy repo id or local path")
-    p.add_argument("--robot-ip", default="192.168.123.164", help="Robot IP (camera + action ports)")
-    p.add_argument("--action-port", type=int, default=ACTION_PORT, help="Onboard ZMQ PULL port for actions")
-    p.add_argument("--camera-port", type=int, default=5555, help="Onboard ZMQ camera PUB port")
-    p.add_argument("--camera-name", default="head_camera", help="Camera name served by run_g1_onboard")
-    p.add_argument("--camera-width", type=int, default=640, help="Camera width")
-    p.add_argument("--camera-height", type=int, default=480, help="Camera height")
-    p.add_argument("--task", default="walk back and forth", help="Language prompt for the VLA")
-    p.add_argument("--fps", type=float, default=30.0, help="Token send rate (matches training inference)")
-    p.add_argument("--device", default="cuda", help="Torch device")
-    p.add_argument("--max-ticks", type=int, default=0, help="Stop after N ticks (0 = run forever)")
-    p.add_argument("--dry-run", action="store_true", help="Run inference but do not PUSH tokens to the robot")
-    args = p.parse_args()
-
-    device = torch.device(args.device)
-
-    # --- Policy + processors (normalization stats baked into the checkpoint) ---
-    logger.info("Loading policy from '%s'...", args.policy_path)
-    policy_cfg = PreTrainedConfig.from_pretrained(args.policy_path)
-    policy_cfg.pretrained_path = args.policy_path
-    policy = get_policy_class(policy_cfg.type).from_pretrained(args.policy_path, config=policy_cfg)
-    policy = policy.to(device)
-    policy.eval()
-    policy.reset()
-
-    preprocessor, postprocessor = make_pre_post_processors(
-        policy_cfg=policy_cfg,
-        pretrained_path=args.policy_path,
-        preprocessor_overrides={"device_processor": {"device": str(device)}},
-    )
-    logger.info("Policy loaded (type=%s, device=%s, chunk=%s)", policy_cfg.type, device,
-                getattr(policy_cfg, "chunk_size", "?"))
-
-    # --- Camera (ZMQ from the robot's onboard image server) ---
-    cam = ZMQCamera(
-        ZMQCameraConfig(
-            server_address=args.robot_ip,
-            port=args.camera_port,
-            camera_name=args.camera_name,
-            width=args.camera_width,
-            height=args.camera_height,
-            fps=int(args.fps),
-        )
-    )
-    logger.info("Connecting camera %s@%s:%d ...", args.camera_name, args.robot_ip, args.camera_port)
-    cam.connect()
-
-    # --- Action PUSH socket to the onboard controller ---
-    import zmq
-
-    ctx = zmq.Context.instance()
-    sock = ctx.socket(zmq.PUSH)
-    sock.setsockopt(zmq.SNDHWM, 2)
-    sock.setsockopt(zmq.LINGER, 0)
-    sock.connect(f"tcp://{args.robot_ip}:{args.action_port}")
-    logger.info("Sending tokens to tcp://%s:%d (dry_run=%s)", args.robot_ip, args.action_port, args.dry_run)
-
-    stop = {"flag": False}
-    signal.signal(signal.SIGINT, lambda *_: stop.__setitem__("flag", True))
-    signal.signal(signal.SIGTERM, lambda *_: stop.__setitem__("flag", True))
-
-    # observation.state = the token currently executing on the robot (last one we sent);
-    # start at the neutral token the decoder holds before the first send, so the very
-    # first inference sees the true executing token (not zeros).
-    prev_token = NEUTRAL_TOKEN.copy()
-    period = 1.0 / args.fps
-    n = 0
-    t_infer_total = 0.0
-    logger.info("Streaming tokens at %.0f Hz. Ctrl-C to stop.", args.fps)
-    try:
-        while not stop["flag"]:
-            t0 = time.time()
-            try:
-                frame = cam.read()  # HxWxC uint8 RGB
-            except Exception as e:  # noqa: BLE001
-                logger.warning("Camera read failed: %s", e)
-                time.sleep(period)
-                continue
-
-            raw_obs = {
-                IMAGE_KEY: np.ascontiguousarray(frame),
-                STATE_KEY: prev_token.copy(),
-            }
-            with torch.inference_mode():
-                obs = prepare_observation_for_inference(raw_obs, device, args.task, "unitree_g1")
-                obs = preprocessor(obs)
-                action = policy.select_action(obs)
-                action = postprocessor(action)
-            token = action.squeeze(0).to("cpu").numpy().astype(np.float32)
-            prev_token = token
-
-            if not args.dry_run:
-                msg = {token_action_key(i): float(token[i]) for i in range(TOKEN_DIM)}
-                with contextlib.suppress(zmq.Again):
-                    sock.send_string(json.dumps(msg), zmq.NOBLOCK)
-
-            n += 1
-            t_infer_total += time.time() - t0
-            if n % 30 == 0:
-                logger.info(
-                    "tick %d | avg %.1f ms/tick | token[:3]=%s",
-                    n, 1000.0 * t_infer_total / 30.0, np.round(token[:3], 3).tolist(),
-                )
-                t_infer_total = 0.0
-
-            if args.max_ticks and n >= args.max_ticks:
-                break
-            time.sleep(max(0.0, period - (time.time() - t0)))
-    finally:
-        logger.info("Stopping sender after %d ticks.", n)
-        with contextlib.suppress(Exception):
-            cam.disconnect()
-        with contextlib.suppress(Exception):
-            sock.close(linger=0)
-
-
-if __name__ == "__main__":
-    main()

src/lerobot/robots/unitree_g1/run_g1_onboard.py

@@ -1,254 +0,0 @@
-#!/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.
-
-"""Run the G1 locomotion / whole-body controller ONBOARD, driven by high-level actions
-from a laptop.
-
-The controller (GR00T / Holosoma / SONIC whole-body) runs on the robot itself against
-local DDS, at full control rate. The laptop ships only the resulting high-level action
-(arm joint targets + joystick axes + gripper flags, or a 64-D SONIC motion token) as
-JSON over ZMQ. This process applies each action via ``UnitreeG1.send_action`` while the
-onboard controller thread keeps the legs balanced / decodes the token.
-
-This is the real-deploy counterpart to running ``lerobot-rollout`` on the laptop with
-``--robot.is_simulation=false`` (the ZMQ *socket bridge*): there the 50 Hz lowcmd
-crosses the network; here only compact high-level actions do, and the control loop stays
-local to the robot. Pair with a laptop client that produces actions (exo teleop, or a
-policy such as ``nepyope/sonic_walk`` emitting ``motion_token.{i}.pos``).
-
-Besides receiving actions, this process publishes ``observation.state`` (29 joint ``.q``)
-on a ZMQ PUB port so a laptop policy client has proprioception.
-
-Safety: type ``e`` then Enter in this terminal to stop immediately (zero-torque + exit).
-Ctrl-C does the normal graceful shutdown (kp ramp).
-
-Examples (on the robot):
-
-    # GR00T locomotion, arm targets from the laptop:
-    python -m lerobot.robots.unitree_g1.run_g1_onboard --controller GrootLocomotionController
-
-    # SONIC whole-body walk policy: laptop ships 64-D tokens, decoder runs here:
-    python -m lerobot.robots.unitree_g1.run_g1_onboard \
-        --controller SonicWholeBodyController --sonic-token-action \
-        --cameras "head_camera:/dev/v4l/by-path/platform-3610000.usb-usb-0:2.1:1.3-video-index0:640x480"
-"""
-
-import argparse
-import contextlib
-import json
-import logging
-import os
-import signal
-import sys
-import threading
-import time
-
-import numpy as np
-import zmq
-
-from lerobot.cameras.zmq.image_server import ImageServer
-from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
-from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
-from lerobot.robots.unitree_g1.run_g1_server import Gripper, build_gripper, parse_camera_specs
-from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
-
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", force=True)
-logger = logging.getLogger("g1_onboard")
-
-ACTION_PORT = 6004
-STATE_PORT = 6005
-
-
-def main() -> None:
-    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    p.add_argument("--controller", default="GrootLocomotionController", help="Controller class name")
-    p.add_argument("--dds-interface", default=None, help="DDS network interface (default: SDK default)")
-    p.add_argument(
-        "--sim",
-        action="store_true",
-        help="Attach to a DDS MuJoCo sim: skip MotionSwitcher + physical remote, default dds-interface 'lo'.",
-    )
-    p.add_argument(
-        "--sonic-token-action",
-        action="store_true",
-        help="SONIC token interface: actions carry a 64-D motion_token.{i}.pos that the decoder consumes.",
-    )
-    p.add_argument("--action-port", type=int, default=ACTION_PORT, help="ZMQ PULL port for laptop actions")
-    p.add_argument("--state-port", type=int, default=STATE_PORT, help="ZMQ PUB port for observation.state")
-    p.add_argument("--state-fps", type=float, default=30.0, help="observation.state publish rate; <=0 disables")
-    p.add_argument("--gravity-compensation", action="store_true", help="Enable arm gravity compensation")
-    # Gripper control (Damiao over CAN).
-    p.add_argument("--grippers", action="store_true", help="Drive Damiao grippers from action L3/R3 flags")
-    p.add_argument("--gripper-port-left", default="can1", help="CAN interface for LEFT gripper")
-    p.add_argument("--gripper-port-right", default="can0", help="CAN interface for RIGHT gripper")
-    p.add_argument("--gripper-send-id", type=lambda x: int(x, 0), default=0x08, help="Motor send CAN id")
-    p.add_argument("--gripper-recv-id", type=lambda x: int(x, 0), default=0x18, help="Motor recv CAN id")
-    p.add_argument("--gripper-motor-type", default="dm4310", help="Damiao motor type")
-    p.add_argument("--gripper-open-deg", type=float, default=-65.0, help="Gripper OPEN position (deg)")
-    p.add_argument("--gripper-close-deg", type=float, default=0.0, help="Gripper CLOSE position (deg)")
-    p.add_argument("--gripper-kp", type=float, default=15.0, help="MIT position gain (stiffness)")
-    p.add_argument("--gripper-kd", type=float, default=0.5, help="MIT damping gain")
-    p.add_argument("--gripper-no-fd", dest="gripper_fd", action="store_false", help="Classic CAN (non-FD)")
-    p.set_defaults(gripper_fd=True)
-    # Optional camera streaming (ZMQ) so the laptop policy client / viewer can connect.
-    p.add_argument("--cameras", default=None, help="Camera spec 'name:device[:WxH[:FOURCC]]', comma-sep")
-    p.add_argument("--camera-fps", type=int, default=30, help="Camera FPS")
-    p.add_argument("--camera-port", type=int, default=5555, help="Camera ZMQ port")
-    p.add_argument("--camera-width", type=int, default=640, help="Default camera width")
-    p.add_argument("--camera-height", type=int, default=480, help="Default camera height")
-    args = p.parse_args()
-
-    dds_interface = args.dds_interface
-    if args.sim and dds_interface is None:
-        dds_interface = "lo"
-
-    cfg = UnitreeG1Config(
-        is_simulation=False,
-        onboard=True,
-        controller=args.controller,
-        dds_interface=dds_interface,
-        gravity_compensation=args.gravity_compensation,
-        release_motion_control=not args.sim,
-        physical_remote=not args.sim,
-        sonic_token_action=args.sonic_token_action,
-        cameras={},
-    )
-
-    # Optional camera server (background thread; independent of DDS/CAN).
-    camera_server = None
-    if args.cameras:
-        cameras = parse_camera_specs(args.cameras, args.camera_width, args.camera_height)
-        camera_server = ImageServer({"fps": args.camera_fps, "cameras": cameras}, port=args.camera_port)
-        threading.Thread(target=camera_server.run, daemon=True).start()
-        cam_summary = ", ".join(f"{name}(dev {c['device_id']})" for name, c in cameras.items())
-        logger.info("Camera server started on :%d: %s", args.camera_port, cam_summary)
-
-    robot = UnitreeG1(cfg)
-    logger.info("Connecting onboard robot (controller=%s, token=%s)...", args.controller, args.sonic_token_action)
-    robot.connect()
-    # Note: with --sonic-token-action the SonicWholeBodyController holds a neutral
-    # (all-zero) token until the first laptop token arrives, then holds the last token
-    # between ticks -- see SonicWholeBodyController.token_mode (set from config).
-
-    grippers: dict[str, Gripper] = {}
-    if args.grippers:
-        for side, port in (("L", args.gripper_port_left), ("R", args.gripper_port_right)):
-            grippers[side] = build_gripper(
-                side, port, args.gripper_send_id, args.gripper_recv_id, args.gripper_motor_type,
-                args.gripper_fd, args.gripper_open_deg, args.gripper_close_deg, args.gripper_kp, args.gripper_kd,
-            )
-        logger.info("Grippers enabled: L3 -> left, R3 -> right")
-
-    ctx = zmq.Context.instance()
-    sock = ctx.socket(zmq.PULL)
-    sock.setsockopt(zmq.CONFLATE, 1)  # only ever act on the freshest command
-    sock.setsockopt(zmq.RCVTIMEO, 200)  # keeps the loop responsive to the stop event
-    sock.bind(f"tcp://0.0.0.0:{args.action_port}")
-    logger.info("Onboard controller live. Waiting for laptop actions on :%d ...", args.action_port)
-    logger.info("Type 'e' then Enter to STOP immediately (or Ctrl-C for graceful shutdown).")
-
-    stop = threading.Event()
-    signal.signal(signal.SIGINT, lambda *_: stop.set())
-    signal.signal(signal.SIGTERM, lambda *_: stop.set())
-
-    def estop_listener() -> None:
-        for line in sys.stdin:
-            if line.strip().lower() == "e":
-                logger.warning("E-STOP ('e'): going passive NOW.")
-                try:
-                    robot._shutdown_event.set()  # stop the controller loop publishing
-                    time.sleep(0.05)
-                    robot._send_zero_torque()  # motors limp; nothing overwrites it now
-                except Exception as e:  # noqa: BLE001
-                    logger.warning("E-stop zero-torque failed: %s", e)
-                os._exit(0)  # immediate hard exit, no slow cleanup
-
-    threading.Thread(target=estop_listener, daemon=True).start()
-
-    # Proprioception feedback: publish observation.state (29 joint .q) so a laptop
-    # inference client can feed it to a policy. DDS stays local; only compact JSON
-    # state crosses the network. (For a token policy the laptop closes the loop on the
-    # token instead, but publishing joint state is harmless and useful for logging.)
-    state_sock = None
-    if args.state_fps > 0:
-        state_sock = ctx.socket(zmq.PUB)
-        state_sock.setsockopt(zmq.SNDHWM, 2)
-        state_sock.setsockopt(zmq.LINGER, 0)
-        state_sock.bind(f"tcp://0.0.0.0:{args.state_port}")
-        logger.info("Publishing observation.state on :%d at %.0f Hz", args.state_port, args.state_fps)
-
-        def publish_state() -> None:
-            period = 1.0 / args.state_fps
-            joint_names = [j.name for j in G1_29_JointIndex]
-            while not stop.is_set():
-                t0 = time.time()
-                obs = robot.get_observation()
-                if obs:
-                    state = {f"{name}.q": float(obs.get(f"{name}.q", 0.0)) for name in joint_names}
-                    with contextlib.suppress(zmq.Again):
-                        state_sock.send_json(state, zmq.NOBLOCK)
-                time.sleep(max(0.0, period - (time.time() - t0)))
-
-        threading.Thread(target=publish_state, daemon=True).start()
-    else:
-        logger.info("observation.state PUB disabled (--state-fps<=0)")
-
-    n = 0
-    try:
-        while not stop.is_set():
-            try:
-                payload = sock.recv()
-            except zmq.Again:
-                continue
-            except zmq.ContextTerminated:
-                break
-
-            try:
-                action = json.loads(payload.decode("utf-8"))
-            except (ValueError, UnicodeDecodeError) as e:
-                logger.warning("Dropping malformed action: %s", e)
-                continue
-
-            robot.send_action(action)
-
-            if grippers:
-                # L3 = remote.button.4 -> left, R3 = remote.button.0 -> right.
-                if "L" in grippers and "remote.button.4" in action:
-                    grippers["L"].apply(bool(action["remote.button.4"]))
-                if "R" in grippers and "remote.button.0" in action:
-                    grippers["R"].apply(bool(action["remote.button.0"]))
-
-            n += 1
-            if n % 60 == 0:
-                axes = {k: round(float(action.get(k, 0.0)), 3) for k in ("remote.lx", "remote.ly", "remote.rx", "remote.ry")}
-                logger.info("Applied %d actions | axes=%s", n, axes)
-    finally:
-        logger.info("Shutting down onboard controller...")
-        stop.set()
-        if state_sock is not None:
-            with contextlib.suppress(Exception):
-                state_sock.close(linger=0)
-        if camera_server is not None:
-            with contextlib.suppress(Exception):
-                camera_server.stop()
-        for g in grippers.values():
-            with contextlib.suppress(Exception):
-                g.bus.disconnect()
-        robot.disconnect()
-
-
-if __name__ == "__main__":
-    main()

src/lerobot/robots/unitree_g1/run_g1_server.py

@@ -28,11 +28,9 @@ import argparse
 import base64
 import contextlib
 import json
-import re
 import threading
 import time
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 import zmq
 from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
@@ -43,9 +41,6 @@ from unitree_sdk2py.utils.crc import CRC
 
 from lerobot.cameras.zmq.image_server import ImageServer
 
-if TYPE_CHECKING:
-    from lerobot.motors.damiao.damiao import DamiaoMotorsBus
-
 # DDS topic names follow Unitree SDK naming conventions
 # ruff: noqa: N816
 kTopicLowCommand_Debug = "rt/lowcmd"  # action to robot
@@ -56,105 +51,6 @@ LOWSTATE_PORT = 6001
 NUM_MOTORS = 35
 
 
-@dataclass
-class Gripper:
-    """A single Damiao gripper that only writes to CAN when the open/close state changes."""
-
-    name: str
-    bus: "DamiaoMotorsBus"
-    open_deg: float
-    close_deg: float
-    _last_cmd: str | None = None  # "open" | "close"
-
-    def apply(self, want_close: bool) -> None:
-        want = "close" if want_close else "open"
-        if want == self._last_cmd:
-            return
-        target = self.close_deg if want_close else self.open_deg
-        self.bus.write("Goal_Position", "gripper", target)
-        self._last_cmd = want
-        print(f"[gripper] {self.name} -> {want.upper()} ({target:.1f} deg)")
-
-
-def build_gripper(
-    name: str,
-    port: str,
-    send_id: int,
-    recv_id: int,
-    motor_type: str,
-    use_can_fd: bool,
-    open_deg: float,
-    close_deg: float,
-    kp: float,
-    kd: float,
-) -> Gripper:
-    from lerobot.motors.damiao.damiao import DamiaoMotorsBus
-    from lerobot.motors.motors_bus import Motor, MotorNormMode
-
-    motors = {
-        "gripper": Motor(
-            id=send_id,
-            model=motor_type,
-            norm_mode=MotorNormMode.DEGREES,
-            motor_type_str=motor_type,
-            recv_id=recv_id,
-        )
-    }
-    bus = DamiaoMotorsBus(port=port, motors=motors, use_can_fd=use_can_fd)
-    print(f"Connecting {name} gripper on {port} (fd={use_can_fd})...")
-    bus.connect(handshake=True)
-    bus.write("Kp", "gripper", kp)
-    bus.write("Kd", "gripper", kd)
-    bus.write("Goal_Position", "gripper", open_deg)  # start open
-    print(f"  {name}: connected, torque enabled, opened.")
-    return Gripper(name, bus, open_deg, close_deg, _last_cmd="open")
-
-
-def parse_camera_specs(spec: str, default_width: int, default_height: int) -> dict[str, dict]:
-    """Parse a multi-camera spec string into an ImageServer ``cameras`` dict.
-
-    Format: comma-separated ``name:device[:WxH[:FOURCC]]`` entries, e.g.
-    ``head_camera:6,left_wrist:0``. ``device`` may be an integer index or an explicit
-    device path (e.g. ``/dev/video6``), including stable ``by-path`` names like
-    ``/dev/v4l/by-path/platform-...:2.1:1.3-video-index0`` which survive USB
-    re-enumeration (unlike bare ``/dev/videoN`` indices). Because a by-path name
-    itself contains colons, the optional ``WxH`` and ``FOURCC`` are parsed from the
-    *right* so the device-path colons are preserved.
-    """
-    wh_re = re.compile(r"\d+x\d+", re.IGNORECASE)
-    fourcc_re = re.compile(r"[A-Za-z0-9]{4}")
-
-    cameras: dict[str, dict] = {}
-    for entry in spec.split(","):
-        entry = entry.strip()
-        if not entry:
-            continue
-        if ":" not in entry:
-            raise ValueError(f"Invalid camera spec '{entry}', expected 'name:device[:WxH[:FOURCC]]'")
-        name, rest = entry.split(":", 1)
-        name = name.strip()
-        tokens = [t.strip() for t in rest.split(":")]
-
-        fourcc = None
-        if len(tokens) >= 3 and wh_re.fullmatch(tokens[-2]) and fourcc_re.fullmatch(tokens[-1]):
-            fourcc = tokens.pop().upper()
-        width, height = default_width, default_height
-        if len(tokens) >= 2 and wh_re.fullmatch(tokens[-1]):
-            w, h = tokens.pop().lower().split("x")
-            width, height = int(w), int(h)
-
-        raw_id = ":".join(tokens).strip()
-        if not raw_id:
-            raise ValueError(f"Invalid camera spec '{entry}', missing device")
-        device_id: int | str = int(raw_id) if raw_id.lstrip("-").isdigit() else raw_id
-        if name in cameras:
-            raise ValueError(f"Duplicate camera name '{name}' in --cameras")
-        cameras[name] = {"device_id": device_id, "shape": [height, width], "fourcc": fourcc}
-    if not cameras:
-        raise ValueError("No cameras parsed from --cameras spec")
-    return cameras
-
-
 def lowstate_to_dict(msg: hg_LowState) -> dict[str, Any]:
     """Convert LowState SDK message to a JSON-serializable dictionary."""
     motor_states = []
@@ -259,11 +155,7 @@ def main() -> None:
     """Main entry point for the robot server bridge."""
     parser = argparse.ArgumentParser(description="DDS-to-ZMQ bridge server for Unitree G1")
     parser.add_argument("--camera", action="store_true", help="Also launch camera server")
-    parser.add_argument("--camera-device", default="4",
-                        help="Camera device: index or /dev/video path or by-path name (default: 4)")
-    parser.add_argument("--cameras", default=None,
-                        help="Multi-camera spec 'name:device[:WxH[:FOURCC]]', comma-separated. Overrides "
-                             "--camera-device; device may be a by-path name to survive USB re-enumeration.")
+    parser.add_argument("--camera-device", type=int, default=4, help="Camera device ID (default: 4)")
     parser.add_argument("--camera-fps", type=int, default=30, help="Camera FPS (default: 30)")
     parser.add_argument("--camera-width", type=int, default=640, help="Camera width (default: 640)")
     parser.add_argument("--camera-height", type=int, default=480, help="Camera height (default: 480)")
@@ -272,20 +164,20 @@ def main() -> None:
 
     # Optionally start camera server in background thread
     camera_thread = None
-    if args.camera or args.cameras:
-        if args.cameras:
-            cameras = parse_camera_specs(args.cameras, args.camera_width, args.camera_height)
-        else:
-            # Single camera; accept an int index or a device/by-path string.
-            dev = args.camera_device
-            dev = int(dev) if str(dev).lstrip("-").isdigit() else dev
-            cameras = {"head_camera": {"device_id": dev, "shape": [args.camera_height, args.camera_width]}}
-        camera_config = {"fps": args.camera_fps, "cameras": cameras}
+    if args.camera:
+        camera_config = {
+            "fps": args.camera_fps,
+            "cameras": {
+                "head_camera": {
+                    "device_id": args.camera_device,
+                    "shape": [args.camera_height, args.camera_width],
+                }
+            },
+        }
         camera_server = ImageServer(camera_config, port=args.camera_port)
         camera_thread = threading.Thread(target=camera_server.run, daemon=True)
         camera_thread.start()
-        cam_summary = ", ".join(f"{n}(dev {c['device_id']})" for n, c in cameras.items())
-        print(f"Camera server started on port {args.camera_port}: {cam_summary}")
+        print(f"Camera server started on port {args.camera_port} (device {args.camera_device})")
 
     # initialize DDS
     ChannelFactoryInitialize(0)

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -33,14 +33,12 @@ from ..robot import Robot
 from .config_unitree_g1 import UnitreeG1Config
 from .g1_kinematics import G1_29_ArmIK
 from .g1_utils import (
-    KEYBOARD_KEYS_FIELD,
     REMOTE_AXES,
+    REMOTE_KEYS,
     G1_29_JointArmIndex,
     G1_29_JointIndex,
     default_remote_input,
-    lowstate_to_obs,
     make_locomotion_controller,
-    obs_to_wb34_state,
 )
 
 if TYPE_CHECKING or _unitree_sdk_available:
@@ -49,12 +47,8 @@ if TYPE_CHECKING or _unitree_sdk_available:
         ChannelPublisher as _SDKChannelPublisher,
         ChannelSubscriber as _SDKChannelSubscriber,
     )
-    from unitree_sdk2py.idl.default import (
-        unitree_hg_msg_dds__HandCmd_ as hg_HandCmd_default,
-        unitree_hg_msg_dds__LowCmd_,
-    )
+    from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
     from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
-        HandCmd_ as hg_HandCmd,
         LowCmd_ as hg_LowCmd,
         LowState_ as hg_LowState,
     )
@@ -64,8 +58,6 @@ else:
     _SDKChannelPublisher = None
     _SDKChannelSubscriber = None
     unitree_hg_msg_dds__LowCmd_ = None
-    hg_HandCmd_default = None
-    hg_HandCmd = None
     hg_LowCmd = None
     hg_LowState = None
     CRC = None
@@ -87,14 +79,6 @@ class LocomotionController(Protocol):
 kTopicLowCommand_Debug = "rt/lowcmd"
 kTopicLowState = "rt/lowstate"
 
-# Wireless-remote button byte layout, mapped to the positional button indices the
-# locomotion controllers expect. Used in onboard mode to read the physical Unitree
-# remote from lowstate (mirrors the exo teleoperator's RemoteController).
-_REMOTE_BUTTON_MAP: list[str] = [
-    "RB", "LB", "start", "back", "RT", "LT", "", "",
-    "A", "B", "X", "Y", "up", "right", "down", "left",
-]
-
 
 @dataclass
 class MotorState:
@@ -138,10 +122,8 @@ class UnitreeG1(Robot):
         # Initialize cameras config (ZMQ-based) - actual connection in connect()
         self._cameras = make_cameras_from_configs(config.cameras)
 
-        # Import channel classes based on mode. Simulation and onboard both talk to a
-        # real (local) DDS via the Unitree SDK; only the laptop-side bridge client uses
-        # the ZMQ socket shim.
-        if config.is_simulation or config.onboard:
+        # Import channel classes based on mode
+        if config.is_simulation:
             self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize
             self._ChannelPublisher = _SDKChannelPublisher
             self._ChannelSubscriber = _SDKChannelSubscriber
@@ -169,100 +151,19 @@ class UnitreeG1(Robot):
         # Lower-body controller loaded dynamically
         self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
 
-        # Token-driven deploy: let a SONIC controller hold a neutral token until the
-        # first real one arrives, then hold the last token between control ticks.
-        if config.sonic_token_action and hasattr(self.controller, "token_mode"):
-            self.controller.token_mode = True
-
         # Controller thread state
         self._controller_thread = None
-        # When set, the controller loop stops publishing low commands so reset() can
-        # drive the joints directly without two publishers fighting (single-publisher).
-        self._controller_paused = threading.Event()
         self._controller_action_lock = threading.Lock()
         self.controller_input = default_remote_input()
         self.controller_output = {}
 
-        # Onboard-only: parser for the physical Unitree wireless remote (read straight
-        # from local lowstate so joystick locomotion works without a laptop round-trip).
-        self._joystick = None
-
-        # Replay-camera state: keep the encoded (raw) cells per camera and decode
-        # frames lazily as the play cursor advances, with a small frame cache, so we
-        # don't materialize gigabytes of decoded RGB at construction time.
-        self._replay_raw: dict[str, list] = {}
-        self._replay_cache: dict[tuple[str, int], np.ndarray] = {}
-        self._replay_cache_cap = 8
-        self._replay_len = 0
-        self._replay_idx = 0
-        if config.replay_camera_parquet and config.replay_camera_map:
-            self._load_replay_frames()
-
-        # Token-mode state: last 64-D SONIC latent token commanded by the policy,
-        # echoed back as ``observation.state`` so a token-output VLA closes the loop
-        # on its own previous token (see ``sonic_token_action``). Seeded to zeros;
-        # the controller's startup blend eases joints in regardless.
-        self._last_token: np.ndarray | None = None
-        if config.sonic_token_action:
-            from .controllers.sonic_whole_body import TOKEN_DIM
-
-            self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32)
-
-    def _load_replay_frames(self) -> None:
-        """Load only the mapped parquet columns (encoded frames); decode on demand."""
-        import pyarrow.parquet as pq
-
-        cols_needed = list(dict.fromkeys(self.config.replay_camera_map.values()))
-        table = pq.read_table(self.config.replay_camera_parquet, columns=cols_needed)
-        self._replay_len = table.num_rows
-        self._replay_raw = {
-            cam_name: table.column(column).to_pylist()
-            for cam_name, column in self.config.replay_camera_map.items()
-        }
-        logger.info(
-            "Loaded %d replay frames (lazy-decode) for cameras %s from %s",
-            self._replay_len,
-            list(self.config.replay_camera_map),
-            self.config.replay_camera_parquet,
-        )
-
-    def _decode_replay_cell(self, cell) -> np.ndarray:
-        import io
-
-        from PIL import Image
-
-        data = cell["bytes"] if isinstance(cell, dict) else cell
-        return np.asarray(Image.open(io.BytesIO(data)).convert("RGB"), dtype=np.uint8)
-
-    def _replay_frame(self, cam_name: str, idx: int) -> np.ndarray:
-        """Decode (and briefly cache) a single replay frame for a camera."""
-        key = (cam_name, idx)
-        cached = self._replay_cache.get(key)
-        if cached is not None:
-            return cached
-        frame = self._decode_replay_cell(self._replay_raw[cam_name][idx])
-        if len(self._replay_cache) >= self._replay_cache_cap:
-            self._replay_cache.pop(next(iter(self._replay_cache)))
-        self._replay_cache[key] = frame
-        return frame
-
     def _subscribe_lowstate(self):  # polls robot state @ 250Hz
         while not self._shutdown_event.is_set():
             start_time = time.time()
 
             # Step simulation if in simulation mode
             if self.config.is_simulation and self.sim_env is not None:
-                try:
-                    self.sim_env.step()
-                except ValueError as e:
-                    # Startup race: the sim thread can step once before reset() has
-                    # written a valid base pose, giving a zero-norm pelvis quaternion
-                    # (scipy>=1.11 raises instead of normalizing). Skip and retry so
-                    # the thread survives instead of dying and freezing the sim.
-                    if "zero norm" not in str(e).lower():
-                        raise
-                    time.sleep(self.control_dt)
-                    continue
+                self.sim_env.step()
 
             msg = self.lowstate_subscriber.Read()
             if msg is not None:
@@ -330,80 +231,15 @@ class UnitreeG1(Robot):
                 features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
         return features
 
-    @property
-    def _wb_state_ft(self) -> dict[str, type]:
-        """34-D whole-body proprio state (``wb_state.{i}.pos``) for dense controllers.
-
-        Exposed only when the controller consumes a dense whole-body command
-        (OpenHLM / pi0.5). These ``.pos`` scalars are aggregated by the rollout
-        pipeline into a single 34-D ``observation.state`` for the policy.
-        """
-        if self.config.sonic_token_action:
-            return {}
-        if not getattr(self.controller, "wb_action", False):
-            return {}
-        from .g1_utils import WB_ACTION_DIM
-
-        return {f"wb_state.{i}.pos": float for i in range(WB_ACTION_DIM)}
-
-    @property
-    def _token_state_ft(self) -> dict[str, type]:
-        """64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
-
-        Exposed only in ``sonic_token_action`` mode; aggregated by the rollout into a
-        64-D ``observation.state`` (the last token the policy commanded).
-        """
-        if not self.config.sonic_token_action:
-            return {}
-        from .controllers.sonic_whole_body import TOKEN_DIM, token_state_key
-
-        return {token_state_key(i): float for i in range(TOKEN_DIM)}
-
-    @property
-    def _empty_cameras_ft(self) -> dict[str, tuple]:
-        """Synthetic zero-image cameras (see ``UnitreeG1Config.empty_cameras``)."""
-        h, w = self.config.empty_camera_hw
-        return dict.fromkeys(self.config.empty_cameras, (h, w, 3))
-
-    @property
-    def _replay_cameras_ft(self) -> dict[str, tuple]:
-        """Replay cameras, shaped from their first (lazily decoded) frame."""
-        if not self._replay_len:
-            return {}
-        return {name: self._replay_frame(name, 0).shape for name in self._replay_raw}
-
     @cached_property
     def observation_features(self) -> dict[str, type | tuple]:
-        return {
-            **self._motors_ft,
-            **self._wb_state_ft,
-            **self._token_state_ft,
-            **self._empty_cameras_ft,
-            **self._replay_cameras_ft,
-            **self._cameras_ft,
-        }
+        return {**self._motors_ft, **self._cameras_ft}
 
     @cached_property
     def action_features(self) -> dict[str, type]:
         if self.controller is None:
             return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
 
-        # Token-output VLA (SONIC decoder): advertise a 64-D latent-token action space
-        # (``motion_token.{i}.pos``) so ``lerobot-rollout`` maps a 64-D policy output
-        # straight onto the decoder, bypassing the encoder.
-        if self.config.sonic_token_action:
-            from .controllers.sonic_whole_body import TOKEN_DIM, token_action_key
-
-            return {token_action_key(i): float for i in range(TOKEN_DIM)}
-
-        # Dense whole-body controllers (SONIC / OpenHLM, pi0.5) consume a single
-        # 34-D command per tick. Expose it as ``wb.{i}.pos`` joint-position features
-        # so ``lerobot-rollout`` maps a 34-D policy output straight onto the robot.
-        if getattr(self.controller, "wb_action", False):
-            from .g1_utils import WB_ACTION_DIM, wb_action_key
-
-            return {wb_action_key(i): float for i in range(WB_ACTION_DIM)}
-
         arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
         remote_features = dict.fromkeys(REMOTE_AXES, float)
         return {**arm_features, **remote_features}
@@ -419,11 +255,6 @@ class UnitreeG1(Robot):
         while not self._shutdown_event.is_set():
             start_time = time.time()
 
-            # Paused during reset() so the reset routine is the sole low-cmd publisher.
-            if self._controller_paused.is_set():
-                time.sleep(control_dt)
-                continue
-
             with self._lowstate_lock:
                 lowstate = self._lowstate
 
@@ -440,13 +271,6 @@ class UnitreeG1(Robot):
                 with self._controller_action_lock:
                     controller_input = dict(self.controller_input)
 
-                # Onboard: the physical Unitree remote (in local lowstate) takes
-                # priority for locomotion when active; otherwise laptop/ZMQ axes stand.
-                if self.config.onboard:
-                    wl = self._wireless_remote_input(lowstate)
-                    if wl is not None:
-                        controller_input.update(wl)
-
                 # Run controller step
                 controller_action = self.controller.run_step(controller_input, lowstate)
 
@@ -469,105 +293,15 @@ class UnitreeG1(Robot):
     def configure(self) -> None:
         pass
 
-    def _wireless_remote_input(self, lowstate) -> dict | None:
-        """Parse the physical Unitree remote from lowstate into controller inputs.
-
-        Onboard only. Returns None when the remote is idle so the laptop-provided
-        (ZMQ) axes keep control; otherwise the physical remote takes priority.
-        """
-        js = self._joystick
-        if js is None:
-            return None
-        wr = getattr(lowstate, "wireless_remote", None)
-        if not wr or len(wr) < 24:
-            return None
-        try:
-            js.extract(wr)
-        except Exception:  # noqa: BLE001
-            return None
-
-        axes = {
-            "remote.lx": float(js.lx.data),
-            "remote.ly": float(js.ly.data),
-            "remote.rx": float(js.rx.data),
-            "remote.ry": float(js.ry.data),
-        }
-        active = any(abs(v) > 1e-2 for v in axes.values())
-        out = dict(axes)
-        for i, name in enumerate(_REMOTE_BUTTON_MAP):
-            if name:
-                val = float(getattr(js, name).data)
-                out[f"remote.button.{i}"] = val
-                if val:
-                    active = True
-        return out if active else None
-
-    def _release_motion_control(self) -> None:
-        """Release the robot's built-in motion services so we can send raw lowcmd.
-
-        Onboard-only. Mirrors run_g1_server.py: on the real robot the factory
-        locomotion/hand services must relinquish control before our controller can
-        write to ``rt/lowcmd``, otherwise commands are ignored or fought.
-        """
-        from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
-
-        msc = MotionSwitcherClient()
-        msc.SetTimeout(5.0)
-        msc.Init()
-        _, result = msc.CheckMode()
-        while result is not None and "name" in result and result["name"]:
-            logger.info("[UnitreeG1] Releasing built-in mode '%s'...", result["name"])
-            msc.ReleaseMode()
-            _, result = msc.CheckMode()
-            time.sleep(1.0)
-
     def connect(self, calibrate: bool = True) -> None:  # connect to DDS
         # Initialize DDS channel and simulation environment
         if self.config.is_simulation:
-            from lerobot.envs.utils import (
-                _download_hub_file,
-                _import_hub_module,
-                _normalize_hub_result,
-            )
+            from lerobot.envs import make_env
 
             self._ChannelFactoryInitialize(0, "lo")
-            # Call the hub env's make_env directly so we can disable the offscreen
-            # head_camera renderer. We drive image-conditioned policies from recorded
-            # frames (see replay_camera_parquet / external obs), never the sim's own
-            # camera, so building a MuJoCo offscreen GL context is pure liability: it
-            # crashes with "Failed to make the EGL context current" when GLFW/SDL
-            # already own a context, killing the sim thread and hanging on
-            # "Waiting for robot state...". publish_images=False -> no renderer.
-            repo_id, _, local_file, _ = _download_hub_file(
-                "lerobot/unitree-g1-mujoco", True, None
-            )
-            hub_mod = _import_hub_module(local_file, repo_id)
-            raw = hub_mod.make_env(n_envs=1, use_async_envs=False, publish_images=False, cameras=[])
-            self._env_wrapper = _normalize_hub_result(raw)
+            self._env_wrapper = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True)
             # Extract the actual gym env from the dict structure
             self.sim_env = self._env_wrapper["hub_env"][0].envs[0]
-        elif self.config.onboard:
-            # Real robot, controller running onboard against local DDS. Initialize the
-            # real SDK channel factory on the robot's DDS interface and take low-level
-            # control from the built-in services before we start writing lowcmd.
-            if self.config.dds_interface:
-                self._ChannelFactoryInitialize(0, self.config.dds_interface)
-            else:
-                self._ChannelFactoryInitialize(0)
-            # Real robot: hand low-level control over from the built-in services.
-            # A DDS sim has no MotionSwitcher, so this is skipped there.
-            if self.config.release_motion_control:
-                self._release_motion_control()
-            # Real robot: read the physical wireless remote from lowstate for
-            # locomotion. A sim has no physical remote, so leave _joystick=None and
-            # let send_action (ZMQ) drive the locomotion axes instead.
-            if self.config.physical_remote:
-                from unitree_sdk2py.utils.joystick import Joystick
-
-                self._joystick = Joystick()
-                for axis in (self._joystick.lx, self._joystick.ly, self._joystick.rx, self._joystick.ry):
-                    axis.smooth = 1.0
-                    axis.deadzone = 0.0
         else:
             self._ChannelFactoryInitialize(0, config=self.config)
 
@@ -577,17 +311,6 @@ class UnitreeG1(Robot):
         self.lowstate_subscriber = self._ChannelSubscriber(kTopicLowState, hg_LowState)
         self.lowstate_subscriber.Init()
 
-        # Dex3 hand command publishers (grasping). Driven by the OpenHLM grip scalars.
-        self._hand_publishers = {}
-        if self.config.publish_hands:
-            self._left_hand_cmd = hg_HandCmd_default()
-            self._right_hand_cmd = hg_HandCmd_default()
-            self._hand_publishers["left"] = self._ChannelPublisher("rt/dex3/left/cmd", hg_HandCmd)
-            self._hand_publishers["right"] = self._ChannelPublisher("rt/dex3/right/cmd", hg_HandCmd)
-            for pub in self._hand_publishers.values():
-                pub.Init()
-            logger.info("Dex3 hand command publishers initialized (rt/dex3/{left,right}/cmd)")
-
         # Start subscribe thread to read robot state
         self.subscribe_thread = threading.Thread(target=self._subscribe_lowstate)
         self.subscribe_thread.start()
@@ -620,9 +343,6 @@ class UnitreeG1(Robot):
 
         self.kp = np.array(self.config.kp, dtype=np.float32)
         self.kd = np.array(self.config.kd, dtype=np.float32)
-        if self.controller is not None and hasattr(self.controller, "kp"):
-            self.kp = np.array(self.controller.kp, dtype=np.float32)
-            self.kd = np.array(self.controller.kd, dtype=np.float32)
 
         for joint in G1_29_JointIndex:
             self.msg.motor_cmd[joint].mode = 1
@@ -630,16 +350,12 @@ class UnitreeG1(Robot):
             self.msg.motor_cmd[joint].kd = self.kd[joint.value]
             self.msg.motor_cmd[joint].q = lowstate.motor_state[joint.value].q
 
-        # Start controller thread if enabled. Skipped when run_controller_thread is
-        # False so a caller can step the controller synchronously (faithful replay).
-        if self.controller is not None and self.config.run_controller_thread:
+        # Start controller thread if enabled
+        if self.controller is not None:
             self._controller_thread = threading.Thread(target=self._controller_loop, daemon=True)
             self._controller_thread.start()
             fps = int(1.0 / self.controller.control_dt)
             logger.info(f"Controller thread started ({fps}Hz)")
-        elif self.controller is not None:
-            logger.info("Controller thread disabled (run_controller_thread=False); "
-                        "caller must drive controller.run_step synchronously.")
 
     def _send_zero_torque(self) -> None:
         """Send a zero-gain command to make joints passive before shutting down."""
@@ -655,59 +371,13 @@ class UnitreeG1(Robot):
         except Exception as e:
             logger.warning(f"Failed to send zero-torque on disconnect: {e}")
 
-    def _graceful_stop(self) -> None:
-        """Soft shutdown: hold the current pose and ramp joint stiffness (kp) to zero
-        over ``graceful_stop_s`` while keeping damping (kd), then go passive.
-
-        Prevents the robot from collapsing the instant control ends (a bare
-        zero-torque command is kp=kd=0 ≈ free-fall). Must run after the controller
-        loop has stopped so the two aren't publishing at once.
-        """
-        if self.config.graceful_stop_s <= 0:
-            self._send_zero_torque()
-            return
-        with self._lowstate_lock:
-            lowstate = self._lowstate
-        if lowstate is None:
-            self._send_zero_torque()
-            return
-        q_hold = {f"{motor.name}.q": lowstate.motor_state[motor.value].q for motor in G1_29_JointIndex}
-        kp = np.array(self.kp, dtype=np.float32)
-        kd = np.array(self.kd, dtype=np.float32)
-        zeros = np.zeros(29, dtype=np.float32)
-        dt = self.controller.control_dt if self.controller is not None else self.config.control_dt
-        steps = max(1, int(self.config.graceful_stop_s / dt))
-        logger.info("Graceful stop: damping down over %.1fs", self.config.graceful_stop_s)
-        for i in range(steps):
-            ratio = (i + 1) / steps
-            self.publish_lowcmd(q_hold, kp=kp * (1.0 - ratio), kd=kd, tau=zeros)
-            time.sleep(dt)
-        self._send_zero_torque()
-
     def disconnect(self):
-        # Stop the controller loop first so it isn't fighting the shutdown ramp.
+        # Put robot in passive mode before stopping threads
+        if not self.config.is_simulation:
+            self._send_zero_torque()
+
+        # Signal thread to stop and unblock any waits
         self._shutdown_event.set()
-        controller_stopped = True
-        if self._controller_thread is not None:
-            # Wait long enough for any in-flight inference tick to finish and the loop
-            # to observe the shutdown flag, so no stray low command is published while
-            # the ramp runs (the shutdown routine must be the single publisher).
-            self._controller_thread.join(timeout=5.0)
-            if self._controller_thread.is_alive():
-                controller_stopped = False
-                logger.error(
-                    "Controller thread did not stop; skipping graceful ramp to avoid "
-                    "concurrent low commands (fail-safe: joints keep last command until exit)"
-                )
-
-        # Soft, damped settle instead of an instant limp (real robot only; the
-        # subscribe thread is still alive here to supply the current pose). Only ramp
-        # once the controller thread has definitely exited.
-        if not self.config.is_simulation and controller_stopped:
-            self._graceful_stop()
-
-        if self.controller is not None and hasattr(self.controller, "shutdown"):
-            self.controller.shutdown()
 
         # Wait for subscribe thread to finish
         if self.subscribe_thread is not None:
@@ -715,6 +385,12 @@ class UnitreeG1(Robot):
             if self.subscribe_thread.is_alive():
                 logger.warning("Subscribe thread did not stop cleanly")
 
+        # Wait for controller thread to finish
+        if self._controller_thread is not None:
+            self._controller_thread.join(timeout=2.0)
+            if self._controller_thread.is_alive():
+                logger.warning("Controller thread did not stop cleanly")
+
         # Close simulation environment
         if self.config.is_simulation and self.sim_env is not None:
             try:
@@ -746,41 +422,44 @@ class UnitreeG1(Robot):
         if lowstate is None:
             return {}
 
-        # Motors + IMU + wireless remote (shared lowstate -> obs mapping)
-        obs = lowstate_to_obs(lowstate)
+        obs = {}
 
-        # Dense whole-body controllers (OpenHLM / pi0.5): expose the 34-D proprio
-        # state as ``wb_state.{i}.pos`` so the rollout aggregates it into
-        # ``observation.state`` for the policy.
-        if self.config.sonic_token_action:
-            # Token mode: echo the last commanded latent token as observation.state
-            # so a token-output VLA closes the loop on its own previous token.
-            from .controllers.sonic_whole_body import token_state_key
+        # Motors - q, dq, tau for all joints
+        for motor in G1_29_JointIndex:
+            name = motor.name
+            idx = motor.value
+            obs[f"{name}.q"] = lowstate.motor_state[idx].q
+            obs[f"{name}.dq"] = lowstate.motor_state[idx].dq
+            obs[f"{name}.tau"] = lowstate.motor_state[idx].tau_est
 
-            token = self._last_token if self._last_token is not None else []
-            for i, v in enumerate(token):
-                obs[token_state_key(i)] = float(v)
-        elif getattr(self.controller, "wb_action", False):
-            wb_state = obs_to_wb34_state(obs)
-            for i, v in enumerate(wb_state):
-                obs[f"wb_state.{i}.pos"] = float(v)
+        # IMU - gyroscope
+        if lowstate.imu_state.gyroscope:
+            obs["imu.gyro.x"] = lowstate.imu_state.gyroscope[0]
+            obs["imu.gyro.y"] = lowstate.imu_state.gyroscope[1]
+            obs["imu.gyro.z"] = lowstate.imu_state.gyroscope[2]
 
-        # Synthetic empty cameras: black frames so image-conditioned policies run
-        # before real cameras are wired.
-        if self.config.empty_cameras:
-            h, w = self.config.empty_camera_hw
-            black = np.zeros((h, w, 3), dtype=np.uint8)
-            for name in self.config.empty_cameras:
-                obs[name] = black
+        # IMU - accelerometer
+        if lowstate.imu_state.accelerometer:
+            obs["imu.accel.x"] = lowstate.imu_state.accelerometer[0]
+            obs["imu.accel.y"] = lowstate.imu_state.accelerometer[1]
+            obs["imu.accel.z"] = lowstate.imu_state.accelerometer[2]
 
-        # Replay cameras: serve the current recorded frame per camera, then advance.
-        if self._replay_len:
-            idx = self._replay_idx
-            if idx >= self._replay_len:
-                idx = self._replay_len - 1 if not self.config.replay_camera_loop else idx % self._replay_len
-            for name in self._replay_raw:
-                obs[name] = self._replay_frame(name, idx)
-            self._replay_idx += 1
+        # IMU - quaternion
+        if lowstate.imu_state.quaternion:
+            obs["imu.quat.w"] = lowstate.imu_state.quaternion[0]
+            obs["imu.quat.x"] = lowstate.imu_state.quaternion[1]
+            obs["imu.quat.y"] = lowstate.imu_state.quaternion[2]
+            obs["imu.quat.z"] = lowstate.imu_state.quaternion[3]
+
+        # IMU - rpy
+        if lowstate.imu_state.rpy:
+            obs["imu.rpy.roll"] = lowstate.imu_state.rpy[0]
+            obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1]
+            obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2]
+
+        # Wireless remote (raw bytes for teleoperator)
+        if lowstate.wireless_remote:
+            obs["wireless_remote"] = lowstate.wireless_remote
 
         # Cameras - read images from ZMQ cameras
         for cam_name, cam in self._cameras.items():
@@ -794,19 +473,9 @@ class UnitreeG1(Robot):
     def send_action(self, action: RobotAction) -> RobotAction:
         action_to_publish = action
         if self.controller is not None:
-            if self.config.sonic_token_action:
-                from .controllers.sonic_whole_body import _extract_token_from_action
-
-                token = _extract_token_from_action(action)
-                if token is not None:
-                    self._last_token = token
-            self._update_controller_action(action)
-            if self.config.publish_hands and getattr(self.controller, "wb_action", False):
-                self._publish_hand_cmds(action)
-            if getattr(self.controller, "full_body", False):
-                return action
             # Controller thread owns legs/waist. Here we only update joystick inputs
             # and publish arm targets from the teleoperator.
+            self._update_controller_action(action)
             arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
             action_to_publish = {
                 key: value
@@ -834,67 +503,11 @@ class UnitreeG1(Robot):
         return action
 
     def _update_controller_action(self, action: RobotAction) -> None:
-        """Update controller input state from an incoming teleop action.
-
-        Controller-agnostic: every value-carrying key is forwarded verbatim into
-        ``controller_input`` (whole-body ``wb.{i}.pos`` from a 34-D VLA, or whatever a
-        future controller expects), and each controller extracts only the keys it
-        understands. The robot deliberately does not enumerate any controller's key
-        schema here.
-
-        KeyboardTeleop is the one special case: it emits the currently-pressed keys as
-        bare action keys with a ``None`` value (``dict.fromkeys(pressed, None)``), so
-        those are collected into a single held-key set under ``KEYBOARD_KEYS_FIELD``,
-        rebuilt each tick so releases clear. Special keys arrive as pynput objects and
-        are normalised to their name ("space", ...).
-        """
+        """Update controller input state from incoming teleop action."""
         with self._controller_action_lock:
-            self.controller_input[KEYBOARD_KEYS_FIELD] = {
-                (k if isinstance(k, str) else getattr(k, "name", str(k)))
-                for k, value in action.items()
-                if value is None
-            }
-            for key, value in action.items():
-                if isinstance(key, str) and value is not None:
-                    self.controller_input[key] = value
-
-    def _publish_hand_cmds(self, action: RobotAction) -> None:
-        """Drive the Dex3 hands from the OpenHLM grip scalars in a 34-D wb action.
-
-        ``wb.7.pos`` is the left grip and ``wb.15.pos`` the right grip. Each scalar in
-        [0, 1] (``hand_open_grip_value`` == fully open) is turned into a curl amount and
-        scaled onto ``hand_closed_pose`` (7 joints), then published as a PD target on
-        ``rt/dex3/{left,right}/cmd`` so the fingers close when the policy grips.
-        """
-        if not self._hand_publishers:
-            return
-        from .g1_utils import wb_action_key
-
-        open_val = float(self.config.hand_open_grip_value)
-        closed_val = float(self.config.hand_closed_grip_value)
-        closed_pose = self.config.hand_closed_pose
-        kp, kd = float(self.config.hand_kp), float(self.config.hand_kd)
-        span = (closed_val - open_val) or 1.0
-
-        def curl_amount(grip: float) -> float:
-            # Fraction of the way from the open scalar to the closed scalar, in [0, 1].
-            return float(min(max((grip - open_val) / span, 0.0), 1.0))
-
-        for side, grip_idx, cmd in (
-            ("left", 7, self._left_hand_cmd),
-            ("right", 15, self._right_hand_cmd),
-        ):
-            grip = action.get(wb_action_key(grip_idx))
-            if grip is None:
-                continue
-            amount = curl_amount(float(grip))
-            for i, closed_q in enumerate(closed_pose):
-                cmd.motor_cmd[i].q = float(closed_q) * amount
-                cmd.motor_cmd[i].dq = 0.0
-                cmd.motor_cmd[i].kp = kp
-                cmd.motor_cmd[i].kd = kd
-                cmd.motor_cmd[i].tau = 0.0
-            self._hand_publishers[side].Write(cmd)
+            for key in REMOTE_KEYS:
+                if key in action:
+                    self.controller_input[key] = action[key]
 
     @property
     def is_calibrated(self) -> bool:
@@ -924,64 +537,43 @@ class UnitreeG1(Robot):
         if default_positions is None:
             default_positions = np.array(self.config.default_positions, dtype=np.float32)
 
-        # Full-body controllers (SONIC / OpenHLM) own the whole 29-DoF command and
-        # ignore ``<joint>.q`` in send_action(), so reset() must publish the default
-        # pose directly. Pause the background controller first so the two aren't both
-        # writing low commands while the robot moves to the default pose.
-        full_body = getattr(self.controller, "full_body", False)
-        paused = False
-        if full_body and self._controller_thread is not None:
-            self._controller_paused.set()
-            paused = True
-            time.sleep(control_dt)  # let any in-flight controller tick settle
+        if self.config.is_simulation and self.sim_env is not None:
+            self.sim_env.reset()
+            self.publish_lowcmd(
+                {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
+            )
+        else:
+            total_time = 3.0
+            num_steps = int(total_time / control_dt)
 
-        try:
-            if self.config.is_simulation and self.sim_env is not None:
-                self.sim_env.reset()
-                self.publish_lowcmd(
-                    {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
-                )
-            else:
-                total_time = 3.0
-                num_steps = int(total_time / control_dt)
+            # get current state
+            obs = self.get_observation()
 
-                # get current state
-                obs = self.get_observation()
+            # record current positions
+            init_dof_pos = np.zeros(29, dtype=np.float32)
+            for motor in G1_29_JointIndex:
+                init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
 
-                # record current positions
-                init_dof_pos = np.zeros(29, dtype=np.float32)
+            # Interpolate to default position
+            for step in range(num_steps):
+                start_time = time.time()
+
+                alpha = step / num_steps
+                action_dict = {}
                 for motor in G1_29_JointIndex:
-                    init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
+                    target_pos = default_positions[motor.value]
+                    interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
+                    action_dict[f"{motor.name}.q"] = float(interp_pos)
 
-                # Interpolate to default position
-                for step in range(num_steps):
-                    start_time = time.time()
+                self.send_action(action_dict)
 
-                    alpha = step / num_steps
-                    action_dict = {}
-                    for motor in G1_29_JointIndex:
-                        target_pos = default_positions[motor.value]
-                        interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
-                        action_dict[f"{motor.name}.q"] = float(interp_pos)
+                # Maintain constant control rate
+                elapsed = time.time() - start_time
+                sleep_time = max(0, control_dt - elapsed)
+                time.sleep(sleep_time)
 
-                    # Full-body controllers no-op in send_action(); publish the pose
-                    # directly (arm-only controllers keep the send_action() path).
-                    if full_body:
-                        self.publish_lowcmd(action_dict)
-                    else:
-                        self.send_action(action_dict)
-
-                    # Maintain constant control rate
-                    elapsed = time.time() - start_time
-                    sleep_time = max(0, control_dt - elapsed)
-                    time.sleep(sleep_time)
-
-            # Reset controller internal state (gait phase, obs history, etc.) before
-            # resuming so its buffers reflect the post-reset pose.
-            if self.controller is not None and hasattr(self.controller, "reset"):
-                self.controller.reset()
-        finally:
-            if paused:
-                self._controller_paused.clear()
+        # Reset controller internal state (gait phase, obs history, etc.)
+        if self.controller is not None and hasattr(self.controller, "reset"):
+            self.controller.reset()
 
         logger.info("Reached default position")

src/lerobot/robots/unitree_go2/__init__.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+# 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.
@@ -14,15 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
+from .config_unitree_go2 import UnitreeGo2Config
+from .unitree_go2 import UnitreeGo2
 
-from .gr00t_locomotion import GrootLocomotionController
-from .holosoma_locomotion import HolosomaLocomotionController
-from .sonic_whole_body import SonicRuntime, SonicWholeBodyController
-
-__all__ = [
-    "GrootLocomotionController",
-    "HolosomaLocomotionController",
-    "SonicRuntime",
-    "SonicWholeBodyController",
-]
+__all__ = ["UnitreeGo2", "UnitreeGo2Config"]

src/lerobot/robots/unitree_go2/config_unitree_go2.py

@@ -0,0 +1,57 @@
+#!/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.
+
+from dataclasses import dataclass, field
+
+from lerobot.cameras import CameraConfig
+
+from ..config import RobotConfig
+
+
+@RobotConfig.register_subclass("unitree_go2")
+@dataclass
+class UnitreeGo2Config(RobotConfig):
+    """Configuration for the Unitree Go2 quadruped (EDU).
+
+    The host machine talks DDS directly to the dog over Ethernet/WiFi via
+    ``unitree_sdk2py`` — no onboard companion computer is required. Actions
+    are high-level sport-mode body velocities; observations are sport-mode
+    odometry plus the dog's built-in front camera.
+    """
+
+    # Network interface on the host that is wired/bridged to the Go2
+    # (the dog lives on 192.168.123.x when connected over Ethernet).
+    network_interface: str = "eth0"
+
+    # DDS domain id (0 for a stock Go2).
+    domain_id: int = 0
+
+    # Safety clamps applied in send_action() before commands reach the dog.
+    # The Go2 accepts far more (vx up to ~3.7 m/s) — keep indoor-sane defaults.
+    max_x_vel: float = 1.0  # m/s, body forward
+    max_y_vel: float = 0.5  # m/s, body left
+    max_theta_vel: float = 1.5  # rad/s, CCW about z-up
+
+    # Built-in front camera, served through the SDK VideoClient.
+    use_front_camera: bool = True
+    front_camera_width: int = 1280
+    front_camera_height: int = 720
+
+    # Send BalanceStand once on connect so the dog is ready to walk.
+    stand_on_connect: bool = True
+
+    # Additional external cameras (standard LeRobot camera configs).
+    cameras: dict[str, CameraConfig] = field(default_factory=dict)

src/lerobot/robots/unitree_go2/unitree_go2.py

@@ -0,0 +1,260 @@
+#!/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.
+
+"""Unitree Go2 quadruped (EDU) — high-level sport-mode integration.
+
+Unlike :class:`~lerobot.robots.unitree_g1.UnitreeG1` (low-level joint
+control at 250 Hz through an on-robot ZMQ bridge), the Go2 is driven with
+sport-mode **body velocity commands** at tens of Hz, which work fine over
+plain DDS from any Linux host on the dog's network — no bridge server, no
+onboard companion computer.
+
+Setup:
+  1. Connect the host to the Go2 via Ethernet (dog is on 192.168.123.x)
+     or put both on the same WiFi network.
+  2. ``pip install unitree_sdk2py`` (Linux only — rides on cyclonedds).
+  3. Find your interface name (``ip link``), then e.g.::
+
+       lerobot-teleoperate \
+           --robot.type=unitree_go2 \
+           --robot.network_interface=enp2s0 \
+           --teleop.type=gamepad
+
+Actions are body-frame velocities ``x.vel`` (forward, m/s), ``y.vel``
+(left, m/s), ``theta.vel`` (CCW yaw, rad/s) — the exact arguments of the
+SDK's ``SportClient.Move``. Observations are planar sport-mode odometry
+(``*.pos`` pose + ``*.vel`` body velocities) and the built-in front
+camera, plus any extra configured cameras.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from functools import cached_property
+from typing import Any
+
+import cv2
+import numpy as np
+
+from lerobot.cameras import make_cameras_from_configs
+from lerobot.types import RobotAction, RobotObservation
+from lerobot.utils.import_utils import require_package
+
+from ..robot import Robot
+from .config_unitree_go2 import UnitreeGo2Config
+
+logger = logging.getLogger(__name__)
+
+# DDS topic names follow Unitree SDK naming conventions
+SPORT_MODE_STATE_TOPIC = "rt/sportmodestate"
+
+
+class UnitreeGo2(Robot):
+    """LeRobot interface to a Unitree Go2 over unitree_sdk2py sport mode."""
+
+    config_class = UnitreeGo2Config
+    name = "unitree_go2"
+
+    def __init__(self, config: UnitreeGo2Config):
+        super().__init__(config)
+        self.config = config
+
+        self._cameras = make_cameras_from_configs(config.cameras)
+
+        # SDK handles — populated in connect(); the SDK import lives there
+        # too so that configs, features and tests work on SDK-less hosts.
+        self._sport = None
+        self._video = None
+        self._state_subscriber = None
+
+        self._state_lock = threading.Lock()
+        self._latest_state = None  # last SportModeState_ message
+        self._connected = False
+
+    # ------------------------------------------------------------------ #
+    # Features
+    # ------------------------------------------------------------------ #
+
+    @cached_property
+    def _odom_ft(self) -> dict[str, type]:
+        return {
+            "x.pos": float,
+            "y.pos": float,
+            "theta.pos": float,
+            "x.vel": float,
+            "y.vel": float,
+            "theta.vel": float,
+        }
+
+    @property
+    def _cameras_ft(self) -> dict[str, tuple]:
+        ft: dict[str, tuple] = {}
+        if self.config.use_front_camera:
+            ft["front"] = (self.config.front_camera_height, self.config.front_camera_width, 3)
+        for name, cam in self._cameras.items():
+            ft[name] = (cam.height, cam.width, 3)
+        return ft
+
+    @property
+    def observation_features(self) -> dict:
+        return {**self._odom_ft, **self._cameras_ft}
+
+    @property
+    def action_features(self) -> dict:
+        return {"x.vel": float, "y.vel": float, "theta.vel": float}
+
+    # ------------------------------------------------------------------ #
+    # Lifecycle
+    # ------------------------------------------------------------------ #
+
+    @property
+    def is_connected(self) -> bool:
+        return self._connected
+
+    def connect(self, calibrate: bool = True) -> None:
+        if self._connected:
+            return
+        require_package("unitree-sdk2py", extra="unitree_go2", import_name="unitree_sdk2py")
+
+        from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelSubscriber
+        from unitree_sdk2py.go2.sport.sport_client import SportClient
+        from unitree_sdk2py.idl.unitree_go.msg.dds_ import SportModeState_
+
+        ChannelFactoryInitialize(self.config.domain_id, self.config.network_interface)
+
+        sport = SportClient()
+        sport.SetTimeout(5.0)
+        sport.Init()
+        self._sport = sport
+
+        subscriber = ChannelSubscriber(SPORT_MODE_STATE_TOPIC, SportModeState_)
+        subscriber.Init(self._on_sport_state, 10)
+        self._state_subscriber = subscriber
+
+        if self.config.use_front_camera:
+            from unitree_sdk2py.go2.video.video_client import VideoClient
+
+            video = VideoClient()
+            video.SetTimeout(3.0)
+            video.Init()
+            self._video = video
+
+        for cam in self._cameras.values():
+            cam.connect()
+
+        if self.config.stand_on_connect:
+            self._sport.BalanceStand()
+
+        self._connected = True
+        self.configure()
+        logger.info(
+            "%s connected (iface=%s, domain=%d)",
+            self,
+            self.config.network_interface,
+            self.config.domain_id,
+        )
+
+    def disconnect(self) -> None:
+        if self._sport is not None:
+            try:
+                self._sport.StopMove()
+            except Exception:
+                logger.exception("StopMove on disconnect failed")
+        for cam in self._cameras.values():
+            try:
+                cam.disconnect()
+            except Exception:
+                logger.exception("camera disconnect failed")
+        self._sport = None
+        self._video = None
+        self._state_subscriber = None
+        self._connected = False
+
+    # Sport mode needs no calibration.
+    @property
+    def is_calibrated(self) -> bool:
+        return True
+
+    def calibrate(self) -> None:
+        pass
+
+    def configure(self) -> None:
+        pass
+
+    # ------------------------------------------------------------------ #
+    # I/O
+    # ------------------------------------------------------------------ #
+
+    def get_observation(self) -> RobotObservation:
+        if not self._connected:
+            raise ConnectionError(f"{self} is not connected.")
+
+        obs: dict[str, Any] = dict.fromkeys(self._odom_ft, 0.0)
+        with self._state_lock:
+            state = self._latest_state
+        if state is not None:
+            obs["x.pos"] = float(state.position[0])
+            obs["y.pos"] = float(state.position[1])
+            obs["theta.pos"] = float(state.imu_state.rpy[2])
+            obs["x.vel"] = float(state.velocity[0])
+            obs["y.vel"] = float(state.velocity[1])
+            obs["theta.vel"] = float(state.yaw_speed)
+
+        if self.config.use_front_camera:
+            obs["front"] = self._read_front_camera()
+
+        for name, cam in self._cameras.items():
+            obs[name] = cam.async_read()
+
+        return obs
+
+    def send_action(self, action: RobotAction) -> RobotAction:
+        if not self._connected:
+            raise ConnectionError(f"{self} is not connected.")
+
+        vx = float(np.clip(action.get("x.vel", 0.0), -self.config.max_x_vel, self.config.max_x_vel))
+        vy = float(np.clip(action.get("y.vel", 0.0), -self.config.max_y_vel, self.config.max_y_vel))
+        vyaw = float(
+            np.clip(action.get("theta.vel", 0.0), -self.config.max_theta_vel, self.config.max_theta_vel)
+        )
+
+        self._sport.Move(vx, vy, vyaw)
+        return {"x.vel": vx, "y.vel": vy, "theta.vel": vyaw}
+
+    # ------------------------------------------------------------------ #
+    # Internals
+    # ------------------------------------------------------------------ #
+
+    def _on_sport_state(self, msg) -> None:
+        with self._state_lock:
+            self._latest_state = msg
+
+    def _read_front_camera(self) -> np.ndarray:
+        """Fetch one frame from the built-in front camera (RGB, HxWx3)."""
+        h, w = self.config.front_camera_height, self.config.front_camera_width
+        code, data = self._video.GetImageSample()
+        if code != 0 or data is None:
+            logger.warning("front camera GetImageSample failed (code=%s)", code)
+            return np.zeros((h, w, 3), dtype=np.uint8)
+        frame = cv2.imdecode(np.frombuffer(bytes(data), dtype=np.uint8), cv2.IMREAD_COLOR)
+        if frame is None:
+            logger.warning("front camera frame failed to decode")
+            return np.zeros((h, w, 3), dtype=np.uint8)
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        if frame.shape[:2] != (h, w):
+            frame = cv2.resize(frame, (w, h), interpolation=cv2.INTER_AREA)
+        return frame

src/lerobot/utils/import_utils.py

@@ -60,18 +60,8 @@ def is_package_available(
                     # If the package can't be imported, it's not available
                     package_exists = False
             else:
-                # The distribution may be published under a name that differs from the
-                # import name (e.g. ``onnxruntime`` imports from ``onnxruntime-gpu`` /
-                # ``onnxruntime-silicon``). Resolve the import name to its actual
-                # distribution(s) and read the version from there before giving up.
-                try:
-                    dists = importlib.metadata.packages_distributions().get(import_name, [])
-                    if dists:
-                        package_version = importlib.metadata.version(dists[0])
-                    else:
-                        package_exists = False
-                except importlib.metadata.PackageNotFoundError:
-                    package_exists = False
+                # For packages other than "torch", don't attempt the fallback and set as not available
+                package_exists = False
         logging.debug(f"Detected {pkg_name} version: {package_version}")
     if return_version:
         return package_exists, package_version
@@ -133,8 +123,6 @@ _pyrealsense2_available = is_package_available("pyrealsense2") or is_package_ava
     "pyrealsense2-macosx", import_name="pyrealsense2"
 )
 _zmq_available = is_package_available("pyzmq", import_name="zmq")
-_onnxruntime_available = is_package_available("onnxruntime")
-_onnx_available = is_package_available("onnx")
 _hebi_available = is_package_available("hebi-py", import_name="hebi")
 _teleop_available = is_package_available("teleop")
 _placo_available = is_package_available("placo")

tests/navigation/test_agent.py

@@ -0,0 +1,248 @@
+#!/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.
+
+"""Unit tests for the C1 deterministic agent + language parser."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from lerobot.navigation.agent import (
+    AgentConfig,
+    DeterministicAgent,
+    HardcodedTaskParser,
+    Task,
+)
+from lerobot.navigation.skills import ExploreResult, GotoResult, LocateResult
+
+# ----- fakes -------------------------------------------------------------
+
+
+@dataclass
+class FakeSkills:
+    """Programmable :class:`SpatialSkills` stand-in. Each method consults a
+    pre-recorded script and bumps a call counter so tests can assert the
+    deterministic policy made the right sequence of calls."""
+
+    locate_script: list[LocateResult] = field(default_factory=list)
+    explore_script: list[ExploreResult] = field(default_factory=list)
+    goto_script: list[GotoResult] = field(default_factory=list)
+
+    locate_calls: list[str] = field(default_factory=list)
+    goto_calls: list[tuple[float, float, float]] = field(default_factory=list)
+    explore_calls: list[str | None] = field(default_factory=list)
+
+    def locate(self, text: str) -> LocateResult:
+        self.locate_calls.append(text)
+        if not self.locate_script:
+            return LocateResult(False, None, -1.0, 0, text)
+        return self.locate_script.pop(0)
+
+    def explore(self, query: str | None = None) -> ExploreResult:
+        self.explore_calls.append(query)
+        if not self.explore_script:
+            return ExploreResult(None, False, 0.0, "no frontier")
+        return self.explore_script.pop(0)
+
+    def goto(self, xyz: tuple[float, float, float], **_: object) -> GotoResult:
+        self.goto_calls.append(xyz)
+        if not self.goto_script:
+            return GotoResult(True, xyz, 0.0, 0, "ok", [])
+        return self.goto_script.pop(0)
+
+    @property
+    def base(self):
+        # Minimal stub: agent's teleport branch isn't exercised by these tests.
+        class _Base:
+            def move(self, *a, **k):
+                pass
+
+            def pose(self):
+                return np.eye(4)
+
+        return _Base()
+
+
+# ----- HardcodedTaskParser ------------------------------------------------
+
+
+def test_parser_simple_go_to():
+    t = HardcodedTaskParser().parse("go to the mug")
+    assert t.targets == ["mug"]
+
+
+def test_parser_strips_punctuation_and_articles():
+    t = HardcodedTaskParser().parse("Find the red lamp.")
+    assert t.targets == ["red lamp"]
+
+
+def test_parser_multi_step():
+    t = HardcodedTaskParser().parse("go to the mug then the chair")
+    assert t.targets == ["mug", "chair"]
+
+
+def test_parser_no_verb_treats_command_as_target():
+    """``parser.parse('couch')`` should still produce a usable Task."""
+    t = HardcodedTaskParser().parse("couch")
+    assert t.targets == ["couch"]
+
+
+def test_parser_empty_string_returns_empty_task():
+    t = HardcodedTaskParser().parse("   ")
+    assert t.targets == []
+
+
+def test_parser_split_by_comma():
+    t = HardcodedTaskParser().parse("go to mug, chair")
+    assert t.targets == ["mug", "chair"]
+
+
+# ----- DeterministicAgent policy -----------------------------------------
+
+
+def _ok_locate(xyz=(1.0, 0.0, 1.0), conf=0.9) -> LocateResult:
+    return LocateResult(True, xyz, conf, 10, "x")
+
+
+def _miss_locate(conf=0.05) -> LocateResult:
+    return LocateResult(False, None, conf, 0, "x")
+
+
+def _ok_goto(xyz=(1.0, 0.0, 1.0)) -> GotoResult:
+    return GotoResult(True, xyz, 0.0, 5, "ok", [])
+
+
+def _failed_goto(xyz=(1.0, 0.0, 1.0)) -> GotoResult:
+    return GotoResult(False, (0.0, 0.0, 0.0), 1.4, 0, "no path", [])
+
+
+def _explore_to(xyz=(2.0, 0.0, 2.0)) -> ExploreResult:
+    return ExploreResult(xyz, True, 2.8, "ok")
+
+
+def test_agent_hit_then_goto():
+    """Found on first call → no explore, single goto."""
+    skills = FakeSkills(
+        locate_script=[_ok_locate()],
+        goto_script=[_ok_goto()],
+    )
+    agent = DeterministicAgent(skills, AgentConfig(max_explore_iters=3))
+    res = agent.execute(Task(targets=["mug"]))
+    assert res.fully_successful
+    assert skills.locate_calls == ["mug"]
+    assert skills.goto_calls == [(1.0, 0.0, 1.0)]
+    assert skills.explore_calls == []
+    assert res.target_results[0].n_explore_iters == 0
+
+
+def test_agent_explore_then_relocate_then_goto():
+    """First locate misses → explore → goto-to-frontier → re-locate finds → final goto."""
+    skills = FakeSkills(
+        locate_script=[_miss_locate(), _ok_locate()],
+        explore_script=[_explore_to((3.0, 0.0, 0.0))],
+        goto_script=[_ok_goto((3.0, 0.0, 0.0)), _ok_goto((1.0, 0.0, 1.0))],
+    )
+    agent = DeterministicAgent(skills, AgentConfig(max_explore_iters=3))
+    res = agent.execute(Task(targets=["mug"]))
+    assert res.fully_successful
+    assert skills.locate_calls == ["mug", "mug"]
+    assert skills.explore_calls == ["mug"]
+    assert skills.goto_calls == [(3.0, 0.0, 0.0), (1.0, 0.0, 1.0)]
+    assert res.target_results[0].n_explore_iters == 1
+
+
+def test_agent_budget_exhaustion():
+    """All N+1 locate calls miss → return budget_exhausted."""
+    skills = FakeSkills(
+        locate_script=[_miss_locate() for _ in range(5)],
+        explore_script=[_explore_to() for _ in range(4)],
+        goto_script=[_ok_goto((2.0, 0.0, 2.0)) for _ in range(4)],
+    )
+    agent = DeterministicAgent(skills, AgentConfig(max_explore_iters=3))
+    res = agent.execute(Task(targets=["mug"]))
+    assert res.fully_successful is False
+    r = res.target_results[0]
+    assert r.reason == "budget_exhausted"
+    assert r.n_explore_iters == 3
+    # 4 locate calls: initial + 3 retries.
+    assert len(skills.locate_calls) == 4
+    assert len(skills.explore_calls) == 3
+
+
+def test_agent_no_frontier_short_circuits():
+    """If explore can't find a frontier, give up immediately — no point looping."""
+    skills = FakeSkills(
+        locate_script=[_miss_locate()],
+        explore_script=[ExploreResult(None, False, 0.0, "no frontier")],
+    )
+    agent = DeterministicAgent(skills, AgentConfig(max_explore_iters=3))
+    res = agent.execute(Task(targets=["mug"]))
+    r = res.target_results[0]
+    assert r.reached is False
+    assert r.reason == "no_frontier"
+    assert len(skills.locate_calls) == 1
+    assert len(skills.explore_calls) == 1
+
+
+def test_agent_failed_goto_does_not_loop_back():
+    """If locate finds the target but goto fails (e.g. no path), report the
+    failure cleanly rather than retrying."""
+    skills = FakeSkills(
+        locate_script=[_ok_locate()],
+        goto_script=[_failed_goto()],
+    )
+    agent = DeterministicAgent(skills)
+    res = agent.execute(Task(targets=["mug"]))
+    r = res.target_results[0]
+    assert r.reached is False
+    assert r.reason == "no path"
+
+
+def test_agent_multi_target_bails_on_first_failure():
+    """The spec says sequential targets stop at the first failure so the
+    caller sees the failure clearly."""
+    skills = FakeSkills(
+        locate_script=[_miss_locate()],
+        explore_script=[ExploreResult(None, False, 0.0, "no frontier")],
+    )
+    agent = DeterministicAgent(skills, AgentConfig(max_explore_iters=0))
+    res = agent.execute(Task(targets=["mug", "chair"]))
+    assert len(res.target_results) == 1  # bailed before chair
+    assert res.target_results[0].target == "mug"
+
+
+def test_agent_swap_parser_does_not_change_policy():
+    """Acceptance from the spec: 'swapping Qwen for a hardcoded target string
+    yields the same spatial behaviour'. Same skills script, same scripted
+    locate/goto, regardless of how the command was parsed."""
+    parser = HardcodedTaskParser()
+    for command in ("mug", "go to the mug", "find the mug"):
+        skills = FakeSkills(locate_script=[_ok_locate()], goto_script=[_ok_goto()])
+        agent = DeterministicAgent(skills)
+        res = agent.execute_command(command, parser)
+        assert res.fully_successful
+        assert skills.goto_calls == [(1.0, 0.0, 1.0)]
+
+
+def test_agent_empty_command_reports_parse_failure():
+    skills = FakeSkills()
+    agent = DeterministicAgent(skills)
+    res = agent.execute_command("", HardcodedTaskParser())
+    assert res.fully_successful is False
+    assert res.target_results[0].reason == "parse_empty"
+    assert skills.locate_calls == []

tests/navigation/test_base_controller.py

@@ -0,0 +1,320 @@
+#!/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.
+
+"""Tests for the navigation base controller.
+
+Hardware-free and SDK-free: the frame math is pure, the stub is
+kinematic, and the robot-backed controller is exercised through a fake
+Robot that records actions and serves canned odometry.
+"""
+
+from __future__ import annotations
+
+import math
+import time
+
+import numpy as np
+import pytest
+
+from lerobot.navigation.base_controller import (
+    BaseController,
+    RobotBaseController,
+    RobotBaseControllerConfig,
+    SafeBaseController,
+    StubBaseController,
+    odometry_to_world_pose,
+    world_velocity_to_body,
+)
+
+# ----- world_velocity_to_body ---------------------------------------------
+
+
+def test_forward_maps_to_body_x():
+    """heading=0, world +z (forward) → (vx>0, 0, 0)."""
+    vx_f, vy_l, vyaw = world_velocity_to_body(0.0, 0.3, 0.0, heading_rad=0.0)
+    assert vx_f == pytest.approx(0.3)
+    assert vy_l == pytest.approx(0.0)
+    assert vyaw == pytest.approx(0.0)
+
+
+def test_world_right_maps_to_negative_left():
+    """heading=0, world +x is the robot's RIGHT → negative y.vel."""
+    vx_f, vy_l, _ = world_velocity_to_body(0.3, 0.0, 0.0, heading_rad=0.0)
+    assert vx_f == pytest.approx(0.0)
+    assert vy_l == pytest.approx(-0.3)
+
+
+def test_world_x_is_forward_after_quarter_turn():
+    vx_f, vy_l, _ = world_velocity_to_body(0.3, 0.0, 0.0, heading_rad=math.pi / 2)
+    assert vx_f == pytest.approx(0.3)
+    assert vy_l == pytest.approx(0.0, abs=1e-9)
+
+
+def test_yaw_rate_sign_flips():
+    _, _, vyaw = world_velocity_to_body(0.0, 0.0, 0.5, heading_rad=0.0)
+    assert vyaw == pytest.approx(-0.5)
+
+
+def test_velocity_magnitude_preserved_under_rotation():
+    vx_f, vy_l, _ = world_velocity_to_body(0.3, 0.4, 0.0, heading_rad=1.234)
+    assert math.hypot(vx_f, vy_l) == pytest.approx(0.5)
+
+
+# ----- odometry_to_world_pose ---------------------------------------------
+
+
+def test_odometry_at_origin_is_identity():
+    pose, heading = odometry_to_world_pose(1.0, 2.0, 0.3, origin=(1.0, 2.0, 0.3))
+    np.testing.assert_allclose(pose, np.eye(4), atol=1e-12)
+    assert heading == pytest.approx(0.0)
+
+
+def test_odometry_forward_maps_to_world_z():
+    pose, heading = odometry_to_world_pose(1.0, 0.0, 0.0, origin=(0.0, 0.0, 0.0))
+    assert pose[0, 3] == pytest.approx(0.0)
+    assert pose[2, 3] == pytest.approx(1.0)
+    assert heading == pytest.approx(0.0)
+
+
+def test_odometry_left_maps_to_world_negative_x():
+    pose, _ = odometry_to_world_pose(0.0, 1.0, 0.0, origin=(0.0, 0.0, 0.0))
+    assert pose[0, 3] == pytest.approx(-1.0)
+    assert pose[2, 3] == pytest.approx(0.0)
+
+
+def test_odometry_yaw_sign_flip():
+    pose, heading = odometry_to_world_pose(0.0, 0.0, 0.5, origin=(0.0, 0.0, 0.0))
+    assert heading == pytest.approx(-0.5)
+    fwd = pose[:3, 2]
+    np.testing.assert_allclose(fwd, [math.sin(-0.5), 0.0, math.cos(-0.5)], atol=1e-12)
+
+
+def test_odometry_origin_yaw_is_derotated():
+    """Motion along the boot-time heading is always world +z, whatever
+    direction the robot faced when odometry started."""
+    origin = (0.0, 0.0, math.pi / 2)
+    pose, heading = odometry_to_world_pose(0.0, 1.0, math.pi / 2, origin=origin)
+    assert pose[0, 3] == pytest.approx(0.0, abs=1e-12)
+    assert pose[2, 3] == pytest.approx(1.0)
+    assert heading == pytest.approx(0.0)
+
+
+# ----- StubBaseController --------------------------------------------------
+
+
+def test_stub_is_basecontroller():
+    assert isinstance(StubBaseController(), BaseController)
+
+
+def test_stub_integrates_forward():
+    c = StubBaseController()
+    c.move(0.0, 0.2, dt=1.0)
+    assert c.position()[2] == pytest.approx(0.2)
+
+
+def test_stub_clamps_velocity():
+    c = StubBaseController(max_lin_speed=0.1)
+    c.move(5.0, 0.0, dt=1.0)
+    assert c.position()[0] == pytest.approx(0.1)
+
+
+# ----- RobotBaseController -------------------------------------------------
+
+
+class FakeRobot:
+    """Minimal Robot stand-in: records actions, serves canned odometry."""
+
+    def __init__(self) -> None:
+        self.actions: list[dict] = []
+        self.obs: dict = {}
+
+    def send_action(self, action: dict) -> dict:
+        self.actions.append(action)
+        return action
+
+    def get_observation(self) -> dict:
+        return self.obs
+
+
+def _robot_controller(**cfg_kwargs) -> tuple[RobotBaseController, FakeRobot]:
+    robot = FakeRobot()
+    cfg = RobotBaseControllerConfig(**cfg_kwargs)
+    return RobotBaseController(robot, cfg), robot
+
+
+def _odom(x=0.0, y=0.0, yaw=0.0) -> dict:
+    return {"x.pos": x, "y.pos": y, "theta.pos": yaw}
+
+
+def test_robot_controller_is_basecontroller():
+    ctl, _ = _robot_controller()
+    assert isinstance(ctl, BaseController)
+
+
+def test_forward_command_reaches_send_action():
+    ctl, robot = _robot_controller()
+    ctl.feed_observation(_odom())  # heading 0
+    ctl.move(vx=0.0, vz=0.3, dt=0.05)
+    assert robot.actions[-1] == {
+        "x.vel": pytest.approx(0.3),
+        "y.vel": pytest.approx(0.0),
+        "theta.vel": pytest.approx(0.0),
+    }
+
+
+def test_command_uses_odometry_heading():
+    """After the robot turns to heading +π/2, a world +x command comes out
+    as pure body-forward. First sample fixes the origin."""
+    ctl, robot = _robot_controller()
+    ctl.feed_observation(_odom())  # origin, heading 0
+    ctl.feed_observation(_odom(yaw=-math.pi / 2))  # turned; heading +π/2
+    ctl.move(vx=0.3, vz=0.0, dt=0.05)
+    assert robot.actions[-1]["x.vel"] == pytest.approx(0.3)
+    assert robot.actions[-1]["y.vel"] == pytest.approx(0.0, abs=1e-9)
+
+
+def test_command_is_clamped_before_send():
+    ctl, robot = _robot_controller(max_lin_speed=0.1)
+    ctl.feed_observation(_odom())
+    ctl.move(vx=0.0, vz=9.0, dt=0.05)
+    assert robot.actions[-1]["x.vel"] == pytest.approx(0.1)
+
+
+def test_pose_comes_from_odometry_not_integration():
+    ctl, _ = _robot_controller()
+    ctl.feed_observation(_odom())
+    ctl.move(0.0, 0.3, dt=1.0)  # would integrate 0.3 m open-loop
+    ctl.feed_observation(_odom(x=0.05))  # ...but odometry says 5 cm forward
+    assert ctl.position()[2] == pytest.approx(0.05)
+
+
+def test_origin_is_first_odometry_sample():
+    ctl, _ = _robot_controller()
+    ctl.feed_observation(_odom(x=3.0, y=-1.0, yaw=0.7))
+    np.testing.assert_allclose(ctl.pose(), np.eye(4), atol=1e-12)
+
+
+def test_open_loop_fallback_without_odometry():
+    """No odometry fed → integrate open-loop like the stub."""
+    ctl, _ = _robot_controller()
+    ctl.move(0.0, 0.2, dt=1.0)
+    assert ctl.position()[2] == pytest.approx(0.2)
+
+
+def test_stop_sends_zero_velocity():
+    ctl, robot = _robot_controller()
+    ctl.stop()
+    assert robot.actions[-1] == {"x.vel": 0.0, "y.vel": 0.0, "theta.vel": 0.0}
+    assert ctl.is_stopped
+
+
+def test_robot_controller_matches_stub_open_loop():
+    """Open-loop pose integration matches StubBaseController for the same
+    command sequence — sim runs must transfer to the real base."""
+    ctl, _ = _robot_controller(max_lin_speed=1.0)
+    stub = StubBaseController()
+    for vx, vz, yaw in [(0.2, 0.0, 0.0), (0.0, 0.3, 0.5), (0.1, 0.1, -0.2)]:
+        ctl.move(vx, vz, yaw, dt=0.5)
+        stub.move(vx, vz, yaw, dt=0.5)
+    np.testing.assert_allclose(ctl.pose(), stub.pose(), atol=1e-9)
+
+
+# ----- SafeBaseController --------------------------------------------------
+
+
+class FakeGrid:
+    """Occupancy stand-in with a single obstacle cell."""
+
+    def __init__(self, obstacle_cell=(5, 6), cell_size=0.1, origin_x=-0.5, origin_z=-0.5):
+        self.obstacle_cell = obstacle_cell
+        self.cell_size = cell_size
+        self.origin_x = origin_x
+        self.origin_z = origin_z
+
+    def world_to_cell(self, x: float, z: float) -> tuple[int, int]:
+        ix = int((x - self.origin_x) / self.cell_size)
+        iz = int((z - self.origin_z) / self.cell_size)
+        return iz, ix
+
+    def is_obstacle(self, iz: int, ix: int) -> bool:
+        return (iz, ix) == self.obstacle_cell
+
+
+def test_safe_passes_normal_moves():
+    inner = StubBaseController()
+    safe = SafeBaseController(inner=inner)
+    safe.feed_watchdog()
+    safe.move(0.0, 0.1, dt=1.0)
+    assert inner.position()[2] == pytest.approx(0.1)
+
+
+def test_safe_clamps_speed():
+    inner = StubBaseController(max_lin_speed=100.0)
+    safe = SafeBaseController(inner=inner, max_lin_speed=0.5)
+    safe.feed_watchdog()
+    safe.move(10.0, 0.0, dt=1.0)
+    assert inner.position()[0] == pytest.approx(0.5)
+
+
+def test_safe_watchdog_latches_on_stale_keyframes():
+    inner = StubBaseController()
+    safe = SafeBaseController(inner=inner, watchdog_timeout_s=0.05)
+    safe.feed_watchdog()
+    time.sleep(0.1)
+    safe.move(0.0, 0.1, dt=1.0)
+    assert safe.e_stop_latched
+    safe.move(0.0, 10.0, dt=1.0)  # refused
+    assert inner.position()[2] == pytest.approx(0.0, abs=1e-6)
+
+
+def test_safe_reset_watchdog_re_enables_motion():
+    inner = StubBaseController()
+    safe = SafeBaseController(inner=inner, watchdog_timeout_s=0.05)
+    safe.feed_watchdog()
+    time.sleep(0.1)
+    safe.move(0.0, 0.1)
+    assert safe.e_stop_latched
+    safe.reset_watchdog()
+    safe.move(0.0, 0.1, dt=1.0)
+    assert inner.position()[2] == pytest.approx(0.1)
+
+
+def test_safe_refuses_move_into_obstacle():
+    inner = StubBaseController()
+    grid = FakeGrid(obstacle_cell=(5, 6))
+    safe = SafeBaseController(inner=inner, occupancy_provider=lambda: grid)
+    safe.feed_watchdog()
+    # +0.15 m in x from origin lands mid-column ix=6 (origin_x=-0.5, cell=0.1).
+    safe.move(vx=0.15, vz=0.0, dt=1.0)
+    assert safe.e_stop_latched
+    assert inner.position()[0] == pytest.approx(0.0, abs=1e-6)
+
+
+def test_safe_allows_move_into_free_cell():
+    inner = StubBaseController()
+    grid = FakeGrid(obstacle_cell=(99, 99))
+    safe = SafeBaseController(inner=inner, occupancy_provider=lambda: grid)
+    safe.feed_watchdog()
+    safe.move(vx=0.1, vz=0.0, dt=1.0)
+    assert inner.position()[0] == pytest.approx(0.1)
+
+
+def test_safe_allows_when_no_map_yet():
+    inner = StubBaseController()
+    safe = SafeBaseController(inner=inner, occupancy_provider=lambda: None)
+    safe.feed_watchdog()
+    safe.move(vx=0.1, vz=0.0, dt=1.0)
+    assert inner.position()[0] == pytest.approx(0.1)

tests/navigation/test_dog_cli.py

@@ -0,0 +1,106 @@
+#!/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.
+
+"""End-to-end dry-run tests for the dog-nav REPL + synthetic scene.
+
+These exercise the whole navigation stack — sim scene → voxel map →
+SigLIP stand-in → skills → agent → controller — with no robot, camera,
+or models.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from lerobot.navigation.dog_cli import DogController, _build_dry_run, main
+from lerobot.navigation.sim import kitchen_scene
+
+
+def test_kitchen_scene_builds_with_all_objects():
+    scene = kitchen_scene()
+    assert {o.name for o in scene.objects} == {"couch", "chair", "lamp", "plant"}
+    assert len(scene.voxel_map) > 0
+    assert scene.voxel_map.feature_dim == scene.feature_dim
+
+
+def test_feature_extractor_matches_object_vectors():
+    scene = kitchen_scene()
+    fx = scene.feature_extractor()
+    couch = scene.object("couch")
+    emb = fx.encode_text("couch")
+    # The couch query should align with the couch's stored basis vector.
+    import numpy as np
+
+    assert float(np.dot(emb, couch.feature_vec / np.linalg.norm(couch.feature_vec))) > 0.9
+
+
+def test_controller_reaches_mapped_object():
+    ctl = _build_dry_run()
+    result = ctl.handle_prompt("couch")
+    assert result.fully_successful
+    tr = result.target_results[0]
+    assert tr.reached
+    # Landed near the couch ground-truth (3.0, _, 2.0).
+    assert tr.final_xyz is not None
+    assert abs(tr.final_xyz[0] - 3.0) < 1.5
+    assert abs(tr.final_xyz[2] - 2.0) < 1.5
+
+
+def test_controller_navigates_to_each_object():
+    for name, (gx, gz) in {
+        "couch": (3.0, 2.0),
+        "chair": (-2.0, -1.5),
+        "plant": (-2.5, 2.5),
+    }.items():
+        ctl = _build_dry_run()
+        result = ctl.handle_prompt(name)
+        assert result.fully_successful, f"failed to reach {name}"
+        fx = result.target_results[0].final_xyz
+        assert abs(fx[0] - gx) < 1.5 and abs(fx[2] - gz) < 1.5
+
+
+def test_controller_abstains_on_absent_object():
+    ctl = _build_dry_run()
+    result = ctl.handle_prompt("banana")  # not in the scene
+    assert not result.fully_successful
+    assert result.target_results[0].reason in {"budget_exhausted", "no_frontier"}
+
+
+def test_idle_tick_explores_or_reports_no_frontier():
+    ctl = _build_dry_run()
+    ex = ctl.idle_tick()
+    # A fully-observed synthetic floor may have no frontier; either way the
+    # call must be well-formed and not raise.
+    assert ex.reason in {"ok", "no frontier"}
+
+
+def test_main_single_command_dry_run_returns_zero():
+    assert main(["--dry-run", "--command", "couch", "--log-level", "WARNING"]) == 0
+
+
+def test_main_absent_object_returns_nonzero():
+    assert main(["--dry-run", "--command", "banana", "--log-level", "WARNING"]) == 1
+
+
+def test_main_live_mode_refuses_until_pipeline_lands():
+    with pytest.raises(SystemExit):
+        main(["--command", "couch"])
+
+
+def test_dogcontroller_stop_is_safe():
+    ctl = _build_dry_run()
+    ctl.stop()
+    assert isinstance(ctl, DogController)

tests/navigation/test_geometry.py

@@ -0,0 +1,120 @@
+#!/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.
+
+"""Tests for geometry runners + the odometry similarity anchor.
+
+Model-free: only the FakeGeometryRunner and the pure-numpy Umeyama fit
+are exercised. LingBotMapRunner is checked for its no-SDK error only.
+"""
+
+# ruff: noqa: N806  — R, U, S, Vt, D: conventional linear-algebra / array-dimension names
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from lerobot.navigation.geometry import (
+    FakeGeometryRunner,
+    GeometryOutput,
+    GeometryRunner,
+    LingBotMapRunner,
+    align_trajectory_to_odometry,
+    umeyama_similarity,
+)
+
+
+def _views(n=2, h=14, w=14) -> np.ndarray:
+    return np.zeros((n, h, w, 3), dtype=np.uint8)
+
+
+def test_fake_runner_satisfies_protocol():
+    assert isinstance(FakeGeometryRunner(), GeometryRunner)
+
+
+def test_fake_runner_output_shapes():
+    out = FakeGeometryRunner(depth=3.0, focal_px=100.0)(_views(2, 14, 14))
+    assert isinstance(out, GeometryOutput)
+    assert out.points.shape == (2, 14, 14, 3)
+    assert out.local_points.shape == (2, 14, 14, 3)
+    assert out.conf.shape == (2, 14, 14)
+    assert out.camera_poses.shape == (2, 4, 4)
+
+
+def test_fake_runner_depth_is_constant():
+    out = FakeGeometryRunner(depth=2.5)(_views())
+    # local_points z channel is the depth everywhere.
+    assert np.allclose(out.local_points[..., 2], 2.5)
+
+
+def test_fake_runner_rejects_bad_shape():
+    with pytest.raises(ValueError, match="N, H, W, 3"):
+        FakeGeometryRunner()(np.zeros((14, 14, 3), dtype=np.uint8))
+
+
+def test_lingbot_runner_raises_without_sdk():
+    runner = LingBotMapRunner(device="cpu")
+    with pytest.raises((RuntimeError, ValueError)):
+        # Either the lazy import fails (no lingbot-map) or shape check trips
+        # first — both are acceptable "did not silently succeed" outcomes.
+        runner(_views())
+
+
+# ----- Umeyama similarity --------------------------------------------------
+
+
+def test_umeyama_recovers_known_similarity():
+    rng = np.random.default_rng(0)
+    src = rng.normal(size=(20, 3))
+    # Known transform: scale 2.5, a rotation about z by 30°, translation.
+    theta = np.deg2rad(30.0)
+    c, s = np.cos(theta), np.sin(theta)
+    R_true = np.array([[c, -s, 0], [s, c, 0], [0, 0, 1.0]])
+    s_true, t_true = 2.5, np.array([1.0, -2.0, 0.5])
+    dst = (s_true * (R_true @ src.T)).T + t_true
+
+    s_fit, R_fit, t_fit = umeyama_similarity(src, dst)
+    assert s_fit == pytest.approx(s_true, rel=1e-6)
+    np.testing.assert_allclose(R_fit, R_true, atol=1e-6)
+    np.testing.assert_allclose(t_fit, t_true, atol=1e-6)
+
+
+def test_umeyama_reconstructs_points():
+    rng = np.random.default_rng(1)
+    src = rng.normal(size=(10, 3))
+    dst = 0.5 * src + np.array([3.0, 0.0, -1.0])
+    s, R, t = umeyama_similarity(src, dst)
+    recon = (s * (R @ src.T)).T + t
+    np.testing.assert_allclose(recon, dst, atol=1e-6)
+
+
+def test_umeyama_rejects_mismatched_shapes():
+    with pytest.raises(ValueError):
+        umeyama_similarity(np.zeros((5, 3)), np.zeros((4, 3)))
+
+
+def test_align_requires_three_points():
+    with pytest.raises(ValueError, match="at least 3"):
+        align_trajectory_to_odometry(np.zeros((2, 3)), np.zeros((2, 3)))
+
+
+def test_align_scale_anchor_makes_metric():
+    """A monocular trajectory at half scale is recovered to metric."""
+    odom = np.array([[0, 0, 0], [1, 0, 0], [1, 0, 1], [0, 0, 1]], dtype=np.float64)
+    cam = odom * 0.5  # model world is half-scale
+    s, R, t = align_trajectory_to_odometry(cam, odom)
+    assert s == pytest.approx(2.0, rel=1e-6)
+    recon = (s * (R @ cam.T)).T + t
+    np.testing.assert_allclose(recon, odom, atol=1e-9)

tests/navigation/test_live_mapper.py

@@ -0,0 +1,115 @@
+#!/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.
+
+"""Exercise the live mapping loop (LiveMapper.tick) with no hardware.
+
+A fake robot serves canned front-camera frames + odometry; FakeGeometryRunner
+supplies planar depth. This validates the perceive → project-through-odometry
+→ integrate path — the logic that runs on the real dog — including the
+frame-consistency fix (voxels land in the odometry world frame).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from lerobot.navigation.base_controller import RobotBaseController
+from lerobot.navigation.dog_cli import LiveMapper
+from lerobot.navigation.geometry import FakeGeometryRunner
+from lerobot.navigation.pipeline import PipelineConfig
+from lerobot.navigation.voxel_map import VoxelMap
+
+
+class FakeGo2:
+    """Robot stand-in: canned front frame + programmable odometry."""
+
+    def __init__(self, h=14, w=14):
+        self.h, self.w = h, w
+        self.odom = {"x.pos": 0.0, "y.pos": 0.0, "theta.pos": 0.0}
+        self.actions = []
+
+    def get_observation(self):
+        return {
+            "front": np.full((self.h, self.w, 3), 120, dtype=np.uint8),
+            **self.odom,
+        }
+
+    def send_action(self, action):
+        self.actions.append(action)
+        return action
+
+
+def _mapper(**pcfg_kwargs):
+    robot = FakeGo2()
+    base = RobotBaseController(robot)
+    vm = VoxelMap(voxel_size=0.05)
+    geom = FakeGeometryRunner(depth=2.0, focal_px=100.0)
+    pcfg = PipelineConfig(focal_px=100.0, **pcfg_kwargs)
+    mapper = LiveMapper(robot, base, geom, siglip=None, voxel_map=vm, pcfg=pcfg)
+    return mapper, robot, base, vm
+
+
+def test_tick_populates_voxel_map():
+    mapper, _, _, vm = _mapper()
+    mapper.tick(0.0)
+    assert len(vm) > 0
+
+
+def test_tick_places_voxels_in_odometry_frame():
+    """With the dog at origin facing +z, the planar floor at depth 2 lands
+    around world z≈2 — i.e. in the odometry world frame, not the model's."""
+    mapper, robot, base, vm = _mapper()
+    mapper.tick(0.0)
+    snap = vm.snapshot()
+    # Median z of the observed floor should be near the camera depth (2 m).
+    assert 1.0 < float(np.median(snap.xyz[:, 2])) < 3.0
+
+
+def test_tick_follows_odometry_translation():
+    """Move the dog forward 5 m in odometry; the new floor voxels shift with
+    it — proof the map tracks the odometry world frame."""
+    mapper, robot, base, vm = _mapper()
+    mapper.tick(0.0)
+    z0 = float(np.median(vm.snapshot().xyz[:, 2]))
+
+    robot.odom = {"x.pos": 5.0, "y.pos": 0.0, "theta.pos": 0.0}  # forward 5 m
+    vm2 = VoxelMap(voxel_size=0.05)
+    mapper.voxel_map = vm2
+    mapper.tick(1.0)
+    z1 = float(np.median(vm2.snapshot().xyz[:, 2]))
+    # Forward odometry (+x_odom → +z_world) shifts the floor ~5 m in world z.
+    assert z1 - z0 > 4.0
+
+
+def test_tick_without_front_frame_is_safe():
+    mapper, robot, _, vm = _mapper()
+    robot.get_observation = lambda: {"x.pos": 0.0, "y.pos": 0.0, "theta.pos": 0.0}
+    mapper.tick(0.0)  # no 'front' → no-op, must not raise
+    assert len(vm) == 0
+
+
+def test_tick_feeds_watchdog_when_present():
+    from lerobot.navigation.base_controller import SafeBaseController
+
+    mapper, _, base, _ = _mapper()
+    safe = SafeBaseController(inner=base)
+    safe.e_stop_latched = True  # will clear on a healthy feed via reset elsewhere
+    mapper.safe = safe
+    # feed_watchdog just refreshes the timer; assert tick calls it (no raise,
+    # and the timestamp advances).
+    before = safe._last_keyframe_walltime
+    mapper.tick(0.0)
+    assert safe._last_keyframe_walltime >= before

tests/navigation/test_occupancy.py

@@ -0,0 +1,266 @@
+#!/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.
+
+"""Unit tests for the B1 occupancy projection + A*."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from lerobot.navigation.occupancy import (
+    NAVIGABLE,
+    OBSTACLE,
+    UNOBSERVED,
+    OccupancyGrid,
+    astar,
+    estimate_ground_y,
+    find_frontier_cells,
+    occupancy_to_rgb,
+    project_voxel_map_to_grid,
+)
+from lerobot.navigation.voxel_map import VoxelMap
+
+
+def _vm_from_points(pts: list[tuple[float, float, float]], voxel_size: float = 0.1) -> VoxelMap:
+    """Build a VoxelMap from a list of world-XYZ points (all at conf=1.0)."""
+    vm = VoxelMap(voxel_size=voxel_size)
+    if not pts:
+        return vm
+    arr = np.asarray(pts, dtype=np.float32).reshape(-1, 1, 3)
+    rgb = np.full((len(pts), 1, 3), 200, dtype=np.uint8)
+    conf = np.ones((len(pts), 1), dtype=np.float32)
+    vm.add(arr, rgb, conf, frame=0, t=0.0)
+    return vm
+
+
+def test_estimate_ground_y_picks_high_percentile():
+    # +y is down in OpenCV; ground = largest y. Numpy's default percentile
+    # interpolation lands at 0.9 here (95% of way from 0.5 to 1.0 of the
+    # last interval = 0.5 + 0.4·1.0 = 0.9), so we just check the estimate
+    # is at least that high and clearly above the median.
+    xyz = np.array(
+        [[0, -1, 0], [0, -0.5, 0], [0, 0.0, 0], [0, 0.5, 0], [0, 1.0, 0]],
+        dtype=np.float32,
+    )
+    g = estimate_ground_y(xyz)
+    assert g >= 0.8
+    median_y = float(np.median(xyz[:, 1]))
+    assert g > median_y
+
+
+def test_projection_makes_ground_navigable_and_obstacles_red():
+    # Ground row at y=1.0 across z=0..1; an obstacle at y=0.2 (above ground).
+    pts: list[tuple[float, float, float]] = []
+    for x in np.linspace(-0.5, 0.5, 11):
+        for z in np.linspace(0.0, 1.0, 11):
+            pts.append((float(x), 1.0, float(z)))  # floor
+    # Standing obstacle column at (x=0, z=0.5)
+    for y in np.linspace(0.2, 0.9, 8):
+        pts.append((0.0, float(y), 0.5))
+    vm = _vm_from_points(pts, voxel_size=0.05)
+    grid = project_voxel_map_to_grid(vm, cell_size=0.1, obstacle_y_range=(-2.0, -0.1))
+    # The obstacle column should yield at least one OBSTACLE cell.
+    assert (grid.classes == OBSTACLE).sum() >= 1
+    # The floor away from the obstacle should be NAVIGABLE.
+    iz, ix = grid.world_to_cell(-0.4, 0.0)
+    assert grid.classes[iz, ix] == NAVIGABLE
+    # Everywhere outside the observed area should still be UNOBSERVED.
+    assert (grid.classes == UNOBSERVED).any()
+
+
+def test_empty_voxelmap_projects_to_single_unobserved_cell():
+    vm = VoxelMap()
+    grid = project_voxel_map_to_grid(vm)
+    assert grid.shape == (1, 1)
+    assert grid.classes[0, 0] == UNOBSERVED
+
+
+def test_world_to_cell_and_back_roundtrip():
+    classes = np.full((4, 6), NAVIGABLE, dtype=np.int8)
+    grid = OccupancyGrid(classes=classes, cell_size=0.5, origin_x=-1.0, origin_z=2.0, ground_y=0.0)
+    iz, ix = grid.world_to_cell(0.25, 3.4)
+    x, z = grid.cell_to_world(iz, ix)
+    # The recovered (x, z) should land inside the same cell.
+    iz2, ix2 = grid.world_to_cell(x, z)
+    assert (iz, ix) == (iz2, ix2)
+
+
+def test_astar_finds_straight_path_when_unblocked():
+    classes = np.full((10, 10), NAVIGABLE, dtype=np.int8)
+    grid = OccupancyGrid(classes=classes, cell_size=0.1, origin_x=0.0, origin_z=0.0, ground_y=0.0)
+    path = astar(grid, (0.05, 0.05), (0.85, 0.85))
+    assert path is not None
+    assert len(path) >= 2
+    # Should land near the goal.
+    assert abs(path[-1][0] - 0.85) < 0.1 and abs(path[-1][1] - 0.85) < 0.1
+
+
+def test_astar_routes_around_obstacle_wall():
+    # Vertical wall at column 5, rows 1..8.
+    classes = np.full((10, 10), NAVIGABLE, dtype=np.int8)
+    classes[1:9, 5] = OBSTACLE
+    grid = OccupancyGrid(classes=classes, cell_size=0.1, origin_x=0.0, origin_z=0.0, ground_y=0.0)
+    path = astar(grid, (0.05, 0.45), (0.95, 0.45))
+    assert path is not None
+    # The path must not pass through any obstacle cell.
+    for x, z in path:
+        iz, ix = grid.world_to_cell(x, z)
+        assert classes[iz, ix] != OBSTACLE
+
+
+def test_astar_returns_none_when_no_path():
+    # Wall that completely separates start from goal.
+    classes = np.full((10, 10), NAVIGABLE, dtype=np.int8)
+    classes[:, 5] = OBSTACLE
+    grid = OccupancyGrid(classes=classes, cell_size=0.1, origin_x=0.0, origin_z=0.0, ground_y=0.0)
+    path = astar(grid, (0.05, 0.5), (0.95, 0.5))
+    assert path is None
+
+
+def test_astar_no_corner_cutting_through_obstacles():
+    # Two diagonal obstacle cells that would let a naive A* squeeze through.
+    classes = np.full((4, 4), NAVIGABLE, dtype=np.int8)
+    classes[1, 2] = OBSTACLE
+    classes[2, 1] = OBSTACLE
+    grid = OccupancyGrid(classes=classes, cell_size=1.0, origin_x=0.0, origin_z=0.0, ground_y=0.0)
+    path = astar(grid, (0.5, 0.5), (2.5, 2.5))
+    assert path is not None
+    # The path should NOT step (1, 1) → (2, 2) since that diagonal cuts the
+    # obstacle corner. Verify by checking no two consecutive cells are a
+    # diagonal move with both perpendicular cells blocked.
+    cells = [grid.world_to_cell(x, z) for x, z in path]
+    for prev, nxt in zip(cells, cells[1:], strict=False):
+        diz = nxt[0] - prev[0]
+        dix = nxt[1] - prev[1]
+        if abs(diz) == 1 and abs(dix) == 1:
+            assert grid.is_navigable(prev[0] + diz, prev[1]) and grid.is_navigable(prev[0], prev[1] + dix), (
+                f"Corner cut detected at step {prev}->{nxt}"
+            )
+
+
+def test_obstacle_inflation_grows_obstacle_class():
+    classes = np.full((5, 5), NAVIGABLE, dtype=np.int8)
+    classes[2, 2] = OBSTACLE
+    # Test the private helper directly — we don't need a voxel map for this.
+    from lerobot.navigation.occupancy import _inflate_obstacles
+
+    inflated = _inflate_obstacles(classes, radius=1)
+    # 3x3 block now obstacle (around the single original cell).
+    assert (inflated[1:4, 1:4] == OBSTACLE).all()
+    # Far corner stays navigable.
+    assert inflated[0, 0] == NAVIGABLE
+
+
+def test_find_frontier_cells_are_navigable_next_to_unobserved():
+    classes = np.full((5, 5), UNOBSERVED, dtype=np.int8)
+    classes[1:4, 1:4] = NAVIGABLE
+    grid = OccupancyGrid(classes=classes, cell_size=0.1, origin_x=0.0, origin_z=0.0, ground_y=0.0)
+    cells = find_frontier_cells(grid)
+    # Every frontier cell is NAVIGABLE itself.
+    for iz, ix in cells:
+        assert classes[iz, ix] == NAVIGABLE
+    # And every frontier cell has at least one UNOBSERVED 4-neighbour.
+    for iz, ix in cells:
+        adj_unobs = (
+            (iz > 0 and classes[iz - 1, ix] == UNOBSERVED)
+            or (iz < 4 and classes[iz + 1, ix] == UNOBSERVED)
+            or (ix > 0 and classes[iz, ix - 1] == UNOBSERVED)
+            or (ix < 4 and classes[iz, ix + 1] == UNOBSERVED)
+        )
+        assert adj_unobs
+
+
+def test_find_frontier_empty_when_no_unobserved():
+    classes = np.full((5, 5), NAVIGABLE, dtype=np.int8)
+    grid = OccupancyGrid(classes=classes, cell_size=0.1, origin_x=0.0, origin_z=0.0, ground_y=0.0)
+    cells = find_frontier_cells(grid)
+    assert cells.shape == (0, 2)
+
+
+def test_occupancy_to_rgb_returns_uint8_image():
+    classes = np.array([[UNOBSERVED, NAVIGABLE, OBSTACLE]], dtype=np.int8)
+    grid = OccupancyGrid(classes=classes, cell_size=0.1, origin_x=0.0, origin_z=0.0, ground_y=0.0)
+    img = occupancy_to_rgb(grid)
+    assert img.shape == (1, 3, 3)
+    assert img.dtype == np.uint8
+    # Obstacle cell is reddish.
+    assert img[0, 2, 0] > img[0, 2, 1] and img[0, 2, 0] > img[0, 2, 2]
+
+
+def test_carving_makes_obstacle_vanish_on_next_projection():
+    """End-to-end: an object voxel becomes an obstacle; after we delete it
+    from the VoxelMap (simulating carving), the next projection no longer
+    flags that cell as OBSTACLE — that's what makes the obstacle map a
+    cheap derived view."""
+    pts = []
+    for x in np.linspace(-0.5, 0.5, 11):
+        for z in np.linspace(0.0, 1.0, 11):
+            pts.append((float(x), 1.0, float(z)))  # floor
+    for y in np.linspace(0.2, 0.9, 8):
+        pts.append((0.0, float(y), 0.5))  # obstacle column
+    vm = _vm_from_points(pts, voxel_size=0.05)
+    grid1 = project_voxel_map_to_grid(vm, cell_size=0.1)
+    assert (grid1.classes == OBSTACLE).sum() >= 1
+
+    # Surgically drop every voxel in the obstacle band.
+    cnt = vm._count.astype(np.float64).reshape(-1, 1)  # noqa: SLF001
+    means = vm._xyz_sum / cnt  # noqa: SLF001
+    keep = (means[:, 1] >= 0.95) | (means[:, 1] <= 0.05)
+    vm._idx = vm._idx[keep]  # noqa: SLF001
+    vm._count = vm._count[keep]  # noqa: SLF001
+    vm._xyz_sum = vm._xyz_sum[keep]  # noqa: SLF001
+    vm._rgb_sum = vm._rgb_sum[keep]  # noqa: SLF001
+    vm._last_frame = vm._last_frame[keep]  # noqa: SLF001
+    vm._last_time = vm._last_time[keep]  # noqa: SLF001
+    vm._lookup = {  # noqa: SLF001
+        (int(vm._idx[i, 0]), int(vm._idx[i, 1]), int(vm._idx[i, 2])): i  # noqa: SLF001
+        for i in range(len(vm._idx))  # noqa: SLF001
+    }
+
+    grid2 = project_voxel_map_to_grid(vm, cell_size=0.1)
+    assert (grid2.classes == OBSTACLE).sum() == 0
+
+
+def test_astar_works_after_inflation():
+    # Without inflation, robot can hug a 1-cell-wide wall; with inflation,
+    # it has to take a wider detour.
+    classes = np.full((11, 11), NAVIGABLE, dtype=np.int8)
+    classes[5, 1:9] = OBSTACLE  # horizontal wall
+    grid_raw = OccupancyGrid(classes=classes.copy(), cell_size=0.1, origin_x=0.0, origin_z=0.0, ground_y=0.0)
+    path_raw = astar(grid_raw, (0.45, 0.0), (0.45, 1.0))
+    assert path_raw is not None  # the wall has an opening at the edges
+    from lerobot.navigation.occupancy import _inflate_obstacles
+
+    inflated = _inflate_obstacles(classes, radius=1)
+    grid_inf = OccupancyGrid(classes=inflated, cell_size=0.1, origin_x=0.0, origin_z=0.0, ground_y=0.0)
+    # After inflation the path is at least as long (often longer).
+    path_inf = astar(grid_inf, (0.45, 0.0), (0.45, 1.0))
+    if path_inf is not None:
+        assert len(path_inf) >= len(path_raw)
+
+
+def test_obstacle_range_outside_voxel_y_produces_only_navigable():
+    """A range well above the actual voxel y values should classify the floor
+    as NAVIGABLE only — no obstacles. (Replaces the prior sub-float32-epsilon
+    test which exercised numpy's float-downcast quirk rather than the
+    occupancy semantics.)"""
+    pts = [(float(x), 1.0, float(z)) for x in [0, 0.1] for z in [0, 0.1]]
+    vm = _vm_from_points(pts, voxel_size=0.05)
+    # ground_y will be 1.0; this range looks for obstacles 5 m–10 m above
+    # the floor, where there are none.
+    grid = project_voxel_map_to_grid(vm, obstacle_y_range=(-10.0, -5.0))
+    assert (grid.classes == OBSTACLE).sum() == 0
+    assert (grid.classes == NAVIGABLE).sum() >= 1

tests/navigation/test_pipeline.py

@@ -0,0 +1,93 @@
+#!/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.
+
+"""Tests for the viz-free keyframe integration loop.
+
+Uses FakeGeometryRunner output fed through integrate_keyframe into a real
+VoxelMap — no models, no viz.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from lerobot.navigation.geometry import FakeGeometryRunner
+from lerobot.navigation.pipeline import (
+    KeyframeContext,
+    PipelineConfig,
+    integrate_keyframe,
+    upsample_features_to_view,
+)
+from lerobot.navigation.voxel_map import VoxelMap
+
+
+def _ctx_from_geometry(frame_idx=0, t=0.0, feat_map=None) -> KeyframeContext:
+    h = w = 14
+    views = np.full((1, h, w, 3), 120, dtype=np.uint8)
+    out = FakeGeometryRunner(depth=3.0, focal_px=100.0)(views)
+    return KeyframeContext(
+        frame_idx=frame_idx,
+        t_sec=t,
+        rgb_uint8=views[0],
+        points_world=out.points[0],
+        local_points=out.local_points[0],
+        conf=out.conf[0],
+        pose=out.camera_poses[0],
+        feat_map=feat_map,
+    )
+
+
+def test_integrate_adds_voxels():
+    vm = VoxelMap(voxel_size=0.05)
+    ctx = _ctx_from_geometry()
+    carve, stats = integrate_keyframe(vm, ctx, PipelineConfig(focal_px=100.0))
+    assert stats.n_added > 0
+    assert len(vm) == stats.n_voxels
+    assert carve.n_removed == 0  # nothing to carve on an empty map
+
+
+def test_integrate_second_frame_updates_not_duplicates():
+    vm = VoxelMap(voxel_size=0.05)
+    pcfg = PipelineConfig(focal_px=100.0)
+    integrate_keyframe(vm, _ctx_from_geometry(frame_idx=0, t=0.0), pcfg)
+    n_after_first = len(vm)
+    _, stats2 = integrate_keyframe(vm, _ctx_from_geometry(frame_idx=1, t=0.5), pcfg)
+    # Same synthetic view → same voxels updated, not a second copy.
+    assert stats2.n_added == 0
+    assert len(vm) == n_after_first
+
+
+def test_integrate_with_features_enables_query():
+    vm = VoxelMap(voxel_size=0.05)
+    h = w = 14
+    d = 8
+    feat = np.zeros((h, w, d), dtype=np.float16)
+    feat[..., 0] = 1.0  # every pixel carries basis vector 0
+    ctx = _ctx_from_geometry(feat_map=feat)
+    integrate_keyframe(vm, ctx, PipelineConfig(focal_px=100.0))
+    assert vm.feature_dim == d
+    q = np.zeros(d, dtype=np.float32)
+    q[0] = 1.0
+    result = vm.query(q, top_k=5)
+    assert result.score.size > 0
+    assert float(result.score.max()) > 0.9  # basis-0 query matches basis-0 voxels
+
+
+def test_upsample_features_to_view_shape():
+    patch = np.zeros((3, 3, 8), dtype=np.float16)
+    up = upsample_features_to_view(patch, view_h=28, view_w=28)
+    assert up.shape == (28, 28, 8)
+    assert up.dtype == np.float16

tests/navigation/test_skills.py

@@ -0,0 +1,268 @@
+#!/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.
+
+"""Unit tests for the unified ``SpatialSkills`` API (B1+B2+B3+B4)."""
+
+# ruff: noqa: N803, N806  — D: conventional feature-dimension name
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import numpy as np
+
+from lerobot.navigation.base_controller import StubBaseController
+from lerobot.navigation.skills import SkillsConfig, SpatialSkills
+from lerobot.navigation.voxel_map import VoxelMap
+
+# ----- fakes / fixtures ---------------------------------------------------
+
+
+@dataclass
+class FakeSiglip:
+    """Tiny stand-in for SiglipFeatureExtractor — text → fixed vector."""
+
+    text_to_vec: dict[str, np.ndarray]
+    feature_dim: int = 4
+
+    def encode_text(self, text: str) -> np.ndarray:
+        v = self.text_to_vec.get(text)
+        if v is None:
+            # Default: random-but-deterministic vector
+            rng = np.random.default_rng(abs(hash(text)) % (2**32))
+            v = rng.normal(size=self.feature_dim).astype(np.float32)
+        v = v.astype(np.float32)
+        v = v / max(np.linalg.norm(v), 1e-6)
+        return v
+
+
+def _vm_with_couch_and_chair(D: int = 4) -> VoxelMap:
+    """Two spatially-separated clusters with distinct unit feature vectors."""
+    vm = VoxelMap(voxel_size=0.1)
+    rgb = np.full((1, 1, 3), 200, dtype=np.uint8)
+    conf = np.ones((1, 1), dtype=np.float32)
+    couch_vec = np.eye(D)[0].astype(np.float16).reshape(1, 1, D)
+    chair_vec = np.eye(D)[1].astype(np.float16).reshape(1, 1, D)
+    # Couch cluster around (5, 1, 3)
+    for x in (4.9, 5.0, 5.1):
+        for z in (2.9, 3.0, 3.1):
+            pts = np.array([[[x, 1.0, z]]], dtype=np.float32)
+            vm.add(pts, rgb, conf, frame=0, t=0.0, feat_map=couch_vec)
+    # Chair cluster around (-3, 1, 1)
+    for x in (-3.1, -3.0, -2.9):
+        for z in (0.9, 1.0, 1.1):
+            pts = np.array([[[x, 1.0, z]]], dtype=np.float32)
+            vm.add(pts, rgb, conf, frame=0, t=0.0, feat_map=chair_vec)
+    return vm
+
+
+# ----- locate() -----------------------------------------------------------
+
+
+def test_locate_returns_centroid_for_matching_query():
+    vm = _vm_with_couch_and_chair()
+    base = StubBaseController()
+    siglip = FakeSiglip(text_to_vec={"couch": np.array([1, 0, 0, 0], dtype=np.float32)})
+    skills = SpatialSkills(vm, base, siglip, SkillsConfig(locate_threshold=0.3))
+    result = skills.locate("couch")
+    assert result.found is True
+    assert result.xyz is not None
+    # Centroid should land near (5, 1, 3).
+    assert abs(result.xyz[0] - 5.0) < 0.2
+    assert abs(result.xyz[2] - 3.0) < 0.2
+    assert result.confidence > 0.5
+
+
+def test_locate_abstains_below_threshold():
+    """Threshold tuned high enough that an unaligned query returns NOT_FOUND
+    rather than picking a "best of the bad" cluster."""
+    vm = _vm_with_couch_and_chair()
+    base = StubBaseController()
+    # Query embedding orthogonal to both clusters' vectors.
+    siglip = FakeSiglip(text_to_vec={"banana": np.array([0, 0, 1, 0], dtype=np.float32)})
+    skills = SpatialSkills(
+        vm,
+        base,
+        siglip,
+        SkillsConfig(locate_threshold=0.5),
+    )
+    result = skills.locate("banana")
+    assert result.found is False
+    assert result.xyz is None
+
+
+def test_locate_distinguishes_two_clusters():
+    """red-cup / blue-cup style: two clusters present, the query should pick
+    the right one rather than averaging across both."""
+    vm = _vm_with_couch_and_chair()
+    base = StubBaseController()
+    siglip = FakeSiglip(
+        text_to_vec={
+            "couch": np.array([1, 0, 0, 0], dtype=np.float32),
+            "chair": np.array([0, 1, 0, 0], dtype=np.float32),
+        }
+    )
+    skills = SpatialSkills(vm, base, siglip, SkillsConfig(locate_threshold=0.3))
+    couch = skills.locate("couch")
+    chair = skills.locate("chair")
+    assert couch.found and chair.found
+    assert abs(couch.xyz[0] - 5.0) < 0.3
+    assert abs(chair.xyz[0] - (-3.0)) < 0.3
+
+
+def test_locate_returns_not_found_without_siglip():
+    vm = _vm_with_couch_and_chair()
+    skills = SpatialSkills(vm, StubBaseController(), siglip=None)
+    assert skills.locate("anything").found is False
+
+
+def test_locate_returns_not_found_without_features():
+    vm = VoxelMap()
+    rgb = np.full((1, 1, 3), 200, dtype=np.uint8)
+    vm.add(np.zeros((1, 1, 3), dtype=np.float32), rgb, np.ones((1, 1), dtype=np.float32), frame=0, t=0.0)
+    skills = SpatialSkills(vm, StubBaseController(), siglip=FakeSiglip({}))
+    assert skills.locate("anything").found is False
+
+
+# ----- goto() -------------------------------------------------------------
+
+
+def _floor_vm(extent: float = 4.0, y_floor: float = 1.0, voxel_size: float = 0.1) -> VoxelMap:
+    """A clear floor of NAVIGABLE cells spanning [-extent, extent] in both x and z.
+
+    Inputs are float64 to avoid float32 precision drift colliding adjacent
+    voxels at exact cell boundaries (Pi3X-shaped outputs are continuous and
+    don't hit this in practice; this fixture deliberately puts points AT
+    voxel boundaries so we'd quietly merge ~25% of them in float32)."""
+    vm = VoxelMap(voxel_size=voxel_size)
+    pts = []
+    # Offset placement by half a voxel so each xz lands at a cell *centre*,
+    # robust to small float drift.
+    half = voxel_size / 2.0
+    for x in np.arange(-extent + half, extent + half, voxel_size):
+        for z in np.arange(-extent + half, extent + half, voxel_size):
+            pts.append((float(x), y_floor, float(z)))
+    arr = np.asarray(pts, dtype=np.float64).reshape(-1, 1, 3)
+    rgb_arr = np.full((len(pts), 1, 3), 200, dtype=np.uint8)
+    conf_arr = np.ones((len(pts), 1), dtype=np.float32)
+    vm.add(arr, rgb_arr, conf_arr, frame=0, t=0.0)
+    return vm
+
+
+def test_goto_reaches_static_goal():
+    vm = _floor_vm()
+    base = StubBaseController()
+    skills = SpatialSkills(
+        vm,
+        base,
+        cfg=SkillsConfig(
+            cell_size=0.1,
+            obstacle_inflate_cells=0,
+            goto_threshold=0.3,
+            goto_max_steps=400,
+            goto_step_size=0.1,
+        ),
+    )
+    result = skills.goto((2.0, 1.0, 2.0))
+    assert result.reached, f"goto did not reach: {result}"
+    assert result.distance_to_target < 0.3
+    # Should have logged the executed path.
+    assert len(result.path_xyz) > 0
+
+
+def test_goto_blocked_with_wall():
+    """A floor with a wall of obstacle voxels splitting the navigable space.
+    The wall extends past the floor on both ends so there is no corner
+    detour — A* must report no-path."""
+    vm = _floor_vm(extent=2.0)
+    # Vertical wall along x=0 at obstacle height, spanning more z than the
+    # floor so neither end of the wall has a navigable bypass cell.
+    wall_pts = [
+        (0.0, float(y), float(z)) for y in np.arange(0.2, 0.9, 0.1) for z in np.arange(-3.0, 3.0, 0.1)
+    ]
+    arr = np.asarray(wall_pts, dtype=np.float64).reshape(-1, 1, 3)
+    rgb_arr = np.full((len(wall_pts), 1, 3), 200, dtype=np.uint8)
+    conf_arr = np.ones((len(wall_pts), 1), dtype=np.float32)
+    vm.add(arr, rgb_arr, conf_arr, frame=0, t=0.0)
+
+    init = np.eye(4)
+    init[0, 3] = -1.0
+    base = StubBaseController(initial_pose=init)
+    skills = SpatialSkills(
+        vm,
+        base,
+        cfg=SkillsConfig(
+            cell_size=0.1,
+            obstacle_inflate_cells=0,
+            goto_threshold=0.2,
+            goto_max_steps=200,
+        ),
+    )
+    result = skills.goto((1.0, 1.0, 0.0))
+    assert result.reached is False
+    assert result.reason == "no path"
+
+
+def test_goto_stops_when_already_at_goal():
+    vm = _floor_vm()
+    init = np.eye(4)
+    init[0, 3] = 0.5
+    base = StubBaseController(initial_pose=init)
+    skills = SpatialSkills(vm, base, cfg=SkillsConfig(goto_threshold=0.5))
+    result = skills.goto((0.5, 0.0, 0.0))
+    assert result.reached and result.n_steps == 0
+
+
+# ----- explore() ----------------------------------------------------------
+
+
+def test_explore_returns_a_frontier_when_one_exists():
+    # Build a small floor and let project_voxel_map_to_grid pad the bbox so
+    # there's UNOBSERVED space around it.
+    vm = _floor_vm(extent=1.0)
+    base = StubBaseController()
+    skills = SpatialSkills(
+        vm,
+        base,
+        cfg=SkillsConfig(
+            cell_size=0.1,
+            obstacle_inflate_cells=0,
+        ),
+    )
+    result = skills.explore()
+    assert result.found_frontier
+    assert result.target_xyz is not None
+
+
+def test_explore_reports_no_frontier_on_empty_voxelmap():
+    vm = VoxelMap()
+    skills = SpatialSkills(vm, StubBaseController(), cfg=SkillsConfig())
+    result = skills.explore()
+    assert result.found_frontier is False
+    assert result.target_xyz is None
+
+
+def test_explore_target_distance_matches_pose():
+    vm = _floor_vm(extent=1.0)
+    init = np.eye(4)
+    init[0, 3] = 0.3
+    init[2, 3] = -0.4
+    base = StubBaseController(initial_pose=init)
+    skills = SpatialSkills(vm, base, cfg=SkillsConfig(cell_size=0.1, obstacle_inflate_cells=0))
+    result = skills.explore()
+    if result.target_xyz is not None:
+        d = math.hypot(result.target_xyz[0] - 0.3, result.target_xyz[2] - (-0.4))
+        assert abs(d - result.distance_to_target) < 1e-3

tests/navigation/test_value_map.py

@@ -0,0 +1,211 @@
+#!/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.
+
+"""Unit tests for the B2-full value-map exploration."""
+
+# ruff: noqa: N803, N806  — D: conventional feature-dimension name
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+from lerobot.navigation.occupancy import (
+    NAVIGABLE,
+    UNOBSERVED,
+    OccupancyGrid,
+    project_voxel_map_to_grid,
+)
+from lerobot.navigation.value_map import (
+    ValueMapConfig,
+    compute_value_maps,
+    pick_best_frontier_cell,
+)
+from lerobot.navigation.voxel_map import VoxelMap
+
+
+def _vm_from(points, *, voxel_size=0.1, t0=0.0, dt=0.0, features=None):
+    vm = VoxelMap(voxel_size=voxel_size)
+    rgb = np.full((1, 1, 3), 200, dtype=np.uint8)
+    conf = np.ones((1, 1), dtype=np.float32)
+    for i, p in enumerate(points):
+        pt = np.array([[[p[0], p[1], p[2]]]], dtype=np.float64)
+        feat = None
+        if features is not None:
+            feat = features[i].reshape(1, 1, -1).astype(np.float16)
+        vm.add(pt, rgb, conf, frame=i, t=t0 + i * dt, feat_map=feat)
+    return vm
+
+
+def _grid_around(vm: VoxelMap, cell_size: float = 0.5) -> OccupancyGrid:
+    return project_voxel_map_to_grid(vm, cell_size=cell_size, inflate_cells=0)
+
+
+def test_recency_high_for_unobserved_cells():
+    """Cells with no voxel projection should default to unknown_value."""
+    vm = _vm_from([(0.0, 1.0, 0.0)])
+    grid = _grid_around(vm)
+    cfg = ValueMapConfig(unknown_value=0.95)
+    vm_values = compute_value_maps(vm, grid, cfg=cfg)
+    # The voxel only fills one cell; the rest should be unknown.
+    unobs = grid.classes == UNOBSERVED
+    assert unobs.any()
+    np.testing.assert_allclose(vm_values.recency[unobs], 0.95, atol=1e-6)
+
+
+def test_recency_drops_for_recent_observation():
+    """A freshly-observed cell scores LOW on V_T (recency)."""
+    vm = _vm_from([(0.0, 1.0, 0.0)], t0=100.0)
+    grid = _grid_around(vm)
+    cfg = ValueMapConfig(recency_mid_s=10.0, recency_scale_s=3.0, unknown_value=1.0)
+    # now_t == t0 → age = 0 → sigmoid((0 - 10) / 3) ≈ 0.04
+    values = compute_value_maps(vm, grid, now_t=100.0, cfg=cfg)
+    # Find the cell that received the voxel.
+    obs_mask = values.last_time > -math.inf
+    assert obs_mask.any()
+    assert values.recency[obs_mask].max() < 0.1
+
+
+def test_recency_grows_with_age():
+    vm = _vm_from([(0.0, 1.0, 0.0)], t0=0.0)
+    grid = _grid_around(vm)
+    cfg = ValueMapConfig(recency_mid_s=10.0, recency_scale_s=3.0)
+    # 30 seconds later — V_T should be near 1.
+    values = compute_value_maps(vm, grid, now_t=30.0, cfg=cfg)
+    obs_mask = values.last_time > -math.inf
+    assert values.recency[obs_mask].max() > 0.9
+
+
+def test_similarity_high_for_matching_query():
+    D = 8
+    feat_couch = np.eye(D)[0]
+    vm = _vm_from(
+        [(0.0, 1.0, 0.0)],
+        features=[feat_couch],
+    )
+    grid = _grid_around(vm)
+    text_emb = np.eye(D)[0]  # same direction as couch
+    cfg = ValueMapConfig(similarity_mid=0.15, similarity_scale=0.05)
+    values = compute_value_maps(vm, grid, text_emb=text_emb, cfg=cfg)
+    assert values.similarity is not None
+    assert values.similarity.max() > 0.95
+
+
+def test_similarity_low_for_orthogonal_query():
+    D = 8
+    feat_couch = np.eye(D)[0]
+    vm = _vm_from([(0.0, 1.0, 0.0)], features=[feat_couch])
+    grid = _grid_around(vm)
+    text_emb = np.eye(D)[3]  # orthogonal
+    values = compute_value_maps(vm, grid, text_emb=text_emb)
+    assert values.similarity is not None
+    # Cells with content but no match: low V_S.
+    has_voxel = values.last_time > -math.inf
+    assert values.similarity[has_voxel].max() < 0.1
+
+
+def test_similarity_is_none_when_no_query():
+    vm = _vm_from([(0.0, 1.0, 0.0)])
+    grid = _grid_around(vm)
+    values = compute_value_maps(vm, grid)
+    assert values.similarity is None
+    np.testing.assert_array_equal(values.combined, values.recency)
+
+
+def test_combined_balances_recency_and_similarity():
+    D = 8
+    # Two voxels with different features: one matches query, one doesn't.
+    feats = [np.eye(D)[0], np.eye(D)[3]]
+    vm = _vm_from(
+        [(0.0, 1.0, 0.0), (2.0, 1.0, 0.0)],
+        features=feats,
+        t0=0.0,
+    )
+    grid = _grid_around(vm, cell_size=0.5)
+    cfg = ValueMapConfig(alpha_similarity=0.7, recency_mid_s=5.0, recency_scale_s=2.0)
+    text_emb = np.eye(D)[0]
+    values = compute_value_maps(vm, grid, text_emb=text_emb, now_t=0.0, cfg=cfg)
+
+    # Cell with matching feature should have HIGHER combined value than the
+    # non-matching observed cell at the same age.
+    snap_xyz = vm.snapshot().xyz
+    iz_m = int((snap_xyz[0, 2] - grid.origin_z) / grid.cell_size)
+    ix_m = int((snap_xyz[0, 0] - grid.origin_x) / grid.cell_size)
+    iz_n = int((snap_xyz[1, 2] - grid.origin_z) / grid.cell_size)
+    ix_n = int((snap_xyz[1, 0] - grid.origin_x) / grid.cell_size)
+    assert values.combined[iz_m, ix_m] > values.combined[iz_n, ix_n]
+
+
+def test_pick_best_frontier_prefers_high_value_cell():
+    classes = np.full((6, 6), UNOBSERVED, dtype=np.int8)
+    classes[1:5, 1:5] = NAVIGABLE
+    grid = OccupancyGrid(classes=classes, cell_size=0.5, origin_x=0.0, origin_z=0.0, ground_y=0.0)
+    frontier_cells = np.array([[1, 1], [4, 4]], dtype=np.int32)
+    from lerobot.navigation.value_map import ValueMaps
+
+    # Make cell (4, 4) more valuable than (1, 1).
+    combined = np.zeros((6, 6), dtype=np.float32)
+    combined[1, 1] = 0.2
+    combined[4, 4] = 0.9
+    values = ValueMaps(
+        last_time=np.full((6, 6), -math.inf),
+        recency=combined.copy(),
+        similarity=None,
+        combined=combined,
+    )
+    best_idx, (x, z), d, score = pick_best_frontier_cell(
+        grid,
+        frontier_cells,
+        values,
+        robot_position_xz=(0.0, 0.0),
+        cfg=ValueMapConfig(distance_discount_per_meter=0.0),
+    )
+    assert best_idx == 1
+    assert score > 0.8
+
+
+def test_distance_discount_prefers_closer_when_values_equal():
+    classes = np.full((6, 6), UNOBSERVED, dtype=np.int8)
+    classes[0:6, 0:6] = NAVIGABLE
+    grid = OccupancyGrid(classes=classes, cell_size=1.0, origin_x=0.0, origin_z=0.0, ground_y=0.0)
+    frontier_cells = np.array([[0, 0], [5, 5]], dtype=np.int32)
+    from lerobot.navigation.value_map import ValueMaps
+
+    same = np.ones((6, 6), dtype=np.float32)
+    values = ValueMaps(
+        last_time=np.full((6, 6), -math.inf),
+        recency=same.copy(),
+        similarity=None,
+        combined=same,
+    )
+    best_idx, _, _, _ = pick_best_frontier_cell(
+        grid,
+        frontier_cells,
+        values,
+        robot_position_xz=(0.0, 0.0),
+        cfg=ValueMapConfig(distance_discount_per_meter=0.5),
+    )
+    # Robot at origin → (0, 0) is closer than (5, 5).
+    assert best_idx == 0
+
+
+def test_compute_value_maps_with_empty_voxel_map_returns_unknown():
+    vm = VoxelMap()
+    classes = np.full((4, 4), UNOBSERVED, dtype=np.int8)
+    grid = OccupancyGrid(classes=classes, cell_size=1.0, origin_x=0.0, origin_z=0.0, ground_y=0.0)
+    values = compute_value_maps(vm, grid)
+    np.testing.assert_allclose(values.recency, 1.0)
+    assert values.similarity is None

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

tests/navigation/test_voxel_map.py

@@ -0,0 +1,149 @@
+#!/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.
+
+"""Unit tests for ``VoxelMap`` — geometry-only behavior (M3)."""
+
+# ruff: noqa: N803, N806  — H, W, D: conventional array-dimension names
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from lerobot.navigation.voxel_map import VoxelMap
+
+
+def _scatter(points: np.ndarray, color: tuple[int, int, int] = (200, 100, 50)) -> tuple[np.ndarray, ...]:
+    """Helper: build (points, rgb, conf) arrays from a list of xyz coords."""
+    rgb = np.tile(np.array(color, dtype=np.uint8), (len(points), 1))
+    conf = np.ones(len(points), dtype=np.float32)
+    return points.astype(np.float32), rgb, conf
+
+
+def test_initially_empty():
+    vm = VoxelMap(voxel_size=0.05)
+    assert len(vm) == 0
+    snap = vm.snapshot()
+    assert snap.xyz.shape == (0, 3)
+    assert snap.rgb.shape == (0, 3)
+
+
+def test_voxel_size_must_be_positive():
+    with pytest.raises(ValueError):
+        VoxelMap(voxel_size=0.0)
+    with pytest.raises(ValueError):
+        VoxelMap(voxel_size=-0.1)
+
+
+def test_single_point_creates_one_voxel():
+    vm = VoxelMap(voxel_size=0.1)
+    pts, rgb, conf = _scatter(np.array([[0.123, 0.456, 0.789]]))
+    stats = vm.add(pts, rgb, conf, frame=0, t=0.0)
+    assert stats == type(stats)(n_voxels=1, n_added=1, n_updated=0)
+    assert len(vm) == 1
+    snap = vm.snapshot()
+    np.testing.assert_allclose(snap.xyz[0], [0.123, 0.456, 0.789], atol=1e-5)
+
+
+def test_points_in_same_voxel_collapse_and_average():
+    vm = VoxelMap(voxel_size=0.1)
+    # Two points inside the voxel [0.0, 0.1) on each axis.
+    pts, rgb, conf = _scatter(np.array([[0.01, 0.01, 0.01], [0.09, 0.09, 0.09]]))
+    stats = vm.add(pts, rgb, conf, frame=0, t=0.0)
+    assert stats.n_voxels == 1
+    assert stats.n_added == 1
+    snap = vm.snapshot()
+    np.testing.assert_allclose(snap.xyz[0], [0.05, 0.05, 0.05], atol=1e-5)
+    assert int(snap.count[0]) == 2
+
+
+def test_second_keyframe_updates_running_mean():
+    vm = VoxelMap(voxel_size=0.1)
+    pts1, rgb1, conf1 = _scatter(np.array([[0.02, 0.02, 0.02]]), color=(100, 100, 100))
+    vm.add(pts1, rgb1, conf1, frame=0, t=0.0)
+    pts2, rgb2, conf2 = _scatter(np.array([[0.08, 0.08, 0.08]]), color=(200, 200, 200))
+    stats = vm.add(pts2, rgb2, conf2, frame=1, t=0.5)
+    assert stats.n_added == 0
+    assert stats.n_updated == 1
+    snap = vm.snapshot()
+    # Mean position = (0.02 + 0.08) / 2 = 0.05; mean color = 150.
+    np.testing.assert_allclose(snap.xyz[0], [0.05, 0.05, 0.05], atol=1e-5)
+    np.testing.assert_allclose(snap.rgb[0], [150, 150, 150], atol=1)
+    assert int(snap.last_frame[0]) == 1
+    assert float(snap.last_time[0]) == pytest.approx(0.5)
+
+
+def test_conf_gate_drops_low_confidence_pixels():
+    vm = VoxelMap(voxel_size=0.1)
+    pts = np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]], dtype=np.float32)
+    rgb = np.array([[10, 20, 30], [40, 50, 60]], dtype=np.uint8)
+    conf = np.array([0.9, 0.1], dtype=np.float32)
+    stats = vm.add(pts, rgb, conf, frame=0, t=0.0, conf_thresh=0.5)
+    assert stats.n_voxels == 1
+    snap = vm.snapshot()
+    np.testing.assert_allclose(snap.xyz[0], [0.0, 0.0, 0.0], atol=1e-5)
+    np.testing.assert_allclose(snap.rgb[0], [10, 20, 30], atol=1)
+
+
+def test_quantization_negative_coordinates():
+    vm = VoxelMap(voxel_size=0.1)
+    pts, rgb, conf = _scatter(np.array([[-0.05, -0.15, -0.25]]))
+    stats = vm.add(pts, rgb, conf, frame=0, t=0.0)
+    assert stats.n_voxels == 1
+    # floor(-0.05/0.1) = floor(-0.5) = -1 (voxel covers [-0.1, 0.0)).
+    # Just check the mean equals the input single point.
+    snap = vm.snapshot()
+    np.testing.assert_allclose(snap.xyz[0], [-0.05, -0.15, -0.25], atol=1e-5)
+
+
+def test_image_shaped_input_is_flattened():
+    """``add`` accepts (H, W, 3) point arrays — typical Pi3X output shape."""
+    vm = VoxelMap(voxel_size=0.5)
+    H, W = 4, 4
+    pts = np.zeros((H, W, 3), dtype=np.float32)
+    pts[..., 0] = np.linspace(0, 5, W)[None, :]  # 16 unique x values? No, 4.
+    rgb = np.full((H, W, 3), 128, dtype=np.uint8)
+    conf = np.ones((H, W), dtype=np.float32)
+    stats = vm.add(pts, rgb, conf, frame=0, t=0.0)
+    # All x values land in voxel slots at 0, 1, 2, 3, 4 (different voxels at
+    # 0.5m size); each row of the image contributes the same 4 unique voxels,
+    # collapsed within the keyframe — but actually voxel indices depend on x.
+    # The point of this test is just that flattening works without raising.
+    assert stats.n_voxels >= 1
+    assert stats.n_voxels <= H * W
+
+
+def test_nonfinite_points_are_dropped():
+    vm = VoxelMap(voxel_size=0.1)
+    pts = np.array(
+        [[0.0, 0.0, 0.0], [np.nan, 0.0, 0.0], [0.0, np.inf, 0.0]],
+        dtype=np.float32,
+    )
+    rgb = np.full((3, 3), 100, dtype=np.uint8)
+    conf = np.ones(3, dtype=np.float32)
+    stats = vm.add(pts, rgb, conf, frame=0, t=0.0)
+    assert stats.n_voxels == 1
+
+
+def test_snapshot_rgb_clipped_to_uint8():
+    """If RGB sums accumulate to > 255 per channel, the snapshot mean is still
+    a clean uint8."""
+    vm = VoxelMap(voxel_size=0.1)
+    pts, rgb, conf = _scatter(np.array([[0.0, 0.0, 0.0]]), color=(255, 255, 255))
+    vm.add(pts, rgb, conf, frame=0, t=0.0)
+    vm.add(pts, rgb, conf, frame=1, t=0.5)
+    snap = vm.snapshot()
+    assert snap.rgb.dtype == np.uint8
+    np.testing.assert_array_equal(snap.rgb[0], [255, 255, 255])

tests/navigation/test_voxel_map_carve.py

@@ -0,0 +1,217 @@
+#!/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.
+
+"""Unit tests for ``VoxelMap.carve`` — DynaMem-style free-space removal (M4)."""
+
+# ruff: noqa: N803, N806  — H, W, D: conventional array-dimension names
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import pytest
+
+from lerobot.navigation.voxel_map import CarveResult, VoxelMap
+
+
+def _identity_pose() -> np.ndarray:
+    """Camera-to-world identity (camera == world)."""
+    return np.eye(4, dtype=np.float64)
+
+
+def _shifted_pose(tx: float = 0.0, ty: float = 0.0, tz: float = 0.0) -> np.ndarray:
+    p = np.eye(4, dtype=np.float64)
+    p[:3, 3] = (tx, ty, tz)
+    return p
+
+
+def _solid_depth_view(H: int, W: int, depth: float, conf: float = 1.0) -> tuple[np.ndarray, np.ndarray]:
+    """A synthetic Pi3X view: every pixel sees a surface at ``depth`` meters."""
+    fov_deg = 90.0
+    focal = W / (2.0 * math.tan(math.radians(fov_deg) / 2.0))
+    cx, cy = (W - 1) / 2.0, (H - 1) / 2.0
+    us, vs = np.meshgrid(np.arange(W), np.arange(H))
+    # local_points[v,u] = depth * (K^-1 @ [u,v,1]); z = depth.
+    x = (us - cx) * depth / focal
+    y = (vs - cy) * depth / focal
+    z = np.full_like(x, depth, dtype=np.float64)
+    local = np.stack([x, y, z], axis=-1).astype(np.float32)
+    return local, np.full((H, W), conf, dtype=np.float32), focal
+
+
+def _seed_voxel(vm: VoxelMap, xyz: tuple[float, float, float], frame: int = 0) -> None:
+    pts = np.asarray([xyz], dtype=np.float32)
+    rgb = np.full((1, 3), 200, dtype=np.uint8)
+    conf = np.ones(1, dtype=np.float32)
+    vm.add(pts, rgb, conf, frame=frame, t=float(frame) * 0.5)
+
+
+# ---------------------------------------------------------------------------
+
+
+def test_empty_map_returns_zero():
+    vm = VoxelMap(voxel_size=0.1)
+    local, conf, focal = _solid_depth_view(32, 32, depth=5.0)
+    result = vm.carve(local, conf, _identity_pose(), focal_px=focal, frame=0, t=0.0)
+    assert isinstance(result, CarveResult)
+    assert result.n_removed == 0
+    assert result.removed_xyz.shape == (0, 3)
+
+
+def test_voxel_in_front_of_surface_is_removed():
+    """Voxel at z=2 m, observed surface at z=5 m, margin 0.05: free space."""
+    vm = VoxelMap(voxel_size=0.1)
+    _seed_voxel(vm, (0.0, 0.0, 2.0))
+    assert len(vm) == 1
+    local, conf, focal = _solid_depth_view(32, 32, depth=5.0)
+    result = vm.carve(local, conf, _identity_pose(), focal_px=focal, frame=1, t=0.5)
+    assert result.n_removed == 1
+    assert len(vm) == 0
+    assert result.removed_xyz.shape == (1, 3)
+
+
+def test_voxel_at_surface_is_kept():
+    """Voxel at z = 5.0 m, surface at 5.0 m — d is NOT < D - margin."""
+    vm = VoxelMap(voxel_size=0.1)
+    _seed_voxel(vm, (0.0, 0.0, 5.0))
+    local, conf, focal = _solid_depth_view(32, 32, depth=5.0)
+    result = vm.carve(local, conf, _identity_pose(), focal_px=focal, frame=1, t=0.5, margin=0.05)
+    assert result.n_removed == 0
+    assert len(vm) == 1
+
+
+def test_voxel_behind_surface_is_kept():
+    """Voxel at z=10 m, surface at z=5 m: voxel is occluded, NOT free space."""
+    vm = VoxelMap(voxel_size=0.1)
+    _seed_voxel(vm, (0.0, 0.0, 10.0))
+    local, conf, focal = _solid_depth_view(32, 32, depth=5.0)
+    result = vm.carve(local, conf, _identity_pose(), focal_px=focal, frame=1, t=0.5)
+    assert result.n_removed == 0
+
+
+def test_voxel_behind_camera_is_kept():
+    """A voxel with z<=0 in the camera frame can't be seen — must not be carved."""
+    vm = VoxelMap(voxel_size=0.1)
+    _seed_voxel(vm, (0.0, 0.0, -1.0))
+    local, conf, focal = _solid_depth_view(32, 32, depth=5.0)
+    result = vm.carve(local, conf, _identity_pose(), focal_px=focal, frame=1, t=0.5)
+    assert result.n_removed == 0
+
+
+def test_voxel_out_of_view_is_kept():
+    """A voxel inside no pixel's frustum should be left alone."""
+    vm = VoxelMap(voxel_size=0.1)
+    # Way off to the side — projects outside the 32x32 image at 90° HFOV.
+    _seed_voxel(vm, (50.0, 0.0, 2.0))
+    local, conf, focal = _solid_depth_view(32, 32, depth=5.0)
+    result = vm.carve(local, conf, _identity_pose(), focal_px=focal, frame=1, t=0.5)
+    assert result.n_removed == 0
+
+
+def test_low_confidence_blocks_carve():
+    """If conf at the projected pixel is below threshold, don't remove."""
+    vm = VoxelMap(voxel_size=0.1)
+    _seed_voxel(vm, (0.0, 0.0, 2.0))
+    local, conf, focal = _solid_depth_view(32, 32, depth=5.0, conf=0.2)
+    result = vm.carve(local, conf, _identity_pose(), focal_px=focal, frame=1, t=0.5, conf_thresh=0.5)
+    assert result.n_removed == 0
+
+
+def test_invalid_depth_blocks_carve():
+    """NaN / non-positive depth at the projected pixel must not trigger carve."""
+    vm = VoxelMap(voxel_size=0.1)
+    _seed_voxel(vm, (0.0, 0.0, 2.0))
+    local, conf, focal = _solid_depth_view(32, 32, depth=5.0)
+    local[:, :, 2] = np.nan
+    result = vm.carve(local, conf, _identity_pose(), focal_px=focal, frame=1, t=0.5)
+    assert result.n_removed == 0
+
+
+def test_margin_protects_near_surface():
+    """A voxel 3 cm in front of a 5 m surface, margin 5 cm: NOT carved."""
+    vm = VoxelMap(voxel_size=0.1)
+    _seed_voxel(vm, (0.0, 0.0, 5.0 - 0.03))
+    local, conf, focal = _solid_depth_view(32, 32, depth=5.0)
+    result = vm.carve(local, conf, _identity_pose(), focal_px=focal, frame=1, t=0.5, margin=0.05)
+    assert result.n_removed == 0
+
+
+def test_margin_zero_carves_near_surface():
+    """Same setup, margin 0: now the 3 cm gap counts as free space."""
+    vm = VoxelMap(voxel_size=0.1)
+    _seed_voxel(vm, (0.0, 0.0, 5.0 - 0.03))
+    local, conf, focal = _solid_depth_view(32, 32, depth=5.0)
+    result = vm.carve(local, conf, _identity_pose(), focal_px=focal, frame=1, t=0.5, margin=0.0)
+    assert result.n_removed == 1
+
+
+def test_carve_compacts_arrays_and_lookup():
+    """After carving some but not all voxels, internal storage stays consistent."""
+    vm = VoxelMap(voxel_size=0.1)
+    # Three voxels: in front (will be carved), at surface (kept), to the side (kept).
+    _seed_voxel(vm, (0.0, 0.0, 2.0))
+    _seed_voxel(vm, (0.0, 0.0, 5.0))
+    _seed_voxel(vm, (50.0, 0.0, 5.0))
+    assert len(vm) == 3
+    local, conf, focal = _solid_depth_view(32, 32, depth=5.0)
+    result = vm.carve(local, conf, _identity_pose(), focal_px=focal, frame=1, t=0.5)
+    assert result.n_removed == 1
+    assert len(vm) == 2
+    # All internal arrays must agree on the new length.
+    assert vm._idx.shape == (2, 3)  # noqa: SLF001
+    assert vm._count.shape == (2,)  # noqa: SLF001
+    assert len(vm._lookup) == 2  # noqa: SLF001
+    # Lookup rows must point at valid indices in the new arrays.
+    for row in vm._lookup.values():  # noqa: SLF001
+        assert 0 <= row < 2
+    # A subsequent `add` to a removed voxel must work cleanly (revives it).
+    pts = np.asarray([[0.0, 0.0, 2.0]], dtype=np.float32)
+    rgb = np.full((1, 3), 100, dtype=np.uint8)
+    conf1 = np.ones(1, dtype=np.float32)
+    add_stats = vm.add(pts, rgb, conf1, frame=2, t=1.0)
+    assert add_stats.n_added == 1
+    assert len(vm) == 3
+
+
+def test_carve_shape_validation():
+    vm = VoxelMap(voxel_size=0.1)
+    _seed_voxel(vm, (0.0, 0.0, 2.0))
+    bad_local = np.zeros((32, 32, 2), dtype=np.float32)
+    conf = np.ones((32, 32), dtype=np.float32)
+    with pytest.raises(ValueError, match=r"H, W, 3"):
+        vm.carve(bad_local, conf, _identity_pose(), focal_px=16.0, frame=0, t=0.0)
+
+    local = np.zeros((32, 32, 3), dtype=np.float32)
+    bad_conf = np.ones((16, 32), dtype=np.float32)
+    with pytest.raises(ValueError, match=r"conf shape"):
+        vm.carve(local, bad_conf, _identity_pose(), focal_px=16.0, frame=0, t=0.0)
+
+    bad_pose = np.eye(3)
+    with pytest.raises(ValueError, match=r"pose must be"):
+        vm.carve(local, conf, bad_pose, focal_px=16.0, frame=0, t=0.0)
+
+
+def test_carve_with_translated_camera():
+    """If the camera has moved, world-space voxels must be transformed correctly
+    before the free-space test."""
+    vm = VoxelMap(voxel_size=0.1)
+    # Voxel at world position (10, 0, 2). With camera at world (10, 0, 0),
+    # the voxel is 2 m in front of the camera's local +Z; surface at 5 m → carve.
+    _seed_voxel(vm, (10.0, 0.0, 2.0))
+    pose = _shifted_pose(tx=10.0, ty=0.0, tz=0.0)
+    local, conf, focal = _solid_depth_view(32, 32, depth=5.0)
+    result = vm.carve(local, conf, pose, focal_px=focal, frame=1, t=0.5)
+    assert result.n_removed == 1

tests/navigation/test_voxel_map_mutation.py

@@ -0,0 +1,99 @@
+#!/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.
+
+"""Unit tests for the C2 scene-mutation helper on VoxelMap."""
+
+from __future__ import annotations
+
+import numpy as np
+
+from lerobot.navigation.voxel_map import VoxelMap
+
+
+def test_remove_voxels_in_box_zero_when_empty():
+    vm = VoxelMap()
+    assert vm.remove_voxels_in_box((-1, -1, -1), (1, 1, 1)) == 0
+
+
+def test_remove_voxels_in_box_only_deletes_inside():
+    vm = VoxelMap(voxel_size=0.1)
+    pts = np.array(
+        [
+            [[0.0, 0.0, 0.0]],  # inside
+            [[0.1, 0.0, 0.0]],  # inside
+            [[2.0, 0.0, 0.0]],  # outside
+        ],
+        dtype=np.float64,
+    )
+    rgb = np.full((3, 1, 3), 200, dtype=np.uint8)
+    conf = np.ones((3, 1), dtype=np.float32)
+    vm.add(pts, rgb, conf, frame=0, t=0.0)
+    assert len(vm) == 3
+
+    n = vm.remove_voxels_in_box((-0.05, -0.05, -0.05), (0.15, 0.05, 0.05))
+    assert n == 2
+    assert len(vm) == 1
+    snap = vm.snapshot()
+    np.testing.assert_allclose(snap.xyz[0], [2.0, 0.0, 0.0], atol=1e-3)
+
+
+def test_remove_voxels_keeps_feature_arrays_aligned():
+    vm = VoxelMap(voxel_size=0.1)
+    pts = np.array([[[0.0, 0.0, 0.0]], [[2.0, 0.0, 0.0]]], dtype=np.float64)
+    rgb = np.full((2, 1, 3), 200, dtype=np.uint8)
+    conf = np.ones((2, 1), dtype=np.float32)
+    feat = np.array([[[1.0, 0.0, 0.0, 0.0]], [[0.0, 1.0, 0.0, 0.0]]], dtype=np.float16)
+    vm.add(pts, rgb, conf, frame=0, t=0.0, feat_map=feat)
+    assert len(vm) == 2
+    assert vm._feat_sum.shape == (2, 4)  # noqa: SLF001
+
+    vm.remove_voxels_in_box((-0.5, -0.5, -0.5), (0.5, 0.5, 0.5))
+    assert len(vm) == 1
+    # Feature arrays now have length 1 — matches _count.
+    assert vm._feat_sum.shape == (1, 4)  # noqa: SLF001
+    snap = vm.snapshot(include_features=True)
+    # Surviving voxel had vector [0, 1, 0, 0]; normalized stays the same.
+    assert snap.feat is not None
+    np.testing.assert_allclose(
+        snap.feat[0].astype(np.float32),
+        [0.0, 1.0, 0.0, 0.0],
+        atol=1e-3,
+    )
+
+
+def test_remove_voxels_compacts_lookup():
+    """After deletion the dict→row map must still point to the right rows."""
+    vm = VoxelMap(voxel_size=0.1)
+    pts = np.array(
+        [[[0.0, 0.0, 0.0]], [[2.0, 0.0, 0.0]], [[4.0, 0.0, 0.0]]],
+        dtype=np.float64,
+    )
+    rgb = np.full((3, 1, 3), 200, dtype=np.uint8)
+    conf = np.ones((3, 1), dtype=np.float32)
+    vm.add(pts, rgb, conf, frame=0, t=0.0)
+
+    vm.remove_voxels_in_box((1.5, -0.5, -0.5), (2.5, 0.5, 0.5))  # delete the middle
+    assert len(vm) == 2
+
+    # Adding a new voxel at one of the surviving positions should UPDATE
+    # (not append), which only works if the lookup row indices are correct.
+    new_pt = np.array([[[0.0, 0.0, 0.0]]], dtype=np.float64)
+    stats = vm.add(
+        new_pt, np.full((1, 1, 3), 50, dtype=np.uint8), np.ones((1, 1), dtype=np.float32), frame=1, t=1.0
+    )
+    assert stats.n_added == 0
+    assert stats.n_updated == 1
+    assert len(vm) == 2

tests/robots/test_unitree_go2.py

@@ -0,0 +1,202 @@
+#!/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.
+
+"""Tests for the Unitree Go2 robot.
+
+The SDK is only imported inside ``UnitreeGo2.connect()``, so everything
+here runs without unitree_sdk2py installed: the sport client, state
+subscriber and video client are replaced with mocks.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import cv2
+import numpy as np
+import pytest
+
+from lerobot.robots.unitree_go2 import UnitreeGo2, UnitreeGo2Config
+
+# ---------------------------------------------------------------------------
+# Config (no SDK needed)
+# ---------------------------------------------------------------------------
+
+
+class TestUnitreeGo2Config:
+    def test_registered_type_name(self):
+        assert UnitreeGo2Config().type == "unitree_go2"
+
+    def test_default_config(self):
+        cfg = UnitreeGo2Config()
+        assert cfg.domain_id == 0
+        assert cfg.use_front_camera is True
+        assert cfg.stand_on_connect is True
+        assert cfg.cameras == {}
+
+    def test_safety_clamps_are_positive(self):
+        cfg = UnitreeGo2Config()
+        assert cfg.max_x_vel > 0
+        assert cfg.max_y_vel > 0
+        assert cfg.max_theta_vel > 0
+
+
+# ---------------------------------------------------------------------------
+# Features (no SDK needed)
+# ---------------------------------------------------------------------------
+
+
+def _make_robot(**cfg_kwargs) -> UnitreeGo2:
+    cfg = UnitreeGo2Config(id="test_go2", **cfg_kwargs)
+    return UnitreeGo2(cfg)
+
+
+class TestFeatures:
+    def test_action_features(self):
+        robot = _make_robot()
+        assert robot.action_features == {"x.vel": float, "y.vel": float, "theta.vel": float}
+
+    def test_observation_features_with_front_camera(self):
+        robot = _make_robot()
+        ft = robot.observation_features
+        assert ft["front"] == (720, 1280, 3)
+        for key in ("x.pos", "y.pos", "theta.pos", "x.vel", "y.vel", "theta.vel"):
+            assert ft[key] is float
+
+    def test_observation_features_without_front_camera(self):
+        robot = _make_robot(use_front_camera=False)
+        assert "front" not in robot.observation_features
+
+    def test_features_available_before_connect(self):
+        robot = _make_robot()
+        assert not robot.is_connected
+        assert robot.observation_features
+        assert robot.action_features
+
+    def test_is_calibrated_always_true(self):
+        assert _make_robot().is_calibrated is True
+
+
+# ---------------------------------------------------------------------------
+# I/O with mocked SDK handles
+# ---------------------------------------------------------------------------
+
+
+def _connected_robot(**cfg_kwargs) -> UnitreeGo2:
+    """A robot with mocked SDK handles, as if connect() had run."""
+    robot = _make_robot(**cfg_kwargs)
+    robot._sport = MagicMock()
+    robot._video = MagicMock()
+    robot._connected = True
+    return robot
+
+
+def _fake_state(x=0.0, y=0.0, yaw=0.0, vx=0.0, vy=0.0, yaw_speed=0.0):
+    return SimpleNamespace(
+        position=[x, y, 0.0],
+        velocity=[vx, vy, 0.0],
+        yaw_speed=yaw_speed,
+        imu_state=SimpleNamespace(rpy=[0.0, 0.0, yaw]),
+    )
+
+
+class TestSendAction:
+    def test_action_reaches_sport_move(self):
+        robot = _connected_robot()
+        sent = robot.send_action({"x.vel": 0.3, "y.vel": -0.1, "theta.vel": 0.5})
+        robot._sport.Move.assert_called_once_with(0.3, -0.1, 0.5)
+        assert sent == {"x.vel": 0.3, "y.vel": -0.1, "theta.vel": 0.5}
+
+    def test_action_is_clamped(self):
+        robot = _connected_robot(max_x_vel=0.5, max_y_vel=0.2, max_theta_vel=1.0)
+        sent = robot.send_action({"x.vel": 9.0, "y.vel": -9.0, "theta.vel": -9.0})
+        robot._sport.Move.assert_called_once_with(0.5, -0.2, -1.0)
+        assert sent == {"x.vel": 0.5, "y.vel": -0.2, "theta.vel": -1.0}
+
+    def test_missing_keys_default_to_zero(self):
+        robot = _connected_robot()
+        sent = robot.send_action({})
+        robot._sport.Move.assert_called_once_with(0.0, 0.0, 0.0)
+        assert sent == {"x.vel": 0.0, "y.vel": 0.0, "theta.vel": 0.0}
+
+    def test_raises_when_not_connected(self):
+        robot = _make_robot()
+        with pytest.raises(ConnectionError):
+            robot.send_action({"x.vel": 0.1})
+
+
+class TestGetObservation:
+    def test_odometry_fields(self):
+        robot = _connected_robot(use_front_camera=False)
+        robot._latest_state = _fake_state(x=1.0, y=2.0, yaw=0.3, vx=0.1, vy=-0.05, yaw_speed=0.2)
+        obs = robot.get_observation()
+        assert obs["x.pos"] == pytest.approx(1.0)
+        assert obs["y.pos"] == pytest.approx(2.0)
+        assert obs["theta.pos"] == pytest.approx(0.3)
+        assert obs["x.vel"] == pytest.approx(0.1)
+        assert obs["y.vel"] == pytest.approx(-0.05)
+        assert obs["theta.vel"] == pytest.approx(0.2)
+
+    def test_odometry_zero_before_first_state(self):
+        robot = _connected_robot(use_front_camera=False)
+        obs = robot.get_observation()
+        assert all(obs[k] == 0.0 for k in robot._odom_ft)
+
+    def test_front_camera_decodes_to_configured_shape(self):
+        robot = _connected_robot(front_camera_width=64, front_camera_height=48)
+        raw = np.full((48, 64, 3), 128, dtype=np.uint8)
+        ok, jpeg = cv2.imencode(".jpg", raw)
+        assert ok
+        robot._video.GetImageSample.return_value = (0, jpeg.tobytes())
+        obs = robot.get_observation()
+        assert obs["front"].shape == (48, 64, 3)
+        assert obs["front"].dtype == np.uint8
+
+    def test_front_camera_resizes_native_frames(self):
+        robot = _connected_robot(front_camera_width=64, front_camera_height=48)
+        native = np.zeros((720, 1280, 3), dtype=np.uint8)
+        ok, jpeg = cv2.imencode(".jpg", native)
+        assert ok
+        robot._video.GetImageSample.return_value = (0, jpeg.tobytes())
+        assert robot.get_observation()["front"].shape == (48, 64, 3)
+
+    def test_front_camera_failure_returns_black_frame(self):
+        robot = _connected_robot(front_camera_width=64, front_camera_height=48)
+        robot._video.GetImageSample.return_value = (1, None)
+        frame = robot.get_observation()["front"]
+        assert frame.shape == (48, 64, 3)
+        assert frame.sum() == 0
+
+    def test_observation_matches_features(self):
+        robot = _connected_robot(front_camera_width=64, front_camera_height=48)
+        raw = np.zeros((48, 64, 3), dtype=np.uint8)
+        _, jpeg = cv2.imencode(".jpg", raw)
+        robot._video.GetImageSample.return_value = (0, jpeg.tobytes())
+        obs = robot.get_observation()
+        assert set(obs.keys()) == set(robot.observation_features.keys())
+
+    def test_raises_when_not_connected(self):
+        robot = _make_robot()
+        with pytest.raises(ConnectionError):
+            robot.get_observation()
+
+
+class TestDisconnect:
+    def test_disconnect_stops_motion(self):
+        robot = _connected_robot()
+        sport = robot._sport
+        robot.disconnect()
+        sport.StopMove.assert_called_once()
+        assert not robot.is_connected