"""The 6 base-policy eval tasks (the tasks the OSCAR 6-task LoRA was trained on).

One entry per task: gym id, the STATELESS base-ACT checkpoint + dataset stats,
the OSCAR text prompt (MUST match oscar-public/finetune/captions.py, the caption
the skeleton LoRA was trained with), and the per-task rollout horizon.

`n_exec` / `max_steps` mirror the per-task `<task>_oscar_wm.yaml` configs:
each WM `imagine` produces wm_chunk=25 frames; we keep the first `n_exec` and
re-plan, so the episode is `ceil(max_steps / n_exec)` WM calls long.
"""
from __future__ import annotations

BASE = "/root/pwm_local/data0/base_policies"

# Captions verbatim from /root/pwm_local/oscar-public/finetune/captions.py
_CAP = {
    "Open-Drawer": "A humanoid robot with two dexterous multi-fingered hands opens a drawer on a wooden table in a bright indoor room.",
    "Open-Laptop": "A humanoid robot with two dexterous multi-fingered hands opens a laptop on a wooden table in a bright indoor room.",
    "Pour-Balls":  "A humanoid robot with two dexterous multi-fingered hands pours small balls from a cup into a bowl on a wooden table in a bright indoor room.",
    "Push-Box":    "A humanoid robot with two dexterous multi-fingered hands pushes a box across a wooden table in a bright indoor room.",
    "Stack-Can":   "A humanoid robot with two dexterous multi-fingered hands stacks a soda can on top of another can on a wooden table in a bright indoor room.",
    "Unload-Cans": "A humanoid robot with two dexterous multi-fingered hands unloads soda cans from a holder onto a wooden table in a bright indoor room.",
}

TASKS = {
    "open_drawer": dict(
        gym="Humanoid-Open-Drawer-v0",
        ckpt=f"{BASE}/open_drawer_30demo_stateless/policy_epoch_2000_seed_0.ckpt",
        stats=f"{BASE}/open_drawer_30demo_stateless/dataset_stats.pkl",
        prompt=_CAP["Open-Drawer"], max_steps=300, n_exec=24,
        table_idx=1, room_idx=1, scene_seed=6,
    ),
    "open_laptop": dict(
        gym="Humanoid-Open-Laptop-v0",
        ckpt=f"{BASE}/open_laptop_30demo_stateless/policy.ckpt",
        stats=f"{BASE}/open_laptop_30demo_stateless/dataset_stats.pkl",
        prompt=_CAP["Open-Laptop"], max_steps=300, n_exec=12,
        table_idx=1, room_idx=1, scene_seed=8,
    ),
    "pour_balls": dict(
        gym="Humanoid-Pour-Balls-v0",
        ckpt=f"{BASE}/pour_balls_50demo_stateless/policy.ckpt",
        stats=f"{BASE}/pour_balls_50demo_stateless/dataset_stats.pkl",
        prompt=_CAP["Pour-Balls"], max_steps=400, n_exec=12,
        table_idx=1, room_idx=1, scene_seed=3,
    ),
    "push_box": dict(
        gym="Humanoid-Push-Box-v0",
        ckpt=f"{BASE}/push_box_30demo_stateless/policy.ckpt",
        stats=f"{BASE}/push_box_30demo_stateless/dataset_stats.pkl",
        prompt=_CAP["Push-Box"], max_steps=300, n_exec=24,
        table_idx=1, room_idx=1, scene_seed=5,
    ),
    "stack_can": dict(
        gym="Humanoid-Stack-Can-v0",
        ckpt=f"{BASE}/stack_can_50demo_stateless/policy.ckpt",
        stats=f"{BASE}/stack_can_50demo_stateless/dataset_stats.pkl",
        prompt=_CAP["Stack-Can"], max_steps=300, n_exec=12,
        table_idx=1, room_idx=1, scene_seed=9,
    ),
    "unload_cans": dict(
        gym="Humanoid-Unload-Cans-v0",
        ckpt=f"{BASE}/unload_cans_50demo_stateless/policy.ckpt",
        stats=f"{BASE}/unload_cans_50demo_stateless/dataset_stats.pkl",
        prompt=_CAP["Unload-Cans"], max_steps=300, n_exec=24,
        table_idx=1, room_idx=1, scene_seed=10,
    ),
}

# init-pose replay set (all 6 tasks share this pickle)
INIT_POSES_PATH = "/root/pwm_local/PWM/EgoVLA_Release/init_poses_fixed_set_100traj.pkl"