code/actionrope/dataset.py

"""训练数据集：预计算 latent + sidecar 相机偏移 + 文本表查表。

一个样本就是一个 clip（81 帧 ⇒ 21 个 latent cell）：

    input_latents  bf16 [48, 21, 30, 52]   VAE latent（data/latent/pool/<clip_id>.pt 里的 "latent"）
    context        bf16 [21, 512, 4096]    逐 cell 的 UMT5 文本嵌入（查表，不挂 T5）
    context_ids    int64 [21]              文本串的全局 id（相同串 ⇒ 相同 id），供 arope_forward 去重
    offset_px      float32 [21, 2]         sidecar 的逐 cell 相机偏移 (dx, dy)，cell 0 = 0
    offset_tok     float32 [21, 2]         = offset_px / 32（token 单位；baseline 臂的主信号，与 ARoPE 同源）
    delta_tok      float32 [21, 2]         逐 cell 增量 offset_tok[k] − offset_tok[k−1]，cell 0 = 0（速度）
    action_idx     int64 [21]              9 类离散动作标签（prompts.ACTIONS 顺序），已按 MIRROR_ACTION 翻回真实屏幕方向
    clip_id        str

后四项就是 baseline/SPEC.md 的 `action_inputs`（collate 后带 batch 维），train.py 原样打包交给动作臂；
arope / plain 臂只用 offset_px。action_idx 来自 .pt 的 `actions`（数据集按镜像的 M 派生的标签，
"moving right" 的 clip 玩家实际向左走），所以要经 MIRROR_ACTION 翻回真实方向，与 offset_px 的符号一致
（sidecar 已用负 gain 修正）；推理侧 infer.py 用 velocity_to_label(mirror=False) 得到同一约定的标签。

为什么逐 cell 给 21 份 context 而不是一份
----------------------------------------
转场语料里 trigger 从句按 cell 变化（前几个 cell 是 A 场景，烧起来之后是
"A burning away into B"，烧完是 B），模型要学的是"文本一变就烧"，所以每帧 token 只能
cross-attend 自己 cell 的文本。普通 clip 的 21 条剥完动作从句后通常完全相同，
context_ids 让 arope_forward 只对唯一串算一次 text_embedding。

text_mode
---------
* "scene"（ARoPE 臂）：prompt 先过 prompts.strip_action 剥掉动作从句，
  再查 text_table_actionrope.pt（361 键）。场景 dropout ⇒ 21 个 cell 全换成 NULL_PROMPT。
* "scene_action"（plain 对照臂）：原 prompt 直接查 text_table_eybx.pt（1,253 键），
  动作从句留在文本里。场景 dropout ⇒ 换成 .pt 自带的 prompts_bossdrop（只剩动作从句，
  这是数据集为该臂准备的 dropout 串；NULL_PROMPT 不在 eybx 表里）。

valid 剔除
----------
sidecar 里 valid=False（日志缺口 / 采样异常）或 sidecar 根本没有该 clip 的，
在 __init__ 时剔除并打印数量。invalid clip 的 offset_px 是实际算出来的值（不是 0），
所以不能靠"偏移为零"兜底，必须在这里剔。

转场 clip 的来源泄漏
--------------------
转场 clip（burn / hardcut）是拿两个普通 clip 拼的：帧 < n_pre 来自 src_a，帧 ≥ n_pre 逐帧
就是 src_b。划分只按 heldout 标志处理转场，没有按 src 归属划，结果 train_eybx 里有 257 条
转场的 src_a / src_b 在 val_eybx（其中 123 条 src_b ⇒ val clip 的帧 n_pre..80 原样在训练集里）。
exclude_src_in=<val 目录> 时按 data/meta/transition.jsonl 的 src_a / src_b 把这些剔掉。

rt_ratio 闸（可选）
------------------
min_rt_ratio > 0 时按 meta 的 rt_ratio（世界实时率）剔除普通 clip。SPEC 0.15 s 的 valid 规则
碰巧兼做了 rt 过滤（invalid 集合里 31% 是 rt<0.8 的降速 clip，valid 集合里 rt<0.8 为 0），
放宽 --max_sample_dt 时必须补上这道闸，否则「世界本来就慢」的 clip 会进来。转场 clip 的
meta 没有 rt_ratio，不受此闸影响。
"""
from __future__ import annotations

import functools
import glob
import json
import os
import random

import torch
from torch.utils.data import Dataset

from actionrope.prompts import MIRROR_ACTION, NULL_PROMPT, mirror_action, strip_action
from baseline import make_action_inputs

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_SIDECAR = os.path.join(ROOT, "data/latent/arope_sidecar.pt")
DEFAULT_META_DIR = os.path.join(ROOT, "data/meta")
TRANSITION_KINDS = ("burn", "hardcut")
TEXT_TABLES = {
    "scene": os.path.join(ROOT, "data/latent/text_table_actionrope.pt"),
    "scene_action": os.path.join(ROOT, "data/latent/text_table_eybx_mirror.pt"),
    # 旧约定：不翻动作词，直接用数据集的镜像标签（mirror_labels=False 时用）
    "scene_action_raw": os.path.join(ROOT, "data/latent/text_table_eybx.pt"),
}

