Add model card
Browse files
README.md
ADDED
|
@@ -0,0 +1,72 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
tags:
|
| 3 |
+
- robotics
|
| 4 |
+
- vla
|
| 5 |
+
- rl-token
|
| 6 |
+
---
|
| 7 |
+
|
| 8 |
+
# pi05-so101-armnetbench-tool-insert-rlt-v50
|
| 9 |
+
|
| 10 |
+
RL Token (RLT) encoder-decoder trained on the SO101 tool insertion task, on top of the **step_24999** checkpoint from [lorenzouttini/pi05-so101-armnetbench-tool-insert-isambard-v50](https://huggingface.co/lorenzouttini/pi05-so101-armnetbench-tool-insert-isambard-v50).
|
| 11 |
+
|
| 12 |
+
## What is this?
|
| 13 |
+
|
| 14 |
+
This model is a lightweight transformer encoder-decoder which takes inputs from a frozen Pi-05 VLA backbone. The encoder compresses the VLA final-layer prefix embeddings into a single RL token via a learned query. The decoder autoregressively reconstructs the original embeddings from only this token, forcing it to act as an information bottleneck. See [Xu et al. (2026), Precise Manipulation with Efficient Online RL](https://www.pi.website/research/rlt) for the method.
|
| 15 |
+
|
| 16 |
+
## Training
|
| 17 |
+
|
| 18 |
+
- **Config:** `pi05_rlt_armnetbench_tool_insert`
|
| 19 |
+
- **VLA backbone:** `lorenzouttini/pi05-so101-armnetbench-tool-insert-isambard-v50` step_24999 (frozen, `rl_vla_loss_weight=0.0`)
|
| 20 |
+
- **Encoder-decoder:** 2-layer transformer, 8 heads, 8192 MLP dim, 2048 embedding dim
|
| 21 |
+
- **Dataset:** `villekuosmanen/armnetbench_tool_insert`
|
| 22 |
+
- **Batch size:** 32
|
| 23 |
+
- **LR:** 2.5e-5 cosine (1k warmup, 20k decay)
|
| 24 |
+
- **Steps:** 20,000
|
| 25 |
+
- **Runtime:** ~4h14m on 4x GH200 (Isambard)
|
| 26 |
+
|
| 27 |
+
No validation split was used — the dataset is too small for a held-out eval split.
|
| 28 |
+
|
| 29 |
+
## Loss progression (train)
|
| 30 |
+
|
| 31 |
+
| Step | Train Loss |
|
| 32 |
+
|------|-----------|
|
| 33 |
+
| 0 | 10754.3 |
|
| 34 |
+
| 1,000 | 873.7 |
|
| 35 |
+
| 5,000 | 552.0 |
|
| 36 |
+
| 10,000 | 430.0 |
|
| 37 |
+
| 15,000 | 377.6 |
|
| 38 |
+
| 19,900 | 356.4 |
|
| 39 |
+
|
| 40 |
+
## Checkpoints
|
| 41 |
+
|
| 42 |
+
| Step | Recommended | Params SHA256 |
|
| 43 |
+
|------|-------------|---------------|
|
| 44 |
+
| 19999 | ✓ | `d9ddbbbefc07b3700f7df5dd161c14de3291bfcf805f71e39ababb902e1501b2` |
|
| 45 |
+
|
| 46 |
+
### Verifying checkpoint hashes
|
| 47 |
+
|
| 48 |
+
```bash
|
| 49 |
+
cd checkpoints/19999 && find params -type f | sort | xargs sha256sum | sha256sum
|
| 50 |
+
```
|
| 51 |
+
|
| 52 |
+
## Repo layout
|
| 53 |
+
|
| 54 |
+
```
|
| 55 |
+
assets/ # Norm stats, valid indices
|
| 56 |
+
checkpoints/19999/params/ # Step 19999 model weights (recommended)
|
| 57 |
+
```
|
| 58 |
+
|
| 59 |
+
## W&B
|
| 60 |
+
|
| 61 |
+
Training curves: https://wandb.ai/pravsels/pi05_rlt_armnetbench_tool_insert/runs/jbdqrmu0
|
| 62 |
+
|
| 63 |
+
## Usage
|
| 64 |
+
|
| 65 |
+
```python
|
| 66 |
+
import openpi.models.model as _model
|
| 67 |
+
import openpi.training.config as _config
|
| 68 |
+
|
| 69 |
+
config = _config.get_config("pi05_rlt_armnetbench_tool_insert")
|
| 70 |
+
params = _model.restore_params("checkpoints/19999/params", restore_type=np.ndarray)
|
| 71 |
+
model = config.model.load(params)
|
| 72 |
+
```
|