Feat/norm stats compatibility report

scripts/check_norm_stats.py

@@ -0,0 +1,60 @@
+"""Compare transformed training samples with configured normalization statistics."""
+
+import compute_norm_stats
+import numpy as np
+import tqdm
+import tyro
+
+import openpi.training.config as _config
+from openpi.training.norm_stats_report import analyze_norm_stats
+from openpi.training.norm_stats_report import format_norm_stats_report
+
+
+def main(config_name: str, max_frames: int = 5000) -> None:
+    """Print a read-only compatibility report for a training configuration."""
+    if max_frames < 1:
+        raise ValueError("max_frames must be positive")
+
+    config = _config.get_config(config_name)
+    data_config = config.data.create(config.assets_dirs, config.model)
+    if data_config.norm_stats is None:
+        raise ValueError(
+            "The selected config does not load normalization statistics. "
+            "Run compute_norm_stats.py first or configure checkpoint assets to compare against."
+        )
+
+    if data_config.rlds_data_dir is not None:
+        data_loader, num_batches = compute_norm_stats.create_rlds_dataloader(
+            data_config, config.model.action_horizon, config.batch_size, max_frames
+        )
+    else:
+        data_loader, num_batches = compute_norm_stats.create_torch_dataloader(
+            data_config,
+            config.model.action_horizon,
+            config.batch_size,
+            config.model,
+            config.num_workers,
+            max_frames,
+        )
+
+    keys = ("state", "actions")
+    samples: dict[str, list[np.ndarray]] = {key: [] for key in keys}
+    for batch in tqdm.tqdm(data_loader, total=num_batches, desc="Checking normalization stats"):
+        for key in keys:
+            if key not in batch:
+                raise KeyError(f"Transformed data does not contain required key '{key}'.")
+            samples[key].append(np.asarray(batch[key]))
+
+    for key, batches in samples.items():
+        if key not in data_config.norm_stats:
+            print(f"{key}\n  result: WARNING (missing configured normalization statistics)\n")
+            continue
+        if not batches:
+            raise ValueError("The data loader produced no samples. Increase max_frames or check the dataset config.")
+        report = analyze_norm_stats(key, np.concatenate(batches), data_config.norm_stats[key])
+        print(format_norm_stats_report(report))
+        print()
+
+
+if __name__ == "__main__":
+    tyro.cli(main)

src/openpi/training/norm_stats_report.py

@@ -0,0 +1,106 @@
+"""Read-only compatibility checks for data and configured normalization statistics."""
+
+import dataclasses
+
+import numpy as np
+
+from openpi.shared.normalize import NormStats
+
+
+@dataclasses.dataclass(frozen=True)
+class NormStatsReport:
+    key: str
+    data_dim: int
+    stats_dim: int
+    num_values: int
+    nonfinite_values: int
+    near_constant_dims: tuple[int, ...]
+    outside_quantile_fraction: float | None
+    largest_mismatch_dim: int | None
+
+    @property
+    def dimensions_match(self) -> bool:
+        return self.data_dim == self.stats_dim
+
+
+def analyze_norm_stats(
+    key: str,
+    values: np.ndarray,
+    stats: NormStats,
+    *,
+    near_constant_threshold: float = 1e-6,
+) -> NormStatsReport:
+    """Compare transformed samples with one configured normalization-stat entry."""
+    values = np.asarray(values)
+    if values.ndim == 0:
+        raise ValueError(f"Expected '{key}' samples to have at least one dimension.")
+
+    values = values.reshape(-1, values.shape[-1])
+    data_dim = values.shape[-1]
+    stats_dim = np.asarray(stats.mean).shape[-1]
+    finite = np.isfinite(values)
+    nonfinite_values = int(values.size - np.count_nonzero(finite))
+
+    finite_values = np.where(finite, values, np.nan)
+    with np.errstate(invalid="ignore"):
+        std = np.nanstd(finite_values, axis=0)
+    near_constant_dims = tuple(np.flatnonzero(np.isfinite(std) & (std <= near_constant_threshold)).tolist())
+
+    outside_fraction = None
+    largest_mismatch_dim = None
+    if stats.q01 is not None and stats.q99 is not None and data_dim == stats_dim:
+        q01 = np.asarray(stats.q01)
+        q99 = np.asarray(stats.q99)
+        valid = finite & np.isfinite(q01)[None, :] & np.isfinite(q99)[None, :]
+        outside = valid & ((values < q01) | (values > q99))
+        valid_counts = valid.sum(axis=0)
+        outside_counts = outside.sum(axis=0)
+        total_valid = int(valid_counts.sum())
+        if total_valid:
+            outside_fraction = float(outside_counts.sum() / total_valid)
+            per_dim = np.divide(
+                outside_counts,
+                valid_counts,
+                out=np.full(data_dim, np.nan, dtype=np.float64),
+                where=valid_counts > 0,
+            )
+            if np.any(np.isfinite(per_dim)):
+                largest_mismatch_dim = int(np.nanargmax(per_dim))
+
+    return NormStatsReport(
+        key=key,
+        data_dim=data_dim,
+        stats_dim=stats_dim,
+        num_values=values.size,
+        nonfinite_values=nonfinite_values,
+        near_constant_dims=near_constant_dims,
+        outside_quantile_fraction=outside_fraction,
+        largest_mismatch_dim=largest_mismatch_dim,
+    )
+
+
+def format_norm_stats_report(report: NormStatsReport, *, outside_warning_threshold: float = 0.05) -> str:
+    """Format a compatibility report for terminal output."""
+    outside = "unavailable" if report.outside_quantile_fraction is None else f"{report.outside_quantile_fraction:.1%}"
+    mismatch = "unavailable" if report.largest_mismatch_dim is None else str(report.largest_mismatch_dim)
+    result = "PASS"
+    if (
+        not report.dimensions_match
+        or report.nonfinite_values
+        or report.outside_quantile_fraction is None
+        or report.outside_quantile_fraction > outside_warning_threshold
+    ):
+        result = "WARNING"
+
+    return "\n".join(
+        [
+            report.key,
+            f"  data dimensions:             {report.data_dim}",
+            f"  norm-stat dimensions:        {report.stats_dim}",
+            f"  non-finite values:           {report.nonfinite_values}",
+            f"  near-constant dimensions:    {list(report.near_constant_dims)}",
+            f"  outside configured q01-q99:  {outside}",
+            f"  largest mismatch dimension:  {mismatch}",
+            f"  result:                      {result}",
+        ]
+    )

