aailabkaist's picture
Link W&B training curves
8c61522 verified
|
Raw
History Blame Contribute Delete
4.59 kB
---
license: gemma
library_name: lerobot
pipeline_tag: robotics
tags:
- lerobot
- robotics
- pi0.5
- pi05
- vla
- so101
- manipulation
- full-finetune
base_model:
- lerobot/pi05_base
datasets:
- aailabkaist/so101_recovery_2task
---
# π0.5 · SO-101 pick & place (2-task) — **full fine-tune**, 5,000 steps
A [π0.5](https://www.physicalintelligence.company/blog/pi05) (`pi05`) vision-language-action policy
fine-tuned on an SO-101 arm for a two-task cup manipulation setup: pick a cup up from the table and
place it onto a blue circle.
**The entire model was trained** — PaliGemma VLM (vision encoder + language model) and the action
expert together (`train_expert_only=false`, `freeze_vision_encoder=false`).
- **Base model:** [`lerobot/pi05_base`](https://huggingface.co/lerobot/pi05_base) — the last
pre-relative-actions revision, matching LeRobot `0.4.x`
- **Dataset:** [`aailabkaist/so101_recovery_2task`](https://huggingface.co/datasets/aailabkaist/so101_recovery_2task)
(749 episodes · 151,255 frames · 30 fps · 8 operators)
- **Checkpoint:** step **5,000** of a 10,000-step run (batch 64 → **320,000** samples ≈ **2.1 epochs**)
- **Framework:** [LeRobot](https://github.com/huggingface/lerobot)
- **Robot:** SO-101, 6-DoF, two RGB cameras (`front`, `wrist`, 480×640)
- **Siblings:** `expert 5K` · `expert 10K` · `full 10K` — see the
[two-task collection](https://huggingface.co/collections/nevertmr/so-101-cup-blue-circle-two-task-6a5f8e25fbd8c0bccda23cd7)
## Task and data
Two instructions, both used at training time verbatim:
```
pick up the cup near the blue circle
place the cup on the blue circle
```
![cup start positions](pick_start_positions_anon.png)
Pick-episode cup start positions (394 episodes, per-operator colors, anonymized). The cups form a
ring around the blue circle rather than covering the workspace uniformly.
![cup transport](place_transport_anon.png)
Place-episode transport vectors, start → placed (350 of 355 place episodes — 5 excluded where the
cup is fully occluded by the gripper in the start frame). Endpoints converge tightly on the
blue circle; the release target is nearly constant across the dataset.
## Training
| | |
|---|---|
| Trainable parameters | **3.617 B (100 %)** |
| Frozen | none |
| Steps | 5,000 (run stopped early at 10K; config said 20K) |
| Batch size | 64 |
| Precision | bfloat16, `gradient_checkpointing=true` |
| Optimizer | AdamW · lr 2.5e-5 · wd 0.01 · betas (0.9, 0.95) · grad-clip 1.0 |
| Schedule | cosine decay with warmup — auto-scaled by LeRobot (20K steps < 30K decay): warmup 1,000→666, decay 30,000→20,000; peak 2.5e-5, LR at this checkpoint ≈2.2e-5 (configured floor 2.5e-6 never reached — run stopped at 10K) |
| Chunk | `chunk_size=50`, `n_action_steps=50`, `n_obs_steps=1` |
| Seed | 1000 |
| Hardware | 1× NVIDIA RTX PRO 6000 Blackwell (96 GB) |
| VRAM / speed | 45.4 GB · 4.93 s/step |
| **Loss at this checkpoint** | **0.059** |
Interactive training curves for this run — loss · LR · grad-norm · epochs ([this variant's W&B run](https://wandb.ai/nevertmr-org/pi05_so101/runs/g676fcay) · [project overview](https://wandb.ai/nevertmr-org/pi05_so101)):
![loss curve](loss_curve.png)
![loss curve, log scale](loss_curve_log.png)
Both runs used identical data, batch size and dtype — the only difference is what was trainable, so
the two curves are directly comparable. `expert-only` is essentially flat after warmup (~step 1K):
0.129 → 0.087 over the remaining ~3.8 epochs; full fine-tuning keeps descending. The run was stopped
at 10K rather than the configured 20K on the strength of these curves.
## Usage
```python
from lerobot.policies.pi05.modeling_pi05 import PI05Policy
policy = PI05Policy.from_pretrained("aailabkaist/pi05_recovery_2task_full_5k")
```
Or serve it for async inference:
```bash
python -m lerobot.async_inference.policy_server --host=0.0.0.0 --port=8080 --fps=30
python -m lerobot.async_inference.robot_client \
--server_address=<host>:8080 \
--policy_type=pi05 --pretrained_name_or_path=aailabkaist/pi05_recovery_2task_full_5k \
--task="pick up the cup near the blue circle" \
--actions_per_chunk=50 --chunk_size_threshold=0.2 \
--aggregate_fn_name=weighted_average --fps=30
```
---
*A single-task variant of this setup — one instruction covering the whole pick-and-place motion,
with a plain white cup — lives in the
[SO-101 · cup → blue circle (single-task) collection](https://huggingface.co/collections/nevertmr/so-101-cup-blue-circle-single-task-6a58a5ef4ebe2dc6c028f093).*