# Created by yams_any4d agent, 2026-07-12 (from rwm4_pp70_scene_lora.py).
# JOINT V4HEAD (Cameron pivot 2026-07-12): continue LoRA-finetuning the SCENE-ONLY
# video model — warm-started from the video-only trained weights (pp70_scene_merged.pt) —
# while training the STACKED-IN-TIME action head (temporal_mode='stack_mlp'): all Tl
# generated-frame features are stacked per patch, run through a res-MLP, upsampled, and
# regressed to ONE volume covering all 41 timesteps (uses every timestep's features, not
# just the last). Scene-only (use_wrist=False). Pure forecast (world_model=0.0) so the
# head cannot shortcut through AdaLN action conditioning. Head trains full-rank (unfrozen
# post-peft); fresh LoRA-B starts at zero so the merged scene model IS the init.

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_v4head_640',
    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(
            # copy from rgb1 (present in the 2-view warm-start), NOT rgb2 (synthesized view
            # absent from ckpt_src -> copy:rgb2 KeyErrors in a4d_surgery).
            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(
        # 9 frames (state_t=3) instead of 41 — sets head t_out=9 (num_lowdim_adaln_tokens).
        action=(0, 9, 0, 20, 'zero/load'),
        action_input_mask=(0, 9, 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,
    data_defaults=dict(
        # v3 = grasp_site (TCP) FK, not flange — keypoints land on the gripper toolpoint
        # (matches DINO model + inference; fixes the ~13.8cm flange offset). 2026-07-12.
        action_format_version='action_format_v3',
    ),
    data_train_overrides=dict(
        subsample=None,
        single_sample='random',
        # 640-RES + 9 frames (Cameron 2026-07-12): resize to 640 wide (-> 40x22 tokens/frame,
        # 2x finer than 320's 20x11) and load a 121-frame window at stride 15 -> 9 frames
        # (3 latent), still spanning ~12.1s. Higher spatial res to break the v4h_vol plateau;
        # fewer frames keeps token budget ~= the 320x176x41 run (~2640 vs 2420).
        resize=[-16, 640],
        length=121,
        frame_stride=15,
    ),
    data_val_overrides=dict(
        resize=[-16, 640],
        length=121,
        frame_stride=15,
        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',   # stacked-in-time single-volume head (Cameron 2026-07-12)
    v4head_use_wrist=False,             # scene-only
    v4head_detach_features=True,        # head on FROZEN video backbone — prevents divergence
                                        # ("plug the video generator in as the backbone")
    v4head_pred_size=128,               # hi-res heatmap (Cameron: ~128 res)
    v4head_n_z=128,
    v4head_n_lat_frames=3,              # 3 latent frames (state_t=3) — stack_mlp input dim
                                        # (was 11 for the 41-frame run); MUST match state_t
    # Inherited joint aux weights (small so RGB stays dominant); bump if we want the LoRA
    # to adapt features more aggressively for the head.
    v4head_loss_weights=dict(v4h_vol=0.02, v4h_grip=0.05, v4h_rot=0.05),
    # T5 not loaded (text_encoder_path=''); use the cached embedding for our single fixed
    # prompt ("pick up the cup and place it on the saucer") -> frees ~22GB for higher res.
    default_text_path='custom/vae/default_text_pp70.pkl',
    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(
    # Warm-start from the VIDEO-ONLY scene-trained weights (Cameron 2026-07-12).
    dit_path='/home/robot-lab/any4d_work/ckpts/pp70_scene_merged.pt',
    text_encoder_path='',  # T5 NOT loaded — cached prompt embedding via default_text_path (~22GB freed)
    vae_path=f'{S3_PRETRAINED_ROOT}/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth',
    fsdp_shard_size=2,  # 2x RTX PRO 6000 96GB
    run_validation=False,  # val windows truncate to non-1+4k on our short episodes
    skip_first_validation='delay3',
    validation_iter=50,
    max_val_iter=2,
    max_iter=6000,
    grad_accum_iter=1,
    context_parallel_size=1,
    device_monitor=0,
    manual_gc_iter=288,
    manual_gc_warm_up=-1,
    state_t=3,               # 9 raw frames -> (9-1)/4+1 = 3 latent
    tokenizer_chunk_duration=9,
)

checkpoint = dict(
    save_iter=500,
    early_sanity_check=-1,
)

optimizer = dict(
    # 3e-5 (was 1e-4): warm-starting from the already-optimized pp70_scene_merged and
    # re-finetuning at 1e-4 OVERSHOOTS after the warmup -> video LoRA degrades, RGB loss
    # climbs 6->10 (Cameron caught this 2026-07-12). Lower LR keeps the re-timing stable.
    lr=3e-5,  # LoRA
)

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/yamyukon_pp70_scene.yaml',
        'custom/config/anydata/web/debug/yamyukon_pp70_scene.yaml',
    ],
    num_workers=2,
    batch_size=2,
)

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)

# NOTE(yams_any4d 2026-07-12): template_any4d_2b (unlike the wan templates) does NOT add a
# grad-clip callback, so joint LoRA+V4Head training had unclipped grads and DIVERGED after
# the 500-step warmup (loss 8 -> 1e23). Add max_norm=1.0 clipping (matches the standalone
# head trainer that was stable).
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,
)