---
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.