autogaze-videomae-stage0-robomme

AutoGaze Stage 0 reconstruction model, fine-tuned on RoboMME

The Stage 0 reward model of the AutoGaze pipeline -- a ViT-MAE-large video reconstruction model that takes a sparse, multi-scale set of gazed patches and reconstructs the frame -- fine-tuned on RoboMME front-camera clips. In the AutoGaze setup this model is the reward: the gaze policy is trained against its reconstruction loss.

Files

Two things that will bite you

Every key carries a module. prefix. The trainer saves from the DDP-wrapped module. Strip it or nothing matches -- the first attempt at loading this reported 962 unexpected and 530 missing keys, i.e. the weights were silently absent and the model was plain vit-mae-large:

sd = torch.load("checkpoint_latest_task.pt", map_location="cpu")
sd = sd.get("model", sd.get("state_dict", sd))
sd = {k.removeprefix("module."): v for k, v in sd.items()}

It was trained with l1+dinov2_reg+siglip2 (weights 1+0.3+0.3), not with a foreground-weighted loss. mae.dinov2_reg.* and mae.siglip2.* are frozen teacher weights baked into the file; they come back as "unexpected" if you load under a different loss, which is harmless. A foreground-only l1_seg loss was later applied to this checkpoint as a scoring function -- for greedy gaze label generation -- but no Stage 0 run has been trained under it.

Setup and result

• Data: RoboMME front camera, 16-frame clips, random window and random stride (max 8), 1,408 train / 192 val episodes. Execution frames only -- the pickle export drops the demonstration prefix.
• Geometry: scales 32+64+112+224 with 16px patches, 265 vision tokens per frame, block-causal over frames.
• Gaze budget during training: exponential over ratio 0.02-0.15 (lambda 10), per-frame allocation Dirichlet.
• Augmentation on; learning rate 5e-5.
• Final validation reconstruction loss 0.5249 at an average gazing ratio of 0.0625.

Trained at commit 2a3f53d of the fork; config.yaml holds every hydra override.

Licence

Weights derive from facebook/vit-mae-large and from training on the RoboMME dataset; those licences govern use. Nothing here is a redistribution of RoboMME data itself.