data: stop letting one bad episode or a rotation channel take a run down

egomimic/rldb/zarr/zarr_dataset_multi.py

@@ -50,9 +50,10 @@
 from egomimic.utils.action_utils import (
     _apply_norm_one,
     _apply_unnorm_one,
+    degenerate_mask,
 )
 from egomimic.utils.env import load_env
-from egomimic.utils.pose_utils import rot6d_channels
+from egomimic.utils.pose_utils import bounds_check_channels, rot6d_channels
 
 
 def create_default_engine():
@@ -990,6 +991,19 @@ class MultiDataset(torch.utils.data.Dataset):
     # normalization must keep normalizing it (see _load_state).
     rot6d_passthrough = True
 
+    # Bounds-check slack per (timestep, channel) cell; see _check_bounds.
+    BOUNDS_ABS_TOL = 1e-6
+    BOUNDS_REL_SLACK = 0.5
+    # Bad-sample retries stay in the failing episode this many times, then
+    # draw from the whole index space so one bad episode can't exhaust them.
+    GLOBAL_FALLBACK_ATTEMPTS = 25
+    # Consecutive bad samples before raising: a systemic-failure detector.
+    MAX_FALLBACK_ATTEMPTS = 100
+    # Retry warnings: the first RETRY_LOG_FIRST per worker, then every RETRY_LOG_EVERY.
+    RETRY_LOG_FIRST = 20
+    RETRY_LOG_EVERY = 500
+    _retry_log_count = 0
+
     def __init__(
         self,
         datasets: (
@@ -1143,6 +1157,19 @@ def _check_bounds(
                 q_low = torch.broadcast_to(q_low, arr.shape)
                 q_high = torch.broadcast_to(q_high, arr.shape)
             except RuntimeError:
+                # Stats were computed for a different layout than this sample
+                # (e.g. a stale precomputed norm_stats.json). Say so once
+                # instead of silently disabling the bounds check for the key;
+                # normalize() raises on the same mismatch anyway.
+                warn_key = f"bounds-shape:{zarr_key}"
+                if warn_key not in self._warned_violations:
+                    self._warned_violations.add(warn_key)
+                    logger.warning(
+                        f"[MultiDataset] bounds check skipped for {zarr_key}: "
+                        f"stats shape {tuple(q_low.shape)} does not broadcast to "
+                        f"sample shape {tuple(arr.shape)} (norm stats computed "
+                        "for a different layout?)"
+                    )
                 continue
 
             if torch.any(torch.isnan(arr)) or torch.any(torch.isinf(arr)):
@@ -1153,8 +1180,37 @@ def _check_bounds(
                     logger.warning(prefix)
                 return prefix
 
-            below = arr < q_low
-            above = arr > q_high
+            # Rotation channels are either Euler ypr (wraps at +-pi) or
+            # continuous 6D columns; quantile bounds on them are meaningless
+            # and reject otherwise-valid frames, so only the translation /
+            # gripper / keypoint channels are bounds-checked. Unrecognized
+            # widths fall through to a full-vector check; NaN/Inf above still
+            # covers the full vector.
+            check_idx = bounds_check_channels(zarr_key, arr.shape[-1])
+            if check_idx is not None:
+                arr_q = arr[..., check_idx]
+                q_low = q_low[..., check_idx]
+                q_high = q_high[..., check_idx]
+            else:
+                arr_q = arr
+
+            # Per-cell any-cell checks against subsampled stats reject valid
+            # frames in bulk, so allow BOUNDS_REL_SLACK of each cell's quantile
+            # range (plus an absolute floor). Corrupt values sit orders of
+            # magnitude outside and are still caught.
+            tol = self.BOUNDS_ABS_TOL + self.BOUNDS_REL_SLACK * (q_high - q_low)
+            below = arr_q < q_low - tol
+            above = arr_q > q_high + tol
+            # Cells normalize() treats as constant (wrist-frame t=0, the kp0 /
+            # kp9 structural zeros) cannot blow up the loss whatever offset
+            # they carry, so don't reject on them. The exemption follows
+            # norm_mode's range stat, the same one normalize() uses.
+            constant = self._degenerate_cells(stats, arr.shape)
+            if constant is not None:
+                if check_idx is not None:
+                    constant = constant[..., check_idx]
+                below = below & ~constant
+                above = above & ~constant
             if torch.any(below) or torch.any(above):
                 prefix = f"Bounds violation in {zarr_key} ep={episode_name} frame={idx}"
                 warn_key = f"bounds:{episode_name}:{zarr_key}"
@@ -1164,11 +1220,20 @@ def _check_bounds(
                     n_above = int(above.sum().item())
                     logger.warning(
                         f"{prefix} | n_below={n_below} n_above={n_above} "
-                        f"arr_range=[{arr.min().item():.4f}, {arr.max().item():.4f}]"
+                        f"arr_range=[{arr_q.min().item():.4f}, {arr_q.max().item():.4f}]"
                     )
                 return prefix
         return None
 
+    def _degenerate_cells(self, stats: dict, shape) -> torch.Tensor | None:
+        """Mask of the cells ``normalize()`` treats as constant, or None when
+        this key's stats do not carry the range ``norm_mode`` reads."""
+        try:
+            mask = degenerate_mask(stats, self.norm_mode)
+        except (KeyError, ValueError):
+            return None
+        return torch.broadcast_to(mask, shape)
+
     def __getitem__(self, idx, _attempts: int | None = None):
         attempts = _attempts
         while True:
@@ -1177,14 +1242,13 @@ def __getitem__(self, idx, _attempts: int | None = None):
             try:
                 data = dataset[local_idx]
             except Exception as e:
-                next_idx, attempts = self._next_after_failure(
-                    idx,
-                    dataset_name,
-                    attempts,
-                    reason=f"Sample failed ({type(e).__name__}: {e}) at "
-                    f"{dataset_name}[{local_idx}]",
+                reason = (
+                    f"Sample failed ({type(e).__name__}: {e}) at "
+                    f"{dataset_name}[{local_idx}]"
+                )
+                idx, attempts = self._next_after_failure(
+                    idx, dataset_name, attempts, reason=reason
                 )
-                idx = next_idx
                 continue
 
             # If this leaf is itself a MultiDataset, it already ran bounds +
@@ -1194,38 +1258,48 @@ def __getitem__(self, idx, _attempts: int | None = None):
 
             violation = self._check_bounds(data, dataset, local_idx, dataset_name)
             if violation is not None:
-                next_idx, attempts = self._next_after_failure(
-                    idx,
-                    dataset_name,
-                    attempts,
-                    reason=violation,
+                idx, attempts = self._next_after_failure(
+                    idx, dataset_name, attempts, reason=violation
                 )
-                idx = next_idx
                 continue
 
-            # Bounds passed — normalize and return.
+            # Bounds passed: normalize and return.
             if self.norm_stats and data.get("embodiment") in self.norm_stats:
                 data = self.normalize(data, data["embodiment"])
             return data
 
     def _next_after_failure(
         self, idx: int, dataset_name: str, attempts: int | None, *, reason: str
     ) -> tuple[int, int]:
-        global_candidates = self._global_indices_by_dataset[dataset_name]
-        next_idx, attempts = get_fallback_idx(
-            idx=idx,
-            candidates=global_candidates,
-            _attempts=attempts,
-            max_attempts=len(global_candidates),
-            exhausted_error=(
-                f"Entire dataset bad (no valid indices): dataset={dataset_name}"
-            ),
-        )
+        """Pick the next index to try after a bad sample at ``idx``.
+
+        Raises once ``MAX_FALLBACK_ATTEMPTS`` consecutive samples failed."""
+        attempts = (attempts or 0) + 1
+        if attempts >= self.MAX_FALLBACK_ATTEMPTS:
+            raise RuntimeError(
+                f"[MultiDataset] {attempts} consecutive bad samples, the last "
+                f"{attempts - self.GLOBAL_FALLBACK_ATTEMPTS} drawn uniformly "
+                "from the whole index space: the store, the keymap or the norm "
+                f"stats are broken, not one episode. Last: {reason}"
+            )
+        local = self._global_indices_by_dataset[dataset_name]
+        if attempts <= self.GLOBAL_FALLBACK_ATTEMPTS and len(local) > 1:
+            next_idx = idx
+            while next_idx == idx:
+                next_idx = random.choice(local)
+        else:
+            next_idx = random.randrange(len(self.index_map))
         next_dataset_name, next_local_idx = self.index_map[next_idx]
-        logger.warning(
-            f"{reason} | attempt {attempts}, "
-            f"trying {next_dataset_name}[{next_local_idx}]"
-        )
+        self._retry_log_count += 1
+        if (
+            self._retry_log_count <= self.RETRY_LOG_FIRST
+            or self._retry_log_count % self.RETRY_LOG_EVERY == 0
+        ):
+            logger.warning(
+                f"{reason} | attempt {attempts}, "
+                f"trying {next_dataset_name}[{next_local_idx}] "
+                f"(retries so far in this worker: {self._retry_log_count})"
+            )
         return next_idx, attempts
 
     @classmethod
@@ -1572,8 +1646,13 @@ def _compute_stats_for_array(X):
             "median": np.median(X, axis=0),
             "quantile_1": np.percentile(X, 1, axis=0),
             "quantile_99": np.percentile(X, 99, axis=0),
-            "quantile_0_01": np.percentile(X, 0.01, axis=0),
-            "quantile_99_99": np.percentile(X, 99.99, axis=0),
+            # Bounds-check quantiles snap to observed samples ("lower" /
+            # "higher") instead of interpolating: with fewer than 10k samples
+            # the default linear percentile lands strictly inside the sample
+            # range, and the extreme frames of the very episodes the stats
+            # were computed on would be rejected at train time.
+            "quantile_0_01": np.percentile(X, 0.01, axis=0, method="lower"),
+            "quantile_99_99": np.percentile(X, 99.99, axis=0, method="higher"),
         }
 
     def cache_stats(self, save_cache_dir: str):
@@ -2115,13 +2194,22 @@ def _build_frame_to_ann_end(self) -> dict[int, int]:
         annotation span. Annotations use half-open ``[start_idx, end_idx)``.
         """
         mapping: dict[int, int] = {}
-        for ann in self._load_annotations():
-            start_idx = int(ann.get("start_idx", -1))
-            end_idx = int(ann.get("end_idx", -1))
-            if start_idx < 0 or end_idx <= start_idx:
-                continue
+        n_spans = 0
+        for start_idx, end_idx in _annotation_spans(self._load_annotations()):
+            n_spans += 1
             for idx in range(start_idx, end_idx):
                 mapping[idx] = end_idx
+        # Visibility into annotation-cutoff usage: spans/frames_covered of 0
+        # means the cutoff is a no-op for this episode. One line per episode
+        # per rank is thousands of lines on a real split, so this is DEBUG and
+        # the resolver's kept N/M line carries the aggregate.
+        logger.debug(
+            "[AnnotationCutoff] ep=%s spans=%d frames_covered=%d/%d",
+            Path(self.episode_path).name,
+            n_spans,
+            len(mapping),
+            self.total_frames,
+        )
         return mapping
 
     def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) -> int:
@@ -2136,11 +2224,75 @@ def _chunk_end_idx(self, start_idx: int, horizon: int, key_type: str | None) ->
         return min(end_idx, ann_end)
 
 
+def _annotation_spans(anns: list[dict]) -> Iterable[tuple[int, int]]:
+    """Usable half-open ``[start_idx, end_idx)`` spans of decoded annotations."""
+    for ann in anns:
+        start_idx = int(ann.get("start_idx", -1))
+        end_idx = int(ann.get("end_idx", -1))
+        if 0 <= start_idx < end_idx:
+            yield start_idx, end_idx
+
+
+def _episode_has_annotation_spans(ds: "ZarrDataset") -> bool:
+    """True if the episode has at least one usable annotation span."""
+    try:
+        anns = ds._load_annotations()
+    except KeyError:
+        return False
+    except Exception:
+        logger.warning(
+            "[AnnotationCutoff] could not read annotations of %s",
+            getattr(ds, "episode_path", ds),
+            exc_info=True,
+        )
+        return False
+    return next(_annotation_spans(anns), None) is not None
+
+
 class S3AnnotationCutoffEpisodeResolver(S3EpisodeResolver):
-    """S3EpisodeResolver that loads ZarrAnnotationCutoffDataset instances."""
+    """S3EpisodeResolver that loads ZarrAnnotationCutoffDataset instances.
+
+    When ``require_annotations`` is set (default), episodes whose zarr
+    ``annotations`` array has no usable span are dropped; otherwise the
+    annotation cutoff would silently no-op on them.
+    """
 
     _dataset_class = ZarrAnnotationCutoffDataset
 
+    def __init__(self, *args, require_annotations: bool = True, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.require_annotations = require_annotations
+
+    def resolve(self, filters=None, expected_embodiment=None):
+        datasets = super().resolve(
+            filters=filters, expected_embodiment=expected_embodiment
+        )
+        if not self.require_annotations:
+            return datasets
+        kept = {
+            h: ds for h, ds in datasets.items() if _episode_has_annotation_spans(ds)
+        }
+        dropped = sorted(set(datasets) - set(kept))
+        if dropped:
+            logger.warning(
+                "[AnnotationCutoff] dropped %d/%d episodes with no usable "
+                "annotation spans (e.g. %s)",
+                len(dropped),
+                len(datasets),
+                dropped[:5],
+            )
+        logger.info(
+            "[AnnotationCutoff] kept %d/%d episodes with usable annotations",
+            len(kept),
+            len(datasets),
+        )
+        if not kept:
+            raise ValueError(
+                "[AnnotationCutoff] no resolved episodes contain usable annotation "
+                "spans; check the filter / annotation injection for this dataset."
+            )
+        return kept
+
 
 class LocalAnnotationCutoffEpisodeResolver(LocalEpisodeResolver):
     """LocalEpisodeResolver that loads ZarrAnnotationCutoffDataset instances."""

egomimic/utils/action_utils.py

@@ -116,6 +116,14 @@ def _degenerate(scale: torch.Tensor, norm_mode: str) -> torch.Tensor:
     return full_range < NORM_MIN_RANGE
 
 
+def degenerate_mask(stats: dict[str, Any], norm_mode: str) -> torch.Tensor:
+    """Channels ``_apply_norm_one`` treats as constant under ``norm_mode``.
+
+    Raises ``KeyError`` when ``stats`` lacks the range ``norm_mode`` reads."""
+    _, scale = _norm_center_scale(stats, norm_mode, torch.empty(0))
+    return _degenerate(scale, norm_mode)
+
+
 def _passthrough(
     out: torch.Tensor, tensor: torch.Tensor, channels: Sequence[int] | None
 ) -> torch.Tensor:

egomimic/utils/pose_utils.py

@@ -289,6 +289,20 @@ def rot6d_channels(zarr_key: str, width: int) -> tuple[int, ...] | None:
     return None
 
 
+def bounds_check_channels(zarr_key: str, width: int) -> list[int] | None:
+    """Sorted non-rotation channels (translation, gripper, keypoints) of
+    ``zarr_key`` at ``width``, or ``None`` for an unknown key or width."""
+    if zarr_key in _CARTESIAN_KEYS:
+        layout = bimanual_cartesian_layout(width)
+        if layout is not None:
+            return sorted(layout["xyz"] + layout["grip"])
+    elif zarr_key in _KEYPOINT_KEYS:
+        layout = bimanual_keypoint_layout(width)
+        if layout is not None:
+            return sorted(layout["wrist_xyz"] + layout["keypoints"])
+    return None
+
+
 def _matrix_to_xyzwxyz(mats: np.ndarray) -> np.ndarray:
     """
     args:

tests/fixtures/bounds.py

@@ -0,0 +1,16 @@
+"""Stats-only MultiDataset shells for exercising ``_check_bounds``."""
+
+from __future__ import annotations
+
+
+def bounds_dataset(stats_by_key: dict, norm_mode: str = "quantile", embodiment=0):
+    """A MultiDataset with only the fields ``_check_bounds`` reads: ``stats_by_key``
+    maps each zarr key (also its key name) to its norm stats."""
+    from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset
+
+    md = MultiDataset.__new__(MultiDataset)
+    md.norm_mode = norm_mode
+    md.norm_stats = {embodiment: dict(stats_by_key)}
+    md.zarr_keys = {embodiment: {key: key for key in stats_by_key}}
+    md._warned_violations = set()
+    return md

tests/unit/test_eva_wristframe_norm_pool.py

@@ -238,3 +238,19 @@ def test_state_round_trip_keeps_the_rot6d_convention(saved, expect_identity):
     assert torch.equal(n[:, rot], x[:, rot]) == expect_identity
     back = restored._apply_unnorm_one(n, _stats_for(20), "actions_cartesian")
     torch.testing.assert_close(back, x, atol=1e-5, rtol=0)
+
+
+def test_bounds_check_ignores_constant_cells():
+    from fixtures.bounds import bounds_dataset
+
+    key = "actions_cartesian"
+    q1 = np.full((4, 14), -1.0, dtype=np.float32)
+    q99 = np.full((4, 14), 1.0, dtype=np.float32)
+    q1[0, :3] = q99[0, :3] = 0.0  # wrist-frame t=0: xyz exactly 0
+    md = bounds_dataset({key: {"quantile_1": q1, "quantile_99": q99}})
+
+    arr = np.zeros((4, 14), dtype=np.float32)
+    arr[0, 0] = 1e-3  # off-convention offset at a constant cell: admitted
+    assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is None
+    arr[1, 0] = 50.0  # corrupt value at a regular cell: still rejected
+    assert md._check_bounds({"embodiment": 0, key: arr}, None, 0, "ep") is not None