src/openpi/training/norm_stats_report_test.py

@@ -0,0 +1,74 @@
+import numpy as np
+import pytest
+
+from openpi.shared.normalize import NormStats
+from openpi.training.norm_stats_report import analyze_norm_stats
+from openpi.training.norm_stats_report import format_norm_stats_report
+
+
+def _stats(dim: int = 2) -> NormStats:
+    return NormStats(
+        mean=np.zeros(dim),
+        std=np.ones(dim),
+        q01=np.full(dim, -1.0),
+        q99=np.full(dim, 1.0),
+    )
+
+
+def test_matching_samples_report_outside_fraction_and_largest_dimension():
+    values = np.array([[0.0, 0.0], [2.0, 0.0], [3.0, 2.0]])
+
+    report = analyze_norm_stats("actions", values, _stats())
+
+    assert report.dimensions_match
+    assert report.outside_quantile_fraction == pytest.approx(3 / 6)
+    assert report.largest_mismatch_dim == 0
+
+
+def test_dimension_mismatch_skips_quantile_comparison():
+    report = analyze_norm_stats("state", np.ones((4, 3)), _stats(dim=2))
+
+    assert not report.dimensions_match
+    assert report.outside_quantile_fraction is None
+    assert report.largest_mismatch_dim is None
+
+
+def test_counts_nonfinite_values_without_corrupting_other_checks():
+    values = np.array([[0.0, np.nan], [np.inf, 0.0], [0.5, 0.0]])
+
+    report = analyze_norm_stats("actions", values, _stats())
+
+    assert report.nonfinite_values == 2
+    assert report.outside_quantile_fraction == 0.0
+
+
+def test_identifies_near_constant_dimensions():
+    values = np.array([[1.0, 0.0], [1.0, 1.0], [1.0, 2.0]])
+
+    report = analyze_norm_stats("state", values, _stats())
+
+    assert report.near_constant_dims == (0,)
+
+
+def test_rejects_scalar_samples():
+    with pytest.raises(ValueError, match="at least one dimension"):
+        analyze_norm_stats("state", np.asarray(1.0), _stats(dim=1))
+
+
+def test_format_includes_warning_for_dimension_mismatch():
+    report = analyze_norm_stats("state", np.ones((4, 3)), _stats(dim=2))
+
+    output = format_norm_stats_report(report)
+
+    assert "state" in output
+    assert "norm-stat dimensions:        2" in output
+    assert "result:                      WARNING" in output
+
+
+def test_format_warns_when_too_many_values_are_outside_quantiles():
+    report = analyze_norm_stats("actions", np.array([[2.0, 0.0], [0.0, 0.0]]), _stats())
+
+    output = format_norm_stats_report(report, outside_warning_threshold=0.05)
+
+    assert "outside configured q01-q99:  25.0%" in output
+    assert "result:                      WARNING" in output