Fine-Tuning π0 (Pi-Zero) for robotic manipulation on Amazon SageMaker HyperPod EKS

Introduction

This post demonstrates how to fine-tune π0 (Pi-Zero), a 3-billion parameter flow matching Vision-Language-Action model from Physical Intelligence, on robot manipulation datasets using Amazon SageMaker HyperPod with Amazon Elastic Kubernetes Service (Amazon EKS) orchestration. Two standard robotics benchmarks provide the training data: DROID for real-world manipulation and LIBERO-10 for simulated tabletop tasks, with a strict episode split that keeps the evaluation set out of training entirely. The fine-tuned models reduce action-prediction MSE by 89.7% on DROID and 88.9% on LIBERO relative to the base checkpoint. Training runs on a single ml.p5.48xlarge node (8× H100 80GB) with FSDP FULL_SHARD sharding, and π0’s flow matching architecture generates a full 50-step action chunk in a single ODE step at under 200 ms. The complete pipeline, Kubernetes manifests for training and evaluation, plus an optional container image, is provided as a reusable test case for Amazon SageMaker HyperPod EKS clusters.

For a long time, robotics had a scaling problem: every new task and every new robot required training a policy from scratch on thousands of demonstrations. Physical AI changed this by applying the same recipe that transformed NLP and computer vision, pre-train a large model on internet-scale data, then fine-tune it on a small dataset for your specific problem. The underlying insight is that most of what a robot needs to know doesn’t come from robot data. It already exists in models trained on images, text, and video.

Which kind of pre-trained model makes the best foundation for robot policies is still an open question. Vision-language-action models (VLAs) such as π0, OpenVLA, and GR00T start from a pre-trained vision-language model and fine-tune it to emit actions. World action models (WAMs) start instead from a video model trained to predict how scenes evolve. The trade-off comes down to what each backbone already knows: language pre-training gives a better prior for following instructions about objects that never appeared in a demonstration, while video pre-training gives a better prior for physical dynamics.

π0, from Physical Intelligence, is a leading example of the VLA approach. It pairs a VLM backbone with an action expert that generates continuous motor commands using flow matching, a faster alternative to diffusion. If you are building a robotic manipulation pipeline and want to adapt a foundation model to your specific robot and task, this post is the workflow to start from. You’ll learn how to fine-tune π0 using Hugging Face LeRobot on Amazon SageMaker HyperPod, and you’ll see that fine-tuning the base model on domain-specific data yields an 87-90% improvement in action prediction accuracy across two established robotics benchmarks. You get an end-to-end recipe that runs from an empty cluster to a measured result, including the operational details that are easy to get wrong — dataset splits, checkpoint sizing on shared storage, gated model licenses, and CUDA version mismatches in the base container.

What is π0?

π0 (Pi-Zero) is a 3-billion parameter Vision-Language-Action (VLA) model developed by Physical Intelligence. It was designed as a generalist robot foundation model capable of controlling diverse robot embodiments.

• Model: lerobot/pi0_base
• Architecture: PaliGemma vision-language backbone + flow matching action expert (shared transformer)
• Parameters: ~3 billion total
• Pre-training: 10,000+ hours of robot demonstrations across multiple platforms
• Action generation: Flow matching — learns a velocity field that transports noise to actions via ODE integration
• Action prediction: Action chunks — 50 future timesteps predicted simultaneously
• Inputs: RGB images (multi-view) + robot proprioceptive state + language instruction
• License: Apache 2.0 (LeRobot port of openpi)

Flow Matching Training Objective

π0 generates actions by integrating a learned ODE rather than running iterative denoising:

x_t = (1 - t) * noise + t * action       # linear interpolation path
target_velocity = action - noise          # constant along the path
loss = ||v_predicted(x_t, t, obs) - target_velocity||²

At inference, π0 starts from Gaussian noise and integrates dx/dt = v(x, t, obs) from t=0 to t=1 using Euler steps. The linear interpolation used during training produces straight paths that require fewer integration steps than diffusion’s curved noise schedules.

Flow Matching vs Diffusion

Why Amazon SageMaker HyperPod with Amazon EKS?

