| --- |
| license: mit |
| tags: |
| - robotics |
| - reinforcement-learning |
| - mimicgen |
| - visual-encoder |
| --- |
| |
| # IMPALA pixel encoder — MimicGen `threading_d0` |
| |
| The visual encoder from a DICE-RL pixel-Q critic that reaches **90% success** on MimicGen |
| `threading_d0` (needle-through-tripod insertion), trained from 224x224 RGB with DrQ random-shift |
| augmentation. |
|
|
| ## What this is |
|
|
| Two IMPALA conv towers, one per camera, whose outputs are concatenated with 9 proprio dims: |
|
|
| ``` |
| images uint8 [B, 2, 224, 224, 3] + proprio [B, 9] |
| -> per-camera ImpalaTower (channels 16/32/32, residual blocks, 3x stride-2 pools) |
| -> 32 x 28 x 28 feature map, flattened -> Linear -> LayerNorm -> 128 |
| -> concat([tower0, tower1, proprio]) = 265 |
| ``` |
|
|
| 6.6M parameters. Proprio is `eef_pos(3) + eef_quat(4) + gripper_qpos(2)`, StateScaler-scaled. |
|
|
| ## Loading |
|
|
| ```python |
| import torch |
| from online_rl_mimicgen.pixel import PixelEncoder # PIXEL_IMG=224 must be set |
| |
| enc = PixelEncoder(out_dim=128) |
| enc.load_state_dict(torch.load("impala_encoder.pt", map_location="cpu")) |
| enc.eval() |
| |
| feats = enc(images_uint8, proprio_scaled) # -> [B, 265] |
| visual_only = feats[:, :256] # drop the proprio block |
| ``` |
|
|
| `PIXEL_IMG` is read from the environment at import time and fixes the head's input width, so it |
| must be `224` before `online_rl_mimicgen.pixel` is imported. At the library default of 84 the |
| `Linear` shapes will not match these weights. |
|
|
| ## What it encodes — measured, in millimetres |
|
|
| Held-out state read-off (200 episodes, 20% held out **by episode**, 2-layer MLP head), against the |
| same protocol applied to frozen mean-pooled DINOv2 features: |
|
|
| | target | this encoder | frozen DINOv2 (768) | |
| |---|---|---| |
| | `eef_pos` | 33.3 mm | **15.8 mm** | |
| | `obj_pos` (needle) | 30.2 mm | **17.3 mm** | |
| | `goal_pos` (tripod) | **46.4 mm** | 73.7 mm | |
| | offset (needle − tripod) | **47.9 mm** | 97.1 mm | |
| | ‖offset‖ (scalar distance) | **31.7 mm** | 79.8 mm | |
|
|
| The pattern is the useful part: this encoder is **worse at absolute positions and better at every |
| relative quantity**. Task training reallocates precision toward the needle-hole relationship that |
| TD learning needs, at the cost of world-coordinate localisation it does not — with a third of the |
| feature dimensions. |
|
|
| Note it resolves the offset to ~48 mm while solving a task whose tolerance is far tighter, so |
| millimetre-accurate geometry is evidently not required for 90% success. |
|
|
| ## Provenance and caveats |
|
|
| - Run `thrpix2_q_s42`, 300k env steps, DICE-RL model-free pixel Q (`--pixel-obs`, no world model). |
| - Trajectory: 57% at 0k (BC prior) → 86% at 50k → 89% at 100k → **90% at 300k**. |
| - **Trained WITH DrQ augmentation** (`random_shift`, pad=4 → ±1.79% of a 224 frame). This matters: |
| otherwise-identical runs **without** augmentation reach only **63.4%** (n=5) — augmentation is |
| worth ~+28 points here, so these weights should not be read as "what an IMPALA encoder learns" |
| independent of that regulariser. |
| - A sibling seed reached 96%, but its checkpoint predates encoder saving and is unrecoverable; |
| this is the best encoder that actually exists on disk. |
| - Single task, single simulator, two fixed camera poses. No claim of transfer. |
|
|
