# Created by yams_any4d agent, 2026-07-13 (from rwm4_pp70_640_headfrozen.py). # LIVE HEAD-ON-FROZEN 640ft (Cameron 2026-07-13: "instead of doing extraction can we just # run the video model at a random timestep without gradient"). The FULLY-FINETUNED 640 video # model (pp70_640ft_merged.pt) runs its normal training forward each step, which samples a # RANDOM diffusion timestep (sigma) -> the scene-only V4Head reads those pre-final-layer # features and regresses the volume. Backbone FROZEN + features DETACHED (no grad into the # DiT), so this is exactly "run the video model at a random timestep without gradient" — no # feature cache / extraction pass. Online per-channel standardization (input_channel_norm) # replaces the offline feat-norm stats the cache path needed. 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_headlive', 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=True, # no grad into the frozen DiT 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=True, # FREEZE video; train ONLY the head on live low-sigma feats v4head_input_channel_norm=True, # ONLINE per-channel feat-norm (no extraction pass) # Override the diffusion timestep to a tight LOW band (the "last generation step") instead of # the full EDM schedule (Cameron 2026-07-13). Full-range random sigma fed the head ~50% # near-pure-noise features -> huge gradient variance + v4h_vol climbing after warmup. Clamping # low = consistently informative features (like the cache's fixed 0.05) with no extraction. v4head_train_sigma_band=(0.02, 0.12), skip_video_visuals=True, # frozen backbone -> skip RGB video gallery/psnr; keep 2D head panels # Head is the ONLY trainable path (RGB has no grad), so weight the head losses 1:1:1 like # the standalone cached trainer (small joint weights would starve the head of gradient). v4head_loss_weights=dict(v4h_vol=1.0, v4h_grip=1.0, v4h_rot=1.0), 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 (Cameron 2026-07-13): the scheduled run's best v4h_vol (8.9) hit at effective # lr ~1.4e-4 mid-warmup, then climbed as lr peaked at 2e-4. Drop the schedule entirely and # hold a steady, moderate rate. Larger batch (16) + low-sigma features keep it stable. lr=1e-4, ) 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=8, # 2x the original batch-4 averaging. Batch 16 pushed reserved VRAM to 93/96GB # (OOM risk overnight); batch 8 (~22GB alloc/rank) + expandable_segments is # safe, and the low-sigma band is the real gradient-variance fix (2026-07-13) ) 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, )