# Created by yams_any4d agent, 2026-07-13 (from rwm4_pp70_640ft_headlive.py). # JOINT LoRA + HEAD on the 640ft base (Cameron 2026-07-13: "take the current fully finetuned # model and add the action head on top but make the video model lora finetuned"). Same base # (pp70_640ft_merged.pt) + same scene-only V4Head, but instead of FREEZING the video we now # LoRA-finetune it JOINTLY with the head: the head's gradients flow into the LoRA adapters # (detach_features=False, freeze_backbone=False) so the video features ADAPT to be more # predictive for actions — lifting the ceiling the frozen features hit. LoRA-B is zero-init so # step 0 == the frozen 640ft features, then the adapters gently move. Grad clip (added below) # + low constant lr + zero-init + low-sigma band keep it stable (the earlier joint run diverged # only because it lacked grad clip / used the full sigma schedule / a fresh un-finetuned base). import os from hydra.core.config_store import ConfigStore from custom.experiment.template import template_any4d_2b, S3_PRETRAINED_ROOT, S3_OUTPUT_ROOT ################################################################# job = dict( project='a4d2', group='pickplace', name='pp70_640ft_headlora', prepend_datetime=True, local_root='/home/robot-lab/any4d_work/local_root', s3_root='', ) wandb = dict( enabled=True, entity='cameronsmithbusiness', project='yam_any4d', num_validation_logs=2, ) any4d_config = dict( dataloader='unified_anyact', vae='a4d_vae', video_entries=[ dict( rgb0=(0, 16, 'load', 'load'), rgb0_input_mask=(16, 17, 'load', None), rgb0_output_mask=(17, 18, 'load', None), ), dict( rgb1=(0, 16, 'copy:rgb0/load', 'copy:rgb0/load'), rgb1_input_mask=(16, 17, 'copy:rgb0_input_mask/load', None), rgb1_output_mask=(17, 18, 'copy:rgb0_output_mask/load', None), ), dict( rgb2=(0, 16, 'copy:rgb1/load', 'copy:rgb1/load'), rgb2_input_mask=(16, 17, 'copy:rgb1_input_mask/load', None), rgb2_output_mask=(17, 18, 'copy:rgb1_output_mask/load', None), ), dict( rgb3=(0, 16, 'copy:rgb1/load', 'copy:rgb1/load'), rgb3_input_mask=(16, 17, 'copy:rgb1_input_mask/load', None), rgb3_output_mask=(17, 18, 'copy:rgb1_output_mask/load', None), ), ], lowdim_adaln_entries=dict( # 17 frames (state_t=5) -> head t_out=17 (num_lowdim_adaln_tokens); matches 640ft regime. action=(0, 17, 0, 20, 'zero/load'), action_input_mask=(0, 17, 20, 21, 'zero/load'), ), num_views=4, video_concat_mode='view', video_proj_mode='per_view', view_timestep_mode='per_view', block_gate_fix=True, disable_risky_sharding=True, train_cond_aug_sigma_range=(0.001, 0.01), val_cond_aug_sigma=0.001, # FROZEN VIDEO via v4head_freeze_backbone=True. RGB weights stay 1.0 so EDM loss averaging # has active entries (all-zero -> "No active entries for EDM loss averaging" crash); the RGB # path just has no trainable params so it does not update the frozen video. loss_weights=dict( rgb0=1.0, rgb1=1.0, rgb2=1.0, rgb3=1.0, ), harmonize_streams=True, harmonize_frames=True, data_defaults=dict( action_format_version='action_format_v3', # grasp_site (TCP) FK ), data_train_overrides=dict( subsample=None, single_sample='random', # Match the 640ft finetune regime EXACTLY: 640 wide (40x22 tokens), 17 frames (5 latent), # stride 2 (sharp). Running the frozen backbone in-distribution -> best features. resize=[-16, 640], length=33, frame_stride=2, ), data_val_overrides=dict( resize=[-16, 640], length=33, frame_stride=2, subsample=20, single_sample='first', ), use_views='all', shuffle_cams2views=False, task_probs=dict( cross_modal=0.0, dyn_view_synth=0.0, forecast=1.0, pose_est=0.0, inv_dyn=0.0, policy=0.0, world_model=0.0, # UNCONDITIONAL: pure forecast, no action conditioning ), train_directives=dict(), val_directives=dict(), action_multiplier=3.0, v4head_enabled=True, v4head_prep_path='/home/robot-lab/any4d_work/v4head_prep_grasp.npz', v4head_temporal_mode='stack_mlp', v4head_use_wrist=False, # scene-only v4head_detach_features=False, # JOINT: head gradients flow INTO the LoRA adapters so the # video features adapt to be predictive for actions v4head_pred_size=128, v4head_n_z=128, v4head_n_lat_frames=5, # 5 latent frames (state_t=5) — stack_mlp input dim v4head_freeze_backbone=False, # LoRA-FINETUNE the video jointly (not frozen) — LoRA + head both train v4head_input_channel_norm=True, # ONLINE per-channel feat-norm (no extraction pass) # MODERATE sigma band (Cameron 2026-07-13, v2 after joint DIVERGED at step ~196): the ultra-low # band (0.02,0.12) is fine for a FROZEN backbone but ill-conditions a TRAINABLE LoRA — the EDM # loss weight ~1/sigma^2 explodes as sigma->0, so training the video there blows up. Center the # band near sigma_data=1.0 (the EDM training sweet spot): well-conditioned RGB gradients + still # reasonably informative, lower-variance features for the head than the full schedule. v4head_train_sigma_band=(0.2, 1.2), skip_video_visuals=True, # skip RGB video gallery/psnr; keep 2D head panels # JOINT weights reduced to 0.5 (v2 after divergence): gentler head gradient into the LoRA so # it nudges rather than corrupts the shared video features, while RGB (1.0, moderate sigma) # anchors the video. Bump back up if stable and the LoRA isn't lifting the head. v4head_loss_weights=dict(v4h_vol=0.5, v4h_grip=0.5, v4h_rot=0.5), default_text_path='custom/vae/default_text_pp70.pkl', # cached T5 embedding (T5 not loaded) track_metrics=dict( rgb0=['psnr', 'ssim'], rgb1=['psnr', 'ssim'], rgb2=['psnr', 'ssim'], rgb3=['psnr', 'ssim'], ), train_visuals_interval=99, train_visuals_detail=1, val_visuals_detail=1, visuals_quality=8, viz_input_blacklist=[], viz_extra_modes=['anyact1'], viz_mask_border_width=2, val_num_steps=35, ) ################################################################# model = dict( # The FULLY-FINETUNED 640 video model (full ft from cosmos base; no LoRA merge needed). dit_path='/home/robot-lab/any4d_work/ckpts/pp70_640ft_merged.pt', text_encoder_path='', # T5 NOT loaded — cached prompt embedding via default_text_path vae_path=f'{S3_PRETRAINED_ROOT}/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth', fsdp_shard_size=2, run_validation=False, skip_first_validation='delay3', validation_iter=50, max_val_iter=2, max_iter=15000, grad_accum_iter=1, context_parallel_size=1, device_monitor=0, manual_gc_iter=288, manual_gc_warm_up=-1, state_t=5, # 17 raw frames -> (17-1)/4+1 = 5 latent tokenizer_chunk_duration=17, ) checkpoint = dict( save_iter=1000, early_sanity_check=-1, ) optimizer = dict( # CONSTANT lr, lowered to 2e-5 (v2): 5e-5 diverged at step ~196 (ill-conditioned low-sigma # LoRA training). 2e-5 + moderate sigma + reduced head weights + grad clip for stability. lr=2e-5, ) scheduler = dict( # CONSTANT: no warmup, no decay (f_start=f_max=f_min=1.0 -> lr held at optimizer.lr). warm_up_steps=[1], cycle_lengths=[model['max_iter']], f_start=[1.0], f_max=[1.0], f_min=[1.0], ) dataset_train = dict( config=[ 'custom/config/anydata/web/debug/yamyukon_pp70_scene.yaml', 'custom/config/anydata/web/debug/yamyukon_pp70_scene.yaml', ], num_workers=4, batch_size=4, # JOINT (detach=False) backprops through the full 2B DiT -> all activations # retained; batch 8 hit 90/96GB. Batch 4 (~50GB) is safe overnight; the # low-sigma band handles gradient variance so we don't need a big batch. ) dataset_val = dict( config=dict( YAMYUKON='custom/config/anydata/web/debug/yamyukon_pp70_scene.yaml', ), num_workers=0, batch_size=1, ) metrics = None ################################################################# cs = ConfigStore.instance() this_config = template_any4d_2b( job, wandb, any4d_config, model, checkpoint, optimizer, scheduler, metrics, dataset_train, dataset_val) # grad-clip (template_any4d_2b omits it; unclipped grads diverged the earlier joint run). from imaginaire.lazy_config import LazyCall as _L # noqa: E402 from cosmos_predict2.callbacks.grad_clip import GradClipCallback as _GradClip # noqa: E402 this_config['trainer']['callbacks']['grad_clip'] = _L(_GradClip)(max_norm=1.0, log_every_n=10) experiment_name = 'any4d_' + os.path.splitext(os.path.basename(__file__))[0] cs.store( group='experiment', package='_global_', name=experiment_name, node=this_config, )