pravsels commited on
Commit
bba62c2
·
verified ·
1 Parent(s): c9ef343

Add model card

Browse files
Files changed (1) hide show
  1. README.md +72 -0
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
+ ```