| # Stage 3 — UR5 teleop, CURRENT-VIEW variant (GPU 0). | |
| # Identical to stage3-ur5-qwen25vl-actheavy.yaml except BOTH V-JEPA branches | |
| # read the CURRENT view (context frame repeated to fill the 8-frame clip) | |
| # instead of the future frames. The objective keeps its full shape — align z_a | |
| # to the stage-1 task space, decorrelate z_b from the stage-1 domain space — | |
| # but no future frame enters it anywhere. | |
| # | |
| # loss = 1.0*L1(action) | |
| # + 0.02*InfoNCE(z_a, z_task_CURRENT) | |
| # + 0.002*SIGReg([z_domain_CURRENT ; z_b]) | |
| # | |
| # Pairs with act-heavy, which is the same objective computed on future frames. | |
| # The two isolate one question: does the V-JEPA target have to look AHEAD, or | |
| # is shaping the representation against the current frame enough? | |
| data_root: /data3/ngseo/lerobot/ur5_teleop_multitask | |
| inmem_frame_dir: /data3/ngseo/cache/frames_ur5_teleop | |
| vjepa_cache_dir: null # null -> run V-JEPA online every step | |
| cams: cam0,cam1,cam2 | |
| # Stage 1 frozen target encoder (online mode) | |
| vjepa_checkpoint: /data3/ngseo/experiments/disentangle-stage1/hf-Ngseo-stage1/e4.pt | |
| vjepa_arch: vit_large | |
| vjepa_img_size: 256 | |
| vjepa_pooler_depth: 2 | |
| vjepa_num_queries: 8 | |
| vjepa_mlp_proj: true | |
| vjepa_lora_rank: 32 | |
| vjepa_lora_alpha: 32 | |
| vjepa_lora_dropout: 0.05 | |
| aug_future: true # independent aug on future frames (online only) | |
| vjepa_task_input: current # task branch sees the current view, not the future | |
| vjepa_domain_input: current # domain branch too — nothing in the loss looks ahead | |
| # Must match the caches | |
| chunk_size: 30 # 1.0 s at 30 fps | |
| num_future_frames: 8 | |
| future_stride: 4 # 8x4 = 32 frames = 1.07 s lookahead | |
| # Backbone | |
| vlm_model: Qwen/Qwen2.5-VL-3B-Instruct | |
| image_size: 224 | |
| pooling_type: last | |
| lora_rank: 32 | |
| # Qwen2.5-VL splits attention differently per tower: the LLM has q_proj/v_proj, | |
| # the vision tower has a single fused qkv. Listing all three puts LoRA on BOTH | |
| # (LLM 7.37M + vision 5.24M). Drop `qkv` and the vision encoder trains nothing. | |
| lora_target_modules: q_proj,v_proj,qkv | |
| # Heads (match stage1 e4.pt: proj_dim 4096, 8 queries, depth 2) | |
| stage2_proj_dim: 4096 | |
| stage2_pooler_num_queries: 8 | |
| stage2_pooler_depth: 2 | |
| stage2_pooler_num_heads: 16 | |
| stage2_action_head_hidden_dim: 2048 | |
| stage2_action_head_num_blocks: 2 | |
| # Loss | |
| lambda_action: 1.0 | |
| latent_loss_type: none # InfoNCE carries the alignment; no L1/cosine term | |
| latent_infonce_weight: 0.02 | |
| latent_infonce_temperature: 0.1 | |
| latent_infonce_negatives: teacher | |
| dual_head: true | |
| sigreg_weight: 0.002 | |
| # Inputs | |
| use_state: true # state[t] 1 frame in; the 30-step chunk still has to be predicted | |
| no_aug: false | |
| # Training | |
| batch_size: 32 # throughput is flat vs bs; 31 InfoNCE negatives per anchor | |
| lr: 1.0e-4 | |
| steps_per_epoch: 10000 | |
| num_epochs: 6 # 60k steps = 1.92M samples = 15 passes over the data | |
| num_workers: 10 | |
| seed: 42 | |
| log_freq: 50 | |
| lr_scaling_mode: none | |
| output_dir: /data3/ngseo/experiments/stage3_ur5_qwen25vl_taskcurrent | |
| resume_checkpoint: null | |