diff --git a/scripts/check_norm_stats.py b/scripts/check_norm_stats.py new file mode 100644 index 0000000000..f90d2caf44 --- /dev/null +++ b/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) diff --git a/src/openpi/training/norm_stats_report.py b/src/openpi/training/norm_stats_report.py new file mode 100644 index 0000000000..6dc01d9990 --- /dev/null +++ b/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}", + ] + ) diff --git a/src/openpi/training/norm_stats_report_test.py b/src/openpi/training/norm_stats_report_test.py new file mode 100644 index 0000000000..e9dd572472 --- /dev/null +++ b/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