# Created by BVH/claude, 2026-06-22.
# 4-VIEW YAMxdof train + val: uses ALL 4 webbed cameras (top_left + top_right + both wrists).
# Variant of rwm3_yamxdof.py: num_views 3->4 (+rgb3 entry), uses yamxdof_ila4 (cameras_aux_sample
# commented out so both wrists load), and world_model task @ 0.5 sampled at train AND val
# (forecast always; val_directives left empty so val mirrors the train task mix).

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='raiden',
    name='yamraiden_lora',
    prepend_datetime=True,
    local_root='/home/robot-lab/any4d_work/local_root',
    s3_root='',  # no S3 sync from russet
)

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(
            # NOTE: copy from rgb1 (view 1, present in the 2-view rwm2 warm-start ckpt), NOT rgb2 --
            # rgb2 is a synthesized view (x_embedder_newviews.1) absent from ckpt_src, so copy:rgb2
            # KeyErrors in a4d_surgery (copies read from the ORIGINAL ckpt, not progressively).
            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(
        action=(0, 41, 0, 20, 'zero/load'),
        action_input_mask=(0, 41, 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,  # keep final_layer / t_embedder un-sharded to avoid multi-view hangs
    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,
    ),
    harmonize_streams=True,
    harmonize_frames=True,
    # v2 = YAM joints -> EE pose via FK (yam_fk.py); partial dict merges with the dataclass defaults.
    data_defaults=dict(
        action_format_version='action_format_v2',
    ),
    data_train_overrides=dict(
        subsample=None,
        single_sample='random',
    ),
    data_val_overrides=dict(
        frame_stride=1,
        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.6,
    ),
    train_directives=dict(),
    val_directives=dict(
        # empty -> val samples tasks from task_probs (forecast always, world_model @ 0.5,
        # same as train); do NOT force forecast+world_model at 100% every val step.
    ),
    action_multiplier=3.0,
    # NOTE(bvh/claude): lpips skipped (slow VGG16 path, see rwm3_droid).
    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(
    # AnyAct v2:
    dit_path='/home/robot-lab/any4d_work/ckpts/s52h_yam4_iter_44000.pt',
    text_encoder_path='disable',  # T5-11B does not fit next to DiT on 24 GB
    vae_path=f'{S3_PRETRAINED_ROOT}/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth',
    fsdp_shard_size=1,  # single 4090
    run_validation=False,  # val windows truncate to non-1+4k on our short episodes
    skip_first_validation='delay3',
    validation_iter=50,  # frequent val for fast debugging feedback
    max_val_iter=2,
    max_iter=2000,  # LoRA sanity run on 2 episodes
    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=500,  # save mid-run; max_iter loop ends at 1999 and never hits a 2000 trigger
    early_sanity_check=-1,
)

optimizer = dict(
    lr=1e-4,  # LoRA
)

scheduler = dict(
    warm_up_steps=[500],
    cycle_lengths=[model['max_iter']],
    f_start=[0.01],
    f_max=[1.0],
    f_min=[0.05],
)

# YAMxdof-only train: 4 entries (one per rank) so all 4 GPUs train YAM.
dataset_train = dict(
    config=[
        'custom/config/anydata/web/debug/yamraiden_ego3.yaml',
    ],
    num_workers=2,
    batch_size=1,
)

# YAMxdof-only val.
dataset_val = dict(
    config=dict(
        YAMRAIDEN='custom/config/anydata/web/debug/yamraiden_ego3.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,
)