# SPEC 的形状契约，__getitem__ 里逐项核对：latent 编码脚本若哪天换了分辨率，这里第一时间报错
LATENT_SHAPE = (48, 21, 30, 52)
CELLS = 21
CONTEXT_SHAPE = (512, 4096)


def load_sidecar(sidecar) -> dict:
    """sidecar 可以是路径或已加载的 dict（多个 Dataset 共用一份时不必重复读 32 MB）。"""
    if isinstance(sidecar, dict):
        return sidecar
    return torch.load(sidecar, map_location="cpu", weights_only=False)


def load_text_table(text_table) -> dict[str, torch.Tensor]:
    if isinstance(text_table, dict):
        return text_table
    return torch.load(text_table, map_location="cpu", weights_only=False)


@functools.lru_cache(maxsize=4)
def load_meta_index(meta_dir: str) -> dict[str, dict]:
    """data/meta/*.jsonl 里训练侧用得着的几个字段：src_a / src_b（转场）、rt_ratio（普通）。
    65 MB 的 jsonl 解析约 2 s，train / val 两个 Dataset 共用一份，所以带 lru_cache。"""
    out = {}
    for f in sorted(glob.glob(os.path.join(meta_dir, "*.jsonl"))):
        with open(f, encoding="utf-8") as fh:
            for line in fh:
                r = json.loads(line)
                out[r["clip_id"]] = {k: r[k] for k in ("src_a", "src_b", "rt_ratio", "kind") if k in r}
    return out


def clip_ids_in(split_dir: str) -> set[str]:
    return {os.path.splitext(f)[0] for f in os.listdir(split_dir) if f.endswith(".pt")}


