# 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