Amazon SageMaker HyperPod with Amazon EKS orchestration provides a persistent GPU cluster managed by Amazon SageMaker with Kubernetes-native job submission. Compared to Amazon SageMaker training jobs (ephemeral, fully managed), HyperPod offers a different set of trade-offs:

For iterative experimentation — trying different hyperparameters, datasets, or model variants — HyperPod’s persistent storage and instant job submission eliminate the data-transfer overhead that dominates short training runs. It also means storage is a shared resource you have to manage, which matters for π0 (see the checkpoint sizing note below).

Datasets and Splits

DROID (Distributed Robot Interaction Dataset) is a large-scale real-world manipulation dataset collected across multiple institutions. This post uses lerobot/droid_100, a 100-episode subset with 7-DoF joint actions recorded from a real Franka arm performing diverse manipulation tasks.

LIBERO-10 is a simulated tabletop manipulation benchmark with 10 task categories. This post uses lerobot/libero_10 with 7-DoF actions (6 joints + gripper) from a simulated Franka Panda arm. Note that π0’s pre-training mixture includes LIBERO data.

Both datasets are already in LeRobot V2 format (Parquet + MP4 + metadata), so no preprocessing or action synthesis is required.

Train and evaluation episodes are split explicitly, and the split is enforced on both sides: training passes --dataset.episodes with train-only indices, and evaluation passes --eval-episodes with held-out indices. This matters — an earlier iteration of this work reported much higher improvements because evaluation had been run on episodes the model had already trained on.

lerobot/libero_10 contains 379 episodes, not the 500 the benchmark description might suggest. Hardcoding 400/500 produces an index error partway into the run, so the episode count is worth confirming from meta/info.json before launching a 6-hour job.

Solution Architecture

The pipeline runs on an Amazon SageMaker HyperPod EKS cluster with:

• Compute: 1× ml.p5.48xlarge (8× H100 80GB, 32× EFA)
• Storage: Amazon FSx for Lustre (1.2 TiB) mounted at /fsx via a PVC named fsx-claim
• Orchestration: Kubeflow Training Operator v1.9.1 (PyTorchJob CRD)
• Container: AWS Deep Learning Containers (PyTorch 2.9, CUDA 13, Python 3.12) with LeRobot installed at pod startup

The training pod requests the whole node — nvidia.com/gpu: 8, vpc.amazonaws.com/efa: 32, 90 vCPU, 700 GiB memory — plus a 128 GiB /dev/shm emptyDir backed by medium: Memory, since PyTorch’s dataloader workers move decoded video frames through shared memory and the container default is far too small for four workers reading MP4. The evaluation Job is much smaller: 1 GPU, 8 vCPU, 64 GiB, 32 GiB /dev/shm.

Two Kubernetes manifests per dataset:

1. Training PyTorchJob — downloads the dataset to Amazon FSx for Lustre on first run, then runs FSDP training on 8 GPUs for 20,000 steps
2. Evaluation Job — loads the base and fine-tuned policies on 1 GPU and compares them on the held-out episodes, with an ODE step-count sweep

There is no separate download job. The training manifest checks for meta/info.json on Amazon FSx for Lustre and pulls the dataset from the Hugging Face Hub only if it is missing, which keeps the happy path to a single kubectl apply.

Container Strategy

The manifests point directly at the AWS Deep Learning Containers and install LeRobot in the container’s entrypoint:

763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.9.0-gpu-py312-cu130-ubuntu22.04-sagemaker-v1.9

pip install --quiet \
  "lerobot[pi,dataset]@git+https://github.com/huggingface/lerobot.git@ddc2aa7a27ba725ae527959c7e4814aed550e452" av
pip uninstall -y flash-attn 2>/dev/null || true

This costs about 3-5 minutes of pod startup and removes the need to build, push, and version a custom image — a reasonable trade for a 6-hour training job. A Dockerfile and buildspec.yml are included for teams that prefer a pre-baked image (roughly 30 seconds to start instead of 5 minutes); in that case only the image: field in the manifests changes. The buildspec is written for AWS CodeBuild.

Three details in that install command are load-bearing:

LeRobot is pinned to a commit, not a tag. The released v0.5.0 has a dataclass definition in the GR00T policy that fails to import on Python 3.12, which is what the AWS Deep Learning Containers (DLC) ships. Commit ddc2aa7a fixes it. The pin also keeps the recipe reproducible as LeRobot’s main branch moves.

The `dataset` extra only exists on that commit. lerobot[pi,dataset] resolves on ddc2aa7a but not on v0.5.0, so the extras and the pin have to move together. av is installed explicitly because the dataloader reads MP4 video with the PyAV backend (--dataset.video_backend=pyav).

`flash_attn` is uninstalled at startup. It is pulled in transitively by LeRobot’s XVLA policy and is compiled against CUDA 12, so importing it inside the CUDA 13 DLC fails on libcudart.so.12. π0 does not use flash attention, so removing the package is the cheapest fix. Rebuilding flash_attn against CUDA 13 in a custom image would be cleaner.

Earlier iterations of this recipe used a usercustomize.py monkey-patch to stop LeRobot from calling the Hugging Face Hub for local datasets. That is no longer needed: passing --dataset.root pointing at a local directory is handled natively.

Training Configuration

Training Notes

FSDP FULL_SHARD. Shards model parameters, gradients, and optimizer states across the 8 GPUs. This leaves headroom for activations, which are the dominant memory cost with a 50-step action chunk and multi-view image inputs.

`–policy.dtype=float32` combined with `–mixed_precision=bf16`. The model must load in fp32 so FSDP’s flat-parameter construction sees a uniform dtype; a mix of bf16 and fp32 layers raises “Must flatten tensors with uniform dtype”. FSDP’s MixedPrecision policy then handles the bf16 cast for compute, so this costs nothing in throughput.

20,000 steps. Flow matching converges more gradually than diffusion, and π0’s 50-step action chunk (versus 16 for diffusion-based GR00T) means each step covers more of the action distribution to learn.

Save only the final checkpoint. A π0 checkpoint is roughly 45 GB — model weights plus sharded optimizer state. With the default 2,000-step save interval, a 20,000-step run writes ten of them and fills a 1.2 TiB Amazon FSx for Lustre filesystem well before training completes. Setting --save_freq equal to --steps writes one checkpoint at the end. If you need intermediate checkpoints for a learning-curve study, size Amazon FSx for Lustre accordingly or prune older checkpoints from a sidecar.

Camera key mapping. π0 expects observation.images.base_0_rgb and observation.images.left_wrist_0_rgb. DROID and LIBERO use different key names, so --rename_map maps them, and --policy.empty_cameras=1 fills the third camera slot the base model was trained with.

Prerequisites

Before you start the walkthrough, make sure you have the following in place:

• An AWS account with sufficient service quota for the GPU instance type you plan to use (this post uses one ml.p5.48xlarge).
• An Amazon SageMaker HyperPod cluster orchestrated by Amazon EKS, with GPU nodes in the us-west-2 Region.
• An Amazon FSx for Lustre filesystem exposed to the cluster as a PersistentVolumeClaim named fsx-claim and mounted at /fsx. The included kubernetes/pvc-fsx-lustre.yaml provisions this.
• The Kubeflow Training Operator (v1.9.1) installed on the cluster, which provides the PyTorchJob custom resource.
• kubectl installed locally and configured against the cluster.
• A Hugging Face account with an access token, and the Gemma license accepted for google/paligemma-3b-pt-224. The π0 checkpoint itself is ungated, but its processor pulls the gated PaliGemma repo, so training fails without this.
• The token stored in a Kubernetes Secret named pi0-lerobot-secrets under the key HF_TOKEN.

Walkthrough

Step 1: Reserve Capacity and Create the Cluster

Reserve an ml.p5.48xlarge through Amazon SageMaker Training Plans, then create an Amazon SageMaker HyperPod EKS cluster. Provision an Amazon FSx for Lustre filesystem and expose it to the cluster as a PVC named fsx-claim. The included kubernetes/pvc-fsx-lustre.yaml does this via dynamic provisioning with reclaimPolicy: Retain, so deleting the PVC does not destroy the filesystem and the datasets on it.

