Add TWIN preprocessing and norm-stats helper scripts

openpi/scripts/compute_norm_stats_repo.py

@@ -0,0 +1,65 @@
+"""Compute normalization statistics for a config with repo/loader overrides.
+
+This exists so multiprocessing workers can respawn from a real on-disk module
+instead of failing when the parent process was launched from stdin.
+"""
+
+import dataclasses
+
+import numpy as np
+import tqdm
+import tyro
+
+import openpi.shared.normalize as normalize
+import openpi.training.config as config_lib
+import compute_norm_stats as base_compute
+
+
+def main(
+    config_name: str,
+    repo_id: str,
+    batch_size: int = 64,
+    num_workers: int = 12,
+    assets_base_dir: str = "./assets",
+    max_frames: int | None = None,
+):
+    config = dataclasses.replace(
+        config_lib.get_config(config_name),
+        batch_size=batch_size,
+        num_workers=num_workers,
+        assets_base_dir=assets_base_dir,
+    )
+    data_factory = dataclasses.replace(config.data, repo_id=repo_id)
+    data_config = data_factory.create(config.assets_dirs, config.model)
+
+    if data_config.rlds_data_dir is not None:
+        data_loader, num_batches = base_compute.create_rlds_dataloader(
+            data_config,
+            config.model.action_horizon,
+            config.batch_size,
+            max_frames,
+        )
+    else:
+        data_loader, num_batches = base_compute.create_torch_dataloader(
+            data_config,
+            config.model.action_horizon,
+            config.batch_size,
+            config.model,
+            config.num_workers,
+            max_frames,
+        )
+
+    keys = ["state", "actions"]
+    stats = {key: normalize.RunningStats() for key in keys}
+
+    for batch in tqdm.tqdm(data_loader, total=num_batches, desc=f"{config_name} :: {repo_id}"):
+        for key in keys:
+            stats[key].update(np.asarray(batch[key]))
+
+    output_path = config.assets_dirs / repo_id
+    print(f"Writing stats to: {output_path}")
+    normalize.save(output_path, {key: value.get_statistics() for key, value in stats.items()})
+
+
+if __name__ == "__main__":
+    tyro.cli(main)