training/VLA/cot_belief_dataset.py

"""Qwen3-VL chat-template dataset for CoT + per-frame BeliefToken SFT.

Two supervision modes (auto-detected per record):

(1) Per-frame POMDP target  — when belief.actions_per_frame is present:
        Scene: {scene}
        Critical: {critical}
        Threat: {threat}
        <|BELIEF|> <|A_0|> </|BELIEF|>
        <|BELIEF|> <|A_1|> </|BELIEF|>
        ...
        <|BELIEF|> <|A_{T-1}|> </|BELIEF|>

(2) Clip-level (legacy)      — when only belief.action is present:
        Scene: {scene}
        Critical: {critical}
        Threat: {threat}
        <|BELIEF|> <|ACTION|> </|BELIEF|>

At SFT time only the assistant tokens receive gradient (prefix masked with -100).
At belief-extraction time we teacher-force the full assistant string and read
`last_hidden_state` at each `<|BELIEF|>` position — T 2560-D vectors per clip.
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, List

import torch
from torch.utils.data import Dataset

from training.VLA.frame_utils import sample_frames, sample_frames_from_mp4


# ───────────────────── special tokens ─────────────────────
BELIEF_OPEN   = "<|BELIEF|>"
BELIEF_CLOSE  = "</|BELIEF|>"
ACTION_ALERT   = "<|ALERT|>"
ACTION_OBSERVE = "<|OBSERVE|>"
ACTION_SILENT  = "<|SILENT|>"
ACTION_TOKENS  = [ACTION_ALERT, ACTION_OBSERVE, ACTION_SILENT]

ALL_SPECIAL = [BELIEF_OPEN, BELIEF_CLOSE] + ACTION_TOKENS

ACTION_MAP = {
    "ALERT":   ACTION_ALERT,
    "OBSERVE": ACTION_OBSERVE,
    "SILENT":  ACTION_SILENT,
}

ACTION_TO_IDX = {"ALERT": 0, "OBSERVE": 1, "SILENT": 2}


# ───────────────────── prompts ─────────────────────
SYSTEM_PROMPT = (
    "You are a driving-safety assistant. Given N dashcam frames (earliest → latest), "
    "produce a short chain-of-thought analysis, then emit one risk action token "
    "per frame wrapped in <|BELIEF|> ... </|BELIEF|>. "
    "Actions: <|ALERT|> (collision < 0.5s), <|OBSERVE|> (threat 0.5-2.5s), "
    "<|SILENT|> (no threat). Keep prose minimal; the belief blocks are mandatory."
)

USER_PROMPT = "Analyze the frames and emit scene analysis + per-frame belief blocks."


def _parse_per_frame_belief(threat: str) -> Dict[int, str]:
    """Parse 'f0: phrase; f1: phrase; ...' into {frame_idx: phrase}."""
    import re
    out = {}
    if not threat: return out
    parts = re.split(r"f(\d+):\s*", threat)
    # parts looks like ['', '0', 'phrase0;', '1', 'phrase1;', ...]
    for i in range(1, len(parts) - 1, 2):
        try:
            idx = int(parts[i])
            phrase = parts[i + 1].strip().rstrip(";").strip()
            if phrase:
                out[idx] = phrase
        except (ValueError, IndexError):
            continue
    return out


def _state_phrase_prefix(state: str) -> str:
    """Prefix that hints the model what kind of belief to encode per state.

    SILENT  → broad scene context (lane / traffic / weather)
    OBSERVE → suspect agent + predicted trajectory
    ALERT   → hazard itself + distance / urgency
    """
    return {
        "SILENT":  "context:",
        "OBSERVE": "watching:",
        "ALERT":   "hazard:",
    }.get(state, "context:")


def format_assistant_v4(beliefs_per_frame: List[str]) -> str:
    """v4 canonical assistant text: one <|BELIEF|> {scene+danger} </|BELIEF|>
    per frame. No action token inside the span (action is emitted by the
    policy head downstream). This matches tools/make_cache_gt_belief.py.
    """
    return "\n".join(
        f"{BELIEF_OPEN} {b.strip()} {BELIEF_CLOSE}"
        for b in beliefs_per_frame
    )


def format_assistant(cot: Dict[str, Any], actions: List[str],
                      state_conditional: bool = False) -> str:
    """Build the exact assistant string the model must produce.

    `actions` is a list of action *names* (e.g. ["OBSERVE","OBSERVE","ALERT",...]).
    Single-element list degenerates to the legacy clip-level format.

    When `state_conditional=True`, emit per-frame state-specific phrases
    extracted from `cot.threat_analysis` *between* `<|BELIEF|>` and the
    action token (Stage A of VLAlert-X plan §B). The phrase content
    forces the BELIEF hidden state to encode different information per
    state.
    """
    scene    = str(cot.get("scene", "")).strip()
    critical = "; ".join(str(x).strip() for x in cot.get("critical_objects", [])
                         if str(x).strip())
    threat   = str(cot.get("threat_analysis", "")).strip()
    lines = [f"Scene: {scene}",
             f"Critical: {critical}",
             f"Threat: {threat}"]

    if state_conditional:
        per_frame = _parse_per_frame_belief(threat)
        for i, a in enumerate(actions):
            phrase = per_frame.get(i, "").strip()
            # truncate phrase to ~15 words for token budget
            phrase = " ".join(phrase.split()[:15])
            prefix = _state_phrase_prefix(a)
            if phrase:
                lines.append(f"{BELIEF_OPEN} {prefix} {phrase} "
                              f"{ACTION_MAP[a]} {BELIEF_CLOSE}")
            else:
                # fallback to legacy if no per-frame phrase available
                lines.append(f"{BELIEF_OPEN} {ACTION_MAP[a]} {BELIEF_CLOSE}")
    else:
        for a in actions:
            lines.append(f"{BELIEF_OPEN} {ACTION_MAP[a]} {BELIEF_CLOSE}")
    return "\n".join(lines)


def build_chat(frames, assistant_text: str | None):
    user_content = [{"type": "image", "image": img} for img in frames]
    user_content.append({"type": "text", "text": USER_PROMPT})
    messages = [
        {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
        {"role": "user",   "content": user_content},
    ]
    if assistant_text is not None:
        messages.append({"role": "assistant",
                         "content": [{"type": "text", "text": assistant_text}]})
    return messages


def _resolve_actions(belief: Dict[str, Any], n_frames: int) -> List[str]:
    """Prefer per-frame POMDP labels; fall back to clip-level repeated T times."""
    pf = belief.get("actions_per_frame")
    if pf is not None and len(pf) > 0:
        if len(pf) < n_frames:
            pf = pf + [pf[-1]] * (n_frames - len(pf))
        elif len(pf) > n_frames:
            pf = pf[:n_frames]
        return list(pf)
    return [belief["action"]] * n_frames


class CoTBeliefDataset(Dataset):
    """Yields Qwen3-VL chat-template tensors with per-token labels.

    Requires the processor's tokenizer to ALREADY have the 5 special tokens added
    (via `add_special_tokens({"additional_special_tokens": ALL_SPECIAL})`).
    """

    def __init__(
        self,
        jsonl_path: str,
        video_dir: str,
        processor,
        n_frames: int = 8,
        resize_short: int = 336,
        max_len: int = 4096,
        per_frame: bool = True,
        state_conditional: bool = False,
        video_root_override: str | None = None,
    ):
        self.video_dir = Path(video_dir)
        self.processor = processor
        self.n_frames = n_frames
        self.resize_short = resize_short
        self.max_len = max_len
        self.per_frame = per_frame
        self.state_conditional = state_conditional
        self.video_root_override = Path(video_root_override) if video_root_override else None

        self.records: List[Dict[str, Any]] = []
        missing = 0
        with open(jsonl_path) as f:
            for line in f:
                rec = json.loads(line)
                if rec.get("cot") is None or rec.get("belief") is None:
                    missing += 1
                    continue
                self.records.append(rec)
        if missing:
            print(f"[CoTBeliefDataset] skipped {missing} records without cot+belief")

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

    def _resolve_video_path(self, rec: Dict[str, Any]) -> Path:
        if rec.get("video_path"):
            return Path(rec["video_path"])
        clip_id = str(rec["id"]).zfill(5)
        return self.video_dir / f"{clip_id}.mp4"

    def __getitem__(self, idx):
        rec = self.records[idx]
        clip_id = str(rec["id"])
        video_path = self._resolve_video_path(rec)
        frame_idx = rec.get("belief", {}).get("frame_indices")
        frames = sample_frames(video_path, n_frames=self.n_frames,
                               resize_short=self.resize_short,
                               frame_indices=frame_idx)

        if self.per_frame:
            actions = _resolve_actions(rec["belief"], self.n_frames)
        else:
            actions = [rec["belief"]["action"]]
        assistant_text = format_assistant(rec["cot"], actions,
                                            state_conditional=self.state_conditional)

        full_msgs   = build_chat(frames, assistant_text=assistant_text)
        prefix_msgs = build_chat(frames, assistant_text=None)

        proc = self.processor
        full_text   = proc.apply_chat_template(full_msgs,   tokenize=False,
                                               add_generation_prompt=False)
        prefix_text = proc.apply_chat_template(prefix_msgs, tokenize=False,
                                               add_generation_prompt=True)

        full = proc(text=[full_text], images=[frames], return_tensors="pt",
                    padding=False, truncation=True, max_length=self.max_len)
        prefix = proc(text=[prefix_text], images=[frames], return_tensors="pt",
                      padding=False, truncation=True, max_length=self.max_len)

        input_ids = full["input_ids"][0]
        labels = input_ids.clone()
        prefix_len = prefix["input_ids"].shape[1]
        labels[:prefix_len] = -100

        action_idx = [ACTION_TO_IDX[a] for a in actions]

        item = {
            "input_ids":      input_ids,
            "attention_mask": full["attention_mask"][0],
            "labels":         labels,
            "pixel_values":   full["pixel_values"],
            "image_grid_thw": full["image_grid_thw"],
            "label":          int(rec["label"]),
            "actions":        actions,
            "action_idx":     torch.tensor(action_idx, dtype=torch.long),
            "id":             clip_id,
        }
        return item


def collate_fn(batch, pad_token_id: int):
    max_len = max(b["input_ids"].size(0) for b in batch)
    input_ids, attn, labels, pixel_values, grid_thw = [], [], [], [], []
    for b in batch:
        pad_n = max_len - b["input_ids"].size(0)
        input_ids.append(torch.cat([b["input_ids"],
                                    torch.full((pad_n,), pad_token_id, dtype=torch.long)]))
        attn.append(torch.cat([b["attention_mask"],
                               torch.zeros(pad_n, dtype=b["attention_mask"].dtype)]))
        labels.append(torch.cat([b["labels"],
                                 torch.full((pad_n,), -100, dtype=torch.long)]))
        pixel_values.append(b["pixel_values"])
        grid_thw.append(b["image_grid_thw"])
    T = max(len(b["actions"]) for b in batch)
    action_idx = torch.full((len(batch), T), -1, dtype=torch.long)
    for i, b in enumerate(batch):
        action_idx[i, :len(b["actions"])] = b["action_idx"]
    return {
        "input_ids":      torch.stack(input_ids),
        "attention_mask": torch.stack(attn),
        "labels":         torch.stack(labels),
        "pixel_values":   torch.cat(pixel_values, dim=0),
        "image_grid_thw": torch.cat(grid_thw, dim=0),
        "_clip_ids":      [b["id"] for b in batch],
        "_actions":       [b["actions"] for b in batch],
        "_action_idx":    action_idx,
        "_cls_labels":    torch.tensor([b["label"] for b in batch], dtype=torch.long),
    }