Step 2: Install Prerequisites

# Connect to the cluster
aws eks update-kubeconfig --name <cluster-eks-name> --region us-west-2

# Install the Kubeflow Training Operator (provides the PyTorchJob CRD)
kubectl apply -k "github.com/kubeflow/training-operator/manifests/overlays/standalone?ref=v1.9.1"

# Hugging Face token (needed for the gated PaliGemma processor)
kubectl create secret generic pi0-lerobot-secrets --from-literal=HF_TOKEN="<your-token>"

lerobot/pi0_base itself is ungated. The gate is the transitive dependency: π0’s processor downloads google/paligemma-3b-pt-224, which requires accepting the Gemma license on your Hugging Face account. If the token is missing or the license has not been accepted, the job runs for roughly 15 minutes and then fails with a GatedRepoError (HTTP 401) — long enough to look like a successful start, so it is worth verifying the token before walking away.

The manifests target HyperPod’s node labels, which carry an ml. prefix:

nodeSelector:
  node.kubernetes.io/instance-type: ml.p5.48xlarge

On plain Amazon EKS (managed node groups or Karpenter) the label is p5.48xlarge without the prefix, so this line needs editing.

Step 3: Stage the Evaluation Script on Amazon FSx for Lustre

The evaluation Job reads evaluate_pi0.py from /fsx/pi0-lerobot/, which keeps the manifests free of embedded Python and avoids fetching code from a URL at runtime. Staging it once is the one manual step in the pipeline.

A ConfigMap is the reliable way to do it. Piping the file into a pod’s stdin with kubectl run -it looks simpler but can produce a zero-byte file if the TTY closes before the write completes:

kubectl create configmap pi0-eval-script --from-file=evaluate_pi0.py=src/evaluate_pi0.py

kubectl run stager --rm --restart=Never --image=busybox:1.37 --overrides='
{"spec":{
  "volumes":[
    {"name":"fsx","persistentVolumeClaim":{"claimName":"fsx-claim"}},
    {"name":"script","configMap":{"name":"pi0-eval-script"}}],
  "containers":[{"name":"stager","image":"busybox:1.37",
    "command":["sh","-c","mkdir -p /fsx/pi0-lerobot && cp /script/evaluate_pi0.py /fsx/pi0-lerobot/"],
    "volumeMounts":[
      {"name":"fsx","mountPath":"/fsx"},
      {"name":"script","mountPath":"/script"}]}]}}'

Verify the file actually landed with a non-zero size before submitting the evaluation Job:

kubectl run check --rm -it --restart=Never --image=busybox:1.37 --overrides='
{"spec":{
  "volumes":[{"name":"fsx","persistentVolumeClaim":{"claimName":"fsx-claim"}}],
  "containers":[{"name":"c","image":"busybox:1.37",
    "command":["sh","-c","wc -l /fsx/pi0-lerobot/evaluate_pi0.py"],
    "volumeMounts":[{"name":"fsx","mountPath":"/fsx"}]}]}}'

Step 4: Submit Training

kubectl apply -f kubernetes/droid/droid-finetune.yaml
kubectl logs -f pi0-lerobot-droid-finetune-worker-0

The pod installs LeRobot, downloads the dataset to Amazon FSx for Lustre if needed, links it into the LeRobot cache, and launches training:

accelerate launch \
    --num_processes=8 --num_machines=1 \
    --use_fsdp --mixed_precision=bf16 \
    --fsdp_sharding_strategy=FULL_SHARD \
    --fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP \
    --fsdp_state_dict_type=SHARDED_STATE_DICT \
    --fsdp_use_orig_params=true \
    --fsdp_cpu_ram_efficient_loading=false \
    --fsdp_sync_module_states=true \
    --fsdp_offload_params=false \
    $(which lerobot-train) \
    --dataset.repo_id=droid_local \
    --dataset.root=/fsx/datasets/droid_100 \
    --dataset.episodes="[0, 1, ..., 79]" \
    --dataset.video_backend=pyav \
    --policy.path=lerobot/pi0_base \
    --policy.dtype=float32 \
    --policy.device=cuda \
    --policy.gradient_checkpointing=true \
    --policy.chunk_size=50 --policy.n_action_steps=50 \
    --policy.empty_cameras=1 \
    --policy.push_to_hub=false \
    --rename_map="$RENAME_MAP" \
    --steps=20000 --batch_size=4 \
    --save_freq=20000 \
    --optimizer.lr=2.5e-5 \
    --num_workers=4 \
    --output_dir=/fsx/runs/pi0-droid/training

The last three flags are the configuration this recipe was validated with: --fsdp_sync_module_states=true broadcasts rank 0’s loaded weights to the other ranks, CPU-RAM-efficient loading is left off, and parameter offload is left off because the 8× 80 GB budget does not need it and offloading would cost throughput.

Training output shows step progress, loss, and throughput:

Training: 100%|██████████| 20000/20000 [6:15:00<00:00, 1.12s/step]
step:20K loss:0.060 grdn:0.800 lr:2.5e-06
Model weights saved to /fsx/runs/pi0-droid/training/checkpoints/020000/pretrained_model!

Measured wall-clock: 6h15m for DROID and 6h11m for LIBERO at 20,000 steps on 8× H100.

LeRobot raises FileExistsError if --output_dir already exists, which is a useful guard against silently overwriting a completed run. Set FORCE_RESTART=1 on the pod to delete and start fresh.

Step 5: Evaluate

Free the GPUs, then submit the evaluation Job:

kubectl delete pytorchjob pi0-lerobot-droid-finetune
kubectl apply -f kubernetes/droid/droid-eval.yaml

python /fsx/pi0-lerobot/evaluate_pi0.py \
    --dataset droid \
    --finetuned-path /fsx/runs/pi0-droid/training/checkpoints/020000/pretrained_model \
    --test-dataset-local /fsx/datasets/droid_100 \
    --eval-episodes 80 81 82 ... 99 \
    --results-out /fsx/runs/pi0-droid/eval_results_droid.json \
    --num-inference-steps 10 5 3 1

The script loads the base and fine-tuned policies sequentially on a single GPU (one at a time, to stay within memory), predicts 50-step action chunks on the held-out episodes, and compares them to ground truth. torch.manual_seed(episode_index) is called before each trajectory so the ODE’s noise draw is reproducible across models and across step counts. The results JSON records a processors_active flag per model so a run where the normalization pipeline silently failed to load can be identified rather than reported as an improvement.

Two details about --num-inference-steps 10 5 3 1 are worth knowing before reading the numbers. The base model is scored once, at the first value in that list, while the fine-tuned model is swept across all of them and the headline comparison uses the largest. Because the list is given in descending order, both sides land on 10 steps and the comparison is properly paired — but that pairing follows from the argument order, so reordering the list would silently compare a 1-step base against a 10-step fine-tuned model. Separately, the script scores the first --num-trajectories of the episodes passed to --eval-episodes, and that flag defaults to 5.

Results

Both models were evaluated on episodes that were excluded from training. The numbers below are averaged over the first five held-out episodes (N=5) of each dataset — episodes 80-84 for DROID and 304-308 for LIBERO. The evaluation Jobs pass the full held-out range via --eval-episodes, but the script scores --num-trajectories of them, which defaults to 5.

DROID — 7-DoF Real Joint Actions

Trained on episodes 0-79, evaluated on held-out episodes.

LIBERO-10 — 7-DoF Tabletop Manipulation

Trained on episodes 0-303, evaluated on held-out episodes.

Headline improvement percentages compare both models at 10 ODE steps, so the two sides are paired at the same step count.

ODE Step-Count Sweep (Flow Matching Speed Dial)

Fine-tuned model, same held-out episodes, varying the number of Euler integration steps.

DROID:

LIBERO:

Key Findings

Fine-tuning cuts prediction error by roughly 9× on held-out data. MSE drops 89.7% on DROID and 88.9% on LIBERO. The two datasets land within a percentage point of each other despite one being real-world and the other simulated, which suggests the gain comes from adapting the action distribution to the target embodiment rather than from anything dataset-specific.

MAE improves less than MSE. DROID’s MAE reduction (72.8%) trails its MSE reduction (89.7%). Since MSE penalizes large errors quadratically, this gap says most of the gain is in eliminating the base model’s occasional large mispredictions, while typical per-step error shrinks by a smaller factor. MSE alone would overstate the improvement.

One ODE step is enough, and it is roughly 2× faster. On DROID, a single Euler step gives the lowest MSE in the sweep (5.151e-02 versus 5.664e-02 at 10 steps). On LIBERO, going from 10 steps to 1 cuts latency from 379 ms to 197 ms with no accuracy cost. The learned velocity field is close to straight after fine-tuning, so additional integration steps buy nothing. The LIBERO sweep is not monotonic — 10 steps scores slightly better than 5 — which is the expected noise floor at N=5 episodes and a reason not to read fine ordering into these numbers.

Sub-200 ms for 50 actions. At 1 ODE step on H100, a full 50-timestep chunk takes under 200 ms, so the policy can re-plan at about 5 Hz while each chunk covers roughly one second of robot motion.

Inference is not memory-bound. Peak VRAM is about 15.7 GB for both base and fine-tuned models — full fine-tuning added no parameters, and a single H100 has ample headroom. The 80 GB GPUs are needed for training, not serving.

Amazon FSx for Lustre removes data-transfer overhead from the iteration loop. With datasets and checkpoints persisting on Amazon FSx for Lustre, resubmitting a job starts training within the pod-startup time. The cost is that storage is finite and shared, which is exactly why the 45 GB checkpoint size has to be managed explicitly.

Practical Interpretation

With 7-DoF actions, the base model’s MSE of 0.55-0.77 reflects predictions that are not tracking the target embodiment’s action distribution. The fine-tuned model reduces per-joint error by 3.7× (DROID) to 5.6× (LIBERO) in MAE terms, and roughly 9× in MSE terms. At about 200 ms per 50-step chunk on H100, the model re-plans at about 5 Hz — a workable rate for manipulation, where each chunk already covers about a second of motion.

Known Limitations

Evaluation coverage is N=5. The reported numbers average the first five held-out episodes, not the full held-out set (20 for DROID, 75 for LIBERO). The direction and magnitude of the improvement are clear at this sample size, but the ODE sweep ordering is within noise. Raising –num-trajectories to cover the full held-out set is a straightforward extension and the right move before quoting these figures as benchmarks.

Open-loop evaluation only. This evaluation measures prediction accuracy against recorded trajectories. Closed-loop deployment, where prediction errors compound across timesteps, requires a simulator or real hardware. Open-loop MSE is a necessary but not sufficient signal for task success.

Single-node training. The recipe uses one node with 8 GPUs. Multi-node FSDP works but needs additional NCCL/EFA configuration and validation that is out of scope here.

Small training sets. 80 episodes for DROID and 304 for LIBERO. Full DROID is roughly 76,000 episodes; scaling up would test generalization rather than just adaptation.

Fixed task instructions. Neither dataset provides diverse phrasings of the same task, so the fine-tuned model’s language robustness is untested.

Asymmetric dataset loading. LIBERO passes the real Hub repo id (lerobot/libero_10) with –dataset.root, while DROID uses a local alias (droid_local) plus a symlink into the LeRobot cache. Both work; the LIBERO form is less dependent on cache-resolution behavior.

`flash_attn` uninstall is a workaround. Removing the package at startup is correct for π0 but leaves the image inconsistent for any policy that does need flash attention. A custom image with flash_attn built against CUDA 13 would resolve it properly.

Cleanup

The cluster, its FSx for Lustre filesystem, and any reserved capacity keep billing while they exist. Remove what you no longer need in this order.

