Skip to content

Add pi0-lerobot test case: fine-tune Pi0 on DROID/LIBERO with HyperPod EKS - #1174

Open
amindm wants to merge 10 commits into
awslabs:mainfrom
amindm:pi0-lerobot-test-case
Open

Add pi0-lerobot test case: fine-tune Pi0 on DROID/LIBERO with HyperPod EKS#1174
amindm wants to merge 10 commits into
awslabs:mainfrom
amindm:pi0-lerobot-test-case

Conversation

@amindm

@amindm amindm commented Jul 9, 2026

Copy link
Copy Markdown

Summary

Adds a test case for fine-tuning Physical Intelligence's π0 (Pi-Zero) 3B VLA model
on DROID and LIBERO robotics benchmarks using SageMaker HyperPod EKS.

Results (held-out evaluation, p5.48xlarge)

Dataset Train Episodes Eval Episodes MSE Reduction Latency (1 ODE step)
DROID 0-79 80-99 89.7% 199 ms
LIBERO 0-303 304-378 88.9% 197 ms

Training: 20K steps, ~6.5 hours per dataset, FSDP on 8× H100.

What's included

  • DLC-direct manifests (no custom image build required)
  • Dockerfile + buildspec (optional, for faster startup)
  • evaluate_pi0.py with ODE step-count sweep
  • Proper train/test split enforced via --dataset.episodes
  • Tested end-to-end on HyperPod EKS (us-west-2, p5.48xlarge)

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 1/8 — Evaluation Methodology

The reported MSE reductions measure the training set, not a held-out set — 🔴 the headline claim doesn't hold as written [confirmed]

Observation: droid-finetune.yaml trains on /fsx/datasets/droid_100 (the full 100-episode dataset — it hf downloads it itself if absent, applying no split), and droid-eval.yaml passes --test-dataset-local /fsx/datasets/droid_100 — the same path. evaluate_pi0.py then scores for ep in range(args.num_trajectories) — episodes 0–4, i.e. the first five training episodes. The download Job's comment concedes the 80/20 split is unimplemented ("For now, copy the full dataset as train") and copies the full set into both _train and _test, and the finetune/eval manifests don't consume those paths anyway. The second independent pass concurred from the code side (s3_suffix naming implies a split that nothing enforces).
Impact: "DROID 3.897e-01 → 1.353e-02 (96.5%)" and "LIBERO 7.214e-01 → 5.003e-02 (93.1%)" are open-loop error on data the model trained on for 20K steps. evaluate_pi0.py's "held-out test set" docstring is inaccurate, and the base-vs-fine-tuned comparison is measuring memorization.
Suggestion: Use LeRobot's native --dataset.episodes for a real split (no copying): train on --dataset.episodes="[0..79]", evaluate LeRobotDataset(..., episodes=list(range(80,100))). This also lets you delete both download Jobs. Even if the deltas stay large, the numbers would then support the README's claim. — repo precedent: DatasetConfig.episodes: list[int] | None in src/lerobot/configs/default.py (verified live, 2026-07-09).
e2e confirmation (2026-07-10, B300): trained a checkpoint for just 30 steps and ran the PR's own eval against droid_100 episodes 0–1 (the same data): base MSE 0.349 → "fine-tuned" 0.104, a 70% reduction after 30 steps. That large a drop from a trivial run is memorization of the eval episodes — direct evidence the headline 96.5%/93.1% (20K steps, same contamination) is not measuring generalization.


Open-loop MSE is a proxy, not the LIBERO metric — 🟡 declare the smoke-test framing, add a seed, fix the results-table pairing

