# Created by BVH, Jul 2025. import os from hydra.core.config_store import ConfigStore from custom.experiment.template import template_any4d_2b ################################################################# job = dict( # this becomes JobConfig project='a4d2', group='debug', name='ipw2', prepend_datetime=True, # becomes {now:%m-%d-%H-%M}_{config.job.name} ) wandb = dict( enabled=True, entity='tri', project='a4d2', num_validation_logs=5, ) any4d_config = dict( # this is part of Any4DConfig, which inherits from Predict2Video2WorldModelConfig transforms='default', # NOTE(bvh): reference_camera now separate per dataset dataloader='basile', 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 # NOTE(bvh): remove cams0 relative to AnyView ), # 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), # NOTE(bvh): remove cams1 relative to AnyView ), ], lowdim_sattn_entries=dict( # key: (seq_start, seq_end, channel_start, channel_end, in_proj_init, out_proj_init) # NOTE: each token here exists only once action=(0, 41, 0, 20, 'zero/load', 'zero/load'), action_input_mask=(0, 41, 20, 21, 'zero/load', None), action_output_mask=(0, 41, 21, 22, 'zero/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 proprio=(0, 41, 0, 20, 'zero/load'), proprio_input_mask=(0, 41, 20, 21, 'zero/load'), ), num_views=2, 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 loss_weights=dict( rgb0=1.0, rgb1=1.0, action=10.0, # NOTE we are varying this ), harmonize_streams=True, harmonize_frames=True, load_modals=['rgb', 'language', 'action'], # NOTE(bvh): cams is needed in vidar to plot actions task_probs=dict( cross_modal=0.0, dyn_view_synth=0.0, forecast=1.0, pose_estimation=0.0, inverse_dynamics=0.2, policy=0.4, world_model=0.4, # NOTE(bvh): ^ actions are never DC (don't care), always input or output or mixed ), use_views='all', shuffle_cams2views=False, # NOTE(bvh): important to have views become consistent & semantically meaningful! normalizer_stats='', # TODO: revisit this relative_actions=False, # TODO: revisit this action_multiplier=3.0, # NOTE(bvh): amplify values to avoid drowning in noise train_directives=dict( ), val_directives=dict( # tasks='forecast,world_model', # apply same random sampling as above for now remove_cond=True, ), data_val_overrides=dict(frame_stride=1), track_metrics=dict( rgb0=['psnr', 'ssim'], rgb1=['psnr', 'ssim'], ), train_visuals_interval=99, train_visuals_detail=1, val_visuals_detail=1, visuals_quality=8, viz_input_blacklist=[], viz_mask_border_width=2, # NOTE(bvh): useful for temporally varying stuff such as forecasting val_num_steps=35, ) ################################################################# S3_PRETRAINED_PREFIX = r's3://tri-ml-sandbox-16011-us-west-2-datasets/cosmos-predict-2/checkpoints' model = dict( # this makes up parts of Predict2Video2WorldModelConfig and Predict2ModelManagerConfig # dit_path=f'{S3_PRETRAINED_PREFIX}/nvidia/Cosmos-Predict2-2B-Video2World/model-480p-10fps.pt', # dit_path=f's3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/09-03-21-11_s14b_wm2_lbm/model/iter_000006000_003072000.pt', # resume from official AnyView HR ckpt (instead of previous WM): # dit_path='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/cvpr/10-31-19-46_officialHR2v/model/iter_000008700_001670400.pt', # resume from latest WM: # dit_path='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/11-26-03-47_s26_rwm2_lbm/model/iter_000047000_012032000.pt', # resume from latest IPW: dit_path='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/a4d2/debug/12-06-22-19_s29c_ipw2_lbm/model/iter_000049000_006272000.pt', text_encoder_path=f'{S3_PRETRAINED_PREFIX}/google-t5/t5-11b', # text_encoder_path='', # disable vae_path=f'{S3_PRETRAINED_PREFIX}/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth', fsdp_shard_size=8, sac_mode='none', # default is predict2_2b_720 but that (and even mm_only) breaks with extra sattn # sac_mode='any4d_21', # TODO ############ run_validation=True, skip_first_validation='delay', # to check VRAM validation_iter=500, # frequency / interval of validation runs max_iter=100000, # total number of training steps # grad_accum_iter=1, grad_accum_iter=2, # grad_accum_iter=4, context_parallel_size=1, device_monitor=0, manual_gc_iter=288, manual_gc_warm_up=-1, # never disable automatic GC to be safe (weird VRAM issue) ############ state_t=11, # for noise level; = latent # frames for now tokenizer_chunk_duration=41, # for VAE; = raw # frames for now ) checkpoint = dict( save_iter=1000, s3_folder='s3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/cosmos-predict2/', early_sanity_check=5, ) optimizer = dict( lr=2e-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/train/lbm40_most_2v.yaml', # only 24 frames # config='custom/config/debug/lbm_ikeal2_40d.yaml', # config='custom/config/train/lbm_ikeal2_40h.yaml', config='custom/config/train/lbm_ikeal2_40l.yaml', num_workers=4, # batch_size=4, batch_size=2, # batch_size=1, ) dataset_val = dict( # config='custom/config/val/lbm40_some_2v.yaml', # only 24 frames # config='custom/config/debug/lbm_ikeal2_40d.yaml', # config='custom/config/val/lbm_ikeal2_40h.yaml', # subsample 128 so up to 16 nodes config='custom/config/val/lbm_ikeal2_40l.yaml', # subsample 128 so up to 16 nodes num_workers=0, # can crash debugger otherwise batch_size=1, ) metrics = None # dict() # disable because fragile ################################################################# cs = ConfigStore.instance() this_config = template_any4d_2b( job, wandb, any4d_config, model, checkpoint, optimizer, scheduler, metrics, dataset_train, dataset_val) # Use the filename (without extension) as the experiment name experiment_name = 'any4d_' + os.path.splitext(os.path.basename(__file__))[0] cs.store( group='experiment', package='_global_', name=experiment_name, node=this_config, )