class AropeLatentDataset(Dataset):
    def __init__(
        self,
        split_dir: str,
        sidecar=DEFAULT_SIDECAR,
        text_table=None,
        text_mode: str = "scene",
        scene_dropout: float = 0.1,
        limit: int | None = None,
        seed: int = 0,
        verbose: bool = True,
        exclude_src_in: str | None = None,
        min_rt_ratio: float = 0.0,
        meta_dir: str = DEFAULT_META_DIR,
        mirror_labels: bool = True,
    ):
        if text_mode not in ("scene", "scene_action"):
            raise ValueError(f"text_mode 只能是 scene / scene_action，得到 {text_mode!r}")
        # 数据集的动作词是镜像的（prompts.mirror_action 的注释）：plain 臂默认翻回真实屏幕方向，
        # 用补齐了镜像串的 text_table_eybx_mirror.pt；mirror_labels=False 走旧约定（原表、原串）
        self.mirror_labels = bool(mirror_labels) and text_mode == "scene_action"
        self.split_dir = split_dir
        self.text_mode = text_mode
        self.scene_dropout = float(scene_dropout)
        self.min_rt_ratio = float(min_rt_ratio)
        # 受保护的 clip：转场 clip 的 src_a / src_b 落在这里 ⇒ 剔除（val 的帧不能原样进训练集）
        protected = clip_ids_in(exclude_src_in) if exclude_src_in else set()
        meta = load_meta_index(meta_dir) if (protected or self.min_rt_ratio > 0) else {}

        sc = load_sidecar(sidecar)
        self.sidecar_meta = sc["meta"]
        clips = sc["clips"]
        table_key = "scene_action_raw" if (text_mode == "scene_action" and not self.mirror_labels) else text_mode
        self.text_table = load_text_table(TEXT_TABLES[table_key] if text_table is None else text_table)
        # 文本串的全局 id：按键排序后的下标，跨样本、跨 batch 都稳定，相同串 ⇒ 相同 id
        self.key_to_id = {k: i for i, k in enumerate(sorted(self.text_table))}

        files = sorted(glob.glob(os.path.join(split_dir, "*.pt")))
        if not files:
            raise FileNotFoundError(f"{split_dir} 下没有 .pt")
        n_missing = n_invalid = n_leak = n_low_rt = 0
        self.items: list[tuple[str, str]] = []  # (clip_id, path)
        for p in files:
            cid = os.path.splitext(os.path.basename(p))[0]
            rec = clips.get(cid)
            if rec is None:
                n_missing += 1
                continue
            if not bool(rec["valid"]):
                n_invalid += 1
                continue
            m = meta.get(cid, {})
            if rec["kind"] in TRANSITION_KINDS:
                if protected and (m.get("src_a") in protected or m.get("src_b") in protected):
                    n_leak += 1
                    continue
            elif self.min_rt_ratio > 0 and float(m.get("rt_ratio", 1.0)) < self.min_rt_ratio:
                n_low_rt += 1
                continue
            self.items.append((cid, p))
        self.n_total, self.n_missing, self.n_invalid = len(files), n_missing, n_invalid
        self.n_leak, self.n_low_rt = n_leak, n_low_rt

        if limit is not None and limit < len(self.items):
            # smoke 用：按 seed 打乱后取前 limit 个，避免只取到排序靠前的同一 session
            rng = random.Random(seed)
            rng.shuffle(self.items)
            self.items = sorted(self.items[:limit])
        # 只保留用得着的 sidecar 条目，DataLoader worker fork 时少复制一点
        self.offsets = {cid: clips[cid]["offset_px"].to(torch.float32).clone() for cid, _ in self.items}
        # offset_tok / delta_tok 只依赖 offset_px，在这里一次算好（与 infer.py 共用 baseline.make_action_inputs 的定义）
        self.action_inputs = {cid: make_action_inputs(off) for cid, off in self.offsets.items()}
        self.kinds = {cid: clips[cid]["kind"] for cid, _ in self.items}

        if verbose:
            print(
                f"[AropeLatentDataset] {split_dir}: {self.n_total} 个 .pt，"
                f"sidecar 缺失 {n_missing}，valid=False 剔除 {n_invalid}，"
                + (f"src 在 {os.path.basename(exclude_src_in.rstrip('/'))} 的转场剔除 {n_leak}，" if exclude_src_in else "")
                + (f"rt_ratio<{self.min_rt_ratio} 剔除 {n_low_rt}，" if self.min_rt_ratio > 0 else "")
                + f"可用 {len(self.items)}"
                + (f"（limit={limit}）" if limit is not None else "")
                + f"；text_mode={text_mode}，文本表 {len(self.text_table)} 键，scene_dropout={self.scene_dropout}"
                + (f"，动作词镜像翻回真实方向={self.mirror_labels}" if text_mode == "scene_action" else "")
            )

    def __len__(self):
        return len(self.items)

    def clip_ids(self) -> list[str]:
        return [cid for cid, _ in self.items]

    # ------------------------------------------------------------------
    def _prompt_keys(self, d: dict) -> list[str]:
        """21 个 cell 的文本表键。场景 dropout 在这里决定（每个样本独立抽签）。"""
        drop = self.scene_dropout > 0 and random.random() < self.scene_dropout
        if self.text_mode == "scene":
            if drop:
                return [NULL_PROMPT] * CELLS
            return [strip_action(p) for p in d["prompts"]]
        prompts = d["prompts_bossdrop"] if drop else d["prompts"]
        if self.mirror_labels:
            return [mirror_action(p) for p in prompts]
        return list(prompts)

    def __getitem__(self, idx: int) -> dict:
        cid, path = self.items[idx]
        d = torch.load(path, map_location="cpu", weights_only=False)
        latent = d["latent"]
        if tuple(latent.shape) != LATENT_SHAPE:
            raise ValueError(f"{cid}: latent 形状 {tuple(latent.shape)} ≠ {LATENT_SHAPE}")
        keys = self._prompt_keys(d)
        if len(keys) != CELLS:
            raise ValueError(f"{cid}: prompts 有 {len(keys)} 条，应为 {CELLS}")
        try:
            context = torch.stack([self.text_table[k] for k in keys], 0)
        except KeyError as e:
            raise KeyError(f"{cid}: 文本表（text_mode={self.text_mode}）里没有 {e.args[0]!r}") from None
        context_ids = torch.tensor([self.key_to_id[k] for k in keys], dtype=torch.int64)
        actions = d.get("actions")
        if actions is None or len(actions) != CELLS:
            raise ValueError(f"{cid}: .pt 里没有 {CELLS} 条 actions 标签（得到 {None if actions is None else len(actions)}）")
        ai = self.action_inputs[cid]
        return {
            "input_latents": latent.to(torch.bfloat16).contiguous(),
            "context": context.to(torch.bfloat16),
            "context_ids": context_ids,
            "offset_px": ai["offset_px"],
            "offset_tok": ai["offset_tok"],
            "delta_tok": ai["delta_tok"],
            # 数据集标签是镜像的，翻回真实屏幕方向（与 offset_px 的符号一致）
            "action_idx": torch.tensor([MIRROR_ACTION[int(a)] for a in actions], dtype=torch.int64),
            "clip_id": cid,
            "kind": self.kinds[cid],
        }


def arope_collate(batch: list[dict]) -> dict:
    """张量按第 0 维堆叠，字符串收成 list。batch_size=1 时也统一成带 batch 维的形状。"""
    out = {}
    for k in batch[0]:
        v = batch[0][k]
        if isinstance(v, torch.Tensor):
            out[k] = torch.stack([b[k] for b in batch], 0)
        else:
            out[k] = [b[k] for b in batch]
    return out