From 7a13410ed441621e902034cbfcae000a52f43617 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Wed, 16 Sep 2026 17:49:00 -0400 Subject: [PATCH] data: stop letting one bad episode or a rotation channel take a run down Three independent ways MultiDataset could fail a whole run on data that was fine, or on data it should have rejected up front. Bounds check. Quantile bounds on a rotation channel are meaningless: euler ypr wraps at +-pi and rot6d columns are matrix entries, so neither respects a quantile interval. _bounds_check_channels uses the layout tables to check only translation, gripper and keypoint channels; unrecognized widths still get the full-vector check, and the NaN/Inf check always covers everything. The bounds themselves get slack -- 1e-6 absolute plus half of each cell's own quantile range -- because the check is any-cell over ~13 200 cells per keypoint chunk. Measured on the 2815-frame flagship episode through the real keypoints_wristframe_6d pipeline (zscore, stats from sample_frac 0.10 over its own frames, 403 probe frames, any-key rejection): 39.5 % of frames rejected at slack 0, 17.9 % at 0.05, 11.4 % at 0.1, 5.2 % at 0.25, 3.2 % at 0.5, 1.0 % at 1.0. The 132-channel action chunk drives all of it. Half a range is the knee, not a cushion -- though the absolute numbers move with which frames the stats were drawn from, so the comment says how to reproduce them. Corrupt values sit orders of magnitude outside the range and are still caught. Cells the normalizer treats as constant are exempt, since a tiny offset there cannot blow up the loss -- and which stat decides that is norm_mode's business (std under zscore, the quantile range under quantile), so the exemption asks _degenerate instead of hardcoding q99 - q1; every shipped config is zscore, and a channel that is 0 in 99.6 % of samples and 1.0 in the rest has q99 - q1 = 0 but std = 0.063, i.e. it is still divided by and must still be checked. quantile_0_01 / quantile_99_99 now snap to observed samples ("lower" / "higher") instead of interpolating: under 10k samples the linear percentile lands strictly inside the sample range, so the extreme frames of the very episodes the stats came from were rejected (3-16 % of a single episode in the overfit runs). A stats/sample shape mismatch is logged once per key instead of silently disabling the check. Sampler. Retries stayed inside the failing episode until they exhausted it, so one wholly-bad episode killed the run and NCCL-timed-out every other rank. Retries now widen to the global index space after 25 attempts, and 100 consecutive failures raise. The cap is a systemic-failure detector, not a tolerance dial: attempts resets per __getitem__, so with p = the bad fraction of the whole index space it fires with probability p**(N - 25), a step function in p. At N = 100 the firing band sits at p >~ 0.8, so everything it can reach is systemic; at N = 1000 it needed ~99 % of the store to be bad and got there through ~975 sequential zarr/S3 reads per call per worker, which reads as a hang rather than a failure. Retry logging was one WARNING per attempt -- tens of thousands of lines per run -- now the first 20 per worker and then one in 500. Annotation cutoff. Many Scale-"completed" episodes have no usable annotation span because the injection step lagged, which makes the cutoff a silent no-op on exactly the runs whose point is to clamp chunks at annotation boundaries. S3AnnotationCutoffEpisodeResolver(require_annotations=True) drops them, logs the count it dropped and the count it kept, and raises if nothing is left. The per-episode span line drops to DEBUG: one line per episode per rank is thousands of lines on a real split, and the resolver's aggregate says the same thing. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/rldb/zarr/zarr_dataset_multi.py | 232 ++++++++++++++++---- egomimic/utils/action_utils.py | 8 + egomimic/utils/pose_utils.py | 14 ++ tests/fixtures/bounds.py | 16 ++ tests/unit/test_eva_wristframe_norm_pool.py | 16 ++ tests/unit/test_infra_port_wiring.py | 88 ++++++++ tests/unit/test_pi05_norm_rot6d.py | 185 ++++++++++++++++ 7 files changed, 519 insertions(+), 40 deletions(-) create mode 100644 tests/fixtures/bounds.py diff --git a/egomimic/rldb/zarr/zarr_dataset_multi.py b/egomimic/rldb/zarr/zarr_dataset_multi.py index 9eba91f18d..26347a84df 100644 --- a/egomimic/rldb/zarr/zarr_dataset_multi.py +++ b/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,16 +1258,12 @@ 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 @@ -1211,21 +1271,35 @@ def __getitem__(self, idx, _attempts: int | None = None): 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.""" diff --git a/egomimic/utils/action_utils.py b/egomimic/utils/action_utils.py index 2d1fad36ec..c6cd38f8f6 100644 --- a/egomimic/utils/action_utils.py +++ b/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: diff --git a/egomimic/utils/pose_utils.py b/egomimic/utils/pose_utils.py index a87148dd79..e874be0bea 100644 --- a/egomimic/utils/pose_utils.py +++ b/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: diff --git a/tests/fixtures/bounds.py b/tests/fixtures/bounds.py new file mode 100644 index 0000000000..0a2387e3c8 --- /dev/null +++ b/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 diff --git a/tests/unit/test_eva_wristframe_norm_pool.py b/tests/unit/test_eva_wristframe_norm_pool.py index d9ad852f0c..0ceba0b3ec 100644 --- a/tests/unit/test_eva_wristframe_norm_pool.py +++ b/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 diff --git a/tests/unit/test_infra_port_wiring.py b/tests/unit/test_infra_port_wiring.py index f02155f1cb..9b59cbfa2a 100644 --- a/tests/unit/test_infra_port_wiring.py +++ b/tests/unit/test_infra_port_wiring.py @@ -4,6 +4,7 @@ from __future__ import annotations from pathlib import Path +from types import SimpleNamespace import numpy as np import pytest @@ -255,3 +256,90 @@ def test_pi_loss_is_reduced_over_the_packed_width(): pi._reduce_loss(torch.tensor([1.0, 3.0]), None, eva, "actions_cartesian").item() == 2.0 ) + + +# --------------------------------------------------- annotation-cutoff filter +def test_episode_has_annotation_spans(): + from egomimic.rldb.zarr.zarr_dataset_multi import _episode_has_annotation_spans + + def ds(anns): + return SimpleNamespace(_load_annotations=lambda: anns) + + assert _episode_has_annotation_spans( + ds([{"text": "a", "start_idx": 0, "end_idx": 5}]) + ) + assert not _episode_has_annotation_spans(ds([])) + assert not _episode_has_annotation_spans(ds([{"text": "a"}])) # span-less + assert not _episode_has_annotation_spans(ds([{"start_idx": 5, "end_idx": 5}])) + + def boom(): + raise OSError("corrupt") + + assert not _episode_has_annotation_spans(SimpleNamespace(_load_annotations=boom)) + + +def test_annotation_cutoff_resolver_defaults_to_requiring_spans(monkeypatch): + import egomimic.rldb.zarr.zarr_dataset_multi as m + + good = SimpleNamespace(_load_annotations=lambda: [{"start_idx": 0, "end_idx": 3}]) + bad = SimpleNamespace(_load_annotations=lambda: []) + episodes = {"g": good, "b": bad} + monkeypatch.setattr(m.S3EpisodeResolver, "__init__", lambda self, *a, **k: None) + monkeypatch.setattr( + m.S3EpisodeResolver, + "resolve", + lambda self, filters=None, expected_embodiment=None: dict(episodes), + ) + r = m.S3AnnotationCutoffEpisodeResolver() + kept = r.resolve(filters=None, expected_embodiment="human_bimanual") + assert set(kept) == {"g"} + r.require_annotations = False + assert set(r.resolve()) == {"g", "b"} + episodes.pop("g") + r.require_annotations = True + with pytest.raises(ValueError, match="no resolved episodes"): + r.resolve() + + +# ------------------------------------------------------- bounds check +def test_bounds_quantiles_include_the_observed_extremes(): + """With fewer than 10k samples a linear 0.01 / 99.99 percentile lands + strictly inside the sample range, so the stats' own extreme frames were + rejected at train time; the bounds quantiles must snap to samples.""" + from fixtures.bounds import bounds_dataset + + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + X = np.random.default_rng(0).normal(size=(500, 4)).astype(np.float32) + st = MultiDataset._compute_stats_for_array(X) + np.testing.assert_array_equal(st["quantile_0_01"], X.min(axis=0)) + np.testing.assert_array_equal(st["quantile_99_99"], X.max(axis=0)) + md = bounds_dataset({"actions_keypoints": st}) + for row in (X.min(axis=0), X.max(axis=0)): + assert ( + md._check_bounds( + {"embodiment": 0, "actions_keypoints": row[None]}, None, 0, "ep" + ) + is None + ) + + +def test_bounds_check_has_relative_slack_but_catches_corrupt_values(): + """Per-cell bounds tolerate frames moderately beyond the stats sample's + range (valid extreme motion) and still reject values orders of magnitude + off (fill constants, wrong-frame data).""" + from fixtures.bounds import bounds_dataset + + lo, hi = np.full(18, -1.0, np.float32), np.full(18, 1.0, np.float32) + md = bounds_dataset({"actions_cartesian": {"quantile_1": lo, "quantile_99": hi}}) + + def check(v): + arr = np.zeros((3, 18), np.float32) + arr[1, 0] = v # an xyz channel + return md._check_bounds( + {"embodiment": 0, "actions_cartesian": arr}, None, 0, "e" + ) + + assert check(1.0 + 0.4 * 2.0) is None # within 50 % of the [-1, 1] range + assert check(1.0 + 0.6 * 2.0) is not None + assert check(1e9) is not None diff --git a/tests/unit/test_pi05_norm_rot6d.py b/tests/unit/test_pi05_norm_rot6d.py index c9e13bc2ad..431dafab71 100644 --- a/tests/unit/test_pi05_norm_rot6d.py +++ b/tests/unit/test_pi05_norm_rot6d.py @@ -335,3 +335,188 @@ def test_fix_left_wrist_convention_flag_prepends_correction(): "left.action_ee_pose", "left.obs_ee_pose", } + + +def _bounds_check_dataset(key: str, width: int): + """Minimal MultiDataset shell exposing _check_bounds with +-1 quantile + bounds on ``key`` for embodiment 0.""" + from fixtures.bounds import bounds_dataset + + return bounds_dataset( + { + key: { + "quantile_1": np.full(width, -1.0, dtype=np.float32), + "quantile_99": np.full(width, 1.0, dtype=np.float32), + } + } + ) + + +@pytest.mark.parametrize("key", ["actions_cartesian", "observations.state.ee_pose"]) +@pytest.mark.parametrize("width,rot_idx,xyz_idx", [(14, 3, 0), (20, 4, 0), (18, 5, 9)]) +def test_bounds_check_ignores_rotation_channels(key, width, rot_idx, xyz_idx): + # Rotation channels (Euler wraps at +-pi; 6D columns are ~[-1, 1]) must be + # excluded from quantile bounds checking, while translation/gripper + # channels are still checked and NaN/Inf still rejects the full vector. + md = _bounds_check_dataset(key, width) + arr = np.zeros((5, width), dtype=np.float32) + + arr[2, rot_idx] = 50.0 # far outside +-1, but a rotation channel + assert md._check_bounds({"embodiment": 0, key: arr.copy()}, None, 0, "ep") is None + + bad = arr.copy() + bad[2, xyz_idx] = 50.0 # translation channel out of bounds -> violation + assert md._check_bounds({"embodiment": 0, key: bad}, None, 0, "ep") is not None + + nan = arr.copy() + nan[2, rot_idx] = np.nan # NaN anywhere (even rotation) -> violation + assert md._check_bounds({"embodiment": 0, key: nan}, None, 0, "ep") is not None + + +@pytest.mark.parametrize("key", ["actions_keypoints", "observations.state.keypoints"]) +@pytest.mark.parametrize("width,rot_idx,kp_idx", [(138, 4, 10), (144, 78, 90)]) +def test_bounds_check_ignores_wrist_rotation_in_keypoint_layout( + key, width, rot_idx, kp_idx +): + md = _bounds_check_dataset(key, width) + arr = np.zeros((5, width), dtype=np.float32) + arr[1, rot_idx] = 50.0 # wrist rotation channel: not bounds-checked + assert md._check_bounds({"embodiment": 0, key: arr.copy()}, None, 0, "ep") is None + bad = arr.copy() + bad[1, kp_idx] = 50.0 # a keypoint coordinate: checked + assert md._check_bounds({"embodiment": 0, key: bad}, None, 0, "ep") is not None + + +def test_bounds_check_full_vector_for_other_keys(): + # Keys without a known layout (or unrecognized widths) keep the + # full-vector check. + md = _bounds_check_dataset("some_other_key", 20) + arr = np.zeros((5, 20), dtype=np.float32) + arr[2, 4] = 50.0 + assert ( + md._check_bounds({"embodiment": 0, "some_other_key": arr}, None, 0, "ep") + is not None + ) + + md16 = _bounds_check_dataset("actions_cartesian", 16) + arr16 = np.zeros((5, 16), dtype=np.float32) + arr16[2, 4] = 50.0 + assert ( + md16._check_bounds({"embodiment": 0, "actions_cartesian": arr16}, None, 0, "ep") + is not None + ) + + +def test_constant_cell_exemption_follows_norm_mode(): + """A channel the zscore normalizer still divides by must still be + bounds-checked, even though its quantile range is zero: 0 in 99.6 % of + samples and 1.0 in the rest gives q99 - q1 = 0 but std = 0.063.""" + from fixtures.bounds import bounds_dataset + + stats = { + "quantile_1": np.zeros(18, dtype=np.float32), + "quantile_99": np.zeros(18, dtype=np.float32), + "mean": np.full(18, 0.004, dtype=np.float32), + "std": np.full(18, 0.063, dtype=np.float32), + } + arr = np.zeros((5, 18), dtype=np.float32) + arr[0, 0] = 50.0 # corrupt value in a quantile-collapsed channel + + for mode, rejected in (("quantile", False), ("zscore", True)): + md = bounds_dataset({"actions_cartesian": stats}, norm_mode=mode) + got = md._check_bounds( + {"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep" + ) + assert (got is not None) is rejected, mode + + +def test_bounds_check_tolerates_roundoff_on_collapsed_bounds(): + """Wrist-frame t=0 cells have bounds [0, 0]; values there must not reject + (normalize() maps such constant cells to 0), while other cells are still + checked.""" + from fixtures.bounds import bounds_dataset + + q99 = np.ones(18, dtype=np.float32) + q99[0] = 0.0 # channel 0 collapsed to [0, 0] + md = bounds_dataset( + { + "actions_cartesian": { + "quantile_1": np.zeros(18, dtype=np.float32), + "quantile_99": q99, + } + } + ) + arr = np.zeros((5, 18), dtype=np.float32) + arr[0, 0] = 1e-9 # a roundoff-scale xyz value at a [0, 0] bound + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") + is None + ) + arr[0, 0] = 1e-3 # off-convention offset at the constant cell: admitted + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") + is None + ) + arr[0, 1] = 50.0 # a real violation on a regular cell is still caught + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") + is not None + ) + + +def test_bounds_check_warns_once_on_stat_shape_mismatch(caplog): + md = _bounds_check_dataset("actions_cartesian", 18) + arr = np.zeros((5, 20), dtype=np.float32) # stats are 18-wide + with caplog.at_level("WARNING"): + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 0, "ep") + is None + ) + assert ( + md._check_bounds({"embodiment": 0, "actions_cartesian": arr}, None, 1, "ep") + is None + ) + msgs = [r.message for r in caplog.records if "bounds check skipped" in r.message] + assert len(msgs) == 1, msgs + + +def _fallback_dataset(): + from egomimic.rldb.zarr.zarr_dataset_multi import MultiDataset + + md = MultiDataset.__new__(MultiDataset) + md.index_map = [("bad", i) for i in range(10)] + [("good", i) for i in range(1000)] + md._global_indices_by_dataset = { + "bad": list(range(10)), + "good": list(range(10, 1010)), + } + return md + + +def test_fallback_widens_to_global_after_local_attempts(): + # A wholly-bad episode must not exhaust the sampler: retries stay inside + # the failing episode for GLOBAL_FALLBACK_ATTEMPTS, then widen to the full + # index space. + md = _fallback_dataset() + attempts = None + seen_local, seen_global = set(), set() + idx = 0 + for _ in range(md.GLOBAL_FALLBACK_ATTEMPTS): + idx, attempts = md._next_after_failure(0, "bad", attempts, reason="r") + seen_local.add(md.index_map[idx][0]) + assert seen_local == {"bad"}, "early retries must stay within the episode" + + for _ in range(50): + idx, attempts = md._next_after_failure(idx, "bad", attempts, reason="r") + seen_global.add(md.index_map[idx][0]) + assert "good" in seen_global, "post-threshold retries must sample globally" + + +def test_fallback_cap_raises(): + # The cap is a systemic-failure detector: everything it can reach has + # already been drawn uniformly from the whole index space and failed. + md = _fallback_dataset() + attempts, idx, n = None, 0, 0 + with pytest.raises(RuntimeError, match="consecutive bad samples"): + for n in range(1, md.MAX_FALLBACK_ATTEMPTS + 2): + idx, attempts = md._next_after_failure(idx, "bad", attempts, reason="boom") + assert n == md.MAX_FALLBACK_ATTEMPTS