From 8793d1a4d53fb3f7555d32aca4093d5d75d64926 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Mon, 20 Jul 2026 18:05:15 +0200 Subject: [PATCH] navigation: port voxel map + occupancy/A* mapping backbone Copied from the dyna360 research stack into lerobot.navigation: - voxel_map.py: sparse-hash VoxelMap (5 cm default), count-weighted running-mean geometry + optional per-voxel semantic feature; carve() for DynaMem-style free-space removal (dynamic updates), query() for top-k cosine text matches, remove_voxels_in_box() for scene mutation. Added an inverse.reshape(-1) guard for numpy 2.x's np.unique. - occupancy.py: derived 3-class top-down grid (UNOBSERVED/NAVIGABLE/ OBSTACLE) projected from the voxel map, A* with no corner-cutting, obstacle inflation, and frontier extraction for exploration. Added OccupancyGrid.is_obstacle() used by SafeBaseController's occupancy gate. Pure numpy, no torch/hardware/SDK. 42 new tests (71 total in tests/navigation/). Next: LingBot-Map geometry runner + integration. Co-Authored-By: Claude Opus 4.8 --- src/lerobot/navigation/__init__.py | 21 + src/lerobot/navigation/occupancy.py | 371 ++++++++++++++ src/lerobot/navigation/voxel_map.py | 508 ++++++++++++++++++++ tests/navigation/test_occupancy.py | 266 ++++++++++ tests/navigation/test_voxel_map.py | 149 ++++++ tests/navigation/test_voxel_map_carve.py | 217 +++++++++ tests/navigation/test_voxel_map_mutation.py | 99 ++++ 7 files changed, 1631 insertions(+) create mode 100644 src/lerobot/navigation/occupancy.py create mode 100644 src/lerobot/navigation/voxel_map.py create mode 100644 tests/navigation/test_occupancy.py create mode 100644 tests/navigation/test_voxel_map.py create mode 100644 tests/navigation/test_voxel_map_carve.py create mode 100644 tests/navigation/test_voxel_map_mutation.py diff --git a/src/lerobot/navigation/__init__.py b/src/lerobot/navigation/__init__.py index 8b1b8c19c..0db735b22 100644 --- a/src/lerobot/navigation/__init__.py +++ b/src/lerobot/navigation/__init__.py @@ -31,12 +31,33 @@ from .base_controller import ( odometry_to_world_pose, world_velocity_to_body, ) +from .occupancy import ( + NAVIGABLE, + OBSTACLE, + UNOBSERVED, + OccupancyGrid, + astar, + find_frontier_cells, + project_voxel_map_to_grid, +) +from .voxel_map import CarveResult, QueryResult, VoxelMap, VoxelSnapshot __all__ = [ + "NAVIGABLE", + "OBSTACLE", + "UNOBSERVED", "BaseController", + "CarveResult", + "OccupancyGrid", + "QueryResult", "RobotBaseController", "SafeBaseController", "StubBaseController", + "VoxelMap", + "VoxelSnapshot", + "astar", + "find_frontier_cells", "odometry_to_world_pose", + "project_voxel_map_to_grid", "world_velocity_to_body", ] diff --git a/src/lerobot/navigation/occupancy.py b/src/lerobot/navigation/occupancy.py new file mode 100644 index 000000000..cd24eca97 --- /dev/null +++ b/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 diff --git a/src/lerobot/navigation/voxel_map.py b/src/lerobot/navigation/voxel_map.py new file mode 100644 index 000000000..e37fa1e1c --- /dev/null +++ b/src/lerobot/navigation/voxel_map.py @@ -0,0 +1,508 @@ +#!/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) + + scores = 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) diff --git a/tests/navigation/test_occupancy.py b/tests/navigation/test_occupancy.py new file mode 100644 index 000000000..565efc21f --- /dev/null +++ b/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 diff --git a/tests/navigation/test_voxel_map.py b/tests/navigation/test_voxel_map.py new file mode 100644 index 000000000..413c271ed --- /dev/null +++ b/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]) diff --git a/tests/navigation/test_voxel_map_carve.py b/tests/navigation/test_voxel_map_carve.py new file mode 100644 index 000000000..8eb433305 --- /dev/null +++ b/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 diff --git a/tests/navigation/test_voxel_map_mutation.py b/tests/navigation/test_voxel_map_mutation.py new file mode 100644 index 000000000..3775d226f --- /dev/null +++ b/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