Prune distill

examples/download_dataset.py

@@ -0,0 +1,3 @@
+from datasets import load_dataset
+
+dataset = load_dataset("openvla/modified_libero_rlds", cache_dir="./")

examples/libero/convert_libero_data_to_lerobot.py

@@ -1,52 +1,60 @@
-"""
-Minimal example script for converting a dataset to LeRobot format.
-
-We use the Libero dataset (stored in RLDS) for this example, but it can be easily
-modified for any other data you have saved in a custom format.
+"""Convert a local LIBERO RLDS mirror into a local LeRobot dataset.
 
 Usage:
-uv run examples/libero/convert_libero_data_to_lerobot.py --data_dir /path/to/your/data
-
-If you want to push your dataset to the Hugging Face Hub, you can use the following command:
-uv run examples/libero/convert_libero_data_to_lerobot.py --data_dir /path/to/your/data --push_to_hub
+uv run examples/libero/convert_libero_data_to_lerobot.py --data_dir /root/pi_train/modified_libero_rlds
 
-Note: to run the script, you need to install tensorflow_datasets:
-`uv pip install tensorflow tensorflow_datasets`
-
-You can download the raw Libero datasets from https://huggingface.co/datasets/openvla/modified_libero_rlds
-The resulting dataset will get saved to the $HF_LEROBOT_HOME directory.
-Running this conversion script will take approximately 30 minutes.
+This writes a LeRobot dataset under HF_LEROBOT_HOME/<repo_name>.
+By default it targets physical-intelligence/libero so the standard pi05_libero
+training config can run fully offline from the local cache.
 """
 
+from collections.abc import Iterator, Sequence
+from io import BytesIO
+from pathlib import Path
 import shutil
 
 from lerobot.common.datasets.lerobot_dataset import HF_LEROBOT_HOME
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-import tensorflow_datasets as tfds
+import numpy as np
+from PIL import Image
+from tfrecord.reader import sequence_loader
 import tyro
 
-REPO_NAME = "your_hf_username/libero"  # Name of the output dataset, also used for the Hugging Face Hub
-RAW_DATASET_NAMES = [
+DEFAULT_REPO_NAME = "physical-intelligence/libero"
+DEFAULT_RAW_DATASET_NAMES = (
     "libero_10_no_noops",
     "libero_goal_no_noops",
     "libero_object_no_noops",
     "libero_spatial_no_noops",
-]  # For simplicity we will combine multiple Libero datasets into one training dataset
+)
+
+
+def _decode_image(value: bytes | np.bytes_) -> np.ndarray:
+    return np.asarray(Image.open(BytesIO(bytes(value))).convert("RGB"))
 
 
-def main(data_dir: str, *, push_to_hub: bool = False):
-    # Clean up any existing dataset in the output directory
-    output_path = HF_LEROBOT_HOME / REPO_NAME
+def _iter_episodes(dataset_root: Path) -> Iterator[dict[str, np.ndarray]]:
+    for shard_path in sorted((dataset_root / "1.0.0").glob("*.tfrecord-*")):
+        for context, _ in sequence_loader(str(shard_path), None):
+            yield context
+
+
+def main(
+    data_dir: str,
+    *,
+    repo_name: str = DEFAULT_REPO_NAME,
+    raw_dataset_names: Sequence[str] = DEFAULT_RAW_DATASET_NAMES,
+    max_episodes_per_dataset: int | None = None,
+):
+    output_path = HF_LEROBOT_HOME / repo_name
     if output_path.exists():
         shutil.rmtree(output_path)
 
-    # Create LeRobot dataset, define features to store
-    # OpenPi assumes that proprio is stored in `state` and actions in `action`
-    # LeRobot assumes that dtype of image data is `image`
     dataset = LeRobotDataset.create(
-        repo_id=REPO_NAME,
+        repo_id=repo_name,
         robot_type="panda",
         fps=10,
+        use_videos=False,
         features={
             "image": {
                 "dtype": "image",
@@ -69,35 +77,44 @@ def main(data_dir: str, *, push_to_hub: bool = False):
                 "names": ["actions"],
             },
         },
-        image_writer_threads=10,
-        image_writer_processes=5,
     )
 
-    # Loop over raw Libero datasets and write episodes to the LeRobot dataset
-    # You can modify this for your own data format
-    for raw_dataset_name in RAW_DATASET_NAMES:
-        raw_dataset = tfds.load(raw_dataset_name, data_dir=data_dir, split="train")
-        for episode in raw_dataset:
-            for step in episode["steps"].as_numpy_iterator():
+    base_dir = Path(data_dir)
+    total_episodes = 0
+    for raw_dataset_name in raw_dataset_names:
+        dataset_root = base_dir / raw_dataset_name
+        saved_for_dataset = 0
+        for context in _iter_episodes(dataset_root):
+            num_steps = len(context["steps/is_last"])
+            actions = np.asarray(context["steps/action"], dtype=np.float32).reshape(num_steps, 7)
+            states = np.asarray(context["steps/observation/state"], dtype=np.float32).reshape(num_steps, 8)
+            main_images = context["steps/observation/image"]
+            wrist_images = context["steps/observation/wrist_image"]
+            tasks = context["steps/language_instruction"]
+
+            for step_idx in range(num_steps):
                 dataset.add_frame(
                     {
-                        "image": step["observation"]["image"],
-                        "wrist_image": step["observation"]["wrist_image"],
-                        "state": step["observation"]["state"],
-                        "actions": step["action"],
-                        "task": step["language_instruction"].decode(),
+                        "image": _decode_image(main_images[step_idx]),
+                        "wrist_image": _decode_image(wrist_images[step_idx]),
+                        "state": states[step_idx],
+                        "actions": actions[step_idx],
+                        "task": bytes(tasks[step_idx]).decode("utf-8"),
                     }
                 )
+
             dataset.save_episode()
+            saved_for_dataset += 1
+            total_episodes += 1
+            print(
+                f"saved {raw_dataset_name} episode {saved_for_dataset} -> total={total_episodes}",
+                flush=True,
+            )
+
+            if max_episodes_per_dataset is not None and saved_for_dataset >= max_episodes_per_dataset:
+                break
 
-    # Optionally push to the Hugging Face Hub
-    if push_to_hub:
-        dataset.push_to_hub(
-            tags=["libero", "panda", "rlds"],
-            private=False,
-            push_videos=True,
-            license="apache-2.0",
-        )
+    print(f"Finished writing {total_episodes} episodes to {output_path}")
 
 
 if __name__ == "__main__":

prune_distill/README.md

@@ -0,0 +1,144 @@
+# Prefix Distillation
+
+This directory contains a low-memory distillation path that trains `gemma_prune` from the frozen `gemma_2b` prefix branch used by the LIBERO `pi05` checkpoint.
+
+The runner keeps memory down by:
+- freezing SigLIP
+- freezing the teacher Gemma-2B
+- dropping the action branch from the distillation graph
+- training only the pruned student on a local LeRobot-format dataset at `/root/flatten_fold_v2`
+- reusing the LIBERO normalization stats from `/root/pi_train/pi05_libero/assets`
+- limiting the loaded training subset by default to about `50_000` examples
+
+The loss is a weighted sum of:
+- hidden-state MSE on all valid prefix tokens
+- cosine distance on the same tokens
+
+The student is warm-started from the teacher by:
+- copying the full token embedder and final norm
+- copying the first 14 attention and norm layers
+- slicing the teacher MLP weights down to the pruned hidden size
+
+Run it from the repo root:
+
+```bash
+.venv/bin/python prune_distill/train_prefix_distill.py \
+  --exp-name gemma_prune_prefix \
+  --dataset-path /root/flatten_fold_v2 \
+  --max-examples 50000 \
+  --overwrite
+```
+
+TensorBoard logs are written to `checkpoints/prune_distill/<exp_name>/tensorboard`.
+
+```bash
+.venv/bin/python -m tensorboard.main --logdir checkpoints/prune_distill
+```
+
+Useful flags:
+
+```bash
+.venv/bin/python prune_distill/train_prefix_distill.py \
+  --exp-name gemma_prune_prefix \
+  --dataset-path /root/flatten_fold_v2 \
+  --batch-size 8 \
+  --max-examples 50000 \
+  --num-train-steps 10000 \
+  --log-interval 20 \
+  --save-interval 500 \
+  --overwrite
+```
+
+If you still hit disk pressure during dataset materialization, lower the subset further:
+
+```bash
+.venv/bin/python prune_distill/train_prefix_distill.py \
+  --dataset-path /root/flatten_fold_v2 \
+  --max-examples 10000
+```
+
+Resume from the latest saved `step_*` checkpoint:
+
+```bash
+.venv/bin/python prune_distill/train_prefix_distill.py \
+  --exp-name gemma_prune_prefix \
+  --dataset-path /root/flatten_fold_v2 \
+  --num-train-steps 10000 \
+  --resume
+```
+
+## PI0.5 Sensitivity Analysis
+
+Use the sensitivity analyzer to score which `pi05` student tensors are most fragile for quantization, pruning, and distillation drift.
+
+It will:
+- rank student layers by distillation gradient / Taylor sensitivity
+- run fake-quant perturbations on the top-ranked tensors
+- run magnitude-prune perturbations on the same tensors
+- save `summary.json`, `candidate_scores.csv`, and `family_summary.csv`
+
+Example:
+
+```bash
+.venv/bin/python prune_distill/analyze_pi05_sensitivity.py \
+  --exp-name pi05_sensitivity \
+  --dataset-path /liujinxin/ZZF/openpi/datasets/piper/flatten_fold_v2 \
+  --student-checkpoint /root/openpi-wr/checkpoints/prune_distill/gemma_prune_prefix/step_0072001/student \
+  --max-examples 2048 \
+  --max-batches 4 \
+  --eval-top-k 24 \
+  --quant-bits 8 4 \
+  --prune-ratios 0.1 0.3 0.5 \
+  --overwrite
+```
+
+Results are written to `checkpoints/prune_distill/sensitivity/<exp_name>`.
+
+## PI0.5 Benchmark
+
+Use the benchmark runner to evaluate:
+- the original full `pi05` checkpoint on offline datasets
+- the distilled pruned prefix checkpoint on offline teacher-agreement metrics
+- the original full `pi05` checkpoint on real LIBERO rollout success
+
+The current distilled student checkpoint is only a pruned prefix model, so it does not support LIBERO action rollouts yet.
+
+Offline benchmark on a custom dataset:
+
+```bash
+.venv/bin/python prune_distill/benchmark_pi05_models.py \
+  --exp-name pi05_dataset_benchmark \
+  --origin-checkpoint-dir /root/pi_train/pi05_libero \
+  --pruned-student-checkpoint /root/openpi-wr/checkpoints/prune_distill/gemma_prune_prefix/step_0072001/student \
+  --dataset-path /liujinxin/ZZF/openpi/datasets/piper/flatten_fold_v2 \
+  --max-examples 50000 \
+  --max-eval-examples 256 \
+  --overwrite
+```
+
+LIBERO rollout success for the original full `pi05` checkpoint:
+
+```bash
+.venv/bin/python prune_distill/benchmark_pi05_models.py \
+  --exp-name pi05_libero_rollout \
+  --origin-checkpoint-dir /root/pi_train/pi05_libero \
+  --no-run-dataset-benchmark \
+  --run-libero-rollout \
+  --libero-task-suite-names libero_spatial libero_object libero_goal libero_10 \
+  --libero-num-trials-per-task 10 \
+  --overwrite
+```
+
+Benchmark outputs are written to `checkpoints/prune_distill/benchmark/<exp_name>`.
+
+If the pruned prefix benchmark hits JAX GPU OOM during model init, rerun it on CPU:
+
+```bash
+JAX_PLATFORMS=cpu .venv/bin/python prune_distill/benchmark_pi05_models.py \
+  --exp-name pi05_dataset_benchmark_cpu \
+  --origin-checkpoint-dir None \
+  --pruned-student-checkpoint /root/openpi-wr/checkpoints/prune_distill/gemma_prune_prefix/step_0072001/student \
+  --dataset-path /liujinxin/ZZF/openpi/datasets/piper/flatten_fold_v2 \
+  --max-eval-examples 64 \
+  --overwrite
+```