stage3-ur5-actheavy

Stage 3 — UR5 teleop VLA (Qwen2.5-VL-3B), action-heavy variant

Training run complete: 60,000 steps. The act-heavy cell from a study on viewpoint↔task spurious correlation in generalist robot policies.

The setup this was trained for

Each of the 7 tasks in Ngseo/ur5_teleop_multitask is only ever shown from one of the 4 cameras, so viewpoint alone almost determines the task. Object colour is deliberately crossed between the two task families so that "colour ⇒ camera" is not a valid shortcut on its own:

A policy that latches onto viewpoint instead of the language instruction will score well in-distribution and fail when a task is shown from the other camera.

Architecture

Qwen2.5-VL-3B + LoRA ─ forward_full_hidden ─┬─ Head A → z_a ─┐
                                            └─ Head B → z_b ─┴─ concat(8192) → ResNetActionHead → 30×7

L = 1.0·L1(action) + 0.02·InfoNCE(z_a, z_task) + 0.002·SIGReg([z_domain ; z_b])

• Head A is pulled toward the task branch of a frozen Ngseo/stage1 disentangled V-JEPA2 ViT-L, evaluated on future frames (8 frames, stride 4 → 1.07 s ahead).
• Head B is free, and SIGReg pushes it to be statistically independent of that model's domain branch.
• Both heads are AttentiveLatentHead (proj 4096, 8 queries, depth 2, 167.8M each).

Final metrics (epoch_6.pt, step 60,000)

Reference points on the same normalised scale, none of which use vision or language:

5.2× below the copy-state baseline. Note that the per-step copy error grows from 0.019 at k=0 to 0.291 at k=29, so most of the chunk genuinely has to be predicted.

Not evaluated on a robot. Everything above is a training-set loss. The question this study is actually about — what happens when a task is requested from a camera it was never trained on — is not answered by these numbers.

The 7-D joint state is also an input here, and a probe on that state alone recovers which of the 7 tasks is running with 94% accuracy (chance 14%). So a viewpoint-OOD failure could not be attributed to the viewpoint shortcut alone in this configuration.

Contents

Each .pt holds model_state_dict (full VLM incl. LoRA), latent_head_state_dict (Head A), free_latent_head_state_dict (Head B), stage2_action_head_state_dict, optimiser state, and the run config. config.yaml is the exact training config.