# Created by BVH, Apr 2026. # Variant of rwm5_ddad_1F1S.py using the ddad5_1F2M format (frames, 2x manifests). 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( # this becomes JobConfig project='a4d2', group='example', name='wm5_ddad', prepend_datetime=True, # becomes {now:%m-%d-%H-%M}_{config.job.name}, local_root='', # empty = _DEFAULT_LOCAL_ROOT (/any4d on DGX, /tmp/a4d2 on SM) s3_root=S3_OUTPUT_ROOT, # default from template.py; set '' to disable S3 sync ) wandb = dict( enabled=True, entity='tri', project='a4d2', num_validation_logs=5, ) any4d_config = dict( # this is part of Any4DConfig, which inherits from Predict2Video2WorldModelConfig dataloader='unified_anydrive', vae='a4d_vae', video_entries=[ # first entry = base / pretrained dict( # key: (channel_start, channel_end, in_proj_init, out_proj_init) # NOTE: there will be T*H*W tokens with this information rgb0=(0, 16, 'load', 'load'), # do not change rgb0_input_mask=(16, 17, 'load', None), # do not change rgb0_output_mask=(17, 18, 'load', None), # do not change # cams0=(18, 50, 'zero/load', 'zero/load'), # cams0_input_mask=(50, 51, 'zero/load', None), # cams0_output_mask=(51, 52, 'zero/load', None), ), # later entries = new viewpoints dict( # key: (channel_start, channel_end, in_proj_init, out_proj_init) 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), # cams1=(18, 50, 'copy:cams0/load', 'copy:cams0/load'), # cams1_input_mask=(50, 51, 'copy:cams0_input_mask/load', None), # cams1_output_mask=(51, 52, 'copy:cams0_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), # cams2=(18, 50, 'copy:cams1/load', 'copy:cams1/load'), # cams2_input_mask=(50, 51, 'copy:cams1_input_mask/load', None), # cams2_output_mask=(51, 52, 'copy:cams1_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), # cams3=(18, 50, 'copy:cams2/load', 'copy:cams2/load'), # cams3_input_mask=(50, 51, 'copy:cams2_input_mask/load', None), # cams3_output_mask=(51, 52, 'copy:cams2_output_mask/load', None), ), dict( rgb4=(0, 16, 'copy:rgb1/load', 'copy:rgb1/load'), rgb4_input_mask=(16, 17, 'copy:rgb1_input_mask/load', None), rgb4_output_mask=(17, 18, 'copy:rgb1_output_mask/load', None), # cams4=(18, 50, 'copy:cams3/load', 'copy:cams3/load'), # cams4_input_mask=(50, 51, 'copy:cams3_input_mask/load', None), # cams4_output_mask=(51, 52, 'copy:cams3_output_mask/load', None), ), ], lowdim_adaln_entries=dict( # key: (seq_start, seq_end, channel_start, channel_end, init) # NOTE: each token here exists only once, and they get flattened into one embedding # NOTE(bvh): driving uses 2D trajectory instead of 20D action traj=(0, 41, 0, 2, 'zero/load'), traj_input_mask=(0, 41, 2, 3, 'zero/load'), ), num_views=5, video_concat_mode='view', video_proj_mode='per_view', view_timestep_mode='per_view', block_gate_fix=True, # should be enabled in every new training run train_cond_aug_sigma_range=(0.001, 0.01), val_cond_aug_sigma=0.001, loss_weights=dict( rgb0=1.0, rgb1=1.0, rgb2=1.0, rgb3=1.0, rgb4=1.0, ), harmonize_streams=True, harmonize_frames=True, data_train_overrides=dict( subsample=None, single_sample='random', ), data_val_overrides=dict( frame_stride=1, subsample=20, # NOTE(bvh): for debugging odd validation set sizes single_sample='first', ), # NOTE(bvh): driving: views should be consistent & semantically meaningful (ego=0, left, right) use_views='all', shuffle_cams2views=False, task_probs=dict( cross_modal=0.0, dyn_view_synth=0.0, forecast=1.0, pose_estimation=0.0, inverse_dynamics=0.0, policy=0.0, world_model=0.9, # NOTE(bvh): leave 10% chance unconditional forecasting ), train_directives=dict( ), val_directives=dict( tasks='forecast,world_model', perturb_traj='basile3', # NOTE(bvh): ^ key flag to generate counterfactual scenarios (at inference only, not training) ), traj_multiplier=1.0, track_metrics=dict( rgb0=['psnr', 'ssim', 'lpips'], rgb1=['psnr', 'ssim', 'lpips'], rgb2=['psnr', 'ssim', 'lpips'], rgb3=['psnr', 'ssim', 'lpips'], rgb4=['psnr', 'ssim', 'lpips'], ), train_visuals_interval=99, train_visuals_detail=1, val_visuals_detail=1, visuals_quality=8, viz_input_blacklist=[], viz_extra_modes=['anydrive1'], viz_mask_border_width=2, val_num_steps=35, ) ################################################################# model = dict( # NOTE(bvh): dit_path = warm-start model weights only (surgery runs, optimizer resets, iter=0). # For full resume (model + optimizer + scheduler + iteration), use checkpoint.resume instead # and set dit_path=None to avoid redundant loading. # pretrained / vanilla cosmos: # dit_path=f'{S3_PRETRAINED_ROOT}/nvidia/Cosmos-Predict2-2B-Video2World/model-480p-10fps.pt', # to warm-start from rwm3_drive (vidar): dit_path='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/12-11-01-56_s27e_rwm3_drive40h/model/iter_000049000_012544000.pt', text_encoder_path='', # disable (driving doesn't use text) vae_path=f'{S3_PRETRAINED_ROOT}/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth', fsdp_shard_size=8, ############ run_validation=True, skip_first_validation='delay3', validation_iter=500, max_val_iter=3, max_iter=60000, grad_accum_iter=1, context_parallel_size=1, device_monitor=0, manual_gc_iter=288, manual_gc_warm_up=-1, ############ state_t=11, tokenizer_chunk_duration=41, ) checkpoint = dict( save_iter=1000, early_sanity_check=5, # NOTE(bvh): full resume (model + optimizer + scheduler + iteration). # Set dit_path=None above when using this to avoid redundant loading. # resume='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/12-11-01-56_s27e_rwm3_drive40h/model/iter_000049000_012544000.pt', ) optimizer = dict( lr=1e-5, ) scheduler = dict( warm_up_steps = [500], cycle_lengths = [model['max_iter']], f_start=[0.01], f_max=[1.0], f_min=[0.05], ) dataset_train = dict( config=[ 'custom/config/anydata/web/debug/misc/formats/ddad5_1F2M.yaml', ], num_workers=4, batch_size=4, ) dataset_val = dict( config=dict( DDAD = 'custom/config/anydata/web/debug/misc/formats/ddad5_1F2M.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) experiment_name = 'any4d_' + os.path.splitext(os.path.basename(__file__))[0] cs.store( group='experiment', package='_global_', name=experiment_name, node=this_config, )