Observation: The standard LIBERO metric (π0 paper, openpi, LeRobot's own lerobot-eval --env.type=libero) is closed-loop success rate; open-loop action MSE weakly correlates with task success. π0 samples actions by integrating an ODE from noise and no seed is set (independent code pass), so metrics vary run-to-run with N=5. The top-level Results table also pairs the 10-step MSE (1.353e-02) with the 1-step latency (199 ms) in one row, and 1-step MSE beats 10-step on both datasets — counter-intuitive for flow matching and likely N=5 noise.
Impact: Readers will read the table as a LIBERO performance claim; the mixed step-count row and unseeded stochasticity make the numbers hard to reproduce.
Suggestion: One sentence in the Results sections stating this validates the end-to-end pipeline and is not a success-rate claim; set a torch/numpy seed; report both MSE and latency at the same step count (or footnote); state N and the raw per-trajectory numbers. A closed-loop lerobot-eval (with [libero] extra + MUJOCO_GL=egl) is a good follow-up, not a blocker — LeRobot's π0.5 LIBERO reproduction anchors ~97.5% avg success. — docs: https://huggingface.co/docs/lerobot/libero

DATASET_PATH="/fsx/datasets/droid_100"
OUTPUT_DIR="/fsx/runs/pi0-droid"
LOCAL_DATASET_NAME="droid_local"
RENAME_MAP='{"observation.images.exterior_image_1_left":"observation.images.base_0_rgb","observation.images.wrist_image_left":"observation.images.wrist_0_rgb"}'

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DROID rename map targets wrist_0_rgb, a key lerobot/pi0_base doesn't declare — 🔴 wrist camera silently dropped [confirmed]

Observation: lerobot/pi0_base's config.json declares exactly three image features: observation.images.base_0_rgb, left_wrist_0_rgb, right_wrist_0_rgb. The DROID map renames wrist_image_leftobservation.images.wrist_0_rgb (none of the three), while the LIBERO map one block over correctly uses left_wrist_0_rgb. Both independent passes flagged this as a typo confirmed by the LIBERO entry.
Impact: The renamed wrist key isn't consumed; empty_cameras=1 pads the wrist slot with zeros, so DROID trains and evaluates on the exterior camera only — half the visual signal discarded, and (if training used this same map) the base-vs-fine-tuned comparison runs on a crippled input pipeline.
Suggestion:

Suggested change
RENAME_MAP='{"observation.images.exterior_image_1_left":"observation.images.base_0_rgb","observation.images.wrist_image_left":"observation.images.wrist_0_rgb"}'
RENAME_MAP='{"observation.images.exterior_image_1_left":"observation.images.base_0_rgb","observation.images.wrist_image_left":"observation.images.left_wrist_0_rgb"}'

Apply the matching change to DATASET_CONFIGS["droid"] in evaluate_pi0.py:40 and re-run DROID. — source: lerobot/pi0_base/config.json (verified live, 2026-07-09).


from transformers import AutoTokenizer
try:
tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224")

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Eval swallows tokenizer/processor/postprocess failures and still emits numbers — 🟡 a green run can be silently meaningless

Observation: AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224") (line 478), make_pre_post_processors (≈458), and preprocess/postprocess (≈266, ≈289) are each wrapped in try/except that warns (only on chunk_start==0) and continues with None/raw values. Under the HF_HUB_OFFLINE=1 the eval sets, an uncached tokenizer load → tokenizer=None → no language conditioning; a failed postprocess leaves predictions in normalized space while ground truth is raw → MSE explodes with no error. (Independent code pass, accepted.)
Impact: The script can write a complete results JSON that is quietly invalid (wrong conditioning, or normalized-vs-raw action comparison). A passing run does not imply a valid measurement.
Suggestion: Make processor/tokenizer load failures fatal, or stamp "degraded": true in the JSON and refuse to emit comparison numbers; print the warnings unconditionally, not just on the first chunk.

# -------------------------------------------------------------------------
# Backfill empty FSDP-frozen params from the base checkpoint
# -------------------------------------------------------------------------
def backfill_empty_params(policy, source_path: str, base_repo_id: str = "lerobot/pi0_base"):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

backfill_empty_params fills partially without erroring — 🟢 (checkpoint-incompleteness concern refuted e2e)

Observation: The eval loads the fine-tuned checkpoint then backfills every numel()==0 param from lerobot/pi0_base. I checked this e2e: a real lerobot-train FSDP save (SHARDED_STATE_DICT) produced a complete checkpoint — 898 params, 0 zero-numel — and the eval's backfill found nothing to fill. So the "checkpoint is not self-contained" worry does not reproduce; the save gathers params correctly. The only residual nit is robustness: load_state_dict(new_state, strict=False) + printing n_filled/len(empty_keys) would silently proceed if a future save/name-scheme change did leave params empty.
Impact: Low — today the backfill is a defensive no-op. It's dead-ish code that could mask a real incomplete-checkpoint bug if one were ever introduced.
Suggestion: Either drop the backfill (the save is complete), or keep it but raise when n_filled != len(empty_keys) so a genuinely incomplete checkpoint fails loudly instead of scoring garbage. — verified live 2026-07-10 (checkpoint inspected on B300).

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 2/8 — Training Configuration Correctness

FSDP FULL_SHARD delivers no real memory sharding for π0 — 🔴 measured e2e, and the obvious fix crashes

Observation: The launch uses --fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP with no --fsdp_transformer_layer_cls_to_wrap. I measured per-GPU HBM (nvidia-smi) on a real run of this exact config (π0 on B300, batch 1): 1 GPU = 16.5 GB, 2 GPU = 79.6 GB, 8 GPU = 73.8 GB. Going 1→2 GPUs increases per-GPU memory ~5× and 2→8 is flat (~7%). That is the opposite of sharding: with no wrappable block, the whole 3B model is a single FSDP unit that gets all-gathered in full every step (plus comm buffers), so FULL_SHARD provides ~no memory scaling. (--fsdp_use_orig_params=true is correctly set — LeRobot builds the optimizer before prepare(), so it's required.) I also verified the naive fix does not work: setting --fsdp_transformer_layer_cls_to_wrap to π0's real block class (_PiGemmaDecoderLayerBase, SiglipEncoderLayer) crashes with RuntimeError: size of tensor a (2048) must match b (0) — π0 is a mixture-of-transformers with joint cross-stream attention, so its decoder layers can't be independently block-wrapped (a shared param shards to size 0 on some ranks).
Impact: It runs on B300 only because ~80 GB fits in 268 GB. On the PR's stated target (p5 / H100-80GB), ~80 GB peak is an OOM risk with no headroom, and adding GPUs won't reduce it — so the recipe likely can't scale on the very hardware its README advertises. The "FSDP FULL_SHARD" label overstates what the config achieves.
Suggestion: Don't claim FULL_SHARD memory scaling without it working. π0 needs a custom FSDP wrap policy that treats each joint attention block (the PaliGemma-layer + action-expert-layer pair that shares attention) as one wrapped unit, or a different strategy (e.g. the openpi/JAX sharding, or per-stream wrapping that respects the shared attention). At minimum, document that this is effectively single-replica-per-rank and size the instance accordingly. — verified live 2026-07-10 on p6-b300; docs: https://huggingface.co/docs/lerobot/en/multi_gpu_training


RUN pip install --no-cache-dir --upgrade pip && \
pip install --no-cache-dir \
"lerobot[pi,dataset] @ git+https://github.com/huggingface/lerobot.git" \

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LeRobot installed from unpinned git HEAD — 🟡 non-reproducible, and π0 has a silent-weight-mismatch history [confirmed]

Observation: Dockerfile:37 and the runtime install in droid-finetune.yaml:70 both pull git+https://github.com/huggingface/lerobot.git with no ref — HEAD at build/run time. The repo's single most-violated rule (pin pip/git clones: SHA > tag > branch). LeRobot is pre-1.0 (latest v0.6.0) and its CLI schema, dataset format, and the transitive transformers pin all move with the ref; π0 weight loading has silently mismatched across transformers bumps before (loads without error, trains from a garbage init). Both passes flagged the unpinned install; the infra pass also noted libero-finetune uses the baked image while droid-finetune installs at runtime, so the two datasets can run different commits.
Impact: Results aren't reproducible; a future HEAD can break the --rename_map/--optimizer.lr flags or silently mis-load π0 weights.
Suggestion: Pin @<tested-sha-or-tag> in the Dockerfile and drop the runtime install (use the custom image like LIBERO does). After pinning, assert the base model's initial eval MSE is in range before the 6.5-hour run as a cheap mismatch guard. The DLC torch 2.9 sits inside LeRobot's torch>=2.7,<2.12, so the pin won't clobber the AWS stack. — upstream: huggingface/lerobot#1406 ; pyproject.toml transformers>=5.4,<5.6 (verified live, 2026-07-09).

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 3/8 — Structure & Repository Hygiene

Build on openvla-oft and cosmos3 rather than a parallel scaffold — 🟡 not a duplicate, but reuse the proven siblings

Observation: π0 + LeRobot + FSDP is a genuinely new model/framework combination (the repo already keeps openvla and openvla-oft separate, so a distinct directory is justified), but it overlaps two siblings closely. openvla-oft is the near-twin scaffold — same LIBERO-on-EKS kubernetes/libero/ recipe — and already solves cleanly every K8s-scaffolding issue flagged in this review: env_vars.example + whitelisted envsubst, HF_TOKEN via Secret only, an image tag pinned to the Dockerfile commit, a verify-tfds-layout.sh sanity check, and the plain-EKS-vs-HyperPod FSx split. cosmos3 post-trains on the same lerobot/droid_100 subset and ships LeRobot-v3 dataset loaders with tests (src/cosmos3_aws/action/lerobot_v3_action_dataset.py).
Impact: The PR reinvents (more weakly) scaffolding a sibling already got right, and hand-rolls a local-dataset loader via monkey-patching when in-repo prior art exists.
Suggestion: Adopt openvla-oft's K8s recipe conventions; look at cosmos3's LeRobot loader before patching the framework. — repo precedent: 3.test_cases/pytorch/openvla-oft/, 3.test_cases/pytorch/cosmos3/.


Download Jobs: broken envsubst, a fake split, and orphaned paths — 🔴 [confirmed]

Observation: (1) $${HF_TOKEN:-} and $${TRAIN_DIR} assume $$$ (Compose/CloudFormation), which GNU envsubst does not do — reproduced live: ${TRAIN_DIR} (unexported) renders to empty (Path('') → CWD, so the copytree is skipped), and $$ then expands to the shell PID, so the token guard is always-true garbage and the real HF_TOKEN secret is never referenced. (2) The "split" copytrees the full dataset into both _train and _test (comment concedes it) — train/test leakage plus 3× FSx usage. (3) Nothing consumes them: they write /data/datasets/droid_100_{train,test} on ${DATA_PVC_NAME} while the finetune/eval Jobs read /fsx/datasets/droid_100 on fsx-claim and re-download themselves; no README invokes them. Both independent passes reached this.
Impact: A user who runs them wastes time/storage and still triggers a second in-Job download; if a gated repo were used, the token guard would silently fail.
Suggestion: Given the --dataset.episodes split (Batch 1) needs no staging Job, delete both download Jobs. If you keep a staging Job, reduce it to the single hf download with paths/PVC aligned to the finetune manifests, and render with an explicit whitelist: envsubst '$NAMESPACE $IMAGE_URI $DATA_PVC_NAME' < file.yaml, referencing ${HF_TOKEN} as the plain env var it is.


Remove the monkey-patch layer — 🟡 unnecessary on current LeRobot, three divergent copies, and the standalone invocation is a no-op

Observation: The Hub-skip patch exists in three divergent copies (standalone lerobot_local_patch.py, a PATCHEOF heredoc in droid-finetune.yaml, and apply_local_dataset_patch() in evaluate_pi0.py), all with bare except: pass. It's patching a problem current LeRobot doesn't have: LeRobotDatasetMetadata.__init__ and LeRobotDataset.__init__ only hit the Hub when the local load fails, so with --dataset.root staged no Hub call happens. And the independent code pass caught that python /opt/.../lerobot_local_patch.py in run_finetuning.sh/evaluate_pi0.sh runs the patch in a throwaway subprocess that exits before accelerate launch — it has zero effect on training; offline behavior rides entirely on HF_HUB_OFFLINE=1.
Impact: Maintenance hazard, and a patch that (for training) does nothing while looking load-bearing. The repo convention is to wrap upstream, not ship divergent runtime monkey-patches.
Suggestion: Pin a LeRobot version with this behavior, pass --dataset.root (already done), and delete lerobot_local_patch.py, the heredoc, the usercustomize.py writers, and apply_local_dataset_patch(). If your pin still hits a Hub lookup on a staged dataset, that's an upstream bug to file, not to patch here.

app: pi0-lerobot-droid-finetune
spec:
nodeSelector:
node.kubernetes.io/instance-type: ml.p5.48xlarge

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Make the manifests run on plain EKS too — 🔴 the ml.-prefixed nodeSelector is HyperPod-only, and blocks the documented p4de option [confirmed]

Observation: All four finetune/eval manifests hardcode node.kubernetes.io/instance-type: ml.p5.48xlarge. The ml. prefix is applied by SageMaker HyperPod; on plain EKS (managed node groups / Karpenter) the AWS cloud-provider sets that label to the raw p5.48xlarge, so these pods stay Pending there. The per-dataset READMEs are titled "…on Amazon EKS" and the parent README's Prerequisites list "p5.48xlarge or p4de.24xlarge" — but the hardcoded selector never schedules p4de either (independent infra pass). Nothing in the training/eval logic is HyperPod-specific.
Impact: Silent unschedulability on plain EKS and on p4de, contradicting the READMEs — and per your steer, no strong reason to restrict it.
Suggestion: Parameterize node.kubernetes.io/instance-type: "${INSTANCE_TYPE}", default INSTANCE_TYPE=p5.48xlarge (no prefix) in an env_vars.example, and note HyperPod users set ml.p5.48xlarge. This is exactly the openvla-oft sibling's pattern (kubernetes/libero/libero-finetune.yaml:83).

Suggested change
node.kubernetes.io/instance-type: ml.p5.48xlarge
node.kubernetes.io/instance-type: "${INSTANCE_TYPE}"

--multi_gpu \
--num_processes="${NUM_GPUS}" \
--mixed_precision=bf16 \
$(python -c "import lerobot; import os; print(os.path.join(os.path.dirname(lerobot.__file__), '..', 'scripts', 'train.py'))" 2>/dev/null || which lerobot-train) \

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

src/run_finetuning.sh is dead code that also can't run — 🔴 wrong CLI namespace, dead entrypoint resolution, and it's DDP not FSDP [confirmed]

Observation: Neither manifest calls this script (both inline their own accelerate launch), the README Layout points readers to it as "Training entrypoint", and as written it cannot work: (a) every flag uses a --training.* namespace that lerobot-train doesn't have (correct forms are --steps, --batch_size, --optimizer.lr, --save_freq, --output_dir); (b) line 69's $(python -c "...os.path.join(...,'scripts','train.py')") only builds a string — the python -c exits 0, so the || which lerobot-train fallback is unreachable, and lerobot/scripts/train.py no longer exists in the package; (c) it launches accelerate launch --multi_gpu = DDP, contradicting the FSDP framing, and passes no --rename_map. Both catches from the independent code pass.
Impact: Anyone following the README runs a script that fails immediately (repo policy: examples must ship working code).
Suggestion: Delete run_finetuning.sh and evaluate_pi0.sh (same class of problem, unreferenced, /data/... paths), letting the manifests be the single source of truth — or make the manifests exec these scripts (already COPYd to /opt/pi0-lerobot/src/) so the command lives in exactly one place, corrected to the working invocation.


# Auto-discover checkpoint if not provided
if args.finetuned_path is None:
job_prefix = f"pi0-finetune-{args.dataset}"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dangling Pi0_SMTJ / sagemaker-user references and dead checkpoint auto-discovery — 🟡

Observation: Manifest headers ("EXACT same invocation as Pi0_SMTJ/scripts/run_finetuning.sh"), the eval docstring ("mirrors 02_training_job.ipynb"), and the ~40-line checkpoint auto-discovery walking /home/sagemaker-user/Pi0_SMTJ/model_artifacts + a sagemaker-{region}-{account} S3 fallback are all artifacts of the SageMaker Training Jobs project this was ported from.
Impact: Dangling references for a repo reader; the auto-discovery can never succeed on an EKS pod (the eval manifests always pass --finetuned-path).
Suggestion: Remove the comments and the auto-discovery/S3-fallback block; make --finetuned-path required — simpler and honest about how this runs.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 4/8 — Deployment Pipeline & Operational Correctness

No env_vars.example; namespace/PVC/instance hardcoded four different ways — 🟡

Observation: The download Jobs use ${IMAGE_URI}/${NAMESPACE}/${DATA_PVC_NAME} while the finetune/eval manifests hardcode namespace: default, claimName: fsx-claim, and the image — the same values spelled inconsistently across six manifests, and pvc-fsx-lustre.yaml provisions ${DATA_PVC_NAME} that the Jobs then reference as the literal fsx-claim (infra pass L1).
Impact: A user who renders the PVC under any other name finds the Jobs can't bind it; config lives in six places.
Suggestion: Add an env_vars.example (mirroring openvla-oft's) defining IMAGE_URI, NAMESPACE, DATA_PVC_NAME, INSTANCE_TYPE, EFA/GPU counts, and render every manifest through a whitelisted envsubst; keep HF_TOKEN in the Secret, out of envsubst.

export PYTHONUNBUFFERED=1

# Clean stale output from previous runs
rm -rf "$OUTPUT_DIR/training" 2>/dev/null || true

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Finetune wipes checkpoints on every start with no resume — 🔴 data loss on any restart

Observation: Both finetune manifests rm -rf "$OUTPUT_DIR/training" at startup (droid line 153, libero line 92) with restartPolicy: Never, despite save_freq=2000 writing checkpoints to persistent FSx (independent infra pass).
Impact: A 6.5-hour run interrupted and re-applied deletes all checkpoints and retrains from step 0; a stray re-apply silently destroys a completed run before eval can read checkpoints/020000/pretrained_model.
Suggestion: Guard the wipe behind an explicit FORCE_RESTART flag, or resume from the latest checkpoint when one is present.

effect: NoSchedule
containers:
- name: eval
image: 783764584149.dkr.ecr.us-east-2.amazonaws.com/pi0-lerobot:v1.0.0

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Eval + LIBERO-finetune hardcode a private ECR account no one else can pull — 🔴 [confirmed]

Observation: droid-eval.yaml:25, libero-eval.yaml:25, and libero-finetune.yaml:30 pin 783764584149.dkr.ecr.us-east-2.amazonaws.com/pi0-lerobot:v1.0.0 — the author's account and a baked region — while droid-finetune.yaml:31 correctly uses the public DLC account 763104351884. So the pattern is inconsistent across the four job manifests. Both passes flagged it.
Impact: Any external user who applies these three (the README Quick Start says to apply them verbatim) gets ImagePullBackOff.
Suggestion: Parameterize as ${IMAGE_URI} (the download Jobs already do) with an <account>.dkr.ecr.<region>... placeholder and a "replace this" cue in the README.

ln -snf /fsx/datasets/droid_100 "$HF_LEROBOT_HOME/droid_local"

# Patch the evaluate script to skip S3 auto-detect
sed -i 's/if args.test_dataset_s3 is None:/if args.test_dataset_s3 is None and not (args.test_dataset_local and args.test_dataset_local.exists()):/' /opt/pi0-lerobot/src/evaluate_pi0.py

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Runtime sed -i self-patch of the eval script is a no-op against the shipped code — 🟡 [confirmed]

Observation: Both eval manifests sed-rewrite evaluate_pi0.py's if args.test_dataset_s3 is None: line at pod start, but the shipped evaluate_pi0.py (lines 404–406) already handles the local case first (args.test_dataset_s3 = "local" when --test-dataset-local exists). Leftover from patching an older baked copy. Independent infra pass also flagged the sed as silently no-op-prone.
Impact: Manifests that rewrite the code they run are a review red flag and drift silently.
Suggestion: Drop both sed lines; if any gap remains, fix it in evaluate_pi0.py (which is in this PR). Repo convention: if a Dockerfile sed is unavoidable, have it grep-verify its own result so an upstream change fails loudly.

containers:
- name: download
image: ${IMAGE_URI}
imagePullPolicy: Always

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

imagePullPolicy: Always on pinned download Jobs — 🟡 [confirmed]

Observation: ${IMAGE_URI} is a pinned tag, so Always just forces a registry round-trip per pod start and breaks air-gapped clusters; the finetune/eval manifests already use IfNotPresent.
Impact / Suggestion:

Suggested change
imagePullPolicy: Always
imagePullPolicy: IfNotPresent

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 5/8 — Infrastructure, NCCL & Container Security

HF token: persisted to shared FSx and interpolated into argv — 🔴 (security carve-out) [confirmed]

Observation: (1) huggingface-cli login / login(token=...) persists the token to $HF_HOME/token, and HF_HOME points at shared FSx (/data/hf-cache, /fsx/hf-cache) — readable by every pod/user on that PVC. (2) run_finetuning.sh:32 and evaluate_pi0.sh:29 interpolate the token into argv (login(token='${HF_TOKEN}')) — visible in ps//proc, leaks under set -x, and breaks on a token containing ' (independent code pass). (3) README Quick Start creates the Secret with --from-literal, leaking into shell history (infra pass, minor).
Impact: Credential exposure on a shared filesystem and in the process table; masked auth errors from 2>/dev/null || true.
Suggestion: Drop the explicit login() calls entirely — huggingface_hub reads HF_TOKEN from the env (the secretKeyRef already provides it). The YAML manifests already do this correctly (os.environ['HF_TOKEN']); make the shell scripts match if they survive the dead-code cleanup.


Training hard-fails without a Gemma-licensed HF token — 🔴 gated PaliGemma, and the README undersells it (found e2e)

Observation: Running the finetune verbatim with no HF token, training crashed ~15 min in (after install + dataset download) with GatedRepoError 401 on google/paligemma-3b-pt-224: π0's processor pipeline (make_pre_post_processors) fetches the gated PaliGemma tokenizer/config at train time. The README lists the token as "for base model download," but lerobot/pi0_base is ungated — the real gate is the transitive PaliGemma dependency, which requires accepting the Gemma license. The manifests mark HF_TOKEN optional: true, so a tokenless user only discovers this deep into the run. (With a Gemma-licensed token, training then ran cleanly to a saved checkpoint.)
Impact: A user following the README without a Gemma-accepted token burns ~15 min before a hard failure that the docs don't predict; optional: true actively misleads.
Suggestion: State in Prerequisites that the token must have accepted the Gemma license for google/paligemma-3b-pt-224 (with the link), make HF_TOKEN required (not optional) for the finetune/eval Jobs, and ideally fail fast with a clear message if it's unset. This also ties to eval finding E3: the eval wraps the same PaliGemma tokenizer load in try/except, so there it fails silently (no language conditioning) instead of loudly. — verified live 2026-07-10.

securityGroupIds: "${FSX_SECURITY_GROUP_ID}"
deploymentType: SCRATCH_2
storageType: SSD
reclaimPolicy: Delete

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

StorageClass reclaimPolicy: Delete (SCRATCH_2) under datasets and checkpoints — 🔴 data loss [confirmed]

Observation: reclaimPolicy: Delete means a kubectl delete pvc (or namespace teardown) destroys the FSx filesystem and every dataset/checkpoint on it — while the READMEs sell "data persists on FSx across jobs." SCRATCH_2 additionally has no replication (infra pass). Both passes flagged it.
Impact: A routine PVC delete silently wipes non-reproducible data, contradicting the recipe's value proposition.
Suggestion:

Suggested change
reclaimPolicy: Delete
reclaimPolicy: Retain

For checkpoints meant to persist, note the SCRATCH_2 durability tradeoff (or use PERSISTENT_2) in the header comment.

ENV HF_HOME=/data/hf-cache
ENV HF_LEROBOT_HOME=/data/lerobot-cache
ENV NCCL_DEBUG=WARN
ENV NCCL_TIMEOUT=1800

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NCCL_TIMEOUT is not an NCCL variable; GLOO_SOCKET_IFNAME=lo is a scale-out trap — 🟡

Observation: NCCL_TIMEOUT=1800 (Dockerfile:59 and both finetune manifests) is a no-op — NCCL has no such variable (collective timeouts come from the framework process-group). NCCL_SOCKET_IFNAME="^lo" is the correct exclusion pattern (good — this is the repo's most-flagged technical issue and it's right). GLOO_SOCKET_IFNAME=lo works only at replicas: 1; scaling the PyTorchJob out makes gloo rendezvous over loopback fail confusingly, and vpc.amazonaws.com/efa: 32 is requested but idle intra-node.
Impact: A dead env var implies behavior it doesn't have; the gloo/loopback default silently blocks the multi-node path the README advertises.
Suggestion: Drop NCCL_TIMEOUT; add a one-line comment that GLOO_SOCKET_IFNAME=lo and the EFA request are single-node defaults to revisit when scaling replicas.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 6/8 — Documentation Consistency

Missing MIT-0 headers on all non-README files — 🟡 lint gate [confirmed]

Observation: Only the three READMEs carry the header; the Dockerfile, buildspec, all seven YAML manifests, and the four src/ files are missing it. This is the repo's single most-recurring miss, and the license lint checks it (YAML included, before any ---). Both passes flagged it.
Suggestion: Prepend the two-line # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. / # SPDX-License-Identifier: MIT-0 to each file.

├── buildspec.yml # AWS CodeBuild spec for image builds
├── README.md # This file
├── src/
│ ├── run_finetuning.sh # Training entrypoint (accelerate + FSDP)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

README Layout documents the broken files; add a Gemma-license note — 🟡

Observation: The Layout tree annotates run_finetuning.sh as "Training entrypoint (accelerate + FSDP)" and lerobot_local_patch.py as part of the flow — but both should be removed or wired in per Batch 3. The Prerequisites also omit the base model's license.
Impact: Sends readers to the broken script; users' legal reviewers have no license cue.
Suggestion: Make the Layout section match the final file list, and add one line: lerobot/pi0_base licensing is ambiguous (LeRobot docs describe the port as Apache-2.0, the HF model card tags license:gemma via PaliGemma) — point users at the model card before commercial use. — source: https://huggingface.co/lerobot/pi0_base

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 7/8 — Minor code-quality nits

link = cache_root / repo_id
if link.is_symlink() or link.exists():
link.unlink() if link.is_symlink() else None
link.symlink_to(test_dir.resolve())

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

setup_lerobot_cache raises FileExistsError when a real directory pre-exists at the link path — 🟢

Observation: link.unlink() if link.is_symlink() else None does nothing when link is a real directory, then link.symlink_to(...) raises (independent code pass). The expression-as-statement is also a smell.
Suggestion: if link.is_symlink() or link.exists(): link.unlink() (with shutil.rmtree for the dir case if that's expected), using a normal if.

tokenizer = None

results = []
for ep in range(args.num_trajectories):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

num_trajectories has no bound against available episodes — 🟢

Observation: for ep in range(args.num_trajectories) (default 5) indexes dataset.meta.episodes[...][ep]; a test set with fewer episodes raises IndexError instead of clamping.
Suggestion: for ep in range(min(args.num_trajectories, n_eps)).

from safetensors.torch import load_file

base_dir = Path(snapshot_download(base_repo_id, allow_patterns=["*.safetensors", "*.json"]))
base_st_files = list(base_dir.glob("*.safetensors")) + list(base_dir.glob("model-*.safetensors"))

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

snapshot_download glob double-counts sharded safetensors — 🟢

Observation: list(base_dir.glob("*.safetensors")) + list(base_dir.glob("model-*.safetensors")) — the first glob already matches the second, so each shard is load_filed twice (the dict update dedups, so it's wasted work, not wrong output).
Suggestion: Drop the second glob.

task = "move object to target"
batch["task"] = task
if isinstance(task, list):
task = task[0] if task else "move object to target"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

build_batch_for_policy leaves batch["task"] a list after unwrapping — 🟢

Observation: task = task[0] if task else ... updates the local var (used correctly by the tokenizer) but not batch["task"], so preprocess/the policy still receive a list (independent code pass).
Suggestion: Assign batch["task"] = task after unwrapping.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 8/8 — Evaluation, Positives & Sources

Things That Look Great

  • It ran end-to-end on real hardware — 20K steps × 2 datasets on p5.48xlarge, an ODE-step sweep, and per-trajectory latency measured with proper torch.cuda.synchronize() bracketing. Far more validation than most first-round test-case PRs.
  • Independently reproduced e2e (reviewer, 2026-07-10, 8× B300): with only environment adaptations, the recipe installed, downloaded, trained, and saved a complete, resumable checkpoint (model + processors + optimizer + RNG). Concretely validated: S1 (adapting the nodeSelector/tolerations makes it schedule on non-HyperPod EKS), and T2 (the unpinned LeRobot HEAD install works today on CUDA-13/Blackwell and preserves the DLC's torch=2.9.0+cu130). The core training path is sound.
  • The FSDP launch is the genuinely supported path: lerobot-train on main is accelerate-integrated (autocast, accelerator.backward, clip_grad_norm_), and --fsdp_use_orig_params=true is correctly set (required because LeRobot builds the optimizer before prepare()).
  • NCCL_SOCKET_IFNAME: "^lo" uses the correct exclusion pattern out of the gate — the repo's single most-flagged technical issue, and this PR got it right.
  • Hyperparameters are sane and consistent between the README table and the manifests (LR 2.5e-5, effective batch 32, chunk 50, gradient checkpointing).
  • --dataset.video_backend=pyav sidesteps the torchcodec/torch wheel-compatibility matrix — the right call for a DLC-pinned torch.
  • NVIDIA_VISIBLE_DEVICES: "void" on the CPU download Jobs stops the device plugin from wasting a GPU slot.
  • Secrets flow through secretKeyRef with optional: true in the manifests (the token nits are about the redundant login() step, not the plumbing), and the HyperPod-EKS-vs-SageMaker-Training-Jobs comparison table plus the symptom→fix troubleshooting section are genuinely useful.

Sources

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Left comments.

…e safety

- Fix DROID wrist camera key: wrist_0_rgb -> left_wrist_0_rgb
- Document ml. nodeSelector prefix (HyperPod vs plain EKS)
- Remove dead src/run_finetuning.sh (never called by manifests)
- Replace rm -rf with resume-aware checkpoint check
- Replace private ECR with public DLC image in eval/libero manifests
- Change FSx reclaimPolicy from Delete to Retain
- Pin LeRobot to v0.5.0 for reproducibility
- Improve eval error messages (explicit warnings for processor failures)
- Remove dangling Pi0_SMTJ/sagemaker-user references
- Remove dead sed patch from eval YAMLs
- Change imagePullPolicy to IfNotPresent on download jobs
- Add NCCL_TIMEOUT comment noting repo convention consistency
@amindm
amindm requested a review from KeitaW July 14, 2026 17:43

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 1/6 — Round-2 Resolution Scoreboard


Thank you for the thorough response round — here's where each round-1 finding stands

Commit d43eee1 ("Address PR review: fix camera key, pin deps, remove dead code, improve eval") genuinely lands several of the round-1 fixes — credit where due. The full ledger, so we're working from the same list:

Resolved (7):

  • run_finetuning.sh dead code — deleted.
  • ✅ LeRobot unpinned git HEAD — now pinned to v0.5.0 (tag resolves, verified live 2026-07-15).
  • ✅ StorageClass reclaimPolicy: Delete — now Retain.
  • imagePullPolicy: Always — now IfNotPresent everywhere.
  • ✅ Runtime sed -i self-patch of evaluate_pi0.py — removed from both eval manifests.
  • ✅ Private ECR account 783764584... hardcoded — removed (but see Batch 2: the replacement introduced a crash).
  • ✅ DROID wrist camera key in the training manifest — left_wrist_0_rgb is now correct (but see Batch 2: the eval copy wasn't fixed).

Partially addressed (6):

  • 🟨 Checkpoint rm -rf wipe — removed ✅, but the replacement "resume" is non-functional and crashes on restart (Batch 2).
  • 🟨 ml.-prefixed nodeSelector — now documented via comments + a README note. Workable; parameterizing via env_vars (Batch 4) would still be better than asking users to hand-edit four manifests.
  • 🟨 NCCL_TIMEOUT — an explanatory comment was added, but it cites a sibling that doesn't actually set the variable (Batch 5).
  • 🟨 Pi0_SMTJ / SageMaker leftovers — header comments removed ✅; the dead checkpoint auto-discovery and the sagemaker-{region}-{account} S3 fallback in evaluate_pi0.py remain (Batch 5).
  • 🟨 HF token in argv — one of two sites removed (with run_finetuning.sh); evaluate_pi0.sh:29 still interpolates the token into argv (Batch 3).
  • 🟨 README documents broken files — run_finetuning.sh gone from the Layout tree ✅; lerobot_local_patch.py still listed, and the Gemma-license note is still missing (Batches 3, 4).

Unresolved (9): eval-on-training-data (the headline numbers), FSDP FULL_SHARD not sharding, Gemma-licensed-token prerequisite, download Jobs (envsubst breakage + fake split + orphaned paths), no env_vars.example, MIT-0 headers outside the READMEs, unseeded eval stochasticity, and the four round-1 code nits in evaluate_pi0.py. Details in Batches 4–5 rather than re-litigating inline threads.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 2/6 — Regressions Introduced by the Round-1 Fixes

# Install patch
USER_SITE=$(python -c "import site; print(site.getusersitepackages())")
mkdir -p "$USER_SITE"
cp /opt/pi0-lerobot/src/lerobot_local_patch.py "$USER_SITE/"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both eval Jobs and the LIBERO finetune now crash on first run — the stock DLC image has no /opt/pi0-lerobot — 🔴 [confirmed]

Observation: Removing the private ECR reference swapped image: in droid-eval.yaml, libero-eval.yaml, and libero-finetune.yaml to the stock DLC (763104351884.dkr.ecr.us-east-2...pytorch-training:2.9.0-...). But all three scripts still reference files that only exist in the custom image this PR's Dockerfile builds: cp /opt/pi0-lerobot/src/lerobot_local_patch.py (droid-eval:42, libero-eval:42, libero-finetune:62) and python /opt/pi0-lerobot/src/evaluate_pi0.py (both evals). libero-finetune.yaml's header even says "custom Docker image — LeRobot pre-installed" while pointing at the DLC, and unlike the DROID finetune it never pip installs lerobot, so lerobot-train wouldn't exist either.
Impact: Under set -e, all three manifests die at the first cp ("No such file or directory") on their first run, as committed. backoffLimit: 1 retries the identical failure once. Only droid-finetune.yaml survives, because it happens to inline everything.
Suggestion: Pick one model and apply it to all four workload manifests: either (a) point them all at the custom image (${IMAGE_URI} via env_vars + envsubst, matching the download Jobs), which also restores the Dockerfile/buildspec to relevance, or (b) go DLC-everywhere and inline the install + scripts as droid-finetune.yaml does. Option (a) is what the README already documents as the primary path. — verified live 2026-07-15: the DLC tag exists and contains no /opt/pi0-lerobot (it's created only by this PR's COPY src/).


# Resume from existing checkpoint if available, otherwise start fresh
if [ -d "$OUTPUT_DIR/training/checkpoints" ] && ls "$OUTPUT_DIR/training/checkpoints"/*/pretrained_model/config.json 1>/dev/null 2>&1; then
echo " Found existing checkpoints — training will resume from latest."

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

"Training will resume from latest" is not true — LeRobot v0.5.0 raises FileExistsError instead — 🔴 [confirmed]

Observation: The round-1 rm -rf "$OUTPUT_DIR/training" is gone (good — that was the data-loss path), replaced in both finetune manifests by a check that only echos "Found existing checkpoints — training will resume from latest." Nothing passes --resume=true, and LeRobot's resume path also requires --config_path pointing at a checkpoint's train_config.json. At the pinned v0.5.0, TrainPipelineConfig.validate() raises FileExistsError: Output directory ... already exists and resume is False whenever output_dir exists without resume (src/lerobot/configs/train.py:122-126, verified live 2026-07-15).
Impact: Any re-run after a node failure, OOM, or a completed run — the exact scenario the message addresses — crashes immediately after the ~2-minute install phase, with a message contradicting what the log just promised. It fails safe (no data loss), but resume simply doesn't exist.
Suggestion: Implement the branch for real:

RESUME_ARGS=""
LATEST_CFG=$(ls -d "$OUTPUT_DIR"/training/checkpoints/*/pretrained_model/train_config.json 2>/dev/null | sort | tail -1)
if [ -n "$LATEST_CFG" ]; then
  echo "  Resuming from $LATEST_CFG"
  RESUME_ARGS="--resume=true --config_path=$LATEST_CFG"
fi

and append $RESUME_ARGS to the lerobot-train invocation (note: on resume, LeRobot takes its config from the checkpoint, so the other CLI flags are ignored by design). Alternatively, drop the echo and document that re-runs need a fresh --output_dir. — https://github.com/huggingface/lerobot/blob/v0.5.0/src/lerobot/configs/train.py#L122

"droid": {
"rename_map": {
"observation.images.exterior_image_1_left": "observation.images.base_0_rgb",
"observation.images.wrist_image_left": "observation.images.wrist_0_rgb",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The camera-key fix missed the eval — training now uses left_wrist_0_rgb, eval still maps to nonexistent wrist_0_rgb — 🔴 [confirmed]

Observation: The commit message says "fix camera key," and droid-finetune.yaml's RENAME_MAP is indeed fixed. But evaluate_pi0.py's DATASET_CONFIGS["droid"]["rename_map"] (line 40) still renames wrist_image_leftobservation.images.wrist_0_rgb — a key lerobot/pi0_base does not declare (its config.json has exactly base_0_rgb, left_wrist_0_rgb, right_wrist_0_rgb; verified live 2026-07-09). The LIBERO map in the same dict is correct.
Impact: The fine-tuned model was trained seeing the wrist camera on left_wrist_0_rgb; at eval it receives that feature under an unknown key, so the input the model most depends on is silently dropped. The DROID eval numbers measure a model deprived of a camera it trained with — on top of the train-set-contamination issue (Batch 4). This is also now a drift between two files that must agree.
Suggestion:

Suggested change
"observation.images.wrist_image_left": "observation.images.wrist_0_rgb",
"observation.images.wrist_image_left": "observation.images.left_wrist_0_rgb",

Longer term, both maps live in two places (manifest env + eval config) — worth deriving one from the other or asserting equality at eval start. — https://huggingface.co/lerobot/pi0_base/blob/main/config.json

aws codebuild start-build --project-name pi0-lerobot-build --region ${AWS_REGION}
```

### Alternative: DLC Direct (no custom image)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The README's image workflow is inverted relative to what the manifests do — 🟡 [confirmed]

Observation: The README presents building/pushing ${REGISTRY}/pi0-lerobot:v1.0.0 as the primary "Container Image" workflow, with "DLC Direct (no custom image)" as an alternative that tells the reader to change the manifests' image: field to the DLC. But every committed manifest already points at the DLC — no manifest consumes the custom image, so the documented build/push (and the buildspec, and CodeBuild section) currently affects nothing a user runs.
Impact: A user who follows the README builds and pushes an image that is never pulled, then hits the Batch-2 crash anyway. The doc describes the intended design; the manifests ship the opposite.
Suggestion: This resolves itself with whichever direction you pick for the crash finding above: option (a) makes the README true (manifests reference ${IMAGE_URI}); option (b) means flipping the README so DLC-direct is the primary path and the Dockerfile/buildspec/CodeBuild sections become the optional fast-startup variant.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 3/6 — Removing the Patch Layer

@@ -0,0 +1,66 @@
"""

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lerobot_local_patch.py and all its injection sites need to go — the maintainers don't want to carry a monkey patch, and at v0.5.0 it's provably unnecessary — 🔴

Observation: Speaking with my maintainer hat on: we don't want to maintain a monkey patch of LeRobot internals in this repo, so this is a merge blocker for me rather than a style preference. Round 1 flagged it at should-fix; it's unchanged in this revision, and the pin to v0.5.0 makes the case stronger — the patch now targets a known version where it demonstrably does nothing useful (all verified live, 2026-07-15, against the v0.5.0 tag):

  1. LeRobotDatasetMetadata.__init__ at v0.5.0 calls load_metadata() first and touches the Hub (get_safe_versionpull_from_repo) only in the except (FileNotFoundError, NotADirectoryError) branch — i.e., only when the local load fails (src/lerobot/datasets/lerobot_dataset.py:106-113; same pattern at 750-752 for the dataset class). Both training and eval pass --dataset.root/root= with the dataset staged, so the patched functions are never reached.
  2. lerobot/datasets/dataset_metadata.py does not exist at v0.5.0 (raw.githubusercontent.com returns 404) — the from lerobot.datasets import dataset_metadata branch is dead code in every copy.
  3. evaluate_pi0.sh:34 runs the patch in a throwaway subprocess (python lerobot_local_patch.py ... || true) — monkey patches are process-local, so this has zero effect on the eval process that follows.
  4. The patch exists in four divergent forms: the standalone file, the PATCHEOF heredoc in droid-finetune.yaml, apply_local_dataset_patch() in evaluate_pi0.py, plus the usercustomize.py writers in four manifests that inject it into every future Python process on the node image.
    Impact: Four copies of dead-at-the-pinned-version code that look load-bearing, will drift from upstream and from each other, and (via usercustomize.py + blanket except: pass) can swallow real import errors invisibly. The repo convention is to wrap or configure upstream, never to ship runtime monkey patches.
    Suggestion: Delete src/lerobot_local_patch.py, the PATCHEOF heredoc, apply_local_dataset_patch() and its call in evaluate_pi0.py, and every usercustomize.py writer. Offline robustness comes from HF_HUB_OFFLINE=1 plus the staged --dataset.root — which the recipe already does. If a staged local dataset still triggers a Hub lookup at v0.5.0, that's an upstream bug to report to huggingface/lerobot, and the report is the artifact we'd want, not a patch. (Also: the Dockerfile's second COPY src/lerobot_local_patch.py and the README Layout entry for it go away with this.) — verified live 2026-07-15: https://github.com/huggingface/lerobot/blob/v0.5.0/src/lerobot/datasets/lerobot_dataset.py#L106-L113


# HuggingFace login for base model download
if [ -n "${HF_TOKEN:-}" ]; then
python -c "from huggingface_hub import login; login(token='${HF_TOKEN}')" 2>/dev/null || true

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

evaluate_pi0.sh is now orphaned dead code with wrong defaults, an argv token leak, and a login-then-offline contradiction — 🟡

Observation: Nothing invokes this wrapper — both eval manifests call evaluate_pi0.py directly (the same situation as round 1's run_finetuning.sh, which was rightly deleted). It also: defaults CHECKPOINT_PATH/TEST_DATA_PATH/RESULTS_PATH to /data/runs/... paths that don't match the manifests' /fsx/runs/... layout; interpolates the HF token into argv (login(token='${HF_TOKEN}') — visible in ps, breaks on ' in a token; the round-1 finding, still present at line 29); logs in and then immediately sets HF_HUB_OFFLINE=1, which blocks the downloads the login was for; and runs the patch subprocess no-op (see above).
Impact: Dead code that a reader will assume is the recipe's entrypoint, with four latent bugs waiting for whoever wires it up.
Suggestion: Delete it (and its README Layout line). If you'd rather keep a wrapper, fix the defaults to the manifest layout, drop the login() (huggingface_hub reads HF_TOKEN from the env on its own), and drop the patch invocation.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 4/6 — Carryover Blockers (unchanged from round 1)


The headline results still measure the training set — the numbers can't support the README's claims — 🔴

Observation: Unchanged from round 1: training consumes all 100 episodes of /fsx/datasets/droid_100; the eval reads the same path (--test-dataset-local /fsx/datasets/droid_100) and scores episodes 0–4 — the first five training episodes. evaluate_pi0.py:392 even constructs the default by stripping _test off the suffix, i.e. it deliberately falls back to the training path. The "held-out test set" docstring and the README's 96.5%/93.1% "Improvement" framing are therefore still unsupported. Round 1 included direct evidence: a 30-step checkpoint already shows a 70% MSE drop on these episodes — that's memorization, not generalization.
Impact: This is the PR's headline claim; as written it will mislead every reader of the results tables.
Suggestion: Same concrete fix as round 1, which needs no staging Job and no dataset copies: train with --dataset.episodes="[0..79]" and evaluate LeRobotDataset(..., episodes=list(range(80, 100))), then re-run and update the tables. Even if the deltas stay strong, they'd then mean what the README says. Alternatively, relabel the tables explicitly as in-sample smoke-test numbers ("open-loop MSE on training episodes — validates the pipeline, not generalization").


FSDP FULL_SHARD still doesn't shard π0, and the recipe still targets 80 GB H100s — 🔴

Observation: The accelerate launch block is unchanged: TRANSFORMER_BASED_WRAP with no wrappable block class means the whole 3B model is one FSDP unit — measured round 1 (per-GPU HBM: 1 GPU = 16.5 GB, 2 GPU = 79.6 GB, 8 GPU = 73.8 GB), and the naive --fsdp_transformer_layer_cls_to_wrap fix crashes on π0's joint cross-stream attention (shared param shards to size 0). ~80 GB peak on the advertised H100-80GB is an OOM cliff with no headroom, and adding GPUs doesn't reduce it.
Impact: The README's "FSDP FULL_SHARD" + "8× H100 80GB" combination overstates what the config achieves on the very hardware it names.
Suggestion: Minimum viable fix is documentation: state that π0 currently trains as effectively one FSDP unit (~80 GB peak per GPU at batch 4), so H100-80GB is at the edge and scaling GPUs adds throughput, not memory headroom. The real fix (custom wrap policy treating each joint PaliGemma+action-expert block pair as a unit) is legitimately follow-up material.


The download Jobs are still broken end-to-end: envsubst corruption, a split that doesn't split, and paths nothing consumes — 🔴

Observation: Unchanged in substance from round 1; the echo text now claims "episodes 0-79 train, 80-99 test" but the Python still copytrees the full dataset to both destinations. And the render pipeline still corrupts the script: GNU envsubst without an allow-list substitutes the unset ${TRAIN_DIR}/${TEST_DIR} inside the Python heredoc to empty strings (Path('') → CWD → both copies silently skipped), doesn't support ${HF_TOKEN:-} syntax (so $${HF_TOKEN:-} degrades to PID garbage and the login guard is always-true), and eats ${TRAIN_DIR} in the verify step (ls "$/meta/"). Net: the Job exits 0 having produced only droid_100_full, which no other manifest reads — finetune/eval use /fsx/datasets/droid_100 on a different PVC (fsx-claim vs ${DATA_PVC_NAME}) and a different mount (/fsx vs /data). Both independent review passes reached the same conclusions this round.
Impact: Users who run the Jobs waste time and 3× FSx capacity, get no split, and the finetune re-downloads anyway. The updated comments now promise a split that the code has never performed — worse than round 1, where the comment at least admitted it.
Suggestion: The --dataset.episodes split (previous finding) removes the need for these Jobs entirely — deleting both remains my recommendation. If you keep a staging Job, reduce it to the single hf download with paths/PVC matching the finetune manifests, and render with an explicit allow-list: envsubst '$NAMESPACE $IMAGE_URI $DATA_PVC_NAME'. — envsubst behavior reproduced live with GNU gettext, 2026-07-09.


Still no env_vars.example, and the render comments now reference a file that doesn't exist — 🟡

Observation: The download/PVC manifests' new render comments say source env_vars && envsubst ..., but no env_vars.example is committed, and the finetune/eval manifests still hardcode namespace: default, claimName: fsx-claim, the instance type, and the image — the same values spelled differently across six files. openvla-oft's kubernetes/libero/env_vars.example remains the in-repo template.
Impact: The one mechanism that would fix the PVC/namespace/image drift (and the ml.-prefix portability note, and the Batch-2 image split) is referenced but not provided.
Suggestion: Commit an env_vars.example defining NAMESPACE, IMAGE_URI, DATA_PVC_NAME, DATA_PVC_SIZE, FSX_SUBNET_ID, FSX_SECURITY_GROUP_ID, INSTANCE_TYPE, and render all manifests through the allow-listed envsubst. — repo precedent: 3.test_cases/pytorch/openvla-oft/kubernetes/libero/env_vars.example.


MIT-0 headers: still only on the three READMEs — 🟡

Observation: Unchanged from round 1: the Dockerfile, buildspec, all seven YAML manifests, and the three src/ files still lack the two-line MIT-0 header. The license lint checks every file, YAML included.
Suggestion: Prepend # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. / # SPDX-License-Identifier: MIT-0 to each.


Eval stochasticity and step-count pairing — still unseeded, and the base/fine-tuned comparison mixes ODE step counts — 🟡

Observation: Round 1's ask (declare the smoke-test framing, set a seed, pair metrics at the same step count) is unaddressed. The second pass this round added a sharper version of the pairing problem: eval_model runs the base policy at args.num_inference_steps[0] but the headline fine-tuned number uses max(finetuned_sweep.keys()) — so --num-inference-steps 5 10 would compare base-at-5 against fine-tuned-at-10 in the summary table and JSON.
Suggestion: Set torch.manual_seed/np.random.seed per trajectory; compare base and fine-tuned at the same step count (evaluate base at each swept count, or fix both to one); state N=5 and report per-trajectory numbers.

- **Note:** HyperPod labels instances with an `ml.` prefix (e.g., `ml.p5.48xlarge`). On plain EKS, the label is `p5.48xlarge`. The manifests use the HyperPod convention by default.
- The [Kubeflow Training Operator](https://github.com/kubeflow/training-operator) installed
- FSx for Lustre PVC (`fsx-claim`) mounted at `/fsx`
- A HuggingFace token (for base model download)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

HF_TOKEN still optional: true and the README still says "for base model download" — the Gemma gate bites 15 minutes in — 🔴

Observation: Unchanged from round 1. lerobot/pi0_base itself is ungated — the real requirement is the transitive gated dependency: π0's processor pipeline fetches google/paligemma-3b-pt-224 at train time, which requires a token that has accepted the Gemma license. Reproduced e2e in round 1: a tokenless run crashes with GatedRepoError 401 only after the ~15-minute install+download phase. All manifests still mark the secret optional: true. In the eval, the same tokenizer load is wrapped in try/except, so there it fails silently instead (language conditioning quietly disabled).
Impact: The docs predict neither the failure nor its timing; optional: true actively misleads.
Suggestion:

Suggested change
- A HuggingFace token (for base model download)
- A HuggingFace token that has **accepted the Gemma license** for [`google/paligemma-3b-pt-224`](https://huggingface.co/google/paligemma-3b-pt-224) — π0's processor downloads this gated tokenizer at train/eval time (the `lerobot/pi0_base` weights themselves are ungated)

and set optional: false on the finetune/eval secretKeyRefs so the failure is immediate and clear.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 5/6 — Smaller Items


CodeBuild path assumes privileged mode and IAM the README never mentions — 🟢

Observation: The buildspec runs docker build/docker push, which needs a CodeBuild project with privileged mode enabled plus IAM for STS + ECR (auth, create-repository, push). The README's three-line CodeBuild section covers none of the project setup.
Suggestion: One sentence + a link to the CodeBuild Docker docs, or drop the section — the buildx path already covers non-x86 hosts.


Round-1 code nits — all four still open — 🟢

Unchanged in evaluate_pi0.py: setup_lerobot_cache raises FileExistsError when a real directory pre-exists at the link path (:126); num_trajectories unbounded against available episodes (:483 area); snapshot_download glob double-counts sharded safetensors (:143); batch["task"] left a list after unwrapping (:209 area). Also still present: the checkpoint auto-discovery prefix pi0-finetune-{dataset} can never match the manifests' /fsx/runs/pi0-{dataset}/training/... layout (falls through to the hardcoded-020000 default — which only works when MAX_STEPS=20000), and the sagemaker-{region}-{account} S3 fallback is a leftover from the SageMaker variant. Fine to sweep these in one pass with the Batch 3 deletions.

value: "WARN"
# NCCL_TIMEOUT is not a native NCCL variable but is read by some
# framework-level process group implementations. Kept for consistency
# with other test cases in this repo (e.g., openvla-oft).

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new NCCL_TIMEOUT comment cites openvla-oft, which doesn't set it — 🟢

Observation: The comment keeps the variable "for consistency with other test cases in this repo (e.g., openvla-oft)" — but grep -rn NCCL_TIMEOUT 3.test_cases/pytorch/openvla-oft/ returns nothing (verified live 2026-07-15). Test cases that do set it: distillation, nemo-rl, vllm/dsv3-uccl-nixl.
Suggestion:

Suggested change
# with other test cases in this repo (e.g., openvla-oft).
# with other test cases in this repo (e.g., distillation, nemo-rl).

(Same line in libero-finetune.yaml.) Dropping the variable outright is equally fine — it's read by neither NCCL nor PyTorch.

# Build (from this directory):
# docker buildx build --platform linux/amd64 \
# --build-arg AWS_REGION=${AWS_REGION} \
# --load -f Dockerfile -t ${REGISTRY}pi0-lerobot:${IMAGE_TAG} .

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dockerfile header example builds an invalid image ref — 🟢 [confirmed]

Observation: The header example tags ${REGISTRY}pi0-lerobot:${IMAGE_TAG} (no slash), while the README and buildspec define REGISTRY without a trailing slash and use ${REGISTRY}/pi0-lerobot:.... Copying the Dockerfile example with the README's REGISTRY yields ...amazonaws.compi0-lerobot.
Suggestion:

Suggested change
# --load -f Dockerfile -t ${REGISTRY}pi0-lerobot:${IMAGE_TAG} .
# --load -f Dockerfile -t ${REGISTRY}/pi0-lerobot:${IMAGE_TAG} .

# ── Patch LeRobot: skip Hub lookups for local datasets ──
# The training script handles this at runtime via the lerobot_local_patch,
# but we pre-copy the patch script into the image for convenience.
COPY src/lerobot_local_patch.py /opt/pi0-lerobot/src/

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Redundant second COPY of the patch file — 🟢 [confirmed]

Observation: COPY src/ /opt/pi0-lerobot/src/ two lines up already includes lerobot_local_patch.py; the dedicated COPY (and its "pre-copy for convenience" comment) is a no-op layer. Moot once the patch is deleted per Batch 3.

RUN pip install --no-cache-dir --upgrade pip && \
pip install --no-cache-dir \
"lerobot[pi,dataset] @ git+https://github.com/huggingface/lerobot.git@v0.5.0" \
accelerate \

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

accelerate, safetensors, huggingface-hub unpinned next to the pinned LeRobot — 🟢

Observation: LeRobot is now pinned (thanks!), but its three siblings float, pip itself is upgraded to latest, and droid-finetune.yaml additionally runs an unpinned pip install --quiet accelerate and --upgrade huggingface-hub at pod start. Rebuilding "v1.0.0" next month can produce a different environment.
Suggestion: Pin all three (e.g. the versions LeRobot v0.5.0's pyproject.toml resolves today) and drop the redundant runtime installs — lerobot[pi,dataset] already pulls accelerate and huggingface-hub as dependencies.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 6/6 — Evaluation, Positives & Sources


Things That Look Great

  • This is a real response round, not a checkbox pass — deleting run_finetuning.sh outright (rather than patching it), pinning LeRobot to a tag, and swapping reclaimPolicy to Retain are exactly the right calls, and the commit message honestly describes the intent.
  • The rm -rf checkpoint wipe is gone — even though the resume replacement needs finishing (Batch 2), removing the data-loss path first was the right priority.
  • The DROID training RENAME_MAP fix is correct, and keeping the LIBERO map's dual-key tolerance (native + re-encoded variants) is thoughtful.
  • The ml.-prefix nodeSelector comments + README note are a genuine portability improvement — every manifest now tells a plain-EKS user exactly what to change.
  • imagePullPolicy: IfNotPresent across the board, and the download Jobs' render comments now at least name the intended envsubst workflow.
  • The Dockerfile's git-lfs smudge-skip block is a nice touch that will save users a multi-GB surprise during pip install.
  • Round 1's positives all still stand: the recipe's core training path is sound (independently reproduced e2e on 8× B300, 2026-07-10, through install → download → train → complete resumable checkpoint), NCCL_SOCKET_IFNAME: "^lo" is the correct exclusion pattern, hyperparameters are consistent across README and manifests, and --dataset.video_backend=pyav remains the right call on a DLC-pinned torch.

Sources

  • LeRobot v0.5.0 (pinned): src/lerobot/configs/train.py L122-126 (FileExistsError without --resume), L92-98 (config_path required to resume); src/lerobot/datasets/lerobot_dataset.py L106-113, L750-752 (Hub contacted only on local-load failure); src/lerobot/datasets/dataset_metadata.py → 404 at this tag — all verified live 2026-07-15. https://github.com/huggingface/lerobot/tree/v0.5.0
  • Tags resolve (verified live 2026-07-15): huggingface/lerobot@v0.5.000b662de; kubeflow/training-operator@v1.9.117077e33.
  • lerobot/pi0_base config.json — three image features (base_0_rgb, left_wrist_0_rgb, right_wrist_0_rgb): https://huggingface.co/lerobot/pi0_base — verified live 2026-07-09.
  • NCCL_TIMEOUT usage in-repo: absent from openvla-oft; present in distillation, nemo-rl, vllm/dsv3-uccl-nixlgrep -rn, verified live 2026-07-15.
  • GNU envsubst ${VAR:-}/unset-variable behavior — reproduced locally, 2026-07-09 (round 1).
  • Gemma gate + FSDP memory measurements — e2e runs on 8× B300, 2026-07-10 (round 1, details in the round-1 threads).
  • Repo precedent: 3.test_cases/pytorch/openvla-oft/kubernetes/libero/env_vars.example.
  • Round 1 review: #1174 (review)

…kers - Fix DROID eval camera key: wrist_0_rgb -> left_wrist_0_rgb (missed in round 1) - Remove lerobot_local_patch.py entirely (merge blocker per maintainer) - Remove dead evaluate_pi0.sh (never called by any manifest) - Rewrite eval YAMLs: self-contained DLC-direct (install + curl script) - Rewrite libero-finetune to not reference /opt/pi0-lerobot paths - Fix resume: fail-fast if output_dir exists (FORCE_RESTART=1 to override) - HF_TOKEN: remove optional:true from finetune manifests (Gemma license required) - README: DLC is primary approach, custom image is optional section - README: add Gemma license prerequisite, FileExistsError troubleshooting - NCCL_TIMEOUT comment: cite distillation/nemo-rl (not openvla-oft) - Dockerfile: pin accelerate/safetensors/huggingface-hub, fix header slash, remove redundant COPY and flash_attn at build time - Pin LeRobot to v0.5.0 in all runtime installs
@amindm
amindm requested a review from KeitaW July 15, 2026 18:18
…ning manifests now pass --dataset.episodes with train-only indices (DROID: 0-79, LIBERO: 0-399) - Eval manifests pass --eval-episodes with held-out indices (DROID: 80-99, LIBERO: 400-499) - evaluate_pi0.py: add --eval-episodes flag, pass to LeRobotDataset(episodes=...) - Remove download Jobs (training downloads data directly; split via episodes flag) - README: replace headline numbers with 'pending re-run' disclaimer (previous numbers were on training data, not held-out)

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 1/6 — Round-3 Resolution Scoreboard


This is the strongest response round yet — net −376 lines, and the hard asks all landed

Commits a63a5fa ("DLC-direct, remove monkey-patch") and ac7bb38 ("enforce 80/20 episode split") resolve 10 of 19 round-2 findings, including every structural blocker:

Resolved (10):

  • Monkey-patch layer deleted entirelylerobot_local_patch.py, the heredoc, apply_local_dataset_patch(), all usercustomize.py writers, and the Dockerfile copy are gone, replaced by an accurate comment ("LeRobot v0.5.0 handles local datasets natively"). Exactly the resolution I hoped for.
  • Train/eval contamination structurally fixed — training now passes --dataset.episodes=[0..79] / [0..399], eval takes --eval-episodes 80–99 / 400–499, and the README honestly retracts the old numbers ("Pending re-run"). One implementation bug remains in how the eval indexes the filtered dataset (Batch 3).
  • ✅ Both download Jobs deleted (and with them the envsubst corruption and the fake split).
  • evaluate_pi0.sh deleted.
  • ✅ Eval DROID rename map fixed (left_wrist_0_rgb).
  • ✅ README image workflow un-inverted — DLC-direct is now the documented default, the custom image is "Optional", and the Dockerfile header says so.
  • ✅ The false "will resume" echo replaced with an explicit fail-fast + FORCE_RESTART=1 escape hatch and a README troubleshooting entry — honest and predictable.
  • ✅ Gemma-license prerequisite documented with both links; optional: true removed from every secretKeyRef.
  • NCCL_TIMEOUT comment now cites siblings that actually set it (distillation, nemo-rl); Dockerfile registry-slash and duplicate-COPY nits fixed.

Partial (2): the three-manifests-crash finding — both evals were reworked (new issues in the delivery mechanism, Batch 2) but libero-finetune.yaml still can't run; the contamination fix — split enforced, eval indexing misaligned (Batch 3).

Regressed (1): the dependency pins — added as asked, but the chosen versions conflict with LeRobot v0.5.0's own requirements, so the Dockerfile no longer builds (Batch 2).

Unresolved (6): FSDP FULL_SHARD claims, MIT-0 headers outside the READMEs, eval seeding/step-pairing, env_vars.example/PVC naming, CodeBuild project prerequisites, and the round-1 code nits in evaluate_pi0.py. Details in Batches 3–5.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 2/6 — Still Won't Run As Committed


mkdir -p "$OUTPUT_DIR"

# --- Quick fixes (pre-installed image, just patch) ---

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

libero-finetune.yaml never installs LeRobot — the "custom image" header is aspirational, the image: is the stock DLC — 🔴 [confirmed]

Observation: This is the surviving half of the round-2 crash finding. The header still says "using custom Docker image — LeRobot pre-installed — startup ~30 sec", and this "Quick fixes (pre-installed image, just patch)" block skips installation accordingly — but image: points at the stock DLC, which ships neither lerobot nor lerobot-train. $(which lerobot-train) expands to empty and the accelerate launch line falls apart. Its sibling droid-finetune.yaml handles the same situation correctly (Steps 1–2 install git deps + lerobot[pi,dataset]@v0.5.0).
Impact: The LIBERO training path — half the PR's payload — cannot run as committed.
Suggestion: Mirror droid-finetune.yaml's Steps 1–2 here (and fix the header comment to match the DLC-direct reality), or if you'd rather keep this manifest as the custom-image demonstration, set image: to a ${IMAGE_URI} placeholder with a loud "replace me" comment and say so in the README. Either is fine — the current halfway state is the only broken option.

EVAL_SCRIPT="/tmp/evaluate_pi0.py"
cp /fsx/pi0-lerobot-src/evaluate_pi0.py "$EVAL_SCRIPT" 2>/dev/null || \
pip show lerobot | grep -q . && \
curl -sL "https://raw.githubusercontent.com/amindm/awsome-distributed-ai/pi0-lerobot-test-case/3.test_cases/pytorch/pi0-lerobot/src/evaluate_pi0.py" -o "$EVAL_SCRIPT"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both eval Jobs fetch their executable from your fork's PR branch at runtime — breaks at merge, and the DROID fallback chain always overwrites — 🔴 [confirmed]

Observation: The evals now source evaluate_pi0.py via cp /fsx/pi0-lerobot-src/evaluate_pi0.py || … curl -sL https://raw.githubusercontent.com/amindm/awsome-distributed-ai/pi0-lerobot-test-case/…. Four problems compound: (1) nothing in the PR or docs ever stages /fsx/pi0-lerobot-src/, so the cp always fails and the curl is the real path; (2) the URL is your fork's PR branch — after this merges the branch will typically be deleted, and curl without -f treats the 404 as success and writes the literal "404: Not Found" body into $EVAL_SCRIPT for Python to "run"; (3) in the DROID variant, cp … || pip show lerobot | grep -q . && curl … parses as (cp || pip show…) && curl — so even when the cp succeeds, curl runs anyway and overwrites the staged copy; (4) operationally, the Job executes mutable remote code with HF_TOKEN in its environment and a writable FSx mount — anyone who can push to that branch controls what runs.
Impact: Both eval Jobs are time bombs that arm themselves the moment this PR merges, and the primary delivery path is dead code today.
Suggestion: Deliver the script through Kubernetes rather than the network: kubectl create configmap pi0-eval-script --from-file=src/evaluate_pi0.py, mount it, and run python /config/evaluate_pi0.py — self-contained, versioned with the user's checkout, no token-adjacent remote fetch. (One added line in each README's Quick Start.) If you keep any curl fallback, point it at awslabs/awsome-distributed-ai's main pinned to a commit SHA and add -f so a 404 fails loudly. — curl(1): without -f/--fail, HTTP error pages are written to output with exit 0.

RUN pip install --no-cache-dir --upgrade pip && \
pip install --no-cache-dir \
"lerobot[pi,dataset] @ git+https://github.com/huggingface/lerobot.git@v0.5.0" \
"accelerate==1.2.1" \

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new Dockerfile pins conflict with LeRobot v0.5.0's own requirements — the image no longer builds — 🔴 [confirmed]

Observation: Round 2 asked for pins on the three floating siblings; the versions chosen predate the LeRobot pin's floors. LeRobot v0.5.0's pyproject.toml requires accelerate>=1.10.0,<2.0.0 and huggingface-hub>=1.0.0,<2.0.0 (verified live 2026-07-16), while this pip install pins accelerate==1.2.1 and huggingface-hub==0.27.1 in the same command — pip's resolver raises ResolutionImpossible and the build dies before the sanity-check layers.
Impact: The optional-image path fails for everyone, at build time.
Suggestion:

Suggested change
"accelerate==1.2.1" \
"accelerate==1.10.1" \

and huggingface-hub==1.0.1 (both exist on PyPI and satisfy v0.5.0's ranges — verified live 2026-07-16); safetensors==0.4.5 is fine. Simplest robust alternative: drop the three explicit pins entirely and let LeRobot's own constraints resolve them — the LeRobot pin is what makes the build reproducible in practice. — https://github.com/huggingface/lerobot/blob/v0.5.0/pyproject.toml#L65-L66

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 3/6 — Held-Out Eval Correctness


Eval robustness bundle — 🟢

Three smaller items in the same file, fine to sweep together: (1) eval_model returns None when every trajectory is shorter than the horizon, and the summary/sweep prints dereference it unconditionally (--action-horizon 500TypeError at the summary); (2) pred[:n] truncation with a full-horizon chunk_start stride silently skips frames when the policy returns fewer than action_horizon actions — worth an assert; (3) the round-1 nits are all still open (setup_lerobot_cache FileExistsError on a pre-existing real directory; the pi0-finetune-{dataset} auto-discovery prefix that can never match /fsx/runs/pi0-{dataset}/… and the SageMaker sagemaker-{region}-{account} S3 fallback — both leftovers worth deleting now that the manifests pass explicit paths; batch["task"] left a list after unwrapping).

):
"""Predict action chunks for one episode and compare to ground truth."""
from_idx = int(dataset.meta.episodes["dataset_from_index"][episode_index])
to_idx = int(dataset.meta.episodes["dataset_to_index"][episode_index])

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The eval reads training episodes' boundaries but held-out episodes' frames — the split is enforced, the segmentation is wrong — 🔴 [confirmed]

Observation: With episodes=eval_episodes passed to LeRobotDataset, v0.5.0 filters hf_dataset and __getitem__(idx) indexes positionally into that subset — but dataset.meta.episodes stays unfiltered (metadata is loaded independently of the episode filter; verified live 2026-07-16 against v0.5.0 lerobot_dataset.py — the filter builds a private _absolute_to_relative_idx map the metadata never consults). This loop reads from_idx/to_idx from meta.episodes[episode_index] for episode_index 0–4 — i.e. training episodes 0–4's absolute boundaries — and applies them as positional offsets into the held-out frames. Symptom you can see immediately: n_eps prints "Loaded 100 test episodes" when only 20 were selected.
Impact: The frames scored are held-out (the contamination fix stands), but "trajectory k" boundaries belong to the wrong episodes: chunks silently span adjacent held-out episodes (a 50-step "chunk" jumps across an episode cut, polluting MSE with discontinuities), per-trajectory numbers don't correspond to real episodes, and a custom --eval-episodes list shorter than ~5 training episodes' span raises IndexError. The pending re-run would publish these numbers.
Suggestion: Derive boundaries positionally from the filtered dataset itself, e.g. once before the loop:

ep_col = dataset.hf_dataset["episode_index"]
bounds = {}
for pos, e in enumerate(ep_col):
    e = int(e)
    bounds.setdefault(e, [pos, pos])[1] = pos + 1
episodes_in_order = list(bounds)  # held-out episode ids, positional [from, to) in bounds[e]

then iterate episodes_in_order[:args.num_trajectories] — this also makes n_eps and the too-short skip honest, and bounds num_trajectories naturally. — https://github.com/huggingface/lerobot/blob/v0.5.0/src/lerobot/datasets/lerobot_dataset.py (_absolute_to_relative_idx construction; __getitem__self.hf_dataset[idx])


# Run base
base_results = eval_model("BASE π0", args.base_policy_id,
num_inference_steps=args.num_inference_steps[0])

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Before the re-run: seed the ODE noise and compare base vs fine-tuned at the same step count — 🟡 [confirmed]

Observation: Carried from rounds 1–2, now load-bearing because the README promises a re-run: flow-matching inference is unseeded (base, fine-tuned, and every sweep entry draw independent noise), and the base is evaluated at num_inference_steps[0] while the headline fine-tuned number uses max() of the sweep — --num-inference-steps 10 5 3 1 happens to align them at 10, but any other ordering compares different solvers. The saved JSON also omits eval_episodes and the base's step count, so a result file can't be tied to the split that produced it.
Impact: The re-run's numbers will be neither reproducible nor strictly apples-to-apples.
Suggestion: torch.manual_seed(s)/np.random.seed(s) per trajectory (same seed for base and fine-tuned); evaluate base at each swept step count or fix both to one; record eval_episodes, seed, and per-model step counts in the JSON.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 4/6 — Documentation Truthfulness


## Results (p5.48xlarge — 8× H100)

| Metric | Base π0 | Fine-Tuned | Improvement |

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The sub-READMEs still publish the results the top README just retracted — 🟡 [confirmed]

Observation: The top-level README now says "Previous numbers were measured on training data and are not shown here" — but kubernetes/droid/README.md and kubernetes/libero/README.md (untouched by the response commits) still present the full 96.5%/93.1% tables and ODE sweeps as valid results. The same two files also kept stale round-2 text: "installs patches" / "patches LeRobot" (the patch is deleted) and the "Container image pushed to ECR" prerequisite (the manifests are DLC-direct now).
Impact: A reader landing on the walkthrough README — the page the Quick Start points them at — sees exactly the numbers the PR retracted, plus instructions for a workflow that no longer exists.
Suggestion: Mirror the top README's "Pending re-run" note in both sub-READMEs (or drop their Results sections until the re-run), and update the prerequisites/step text to the DLC-direct, no-patch reality.


If the training operator crashes with "no matches for MPIJob", install the CRD:
```bash
kubectl delete crd mpijobs.kubeflow.org 2>/dev/null

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

kubectl delete crd mpijobs.kubeflow.org in Troubleshooting deletes every MPIJob on the cluster — 🟡 [confirmed]

Observation: New in this revision: the MPIJob-CRD troubleshooting step now prepends kubectl delete crd mpijobs.kubeflow.org before the re-apply. Deleting a CRD cascades to all resources of that kind cluster-wide — on a shared HyperPod cluster that wipes every other team's running MPIJobs, silently (2>/dev/null).
Impact: A user following the troubleshooting section verbatim can destroy unrelated workloads; this recipe explicitly targets shared clusters.
Suggestion:

Suggested change
kubectl delete crd mpijobs.kubeflow.org 2>/dev/null
kubectl apply -f https://raw.githubusercontent.com/kubeflow/training-operator/v1.9.1/manifests/base/crds/kubeflow.org_mpijobs.yaml

The apply alone is sufficient (it creates the CRD if missing and is a no-op otherwise); reserve deletion for a documented "only if the CRD is corrupted, and only on a cluster you own" caveat if you want to keep it at all.

- The [Kubeflow Training Operator](https://github.com/kubeflow/training-operator) installed
- FSx for Lustre PVC (`fsx-claim`) mounted at `/fsx`
- A HuggingFace token that has accepted **both**:
- [lerobot/pi0_base](https://huggingface.co/lerobot/pi0_base) (Apache 2.0, but gated)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Small factual fixes in Prerequisites/Troubleshooting — 🟢 [confirmed]

Observation: Three factual nits in the otherwise much-improved token documentation: (1) lerobot/pi0_base is described as "Apache 2.0, but gated" — it is not gated; I downloaded its weights with no token at all in the round-1 e2e run (verified live 2026-07-10). The only gate is the transitive PaliGemma dependency, which the next bullet covers correctly. (2) The Troubleshooting section says training fails "with a 403" — the observed failure is GatedRepoError 401. (3) The Gemma-license link points at google/gemma-2b; linking the actually-fetched google/paligemma-3b-pt-224 saves the reader a hop.
Suggestion:

Suggested change
- [lerobot/pi0_base](https://huggingface.co/lerobot/pi0_base) (Apache 2.0, but gated)
- [lerobot/pi0_base](https://huggingface.co/lerobot/pi0_base) (ungated — no acceptance needed, listed for completeness)

plus 403401 (GatedRepoError) and the PaliGemma link in Troubleshooting.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 5/6 — Carryover & Smaller Items


FSDP FULL_SHARD still doesn't shard π0 — the one unresolved must-fix, third round running — 🔴

Observation: Unchanged across all three rounds: TRANSFORMER_BASED_WRAP with no wrappable block means the whole 3B model is one FSDP unit (measured round 1: per-GPU HBM 16.5 GB at 1 GPU → 79.6 GB at 2 → 73.8 GB at 8; the naive block-class fix crashes on π0's joint cross-stream attention). The README still advertises "FSDP FULL_SHARD" on "8× H100 80GB" — an OOM cliff with no headroom that more GPUs won't fix.
Suggestion: I'm not asking for the custom wrap policy in this PR — one honest paragraph in Training Configuration is enough: π0 currently trains as a single FSDP unit (~80 GB peak per GPU at batch 4), so H100-80GB is at the edge and scaling GPUs adds throughput, not memory headroom. Happy to share the measurement methodology for the note.


MIT-0 headers: still only the three READMEs — 🟡

Unchanged from rounds 1–2: Dockerfile, buildspec, all five YAML manifests, and evaluate_pi0.py still lack the two-line header the license lint checks. Now a 7-file, two-minute fix.


Buildspec — orphaned by the README rewrite, plus two setup assumptions — 🟢

Observation: The round-2 README's CodeBuild section (zip + upload + start-build) was dropped in the rewrite, so buildspec.yml is now referenced only by a Layout comment — nothing tells the user how to use it. If kept: it assumes a project with privileged mode + STS/ECR IAM (still undocumented, round-2 carryover), starts in CODEBUILD_SRC_DIR (fine for the old zip flow rooted at this directory; wrong for a repo-connected project, where a cd 3.test_cases/pytorch/pi0-lerobot is needed), and hardcodes AWS_REGION: us-east-2. One short "CodeBuild (optional)" README paragraph — or deleting the buildspec — resolves all of it.

# LeRobot v0.5.0 raises FileExistsError if output_dir exists.
# If FORCE_RESTART=1, remove and start fresh. Otherwise, error early.
if [ -d "$OUTPUT_DIR/training" ]; then
if [ "${FORCE_RESTART:-0}" = "1" ]; then

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ops nits — one bundle — 🟢

Observation: Small things, none blocking: (1) FORCE_RESTART is read here but never declared in the pod env: — a user can't discover or set it without editing the script block; declare it with value: "0" and a comment. (2) pip install --quiet accelerate (line ~131, unpinned) survives just above the launch — redundant since lerobot[pi] already installed it; it can silently float the version the launcher uses. (3) The step banner now runs [2/5][5/5] with no [1/5] after the patch step was removed. (4) The 80- and 400-element --dataset.episodes literals would be less error-prone generated inline: --dataset.episodes="[$(seq -s, 0 79)]". (5) The PVC manifest's ${DATA_PVC_NAME} still has no consumer — every workload hardcodes fsx-claim, and the PVC lands in the caller's current namespace while workloads run in default; the round-2 env_vars.example suggestion (repo precedent: openvla-oft) would close both.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 6/6 — Evaluation, Positives & Sources


Things That Look Great

  • This response round did the hard things. Deleting the entire patch layer (4 sites), deleting both download Jobs, and implementing the episode split via LeRobot's native --dataset.episodes — rather than defending any of them — is exactly how a review round should go. The PR is 376 lines lighter than round 2 and structurally sound.
  • Retracting the results tables instead of leaving them was the right call — "Pending re-run" with the split documented is honest framing that most contributions never manage. The two sub-READMEs just need to catch up (Batch 4).
  • The FORCE_RESTART fail-fast is a genuinely good resolution of the resume saga: explicit, loud, and documented in Troubleshooting — no false promises.
  • The Gemma-license prerequisite is now documented better than most π0 material anywhere: both links, the transitive-dependency explanation, and optional: true removed so the failure is immediate.
  • Round-2 nits were addressed almost line-for-line (registry slash, duplicate COPY, the NCCL_TIMEOUT example swap to distillation/nemo-rl) — thank you for the attention to detail.
  • Still standing from earlier rounds: the core DROID training path is e2e-sound (independently reproduced on 8× B300), NCCL_SOCKET_IFNAME: "^lo" is right, hyperparameters are consistent, --dataset.video_backend=pyav remains the right call.

Sources

  • LeRobot v0.5.0 (all verified live 2026-07-16): pyproject.toml L65-66 (huggingface-hub>=1.0.0,<2.0.0, accelerate>=1.10.0,<2.0.0); src/lerobot/datasets/lerobot_dataset.py__getitem__self.hf_dataset[idx] (positional), _absolute_to_relative_idx built only when episodes is set, meta.episodes loaded unfiltered. https://github.com/huggingface/lerobot/tree/v0.5.0
  • PyPI (verified live 2026-07-16): accelerate 1.10.1 and huggingface-hub 1.0.1 exist and satisfy v0.5.0's ranges; accelerate==1.2.1 and huggingface-hub==0.27.1 fall below the floors.
  • curl(1): without -f/--fail, HTTP 4xx bodies are written to the output file with exit code 0.
  • lerobot/pi0_base downloadable tokenless (ungated); PaliGemma gate returns GatedRepoError 401 — e2e run, 2026-07-10 (round-1 threads).
  • FSDP memory measurements (16.5/79.6/73.8 GB at 1/2/8 GPUs) — e2e on 8× B300, 2026-07-10 (round-1/2 threads).
  • Kubernetes: deleting a CRD garbage-collects all its custom resources cluster-wide — https://kubernetes.io/docs/tasks/extend-kubernetes/custom-resources/custom-resource-definitions/#delete-a-customresourcedefinition
  • Prior rounds: round 1 · round 2

amindm added 2 commits July 16, 2026 09:07
…y Blockers fixed: - Dockerfile: remove conflicting dep pins (let LeRobot v0.5.0 resolve) - Eval Jobs: read evaluate_pi0.py from FSx (user stages once), no curl from fork - libero-finetune: add full pip install block (was missing, pod crashed) - Eval indexing: load full dataset, iterate eval_episodes by index (no episodes= filter) Should-fixes: - Sub-READMEs: retracted old numbers, show train/test split table - Troubleshooting: remove dangerous 'kubectl delete crd', use safe re-apply - README: pi0_base is ungated, error is 401 not 403, fix Gemma link - Ops: FORCE_RESTART in env block, fix step numbering [1/4]-[4/4], episode lists via Python one-liner, remove redundant pip install accelerate
…p5 Results (held-out evaluation, never seen during training): - DROID (train 0-79, eval 80-99): 89.7% MSE reduction, 199ms at 1 ODE step - LIBERO (train 0-303, eval 304-378): 88.9% MSE reduction, 197ms at 1 ODE step Fixes in this commit: - Pin LeRobot to commit ddc2aa7a (HEAD, fixes GR00T dataclass bug on py3.12) - Use lerobot[pi,dataset] everywhere (dataset extra required on HEAD) - LIBERO: 379 total episodes (not 500), train 0-303, eval 304-378 - LIBERO: use real HF repo_id (lerobot/libero_10) to avoid local-name 404 - save_freq=20000 in both manifests (45GB checkpoints fill 1.2TB FSx fast) - All regions us-west-2, all images consistent DLC tag - No stale refs (v0.5.0, private ECR, fork URLs, sagemaker-user, Pi0_SMTJ) - README has real results with proper train/test split documentation - Sub-READMEs updated with matching results and episode counts Tested end-to-end on p5.48xlarge (HyperPod EKS, us-west-2): - DROID: train 6h15m + eval 15min, no errors - LIBERO: train 6h11m + eval 15min, no errors
@amindm
amindm requested a review from KeitaW July 23, 2026 14:24

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 1/5 — Round-4 Scoreboard: the blockers are gone, and I was wrong about one thing


Every round-3 must-fix is resolved, the eval-indexing fix is verified correct against the pinned commit, and the e2e run retires my standing OOM concern

This is the round that lands it. Commits 53bf0dd and 0f5072b resolve all five round-3 must-fixes, and — the part that matters most — the results are now a real held-out evaluation with an end-to-end run behind them. Credit where it's due:

Resolved (all 5 round-3 blockers):

  • libero-finetune now installs LeRobot — full [1/4] install block added; it no longer assumes a pre-baked image. The "custom image" header is gone.
  • Eval scripts no longer curl from your fork branch — both read /fsx/pi0-lerobot/evaluate_pi0.py with a clean fail-fast if unstaged; the supply-chain-via-404 hazard is gone.
  • Dockerfile builds again — the conflicting accelerate==1.2.1/huggingface-hub==0.27.1 pins are removed; LeRobot's own constraints resolve the deps.
  • Eval indexing fixed — and I verified the fix is correct. Loading the full dataset (no episodes= filter) and iterating held-out episode indices means dataset.meta.episodes["dataset_from_index"][ep] and dataset[...] now share the same absolute index space. I checked this against the actual pinned commit ddc2aa7a (dataset_metadata.py load_episodes(root) returns the full unfiltered table with dataset_from_index/dataset_to_index; __getitem__ is absolute-indexed), and the off-by-one is clean: the max frame touched is to_idx − 1, never bleeding into ep+1. The 89.7%/88.9% held-out numbers are trustworthy (see Batch 3 for the N and labeling caveats).
  • Destructive kubectl delete crd removed from Troubleshooting — replaced with the safe re-apply.

And an evidence update on my part: across rounds 1–3 I flagged "FSDP FULL_SHARD doesn't actually shard π0, so ~80 GB peak is an OOM cliff on the H100-80GB the README targets." Your e2e run (DROID train 6h15m + LIBERO 6h11m on p5.48xlarge, no errors) empirically refutes the OOM half of that concern — at batch 4 with gradient checkpointing it fits and trains on 80 GB. The narrow technical point stands (it's one FSDP unit, so memory won't scale down if you later raise batch/sequence), but that's now a footnote, not a blocker, and I'm dropping it as a finding. Thank you for running it to ground.

Retracted-and-replaced results (89.7%/88.9%, down from the contaminated 96.5%/93.1%) with the split documented in three places — exactly the honest resolution. Regions unified to us-west-2 (consistent across Dockerfile, buildspec, all four manifests, and the README export — checked).

What's left is one mechanical merge blocker and a handful of polish/honesty items, below.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 2/5 — The One Blocker + Consistency


MIT-0 license headers are still missing on all seven non-README files — this is the CI lint gate — 🔴 [confirmed]

Observation: Verified on head 0f5072bb: Dockerfile, buildspec.yml, all five YAML manifests (droid-finetune, droid-eval, libero-finetune, libero-eval, pvc-fsx-lustre), and src/evaluate_pi0.py carry no SPDX-License-Identifier: MIT-0 header. Only the three READMEs have it. This has been open since round 1; it's the repo's single most-recurring lint miss and the license check runs on YAML and Dockerfiles too.
Impact: This will fail the license lint in CI — a hard merge blocker, independent of everything else.
Suggestion: Prepend the two-line header to each of the seven files (comment syntax per file type):

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

For evaluate_pi0.py it goes just under the module docstring. This is the last mechanical thing between the PR and a green check.


The PR description still advertises the retracted 96.5%/93.1% numbers — 🟡

Observation: The committed README now shows the real held-out 89.7%/88.9%, but the PR's top-level description still carries the original contaminated table (DROID 96.5% / LIBERO 93.1%, "20K steps, ~6.5 hours"). A maintainer skimming the PR sees the retracted numbers first.
Suggestion: Update the PR description's results table to the held-out numbers (or replace it with a pointer to the README's Results section) so the description and the code agree.

ACTION_HORIZON="${ACTION_HORIZON:-50}"
DATASET_PATH="/fsx/datasets/droid_100"
OUTPUT_DIR="/fsx/runs/pi0-droid"
LOCAL_DATASET_NAME="droid_local"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DROID still loads via the droid_local local-name that LIBERO abandoned to fix a 404 — reconcile the two — 🟡 [confirmed]

Observation: The LIBERO finetune switched LOCAL_DATASET_NAME from libero_local to the real repo_id lerobot/libero_10 and dropped the ln -snf symlink — your commit message says this was "to avoid local-name 404." But the DROID finetune still uses LOCAL_DATASET_NAME="droid_local" + the symlink + --dataset.repo_id=droid_local. Both pass --dataset.root, so both only hit the Hub if the local metadata load fails — at which point a slash-less name like droid_local 404s on the Hub fallback while a real lerobot/... id resolves. That's exactly the asymmetry your LIBERO fix addresses.
Impact: I'm not claiming DROID is broken — your e2e run proves droid_local works today (DROID's local load succeeded, so it never reached the Hub). But it's one failed local-load away (corrupted download, a dataset codebase-version bump) from the same 404 LIBERO hit, and the two siblings now teach two different conventions for the same task.
Suggestion: Mirror the LIBERO fix for symmetry and robustness:

Suggested change
LOCAL_DATASET_NAME="droid_local"
LOCAL_DATASET_NAME="lerobot/droid_100"

and drop the ln -snf "$DATASET_PATH" "$HF_LEROBOT_HOME/$LOCAL_DATASET_NAME" line (the real repo_id + --dataset.root is sufficient, same as LIBERO). — repo precedent: the LIBERO sibling in this same PR.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 3/5 — Evaluation Honesty (the numbers are real, but under-labeled)


results = []
# Evaluate on held-out episodes (by episode index into full dataset)
episodes_to_eval = eval_episodes[:args.num_trajectories]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The headline numbers cover only the first 5 held-out episodes, not the full 80–99 / 304–378 the README implies — 🟡 [confirmed]

Observation: episodes_to_eval = eval_episodes[:args.num_trajectories] with --num-trajectories 5 scores only the first five held-out episodes (80–84 for DROID, 304–308 for LIBERO), even though the manifest passes the full held-out list and the README/eval headers say "held-out episodes 80-99" / "304-378". The print(f"... {eval_episodes[:5]}...{eval_episodes[-1]}") line reinforces the impression the run spans through 99/378. So every headline figure — including the 89.7%/88.9% — is an N=5 average over the first five held-out episodes.
Impact: The frames are correctly held-out (the indexing is right), but "89.7% on held-out episodes 80-99" over-states the coverage; it's 89.7% on episodes 80-84. Not wrong, just under-labeled — and N=5 is a small sample for a headline claim.
Suggestion: Either score the whole held-out set — drop the truncation, or set num_trajectories to len(eval_episodes):

Suggested change
episodes_to_eval = eval_episodes[:args.num_trajectories]
episodes_to_eval = eval_episodes # score the full held-out set, not just the first num_trajectories

or, if 5 is a deliberate smoke-test budget, relabel the README and the print as "first 5 held-out episodes (N=5)" so the coverage is honest.

}

# Run base
base_results = eval_model("BASE π0", args.base_policy_id,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seed the ODE noise and pair base/fine-tuned at the same step count — the LIBERO sweep is already showing the symptom — 🟡 [confirmed]

Observation: Carried from rounds 1–3, and now with visible evidence: flow-matching inference is unseeded (base, fine-tuned, and each sweep point draw independent noise), and the base is evaluated at num_inference_steps[0] while the headline fine-tuned uses max() of the sweep (they coincide at 10 only because the manifest lists 10 5 3 1). The README's LIBERO ODE sweep is non-monotonic — 10-step MSE (8.548e-02) is lower than 5-step (9.236e-02) and 3-step (8.789e-02) — which for a flow-matching solver is the N=5-plus-unseeded-noise signature, not a real step-count trend.
Impact: The sweep can't be read as a step-count relationship as presented, and base-vs-fine-tuned isn't strictly reproducible run to run.
Suggestion: torch.manual_seed(s)/np.random.seed(s) per trajectory (same seed for both models); evaluate the base at the same step count as the reported fine-tuned headline (or fix both to one). Combined with scoring the full held-out set (previous finding), the sweep would become monotone and defensible.

except Exception as e:
print(f" WARNING: pre/post processors failed to load ({type(e).__name__}: {e}).")
print(f" Results will compare raw model output to raw ground truth (no un-normalization).")
preprocess, postprocess = None, None

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Silent processor/tokenizer fallbacks still emit authoritative MSE — 🟡

Observation: Also carried from earlier rounds: if make_pre_post_processors raises, preprocess/postprocess fall back to None and the model is scored on un-normalized inputs vs raw ground truth — a meaningless MSE that is still printed and written to JSON with no flag. Same for the tokenizer load (tokenizer=None → language conditioning silently off). The asymmetric failure mode is the dangerous one: base processors fail but fine-tuned succeed → the "improvement" is a normalization artifact, not training.
Impact: A green eval run can silently produce a fabricated improvement number. Your tested run looks healthy (base ~0.55 → fine-tuned ~0.057 is plausible), so this is a latent hazard rather than a current-run problem — but it's one worth closing before these numbers get cited.
Suggestion: Fail hard if either model's processors don't load, or record a processors_active boolean per model in the JSON and refuse to print the improvement % when either side ran without them.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 4/5 — Nits & Polish

--batch_size=$BATCH_SIZE \
--output_dir=$OUTPUT_DIR/training \
--job_name=pi0_droid_finetune \
--save_freq=20000 \

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SAVE_INTERVAL is now dead, and MAX_STEPS/save_freq/eval-path are silently coupled — 🟢 [confirmed]

Observation: SAVE_INTERVAL="${SAVE_INTERVAL:-20000}" is set but --save_freq=20000 is hardcoded just below, so the env override is a no-op. Meanwhile MAX_STEPS stays env-parametrized while save_freq and the eval's checkpoint path (checkpoints/020000/) are both hardcoded to 20000 — set MAX_STEPS to anything else and you get either no final checkpoint at that step or an eval path that doesn't exist. The single-final-checkpoint / no-resume tradeoff itself is fine and documented (FSx capacity), just make the knobs honest.
Suggestion: Use --save_freq=$SAVE_INTERVAL (and derive the eval checkpoint dir from MAX_STEPS), or drop SAVE_INTERVAL and document that MAX_STEPS is fixed at 20000 for this recipe. (Same in libero-finetune.yaml.)

kubectl create secret generic pi0-lerobot-secrets --from-literal=HF_TOKEN="<your-token>"

# 2. Stage the eval script to FSx (one-time, from any pod that mounts the PVC)
kubectl run stager --rm -it --restart=Never \

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The eval-script staging story is fragile and told two ways — 🟢 [confirmed]

Observation: README step 2 stages the script with kubectl run stager --rm -it --image=busybox:1.37 ... "cat > /fsx/pi0-lerobot/evaluate_pi0.py" < src/evaluate_pi0.py, while the eval YAMLs' header comments say to use kubectl cp src/evaluate_pi0.py <pod>:/fsx/.... Two different methods for one step. The busybox one-liner is also brittle: -it requests a TTY over a redirected file (kubectl warns and drops it), and the --overrides container sets "stdin":true but not "stdinOnce":true, so cat > can block waiting for an EOF that never comes and the pod hangs instead of --rm-ing.
Suggestion: Pick one method. Simplest is to document only the kubectl cp path the YAMLs already reference (it needs a running pod that mounts the PVC — one line to note). If you keep the busybox stager, drop -t and add "stdinOnce":true to the container override.

- The [Kubeflow Training Operator](https://github.com/kubeflow/training-operator) installed
- FSx for Lustre PVC (`fsx-claim`) mounted at `/fsx`
- A HuggingFace token that has accepted:
- [google/paligemma-3b-pt-224](https://huggingface.co/google/paligemma-3b-pt-224) (Gemma license — π0's processor downloads this at train time)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Small stuff — one bundle — 🟢

Observation: A few cosmetic carryovers, none blocking: (1) the Gemma-license link still points at google/gemma-2b here and in both sub-READMEs — link the actually-gated google/paligemma-3b-pt-224 the reader must accept. (2) In both finetune manifests, the comment # Required: π0's PaliGemma processor requires Gemma license acceptance. now sits directly above the newly-inserted FORCE_RESTART env var, so it reads as annotating FORCE_RESTART instead of HF_TOKEN — move one or the other. (3) evaluate_pi0.py's checkpoint auto-discovery prefix pi0-finetune-{dataset} still can never match the real /fsx/runs/pi0-{dataset}/… layout, and the sagemaker-{region}-{account} S3 fallback is dead on EKS — both are harmless now that --finetuned-path/--test-dataset-local are always passed, but they're worth deleting to stop the next reader puzzling over them. (4) The headline MSE row pairs the 10-step MSE with the 1-step latency; for LIBERO the 1-step MSE (7.768e-02) is actually better than the 10-step headline — label the MSE's step count.
Suggestion: Sweep these together whenever you touch the files for the MIT-0 headers.

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 5/5 — Evaluation, Positives & Sources


Things That Look Great

  • Four rounds in, this is a genuinely good test case. You deleted (not defended) the monkey patches, the download Jobs, and the fork-branch fetch; implemented a real held-out split; ran it end-to-end on the target hardware; and retracted the contaminated numbers honestly. That arc is exactly what review is for.
  • The eval-indexing fix is correct, and I verified it against the pinned commit rather than taking it on faith — loading the full dataset and iterating absolute episode indices is the right shape, and the boundaries line up with no off-by-one. The held-out numbers are real.
  • Pinning to an exact commit (ddc2aa7a) with the rationale documented (GR00T dataclass bug on py3.12) is the strongest form of the pin, and it resolves cleanly.
  • The end-to-end run on p5.48xlarge (both datasets, timed, no errors) is more validation than the vast majority of first-time test-case PRs ever provide — and it's what let me retire my OOM concern.
  • Regions unified to us-west-2 consistently; the FORCE_RESTART fail-fast is honest; NCCL_SOCKET_IFNAME: "^lo", --dataset.video_backend=pyav, and the hyperparameters remain sound.
  • The remaining list is short and almost entirely polish: one mechanical CI header fix, one sibling-consistency reconcile, and some eval-labeling honesty. This is close.

Sources

  • LeRobot pinned commit ddc2aa7a27ba (verified live 2026-07-24): src/lerobot/datasets/dataset_metadata.py load_episodes(root) returns the full unfiltered episode table with dataset_from_index/dataset_to_index; src/lerobot/datasets/lerobot_dataset.py __getitem__ is absolute-indexed (get_item(idx)); commit resolves to a real HEAD dated 2026-07-20. https://github.com/huggingface/lerobot/tree/ddc2aa7a27ba725ae527959c7e4814aed550e452
  • MIT-0 header absence verified on head 0f5072bb across Dockerfile, buildspec.yml, the 5 YAML manifests, and evaluate_pi0.py (2026-07-24).
  • Region/image consistency (us-west-2) verified across the Dockerfile ARG, buildspec, four manifest image tags, and the README export (2026-07-24).
  • π0 FSDP memory + earlier e2e reproduction — 8× B300, 2026-07-10 (round-1 threads); the OOM concern is retired by the author's p5.48xlarge run reported in commit 0f5072b.
  • Prior rounds: round 1 · round 2 · round 3

@amindm
amindm requested a review from KeitaW July 24, 2026 16:57

@KeitaW KeitaW left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review Batch 1/1 — Blocker Cleared; Round-4 Should-Fixes Still Open

The MIT-0 blocker is resolved — verified on all eight files, correctly placed — ✅

Commit 22fa516 adds the two-line Copyright … / SPDX-License-Identifier: MIT-0 header to every non-README file — verified on head 22fa5168: Dockerfile, buildspec.yml, all five YAML manifests (droid-finetune, droid-eval, libero-finetune, libero-eval, pvc-fsx-lustre), and src/evaluate_pi0.py. Format and placement are right: the #-comment form is correct for the Dockerfile/YAML/Python files, and on evaluate_pi0.py the header sits above the module docstring — which is fine, since comments aren't statements, so the """…""" is still the module docstring and from __future__ import annotations remains a valid first import. That was the last CI-lint merge blocker, so the license gate should now be green. Thank you for the quick turnaround.

Still open from round 4 (this commit was headers-only) — 🟡/🟢

Since 22fa516 touched only the license headers, everything else from round 4 is unchanged. The inline threads from that round are still live on the same lines — I'm not re-posting them, just consolidating the checklist so nothing gets lost:

Should-fix (🟡):

  • DROID dataset-loading still asymmetric with LIBERO. LIBERO uses the real repo_id lerobot/libero_10 (no symlink) "to avoid the local-name 404"; DROID is still on LOCAL_DATASET_NAME="droid_local" + the symlink. Your e2e run proves DROID works today, but it's one failed local-load from the same 404 — mirror the LIBERO fix (lerobot/droid_100, drop the symlink) for symmetry.
  • Held-out numbers cover only the first 5 episodes. eval_episodes[:num_trajectories] scores episodes 80-84 / 304-308 while the README says "80-99" / "304-378" — the 89.7%/88.9% are N=5 on the first five. Score the full held-out set or relabel as "first 5 held-out (N=5)".
  • Unseeded ODE + base/fine-tuned step-count pairing. No manual_seed; base at num_inference_steps[0] vs headline fine-tuned at max(). The non-monotonic LIBERO sweep (10-step MSE < 5-step) is the visible symptom. Seed per trajectory; pair both models at the same step count.
  • Silent processor/tokenizer fallback still emits authoritative MSE. The except → None paths can publish a normalization-artifact "improvement" on a green run. Fail hard, or record a per-model processors_active flag and suppress the improvement % when either side ran without processors.

Also open (🟢): the PR description still shows the retracted 96.5%/93.1% (update it to the held-out numbers); dead SAVE_INTERVAL + hardcoded save_freq/checkpoint-path coupling to MAX_STEPS; the two-way / fragile eval-script staging (kubectl run -it over piped stdin, missing stdinOnce); Gemma link still google/gemma-2b; the FORCE_RESTART comment misattribution; and the dead checkpoint auto-discovery + SageMaker S3 leftovers in evaluate_pi0.py.

None of these are blockers on their own. Landing the DROID dataset-naming reconcile and the eval-coverage/seed items would make the results section fully defensible; the rest is polish you can sweep in one pass. From my side there are no new findings this round — the PR is one small commit away from an approve.

Sources

  • MIT-0 headers verified present + correctly placed on head 22fa5168 across the Dockerfile, buildspec.yml, the 5 YAML manifests, and evaluate_pi0.py; Python docstring / from __future__ ordering intact (2026-07-24).
  • Round-4 findings (all still open): round 4 review.

…ssors - Eval: default --num-trajectories=None (evaluate ALL held-out episodes) - Seed: torch.manual_seed(episode_index) before each trajectory for reproducibility - Processors: load_policy returns processors_active flag, included in results JSON - Removed --num-trajectories 5 from eval manifests (evaluates full held-out set)
@amindm
amindm requested a review from KeitaW July 27, 2026 18:10
amindm added 2 commits July 30, 2026 10:23
…Jobs pass the full held-out range via --eval-episodes, but evaluate_pi0.py scores only the first --num-trajectories of them (default: 5). The READMEs claimed coverage of 80-99 / 304-378 while the reported metrics actually cover 80-84 / 304-308. Docs-only change; evaluation behavior is unchanged: - README: N=5 note above Results, both result headings labeled, and a 'Scored in reported results' column added to the Dataset Splits table - droid/libero READMEs: Results headings labeled with the scored range, a 'Scored' row added to each Train/Test Split table, and the eval-Job step now states which episodes are scored - Note that ODE-sweep ordering is within noise at N=5, which explains the non-monotonic LIBERO sweep Raise --num-trajectories in the eval Job to score the full held-out set.
… the reviewer's requested change. The trailing clause about ODE-sweep ordering being within noise was my own addition and belongs to a separate open review item (unseeded ODE / base-vs-finetuned step-count pairing), which wants a code fix rather than a README caveat.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants