prlq/pi05-stretch-mobile-manip

pi0.5 — Stretch 3 mobile manipulation (full fine-tune, 50k steps)

Full fine-tune of pi05_base (openpi) on Huggingzhu1/stretch_pickup_nav_v1 (revision d511f5cd, train/ split only: 800 episodes, 420,641 frames, 10 fps).

The dataset covers two language-conditioned tasks on a Hello Robot Stretch 3 in Isaac Sim:

• Pickup: "Pick up the <noun> on the desk and hold it after retracting."
• Navigation: "Drive to the red marker and park."

Trained with the openpi PyTorch trainer on 2× AMD Instinct MI325X (ROCm 7.2), 2026-10-14 to 2026-10-16.

Run inference

The checkpoints are openpi PyTorch checkpoints and run with stock openpi plus the small Stretch config in code/openpi_changes.patch.

1. Install openpi at the commit we used, following openpi's README, including its PyTorch setup step that copies transformers_replace. Then apply our patch:

   git clone https://github.com/Physical-Intelligence/openpi.git && cd openpi && git checkout 215abfb
   git apply /path/to/this_repo/code/openpi_changes.patch
   uv sync && cp -r ./src/openpi/models_pytorch/transformers_replace/* .venv/lib/python3.11/site-packages/transformers/
   

   The patch adds the pi05_stretch config and src/openpi/policies/stretch_policy.py, which inference needs. Its other changes (EMA, the gradient-checkpointing switch) only affect training.

2. Download a checkpoint, for example the final one:

   hf download prlq/pi05-stretch-mobile-manip --include "50000/*" --local-dir ./ckpt
   

   openpi loads model.safetensors. To use the EMA weights instead, make a folder where model.safetensors is the EMA file:

   mkdir ckpt/50000_ema && ln -s ../50000/model_ema.safetensors ckpt/50000_ema/model.safetensors && ln -s ../50000/assets ckpt/50000_ema/assets
   

3. Call the policy:

   from openpi.policies import policy_config
   from openpi.training import config as _config
   policy = policy_config.create_trained_policy(_config.get_config("pi05_stretch"), "ckpt/50000")
   out = policy.infer({
       "observation/head_image":  head_rgb,   # uint8 (1280, 720, 3): head camera as recorded (portrait)
       "observation/wrist_image": wrist_rgb,  # uint8 (480, 640, 3): wrist camera
       "observation/state":       state,      # float32 (11,): same order as the dataset's observation.state
       "prompt": "Pick up the apple on the desk and hold it after retracting.",
   })
   actions = out["actions"]                   # float32 (50, 8): next 50 commands at 10 Hz, dataset action order
   

   • Inputs: images can be any size; they are resized with padding to 224×224 internally. State and actions are normalized with the checkpoint's own assets/train/norm_stats.json.
   • Actions: the 8 dims are, in order: base forward velocity, base yaw velocity, lift velocity, arm extension velocity, wrist-yaw target, head-pan target, gripper close (0/1), wrist-pitch target. These are the dataset's action definition.
   • Ready-to-run example: code/infer_example.py runs the policy on recorded validation frames and compares with the recorded actions:

     python code/infer_example.py --ckpt ckpt/50000 --dataset /path/to/stretch_pickup_nav_v1/validation
     

     Before upload, we ran it on validation episodes (scenes not seen in training) with both raw and EMA weights. The per-dimension mean absolute error against the recorded 50-step chunks was mostly 0.000–0.03, with head-pan up to about 0.35 rad on some frames.
   • Latency: on an AMD MI325X a call takes about 63 ms after warm-up; the first call compiles (about 15 s to 3 min). openpi's PyTorch inference keeps most weights in bf16, as it does for every checkpoint.
   • Serving: for remote inference, openpi's policy server works the same way: uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi05_stretch --policy.dir=ckpt/50000.

Checkpoints

Folder	Contents
10000/ … 40000/	model.safetensors (raw weights), model_ema.safetensors (EMA weights, decay 0.99), assets/train/norm_stats.json, metadata.pt
50000/	Same as above (final step)

Raw and EMA weights are fp32, 14.5 GB each. Both load with openpi's PyTorch PI0Pytorch and the pi05_stretch config below.

Training setup

Setting	Value
Base model	gs://openpi-assets/checkpoints/pi05_base, converted to PyTorch (fp32) with examples/convert_jax_model_to_pytorch.py
Model	Pi0Config(pi05=True, action_horizon=50); action dim 32; discrete state input (pi05 default)
Data	Head camera → base_0_rgb; wrist camera → left_wrist_0_rgb; right_wrist_0_rgb masked. State (11-d) and actions (8-d) zero-padded to 32. Prompt = dataset task. Actions used as recorded (no delta transform).
Normalization	Quantile normalization; statistics from the train split (norm_stats.json)
Steps / batch	50,000 steps, global batch 32 (2 GPUs × 16)
Optimizer	AdamW, gradient clipping 1.0; warmup 1,000 steps to 2.5e-5, cosine decay to 2.5e-6 at 50k
Precision	Full fp32 (pytorch_training_precision="float32"); matmuls with openpi's own float32_matmul_precision("high")
EMA	0.99 (added, see below)
Throughput	≈ 1.64 s/step (≈ 81% of the MI325X fp32 matrix peak by a parameter-count estimate); ≈ 23 h total

Training loss (logged every 100 steps):

Step	0	1k	5k	10k	20k	30k	40k	49.9k
Loss	0.086	0.0065	0.0040	0.0030	0.0021	0.0019	0.0014	0.0012

Interruption: a cluster maintenance event cancelled the run at step 35,894. Training resumed from the 30k checkpoint, restoring the weights, optimizer state and EMA. Over the overlapping steps 30k–35.8k the mean loss was 0.00173 before and 0.00168 after.

No validation or closed-loop evaluation has been run yet.

Training details and changes relative to the original openpi code

This section is only needed to reproduce training.

Code changes

Base: Physical-Intelligence/openpi at commit 215abfb. The full diff is in code/openpi_changes.patch.

1. New dataset config (src/openpi/policies/stretch_policy.py, LeRobotStretchDataConfig, TrainConfig pi05_stretch in src/openpi/training/config.py). This maps the Stretch cameras, state and actions to the pi0.5 inputs as described above. The data-config factory has action_sequence_keys=("action",) and prompt_from_task=True.
2. EMA in the PyTorch trainer (scripts/train_pytorch.py). openpi's PyTorch trainer does not implement EMA, although its JAX trainer enables it by default.
   • We added it with the JAX trainer's update rule, ema = decay·ema + (1 − decay)·params, applied after every optimizer step and initialized from the starting weights. Decay is the config's ema_decay (0.99).
   • It is kept on the main rank in fp32, saved as model_ema.safetensors alongside the raw weights, and reloaded on --resume.
3. Optional gradient checkpointing (scripts/train_pytorch.py). The trainer always enabled gradient checkpointing; OPENPI_NO_GRAD_CKPT=1 now disables it.
   • It was disabled for this run because memory was ample, which took the step time from 1.95 s to 1.64 s.
   • The loss at step 19 was identical with and without it.
4. Normalization statistics. These were computed from the dataset's parquet files (code/norm_stats_parquet.py) instead of scripts/compute_norm_stats.py, so no video had to be decoded.
   • The script uses openpi's own RunningStats, with the same order and batch size of 32 as the original script.
   • Action chunks of 50 are clamped to the episode end, matching LeRobot's padding.

Environment changes (AMD / ROCm, outside the openpi repo)

• Packages: PyTorch 2.14.1+rocm7.2 replaces CUDA torch 2.7.1, and JAX 0.5.3 is the CPU build. JAX is used only for imports and the weight conversion. All other Python packages are pinned to openpi's uv.lock (code/requirements_rocm.txt), with transformers 4.53.2 carrying openpi's transformers_replace patches.
• Video decoding: torchcodec was removed (its CUDA build does not load on ROCm). torchvision ≥ 0.22 no longer has VideoReader, which LeRobot's pinned version calls, so lerobot/common/datasets/video_utils.py now decodes with PyAV using the same keyframe seek. Decoded frames match a full sequential decode exactly (code/lerobot_video_utils.patch).
• OpenCV: opencv-python was replaced by opencv-python-headless 4.11.0.86, because one node lacked libGL. OpenCV is not used in the training data path.
• numpy: during the OpenCV swap, numpy was unintentionally upgraded from 1.26.4 to 2.4.6, so steps 30k–50k ran with numpy 2.4.6. On 50 training samples, the full data pipeline produces bit-identical outputs and dtypes under both versions, covering state, images, actions and tokenized prompt. The environment has since been reverted to 1.26.4.

Differences from openpi's official (JAX) fine-tuning

• Precision: PyTorch full fp32 instead of JAX mixed precision (fp32 weights, bf16 compute). openpi's PyTorch trainer offers only full bf16 or full fp32.
• Numerical detail: the hardware and kernels differ (AMD GPUs vs NVIDIA), so runs are not bitwise reproducible across platforms. They should agree statistically.