# Training config for LingBot-VA on the real-world Franka "fruit_pick" dataset
# (data/fruit_pick_lerobot_v21 -- 100 episodes, 20149 frames, task "Lift the
# lid, put it aside, and pick the black plum.").
#
# Source: downloaded 2026-08-09 from https://huggingface.co/datasets/aabyaneh/fruit_pick
# as LeRobot codebase_version v3.0 (one shared video file per camera, not
# per-episode) -- LingBot-VA's own pipeline (extract_latents.py,
# LatentLeRobotDataset) has no v3.0 awareness (hardcodes revision="v2.1",
# opens per-episode mp4 paths directly), so this repo carries its own genuine
# v2.1 conversion at data/fruit_pick_lerobot_v21, produced by
# scripts/convert_fruit_pick_v3_to_v21.py (decodes both source videos
# sequentially, re-encodes per-episode via the official
# LeRobotDataset.create()/add_frame()/save_episode() writer API -- verified
# byte-identical state/action values and matching frame counts against the
# source for spot-checked episodes) + scripts/backfill_action_config.py (the
# official writer doesn't produce the action_config field extract_latents.py
# needs; same gap hit for place_cube_in_bowl, same fix).
#
# observation.state is already this repo's standard 8-dim eef_pos(3) +
# axisangle(3) + gripper(2) convention (fastwam_train's
# fix_fruit_pick_state_schema.py rewrote it from HF's native 15-dim state
# using the SAME quat2axisangle formula lift2lerobot's own converter uses --
# so this is consistent with lift_new/place_cube_bowl's state convention,
# not a new one). action is 7-dim delta_xyz + delta_rot(always 0) + grasp,
# same schema/mask as every other dataset here.
#
# Gripper convention DIFFERS from place_cube_bowl: fruit_pick's raw action
# column is natively {-1, +1} (verified: np.unique gives exactly [-1., 1.]),
# NOT place_cube_bowl's remapped {0.0=closed, 1.0=open}. This dataset was
# never passed through the (1-g)/2 HDF5-sourced remap (it came pre-converted
# from HF, not from raw HDF5), so the value is whatever the original
# collection tool logged -- which of {-1, +1} means open vs. closed has NOT
# been verified against a physical signal (no gripper_width channel was kept
# after the 15->8 dim state rewrite to cross-check against, unlike
# place_cube_bowl's verification). This does NOT affect training (norm_stat
# below is computed FROM THIS DATASET'S OWN actual range, not reused from
# place_cube_bowl, so the quantile normalization is internally consistent
# regardless of which polarity is "open") -- but get the real mapping
# confirmed before deploying any policy trained from this config.
#
# q01/q99 below are computed from ALL 100 episodes' actual action data
# (data/fruit_pick_lerobot_v21/data/chunk-000/*.parquet, verified byte-copy
# of the source's action column) -- there is no held-out split to exclude
# from the normalizer here (see below).
#
# NO validation split (user decision 2026-08-09, matching the FastWAM
# fruit_pick jobs: use the entire 100-episode dataset for training). Neither
# `num_train_episodes` nor `val_split` is set below -- wan_va/train.py's
# Trainer.__init__ getattr-defaults both to 0/falsy, landing on the
# `else: train_dataset = full_dataset; val_dataset = None` branch (verified
# by reading train.py directly) rather than risking an empty val Subset from
# num_train_episodes=100 (=total episodes).
import os
from copy import deepcopy

from easydict import EasyDict

from wan_va.configs.va_libero_cfg import va_libero_cfg

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

WAN_INIT_PATH = os.environ.get("WAM_WAN_INIT", os.path.join(REPO_ROOT, "ckpts/lingbot_va_wan_init"))
LINGBOT_BASE_PATH = os.environ.get("WAM_LINGBOT_BASE", os.path.join(REPO_ROOT, "ckpts/lingbot-va-base"))
DATASET_PATH = os.environ.get(
    "WAM_FRUIT_PICK_DATASET",
    os.path.join(REPO_ROOT, "data/fruit_pick_lerobot_v21"),
)