1. Delete the training jobs: kubectl delete pytorchjob pi0-lerobot-droid-finetune and kubectl delete pytorchjob pi0-lerobot-libero-finetune.
2. Delete the evaluation jobs: kubectl delete job pi0-lerobot-droid-eval and kubectl delete job pi0-lerobot-libero-eval.
3. Remove the staged evaluation script and the ConfigMap: kubectl delete configmap pi0-eval-script.
4. Free FSx for Lustre space. Checkpoints are roughly 45 GB each, so delete /fsx/runs and /fsx/datasets if you no longer need them. Note that the PVC uses reclaimPolicy: Retain, so deleting the PVC leaves the filesystem, and its cost, in place. Delete the FSx for Lustre filesystem itself in the console or with the AWS Command Line Interface (AWS CLI) when you are finished.
5. Scale the HyperPod cluster down or delete it, which stops the instance charges.
6. Release the Amazon SageMaker Training Plan reservation if you no longer need the capacity.

Conclusion

This post fine-tuned π0, a 3-billion parameter flow matching Vision-Language-Action model, on Amazon SageMaker HyperPod EKS and measured an 89.7% MSE reduction on DROID and 88.9% on LIBERO against strictly held-out episodes. The HyperPod-based approach — Amazon FSx for Lustre for persistent data and checkpoints, PyTorchJob for orchestration, Training Plans for capacity — gives faster iteration than ephemeral training jobs at the same compute performance, with the trade-off that shared storage becomes something you manage.

Two results are worth carrying forward. First, the honest number depends entirely on the split: an earlier version of this work reported ~96% improvement because evaluation episodes overlapped the training set. Enforcing the split on both the training and evaluation side moved the figure to ~89%, and that is the number that means something. Second, flow matching’s adjustable ODE step count is a genuine deployment lever: on H100, a single Euler step gives the best accuracy in the sweep at roughly half the latency of 10 steps, which puts a 3B-parameter VLA inside a real-time control budget without quantization or distillation.

Get started by cloning the repository and running the DROID walkthrough on your own cluster. To learn more, see the Amazon SageMaker HyperPod documentation. If you try this on a different robot embodiment or dataset, share your results in the comments.

Next Steps

1. Full held-out evaluation: score all 20 DROID and 75 LIBERO held-out episodes to tighten the confidence interval on the reported improvements.
2. Closed-loop evaluation: deploy the fine-tuned policy in LIBERO simulation and measure task success rate, not just open-loop action error.
3. Scale the training set: fine-tune on the full DROID dataset to test cross-task generalization.
4. Multi-node FSDP: scale to 2-4 nodes with EFA for larger effective batch sizes.
5. Serving path: package the fine-tuned policy behind a real-time inference service and measure end-to-end latency including observation preprocessing.
6. Diffusion comparison: run a diffusion-based policy on the same splits for a direct flow-matching-versus-diffusion comparison.

Appendix

Repository Structure

pi0-lerobot/
├── Dockerfile                        # Optional: pre-built image (faster startup)
├── buildspec.yml                     # AWS CodeBuild spec for the optional image
├── README.md
├── src/
│   └── evaluate_pi0.py               # Evaluation: MSE/MAE + ODE sweep
└── kubernetes/
    ├── pvc-fsx-lustre.yaml           # FSx PVC (dynamic, reclaimPolicy: Retain)
    ├── droid/
    │   ├── README.md
    │   ├── droid-finetune.yaml       # Training PyTorchJob (8× H100, FSDP)
    │   └── droid-eval.yaml           # Evaluation Job (1× GPU, ODE sweep)
    └── libero/
        ├── README.md
        ├── libero-finetune.yaml
        └── libero-eval.yaml

Resource Usage

Cost depends on whether the node comes from a Training Plan reservation or on-demand capacity; see the Amazon SageMaker pricing page for current rates. Storage sizing is the constraint to plan for: a single π0 checkpoint is about 45 GB.

Troubleshooting

References

• π0 paper — Physical Intelligence, 2024
• LeRobot pi0_base checkpoint
• Hugging Face LeRobot
• openpi (JAX reference implementation)
• DROID dataset
• LIBERO benchmark
• Amazon SageMaker HyperPod
• Amazon SageMaker Training Plans
• Kubeflow Training Operator
• AWS Deep Learning Containers
• GitHub repo: https://github.com/awslabs/awsome-distributed-ai/tree/main/3.test_cases/pytorch/pi0-lerobot