chomeed/threading_d0_impala_pixel_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

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.