---
library_name: lerobot
pipeline_tag: robotics
license: cc-by-sa-4.0
base_model: lerobot/pi05_base
tags:
  - robotics
  - lerobot
  - so101
  - so-101
  - vision-language-action
  - imitation-learning
  - pi0
  - pi05
  - flow-matching
datasets:
  - Project-IRA/TPSoSe2026_Dataset_Lego_LeRobot_SO101
---

# Pi0.5 SO-101 — V4 Lego (single task)

Pi0.5 fine-tuned on the **Lego colour-sorting task only**. This was the first Pi0.5 run and it demonstrated the model was worth scaling to all four tasks. For multi-task use, prefer [V7 Full V2](https://huggingface.co/Project-IRA/TPSoSe2026_Pi05_LeRobot_SO101_Finetuning_V7_Full_V2).

## Demo

<video controls src="https://huggingface.co/Project-IRA/TPSoSe2026_Pi05_LeRobot_SO101_Finetuning_V4_Lego/resolve/main/assets/pi05_lego_demo.mp4"></video>

Pi0.5 sorting lego bricks onto colour-matched plates on the physical SO-101.



Part of **[Project-IRA](https://huggingface.co/Project-IRA)** — Interactive Robotic Arm.
Code: https://github.com/Project-IRA/interactive-robotic-arm

| | |
|---|---|
| Base model | [`lerobot/pi05_base`](https://huggingface.co/lerobot/pi05_base) |
| Robot | SO-101 follower (6-DOF) |
| Training data | [`Project-IRA/TPSoSe2026_Dataset_Lego_LeRobot_SO101`](https://huggingface.co/datasets/Project-IRA/TPSoSe2026_Dataset_Lego_LeRobot_SO101) |
| Recommended checkpoint | **030000** |
| Inputs | `desk_view` + `wrist_left` camera images, 6-dim joint state, English instruction |
| Outputs | 6-dim continuous action chunks |

## Quality

Works **really well** on Lego colour sorting. Recommended checkpoint **030000**.

This model only knows the Lego task — it has never seen desk cleanup, dice, or ball handover data, and will not do them.

## Training

**Single L40S GPU — no FSDP, no multi-GPU.** Pi0.5 fits on one 48 GB L40S at batch 24
with gradient checkpointing and bfloat16, as long as it is not sharded.

| Setting | Value |
|---|---|
| Base | `lerobot/pi05_base` |
| Dataset | `TPSoSe2026_Dataset_Lego_LeRobot_SO101` (460 episodes, 46 prompts) |
| Steps | **30000, completed** (`--save_freq=5000`) |
| Recommended checkpoint | **030000** (the final one) |
| Batch size | 24 |
| Precision | `--policy.dtype=bfloat16` |
| Gradient checkpointing | on |
| Vision encoder | unfrozen |
| Parameters | 4,143,404,816 — **all trainable** (`num_learnable_params == num_total_params`) |
| Optimizer | AdamW, betas (0.9, 0.95), grad clip 1.0 |
| LR schedule | peak `2.5e-05`, 1000 warmup steps, cosine decay to `2.5e-06` over 30000 steps |
| Normalization | `ACTION: MEAN_STD`, `STATE: MEAN_STD`, `VISUAL: IDENTITY` |
| Camera keys | native `wrist_left` / `desk_view` (**no** rename_map) |
| SLURM | `--gres=gpu:L40S:1 --cpus-per-task=16 --mem=92G --time=96:00:00` |

```bash
lerobot-train \
    --dataset.repo_id=TPSoSe2026_Dataset_Lego_LeRobot_SO101 \
    --policy.type=pi05 \
    --policy.pretrained_path=lerobot/pi05_base \
    --policy.normalization_mapping='{"ACTION": "MEAN_STD", "STATE": "MEAN_STD", "VISUAL": "IDENTITY"}' \
    --policy.compile_model=false \
    --policy.gradient_checkpointing=true \
    --policy.dtype=bfloat16 \
    --policy.freeze_vision_encoder=false \
    --policy.train_expert_only=false \
    --policy.device=cuda \
    --policy.push_to_hub=false \
    --batch_size=24 --steps=30000 --save_freq=5000 \
    --num_workers=16 --tolerance_s=0.01 \
    --output_dir=outputs/train/pi05_lego \
    --job_name=pi05_lego \
    --wandb.enable=false
```

Note the contrast with the multi-task V7 runs, where the best checkpoint is very early
(8k–10k): this single-task run is best at its **final** checkpoint. One task means less
overfitting pressure at equivalent step counts.

Pi0.5 requires `lerobot[pi]` installed from GitHub `main`, not the PyPI release.

## Usage

> [!IMPORTANT]
> **Model files are nested under `outputs_V4/`**, so `from_pretrained("Project-IRA/...")`
> on the repo ID will not work:
>
> ```
> outputs_V4/train/pi05_lego/checkpoints/<step>/pretrained_model/   <- weights
> outputs_V4/train/pi05_lego/checkpoints/<step>/training_state/     <- resume only
> ```
>
> **Checkpoints present:** every 5000 steps from `005000` to `030000`.
> Use **`030000`** (the final one). Repo total ~147 GB.
>
> Fetch just that checkpoint's weights (~13 GB instead of 147 GB):
> ```bash
> hf download Project-IRA/TPSoSe2026_Pi05_LeRobot_SO101_Finetuning_V4_Lego \
>   --include 'outputs_V4/train/pi05_lego/checkpoints/030000/pretrained_model/*' \
>   --local-dir ./pi05_lego
> ```

```python
from lerobot.policies.pi0.modeling_pi0 import PI0Policy   # pi05 shares the PI0 module

policy = PI0Policy.from_pretrained("<local path to the checkpoint's pretrained_model/>")
policy = policy.to("cuda").eval()
```

> **Inference dtype:** checkpoints are saved from a `bfloat16` training run. If you hit
> GPU OOM at inference, confirm the loaded policy is in `bfloat16` and not silently
> upcast to `float32`.

Pi0.5 is ~4B parameters. On-robot inference from the robot PC is impractical; we served
it over the **asynchronous gRPC inference server** shipped in the code repo
(https://github.com/Project-IRA/interactive-robotic-arm) and ran the robot as a thin client.

## Robot setup

| | |
|---|---|
| Robot | SO-101 follower arm (6-DOF), `robot_type: so_follower` |
| Teleoperation | SO-101 leader arm |
| Control frequency | 30 fps |
| State / action space | 6-dim: `shoulder_pan.pos`, `shoulder_lift.pos`, `elbow_flex.pos`, `wrist_flex.pos`, `wrist_roll.pos`, `gripper.pos` |
| Camera `observation.images.desk_view` | 800x600, h264 (recording) |
| Camera `observation.images.wrist_left` | 640x480, h264 (recording) |

> **Inference note:** both cameras are run at **640x480 during inference**, not at their
> recording resolutions, to reduce the payload sent to the inference server.

## Environment notes

All training ran on a SLURM cluster with L40S GPUs. Two environment details were required
and are easy to miss when reproducing:

- **ffmpeg libraries for torchcodec.** A minimal conda env supplies the shared libraries
  that `torchcodec` discovers at runtime:
  `export LD_LIBRARY_PATH=$CONDA_PREFIX/envs/ffmpeg_libs_v8/lib:<venv>/lib/python3.12/site-packages/nvidia/npp/lib:$LD_LIBRARY_PATH`
- **`--tolerance_s=0.01`** on every run, to accommodate timestamp jitter in the recorded
  episodes.

Multi-GPU runs additionally set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.
Datasets and the virtualenv were copied to node-local `/scratch` before training rather
than read from shared storage.

No Weights & Biases logging was enabled for any run (`--wandb.enable=false`), so there are
no public training curves — the `job.*.err` SLURM logs are the record.

## Tasks and prompts

The model is conditioned on English natural-language instructions. Prompt phrasing was
varied roughly every 10 episodes during recording, giving 93 distinct prompts in the
merged dataset. **Use one of the training prompts verbatim** for best results — the full
lists are on the [dataset card](https://huggingface.co/datasets/Project-IRA/TPSoSe2026_Dataset_Lego_LeRobot_SO101).

## Limitations

- **Behaviour cloning.** The policy imitates teleoperated demonstrations and has no notion
  of recovery beyond what was demonstrated. It is susceptible to covariate shift and can
  fail to recover from states outside the demonstration distribution.
- **Recovery data is incidental, not systematic.** Recovery behaviour appears in the data
  only where the operator happened to make and correct a mistake during recording; no
  recovery episodes were scripted deliberately.
- **Single environment.** All data comes from one lab desk with one lighting setup, one
  camera geometry, and one set of physical objects. Expect degradation elsewhere.
- **Prompt sensitivity.** Language conditioning was trained on a fixed set of phrasings
  (listed in the dataset card). Prompts far from those phrasings may behave unpredictably.
- **No formal evaluation.** Quality assessments below are qualitative, from operators
  observing rollouts on the physical arm. There are no success-rate numbers.
- **Not safety-rated.** Supervise all physical execution and keep the workspace clear.

## Upstream licensing & attribution

This model is a derivative work of Apache-2.0 licensed components:

| Component | Upstream | License |
|---|---|---|
| LeRobot framework | https://github.com/huggingface/lerobot | Apache-2.0 |
| `lerobot/pi05_base` (Physical Intelligence, openpi) | https://github.com/Physical-Intelligence/openpi | Apache-2.0 |

Apache-2.0 permits relicensing derivative works. We retain the upstream copyright
notices, license text, and NOTICE files for the incorporated material, as Apache-2.0
Section 4 requires. The upstream components remain under Apache-2.0 — only this
project's own contributions (the fine-tuned weights and training configuration) are
offered under **CC BY-SA 4.0**.

**CC BY-SA 4.0** was chosen because it is share-alike: derivatives must be released under
the same licence, so this work cannot be taken closed-source. The project's *source code*
lives in a separate repository under its own licence — see
https://github.com/Project-IRA/interactive-robotic-arm.

## Citation

```bibtex
@misc{project_ira_2026,
  title        = {Project-IRA: Interactive Robotic Arm},
  author       = {Baten, Cleo and Keppler, Bela and Sapper, Jonas},
  year         = {2026},
  howpublished = {\url{https://huggingface.co/Project-IRA}},
  note         = {Code: \url{https://github.com/Project-IRA/interactive-robotic-arm}}
}
```