| --- |
| license: other |
| tags: |
| - robotics |
| - rlt |
| - rl-token |
| - molmoact2 |
| - vla |
| --- |
| |
| # RLT Stage-1 RL Token Encoder (MolmoAct2 / YAM stack-cube) |
|
|
| Backup of the **RL Token (RLT) Stage-1 encoder** for the frozen MolmoAct2-BimanualYAM |
| stack-cube fine-tune. Faithful PyTorch port of openpi's `pi0_rl.py` (Xu et al. 2025): |
| a learned `<rl>` query compresses the VLA's `(M=690, 2560)` prefix hidden states into a |
| single **`z_rl`** token; a causal AR decoder reconstructs the prefix (per-token squared-L2, |
| stop-grad targets, α=0 / frozen VLA). `z_rl` is the state for the downstream SAC actor-critic. |
| |
| ## Chosen encoder |
| **`checkpoints/rl_token_encoder_ctxdrop09_best.pt`** (load `["ema"]`). Trained with the |
| openpi/paper knobs (AdamW 5e-5, 1k warmup, grad-clip 1.0, EMA 0.999, 10k steps) **plus |
| `context_dropout=0.9`** — zeroing 90% of the decoder's teacher-forced context, which fixes |
| the AR-leak that otherwise leaves `z_rl` diffuse (the bare α=0 reconstruction lets the decoder |
| ignore the token). |
| |
| ## Validation |
| | | baseline (α=0) | **dropout-0.9 (chosen)** | |
| |---|---|---| |
| | PCA top-10 var | 15% | **28%** | |
| | temporal smoothness (↓) | 0.72 | **0.69** | |
| | **success-vs-failure** LogReg CV acc | — | **99.2%** (silhouette 0.13) | |
|
|
| `z_rl` cleanly separates success (44 teleop demos) from failure (7 baseline rollouts, SR≈0) |
| in t-SNE — see `outputs/gate_success_fail.png`. Caveat: success/failure are from different |
| sessions, so part of the 99% is domain shift, not pure task semantics — strong upper bound. |
|
|
| ## Data |
| Trained on 9,668 `(690,2560)` prefix shards from the 44 `atharva-pantheon/yam-stack-cube` |
| demos (~1.3 h teleop @ 10 Hz). Matches the RL Token paper's "small per-task demo set" (1–10 h). |
|
|
| ## Files |
| - `code/` — `rl_token_encoder.py` (model), `train_encoder.py`, `collect_prefix.py` (demo→prefix |
| collector), `collect_fail_replay.py` (karma-rollout→prefix), `tsne_gate.py`, `gate_success_fail.py`. |
| - `checkpoints/` — `ctxdrop09_best/final` (chosen), `nodrop_best/final` (baseline), `ctxdrop05_best`. |
| - `plots/` — `tsne_final.png` (phase structure), `gate_success_fail.png` (success/fail), others. |
|
|
| ## Use (Phase-4 actor-critic) |
| ```python |
| import torch |
| from rl_token_encoder import RLTokenAutoencoder, RLTokenConfig |
| ae = RLTokenAutoencoder(RLTokenConfig(dim=2560)) |
| ae.load_state_dict(torch.load("rl_token_encoder_ctxdrop09_best.pt", map_location="cpu")["ema"]) |
| ae.eval() |
| z_rl = ae.encode(prefix, mask) # (b, M, 2560) -> (b, 2560); SAC state x = (z_rl, proprio) |
| ``` |
|
|
| **Gotcha:** validate `z_rl` via `tsne_gate.py` / `gate_success_fail.py`, NOT a first-token |
| ablation — the first prefix token is a constant special id (151645), making that test vacuous. |
|
|