def _make(init_path: str) -> EasyDict:
    cfg = EasyDict(__name__="Config: VA fruit_pick")
    cfg.update(deepcopy(va_libero_cfg))
    cfg.wan22_pretrained_model_name_or_path = init_path

    cfg.dataset_path = DATASET_PATH
    cfg.empty_emb_path = os.path.join(cfg.dataset_path, "empty_emb.pt")

    cfg.height = 256
    cfg.width = 256
    cfg.obs_cam_keys = ["observation.images.image", "observation.images.wrist_image"]

    cfg.action_dim = 30
    cfg.used_action_channel_ids = list(range(0, 7))
    inverse_ids = [len(cfg.used_action_channel_ids)] * cfg.action_dim
    for i, j in enumerate(cfg.used_action_channel_ids):
        inverse_ids[j] = i
    cfg.inverse_used_action_channel_ids = inverse_ids

    cfg.action_norm_method = "quantiles"
    # Computed 2026-08-09 from ALL 100 episodes of data/fruit_pick_lerobot_v21
    # (np.quantile(actions, 0.01/0.99, axis=0)) -- fruit_pick's own actual
    # range, NOT place_cube_bowl's numbers (different task, different scale).
    cfg.norm_stat = {
        "q01": [
            -0.0054, -0.01794857, -0.01162286, 0.0, 0.0, 0.0, -1.0,
        ] + [0.0] * 23,
        "q99": [
            0.0087696, 0.01789714, 0.0120096, 0.0, 0.0, 0.0, 1.0,
        ] + [0.0] * 23,
    }

    cfg.action_channel_names = ["dx", "dy", "dz", "drx", "dry", "drz", "gripper"]
    cfg.action_channel_groups = {"xyz": [0, 1, 2], "rotation": [3, 4, 5], "gripper": [6]}
    cfg.gripper_channel_index = 6  # raw units are {-1.0, 1.0}; polarity meaning unverified, see header

    # Same phase-2 fine-tune recipe as lift_new_base/place_cube_bowl_base (arXiv:2601.21998
    # "Post-Training Details" for adapting a robot-co-pretrained init to a real-world
    # platform): this is the same Franka platform, just a different task, so the same
    # recipe applies rather than the scratch-init 8000-step floor.
    cfg.enable_wandb = True
    cfg.load_worker = 8
    cfg.save_interval = 200
    cfg.val_interval = cfg.save_interval
    cfg.gc_interval = 50
    cfg.cfg_prob = 0.1
    cfg.learning_rate = 1e-5
    cfg.beta1, cfg.beta2 = 0.9, 0.95
    cfg.weight_decay = 1e-1
    cfg.warmup_steps = 10
    cfg.batch_size = 1
    # Global batch is the recipe-critical quantity (maintainers' guidance is 32-64 episodes;
    # place_cube_bowl/lift_new jobs ran at 80) -- derive grad_accum from the launch's GPU
    # count instead of hardcoding it, same as va_place_cube_bowl_cfg.py.
    _global_batch = int(os.environ.get("WAM_GLOBAL_BATCH", 32))
    _world_size = int(os.environ.get("WORLD_SIZE", 1))
    if _global_batch % (cfg.batch_size * _world_size) != 0:
        raise ValueError(
            f"WAM_GLOBAL_BATCH={_global_batch} is not divisible by "
            f"batch_size({cfg.batch_size}) x WORLD_SIZE({_world_size})"
        )
    cfg.gradient_accumulation_steps = _global_batch // (cfg.batch_size * _world_size)
    cfg.num_steps = int(
        os.environ.get("WAM_NUM_STEPS", 8000 if init_path == WAN_INIT_PATH else 3000)
    )

    # Optional weight warm-start; see va_place_cube_bowl_cfg.py's header for the exact
    # semantics (WEIGHTS ONLY, not a bit-exact resume).
    _resume_from = os.environ.get("WAM_RESUME_FROM")
    if _resume_from:
        cfg.resume_from = _resume_from

    return cfg


FRUIT_PICK_CONFIGS = {
    "fruit_pick_scratch": _make(WAN_INIT_PATH),
    "fruit_pick_base": _make(LINGBOT_BASE_PATH),
}