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
| file | size | what it is |
|---|---|---|
checkpoint_latest_task.pt |
1.9 GB | model weights -- this is the one to load |
checkpoint_latest_train.pt |
2.5 GB | optimizer and scheduler state, for resuming |
config.yaml |
3 KB | the full training config |
training.log |
11 KB | the training log |
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.
- Downloads last month
- 8