training/Policy/make_belief_cache_v2.py

#!/usr/bin/env python3
"""
make_belief_cache_v2.py
═══════════════════════════════════════════════════════════════════════════════
Cache pre-VLM features for ablation matrix M0–M14 (CoT-Pool plan, Phase 0).

Modes
─────
  --cache_mode mean_pool      (legacy, sanity-equivalent to v1)
      output:  beliefs       [N, D]      fp16
  --cache_mode dual_pool      (M1: image vs text mean, separately)
      output:  beliefs_img   [N, D]      fp16
               beliefs_text  [N, D]      fp16
  --cache_mode per_frame      (M3-M5: time-axis preserved, spatial pooled)
      output:  beliefs_frame [N, F, D]   fp16     (F = MAX_FRAMES = 8)
               valid_frames  [N, F]      bool
               beliefs_text  [N, D]      fp16     (auxiliary text pool)
  --cache_mode spatial4x4     (M6-M11: time + 4×4 spatial per frame)
      output:  beliefs_grid  [N, F, 16, D] fp16   (16 = 4×4 spatial pooled)
               valid_frames  [N, F]        bool
               beliefs_text  [N, D]        fp16

All modes additionally save: tta_means [N] fp32, tta_vars [N] fp32,
                             schema_version=2, cache_mode, hidden_dim, n_frames.

Why fp16?
  • Belief vectors come from a bf16/fp16 forward; fp32 storage is wasteful.
  • Halves disk + IO; trainer can promote to fp32 at use-time if needed.

Storage budget (217k samples, D=2048, F=8)
  mean_pool     ≈   1.7 GB
  dual_pool     ≈   3.4 GB
  per_frame     ≈  13.5 GB
  spatial4x4    ≈ 113   GB        (use mmap; do NOT load fully into RAM)

Index invariant (same as v1)
  cache[i] corresponds to manifest sample i in
  data/policy_labels/{split}.json["samples"][i].

Usage
─────
  cd PROJECT_ROOT
  python -m training.Policy.make_belief_cache_v2 \\
      --sft_checkpoint checkpoints/SFT/sft_v2/best \\
      --cache_mode     spatial4x4 \\
      --label_dir      data/policy_labels \\
      --out_dir        data/belief_cache_v2 \\
      --batch_size     4
"""

from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn.functional as F
from torch.amp import autocast
from torch.utils.data import DataLoader
from tqdm import tqdm

import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from training.Policy.policy_model import PolicyModel
from training.Policy.policy_dataset import PolicyDataset, policy_collate_fn, MAX_FRAMES

logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("Policy.make_cache_v2")


SCHEMA_VERSION = 2

# ─────────────────────────────────────────────────────────────────────────────
# Helpers — per-image token slicing
# ─────────────────────────────────────────────────────────────────────────────

def _get_spatial_merge_size(model: PolicyModel) -> int:
    """Read spatial_merge_size from VLM vision config.  Qwen2.5-VL = 2."""
    base = model.sft.get_base_model()
    cfg = getattr(base, "config", None)
    vc = getattr(cfg, "vision_config", None) if cfg is not None else None
    sms = getattr(vc, "spatial_merge_size", None) if vc is not None else None
    if sms is None:
        logger.warning("Could not read vision_config.spatial_merge_size; "
                       "defaulting to 2 (Qwen2.5-VL).")
        sms = 2
    return int(sms)


def _per_image_token_counts(image_grid_thw: torch.Tensor,
                            spatial_merge_size: int) -> List[int]:
    """
    For each image i in this batch, how many LLM-visible visual tokens it emits.
      count_i = t_i * h_i * w_i // (spatial_merge_size**2)
    """
    counts: List[int] = []
    sms2 = spatial_merge_size * spatial_merge_size
    for row in image_grid_thw.tolist():
        t, h, w = row[0], row[1], row[2]
        c = (t * h * w) // sms2
        counts.append(int(c))
    return counts


def _spatial_pool_image(tokens: torch.Tensor,
                        h_post: int,
                        w_post: int,
                        out_hw: int = 4) -> torch.Tensor:
    """
    tokens   : [n_tok, D] flattened post-merger spatial sequence for ONE image
    h_post   : post-merger height = h // spatial_merge_size
    w_post   : post-merger width  = w // spatial_merge_size
    out_hw   : target spatial side (4 → 4×4 = 16 outputs)

    Returns  : [out_hw*out_hw, D]
    """
    n_tok, D = tokens.shape
    assert n_tok == h_post * w_post, \
        f"token count {n_tok} != h_post*w_post={h_post * w_post}"
    # → [1, D, h_post, w_post]
    grid = tokens.transpose(0, 1).reshape(1, D, h_post, w_post)
    pooled = F.adaptive_avg_pool2d(grid.float(), (out_hw, out_hw))   # promote to fp32 for AAP
    # → [out_hw*out_hw, D]
    pooled = pooled.reshape(D, out_hw * out_hw).transpose(0, 1)
    return pooled.to(tokens.dtype)


# ─────────────────────────────────────────────────────────────────────────────
# Per-sample feature extraction
# ─────────────────────────────────────────────────────────────────────────────

def _split_sample_visual_tokens(
    hidden_states_b: torch.Tensor,            # [L, D]   one sample's tokens
    input_ids_b:     torch.Tensor,            # [L]
    attention_mask_b: torch.Tensor,           # [L]
    image_grid_thw_b: torch.Tensor,           # [n_img_in_sample, 3]
    image_token_id:   int,
    spatial_merge_size: int,
) -> Tuple[List[torch.Tensor], List[Tuple[int, int]]]:
    """
    Split a single sample's hidden states into per-image chunks.

    Returns
    -------
    chunks  : list of length n_img,  each [count_i, D]   (image-token hiddens)
    shapes  : list of (h_post, w_post) per image
    """
    # 1. Find positions of image_token_id within VALID region.
    valid = attention_mask_b > 0
    is_img = (input_ids_b == image_token_id) & valid
    img_positions = torch.nonzero(is_img, as_tuple=False).squeeze(-1)
    n_img_tokens = int(img_positions.numel())

    counts = _per_image_token_counts(image_grid_thw_b, spatial_merge_size)
    expected_total = sum(counts)

    if n_img_tokens != expected_total:
        raise RuntimeError(
            f"Visual-token count mismatch: input_ids has {n_img_tokens} "
            f"image-token positions, but image_grid_thw expects {expected_total}. "
            f"image_grid_thw rows: {image_grid_thw_b.tolist()}"
        )

    # 2. Slice hidden_states at those positions (already contiguous per Qwen layout).
    img_hidden = hidden_states_b[img_positions]   # [n_img_tokens, D]

    # 3. Partition into per-image chunks; remember (h_post, w_post).
    chunks: List[torch.Tensor] = []
    shapes: List[Tuple[int, int]] = []
    cursor = 0
    for i, c in enumerate(counts):
        chunks.append(img_hidden[cursor:cursor + c])
        t = int(image_grid_thw_b[i, 0].item())
        h = int(image_grid_thw_b[i, 1].item())
        w = int(image_grid_thw_b[i, 2].item())
        # Qwen2.5-VL still images: t==1, post-merger spatial = (h//sms, w//sms).
        # If t > 1 (rare for our pipeline of single frames), we collapse t into
        # the "n_tok" sequence and re-derive spatial as h_post*w_post*t per image.
        # For our use case t=1 always — assert and proceed.
        if t != 1:
            raise RuntimeError(
                f"Unexpected image_grid_thw t={t} (>1). This pipeline assumes "
                f"per-frame image inputs, not video tensors."
            )
        h_post = h // spatial_merge_size
        w_post = w // spatial_merge_size
        shapes.append((h_post, w_post))
        cursor += c

    return chunks, shapes


def _extract_features_for_batch(
    model:              PolicyModel,
    inputs:             Dict[str, torch.Tensor],
    cache_mode:         str,
    spatial_merge_size: int,
    image_token_id:     int,
    n_frames:           int,
) -> Dict[str, torch.Tensor]:
    """
    Run one VLM forward and return (CPU, fp16 where appropriate) tensors
    for the requested cache_mode. All outputs have leading dim B.

    Returns dict with keys depending on cache_mode (see file header).
    """
    # Move tensors to device
    moved: Dict[str, torch.Tensor] = {}
    for k, v in inputs.items():
        if not isinstance(v, torch.Tensor):
            moved[k] = v
            continue
        if k == "pixel_values":
            moved[k] = v.to(model.device, dtype=model.sft.dtype, non_blocking=True)
        else:
            moved[k] = v.to(model.device, non_blocking=True)

    base = model.sft.get_base_model()
    core = getattr(base, "model", None)

    # Run base text+vision encoder; get last hidden state
    with autocast(device_type="cuda", dtype=model._amp_dtype, enabled=True):
        if core is not None:
            out = core(
                input_ids       = moved["input_ids"],
                attention_mask  = moved.get("attention_mask"),
                pixel_values    = moved.get("pixel_values"),
                image_grid_thw  = moved.get("image_grid_thw"),
                use_cache       = False,
                return_dict     = True,
            )
            hs = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
        else:
            out = base(
                input_ids       = moved["input_ids"],
                attention_mask  = moved.get("attention_mask"),
                pixel_values    = moved.get("pixel_values"),
                image_grid_thw  = moved.get("image_grid_thw"),
                use_cache       = False,
                return_dict     = True,
                output_hidden_states = True,
            )
            hs = out.hidden_states[-1]

        # TTA for downstream compatibility — uses the canonical pooled belief.
        belief_canon = model.sft.belief_aggregator(
            hs,
            moved.get("attention_mask"),
            moved.get("input_ids"),
        )
        # belief_aggregator may produce 2D for dual_pool — but we use the
        # ORIGINAL training strategy here (whatever the SFT ckpt has).  The
        # tta_head was trained against THAT strategy, so feed it the canonical.
        tta_mean, tta_logvar = model.sft.tta_head(belief_canon)

    tta_var = torch.exp(tta_logvar.float().clamp(-20.0, 20.0))
    tta_mean = tta_mean.float()

    B = hs.shape[0]
    D = hs.shape[-1]
    attn = moved.get("attention_mask")
    ids  = moved.get("input_ids")
    igt  = moved.get("image_grid_thw")    # [total_images_in_batch, 3]

    out_dict: Dict[str, torch.Tensor] = {
        "tta_means": tta_mean.detach().cpu(),
        "tta_vars":  tta_var.detach().cpu(),
    }

    # ── mean_pool (legacy) ────────────────────────────────────────────────────
    if cache_mode == "mean_pool":
        if attn is not None:
            m = attn.unsqueeze(-1).to(hs.dtype)
            beliefs = (hs * m).sum(dim=1) / m.sum(dim=1).clamp(min=1e-6)
        else:
            beliefs = hs.mean(dim=1)
        out_dict["beliefs"] = beliefs.detach().to(torch.float16).cpu()
        return out_dict

    # ── dual_pool (image-mean, text-mean) ─────────────────────────────────────
    if cache_mode == "dual_pool":
        is_img = (ids == image_token_id)
        if attn is not None:
            valid = attn > 0
            is_img = is_img & valid
            is_text = (~is_img) & valid
        else:
            is_text = ~is_img

        def _mm(mask_b: torch.Tensor) -> torch.Tensor:
            m = mask_b.unsqueeze(-1).to(hs.dtype)
            s = (hs * m).sum(dim=1)
            denom = m.sum(dim=1).clamp(min=1e-6)
            return s / denom

        b_img = _mm(is_img)
        b_txt = _mm(is_text)
        out_dict["beliefs_img"]  = b_img.detach().to(torch.float16).cpu()
        out_dict["beliefs_text"] = b_txt.detach().to(torch.float16).cpu()
        return out_dict

    # ── per_frame / spatial4x4 — both need per-image splitting ────────────────
    if cache_mode in ("per_frame", "spatial4x4"):
        if igt is None:
            raise RuntimeError(
                f"cache_mode={cache_mode} requires image_grid_thw, but the "
                f"processor did not emit it (no images in batch?)."
            )

        # We need to know which (sample, frame) slot each row of image_grid_thw
        # belongs to. The processor concatenates images in batch order; per
        # sample the count equals number of frames passed in. Recover via the
        # number of distinct image-token RUNS in that sample's input_ids.
        # Simpler & more robust: per sample count = number of PIL images we
        # passed.  But here we no longer have access to that; recover from
        # contiguous groups in input_ids.
        #
        # For Qwen2.5-VL each image's tokens form a contiguous run prefixed
        # and suffixed by special <|vision_start|>/<|vision_end|> tokens. We
        # only need image_token_id runs to count images per sample.

        igt_cursor = 0
        beliefs_frame: Optional[torch.Tensor] = None
        beliefs_grid:  Optional[torch.Tensor] = None
        if cache_mode == "per_frame":
            beliefs_frame = torch.zeros(B, n_frames, D, dtype=torch.float16)
        else:  # spatial4x4
            beliefs_grid = torch.zeros(B, n_frames, 16, D, dtype=torch.float16)
        valid_frames = torch.zeros(B, n_frames, dtype=torch.bool)
        beliefs_text = torch.zeros(B, D, dtype=torch.float16)

        for b in range(B):
            ids_b   = ids[b]
            attn_b  = attn[b] if attn is not None else torch.ones_like(ids_b)
            hs_b    = hs[b]

            # Count contiguous runs of image_token_id (= number of images in this sample)
            valid = attn_b > 0
            is_img_b = (ids_b == image_token_id) & valid
            # diff to find run boundaries
            x = is_img_b.to(torch.int8)
            diff = torch.cat([x.new_zeros(1), x[1:] - x[:-1]])
            n_runs = int((diff == 1).sum().item())

            if n_runs == 0:
                # No images for this sample — leave zeros, valid_frames stays False
                # Still compute text mean.
                m_text = valid.unsqueeze(-1).to(hs_b.dtype)
                t_mean = (hs_b * m_text).sum(dim=0) / m_text.sum(dim=0).clamp(min=1e-6)
                beliefs_text[b] = t_mean.detach().to(torch.float16).cpu()
                continue

            # Slice this sample's image_grid_thw rows
            igt_b = igt[igt_cursor:igt_cursor + n_runs]
            igt_cursor += n_runs

            chunks, shapes = _split_sample_visual_tokens(
                hs_b, ids_b, attn_b, igt_b,
                image_token_id, spatial_merge_size,
            )

            n_imgs_use = min(len(chunks), n_frames)
            for f in range(n_imgs_use):
                tok_f = chunks[f]
                h_post, w_post = shapes[f]
                if cache_mode == "per_frame":
                    pooled = tok_f.float().mean(dim=0).to(torch.float16)
                    beliefs_frame[b, f] = pooled.detach().cpu()
                else:  # spatial4x4
                    grid = _spatial_pool_image(tok_f, h_post, w_post, out_hw=4)
                    beliefs_grid[b, f] = grid.detach().to(torch.float16).cpu()
                valid_frames[b, f] = True

            # text mean (non-image valid tokens)
            is_text_b = (~is_img_b) & valid
            m_text = is_text_b.unsqueeze(-1).to(hs_b.dtype)
            denom = m_text.sum(dim=0).clamp(min=1e-6)
            t_mean = (hs_b * m_text).sum(dim=0) / denom
            beliefs_text[b] = t_mean.detach().to(torch.float16).cpu()

        if cache_mode == "per_frame":
            out_dict["beliefs_frame"] = beliefs_frame
        else:
            out_dict["beliefs_grid"]  = beliefs_grid
        out_dict["valid_frames"] = valid_frames
        out_dict["beliefs_text"] = beliefs_text
        return out_dict

    raise ValueError(f"Unknown cache_mode: {cache_mode}")


# ─────────────────────────────────────────────────────────────────────────────
# Cache builder
# ─────────────────────────────────────────────────────────────────────────────

def _flush_chunk(accumulators: Dict[str, List[torch.Tensor]],
                 chunk_dir: Path, chunk_idx: int) -> int:
    """Concat the in-memory batches and atomically save one chunk file.
    Returns number of samples in the chunk."""
    if not accumulators:
        return 0
    part = {k: torch.cat(v, dim=0) for k, v in accumulators.items()}
    n = next(iter(part.values())).shape[0]
    tmp = chunk_dir / f"chunk_{chunk_idx:05d}.pt.tmp"
    fin = chunk_dir / f"chunk_{chunk_idx:05d}.pt"
    torch.save(part, tmp)
    tmp.rename(fin)
    return int(n)


def _scan_chunks(chunk_dir: Path) -> Tuple[int, int]:
    """Return (n_chunks, n_samples_total) present on disk (sorted)."""
    if not chunk_dir.exists():
        return 0, 0
    files = sorted(chunk_dir.glob("chunk_*.pt"))
    # Drop stray .tmp
    for t in chunk_dir.glob("*.tmp"):
        t.unlink(missing_ok=True)
    n_samples = 0
    for f in files:
        try:
            d = torch.load(f, map_location="cpu", weights_only=True)
            n_samples += int(next(iter(d.values())).shape[0])
        except Exception as e:
            logger.warning(f"  [resume] chunk {f.name} unreadable ({e}); dropping")
            f.unlink(missing_ok=True)
    return len(list(chunk_dir.glob("chunk_*.pt"))), n_samples


def _merge_chunks(chunk_dir: Path) -> Dict[str, torch.Tensor]:
    """Load all chunks in order and concatenate into a single cache dict."""
    files = sorted(chunk_dir.glob("chunk_*.pt"))
    if not files:
        return {}
    acc: Dict[str, List[torch.Tensor]] = {}
    for f in files:
        d = torch.load(f, map_location="cpu", weights_only=True)
        for k, v in d.items():
            acc.setdefault(k, []).append(v)
    return {k: torch.cat(lst, dim=0) for k, lst in acc.items()}


@torch.no_grad()
def build_cache(
    model:              PolicyModel,
    loader:             DataLoader,
    split_name:         str,
    cache_mode:         str,
    spatial_merge_size: int,
    image_token_id:     int,
    n_frames:           int,
    chunk_dir:          Optional[Path] = None,
    chunk_size:         int = 200,
    expected_n:         Optional[int] = None,
) -> Dict[str, torch.Tensor]:
    """
    If chunk_dir is provided, save a chunk every `chunk_size` batches and resume
    by scanning existing chunks. `expected_n` is the total sample count (used to
    sanity-check resume alignment).
    """
    model.eval()
    batch_size = loader.batch_size or 1

    # ── Resume detection ────────────────────────────────────────────────────
    start_batch = 0
    chunk_idx = 0
    if chunk_dir is not None:
        chunk_dir.mkdir(parents=True, exist_ok=True)
        n_chunks, n_done = _scan_chunks(chunk_dir)
        if n_chunks > 0:
            # Each chunk (except possibly the last from a previous partial run)
            # contains `chunk_size * batch_size` samples. We skip exactly that
            # many batches so the DataLoader resumes at the next untouched one.
            start_batch = n_chunks * chunk_size
            chunk_idx = n_chunks
            logger.info(
                f"  [resume] found {n_chunks} chunk(s) with {n_done} samples; "
                f"skipping first {start_batch} batches"
            )
            if expected_n is not None and n_done >= expected_n:
                logger.info(f"  [resume] chunks already cover all {expected_n} "
                            f"samples; merging")
                return _merge_chunks(chunk_dir)

    accumulators: Dict[str, List[torch.Tensor]] = {}
    batches_since_flush = 0
    processed_batches = 0

    pbar = tqdm(loader, desc=f"cache[{cache_mode}]{split_name}", ncols=80, leave=True)
    for bi, batch in enumerate(pbar):
        if bi < start_batch:
            # Still need to let DataLoader workers produce the item (cheap — CPU
            # image load only — and keeps ordering deterministic).
            continue
        inputs = model._build_inputs(batch["images"], batch["metadata"])
        feats = _extract_features_for_batch(
            model, inputs, cache_mode,
            spatial_merge_size, image_token_id, n_frames,
        )
        for k, v in feats.items():
            accumulators.setdefault(k, []).append(v)
        batches_since_flush += 1
        processed_batches += 1

        if chunk_dir is not None and batches_since_flush >= chunk_size:
            n_flush = _flush_chunk(accumulators, chunk_dir, chunk_idx)
            pbar.set_postfix_str(f"chunk={chunk_idx} +{n_flush}")
            accumulators = {}
            batches_since_flush = 0
            chunk_idx += 1

    # Final partial chunk
    if chunk_dir is not None and accumulators:
        n_flush = _flush_chunk(accumulators, chunk_dir, chunk_idx)
        logger.info(f"  [chunk] final partial flushed (+{n_flush})")
        accumulators = {}
        chunk_idx += 1

    # ── Assemble final cache ────────────────────────────────────────────────
    if chunk_dir is not None:
        cache = _merge_chunks(chunk_dir)
    else:
        cache = {k: torch.cat(lst, dim=0) for k, lst in accumulators.items()}

    # NaN/Inf sanity
    for k, t in cache.items():
        if t.dtype.is_floating_point:
            n_nan = int(torch.isnan(t).sum().item())
            n_inf = int(torch.isinf(t).sum().item())
            if n_nan or n_inf:
                logger.warning(
                    f"  {split_name}/{k}: {n_nan} NaN, {n_inf} Inf "
                    f"(out of {t.numel()} elems)"
                )

    n = next(iter(cache.values())).shape[0]
    nbytes = sum(t.element_size() * t.numel() for t in cache.values())
    logger.info(
        f"  {split_name}: cached {n} samples  "
        f"keys={list(cache.keys())}  size={nbytes / 1e9:.2f} GB"
    )
    return cache


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser("make_belief_cache_v2")
    ap.add_argument("--sft_checkpoint", default="checkpoints/SFT/sft_v2/best")
    ap.add_argument("--label_dir",      default="data/policy_labels")
    ap.add_argument("--out_dir",        default="data/belief_cache_v2")
    ap.add_argument("--cache_mode",     required=True,
                    choices=["mean_pool", "dual_pool", "per_frame", "spatial4x4"])
    ap.add_argument("--batch_size",     type=int, default=4,
                    help="Smaller for spatial4x4 (more GPU memory for hidden states)")
    ap.add_argument("--num_workers",    type=int, default=2)
    ap.add_argument("--splits",         nargs="+", default=["train", "val"])
    ap.add_argument("--split",          default=None,
                    help="Shortcut for a single split; overrides --splits when set")
    ap.add_argument("--manifest",       default=None,
                    help="Explicit manifest path; overrides label_dir/{split}.json")
    ap.add_argument("--out",            default=None,
                    help="Explicit output .pt path; overrides out_dir/cache_mode/{split}.pt")
    ap.add_argument("--n_frames",       type=int, default=MAX_FRAMES,
                    help="Number of frames per clip (8, 16, 24, ...)")
    ap.add_argument("--sampling",       default="original",
                    choices=["original", "uniform", "last_biased", "last_2s"],
                    help="Frame-index resampling scheme (cf. plan Stage K)")
    ap.add_argument("--source_filter",  default="all",
                    choices=["all", "nexar", "multisrc", "dada", "dad"],
                    help="Restrict samples to a data source (Stage K multi-source variants)")
    ap.add_argument("--debug",          action="store_true",
                    help="Smoke-test on 16 samples per split")
    ap.add_argument("--debug_samples",  type=int, default=16)
    ap.add_argument("--overwrite",      action="store_true")
    ap.add_argument("--chunk_size",     type=int, default=200,
                    help="Flush a chunk to disk every N batches (resume-safe). "
                         "0 disables chunked save.")
    ap.add_argument("--keep_chunks",    action="store_true",
                    help="Keep {out}.chunks/ dir after successful merge "
                         "(default: delete on success).")
    args = ap.parse_args()

    if args.split is not None:
        args.splits = [args.split]

    odir = Path(args.out_dir) / args.cache_mode
    odir.mkdir(parents=True, exist_ok=True)

    # Monkey-patch module-level MAX_FRAMES so _extract_features_for_batch sees it
    # (per_frame / spatial4x4 preallocate buffers based on this).
    import training.Policy.policy_dataset as pds
    pds.MAX_FRAMES = args.n_frames

    logger.info("Loading SFTModel (frozen) for feature extraction...")
    model = PolicyModel(args.sft_checkpoint, use_bf16=True)
    sms = _get_spatial_merge_size(model)
    img_tok_id = model.sft.belief_aggregator.image_token_id
    if img_tok_id is None:
        img_tok_id = 151655
    logger.info(f"  spatial_merge_size = {sms}")
    logger.info(f"  image_token_id     = {img_tok_id}")
    logger.info(f"  hidden_dim         = {model.hidden_dim}")
    logger.info(f"  cache_mode         = {args.cache_mode}")
    logger.info(f"  n_frames           = {args.n_frames}")
    logger.info(f"  sampling           = {args.sampling}")
    logger.info(f"  source_filter      = {args.source_filter}")

    for split in args.splits:
        if args.manifest is not None:
            label_path = Path(args.manifest)
        else:
            label_path = Path(args.label_dir) / f"{split}.json"
        if not label_path.exists():
            logger.warning(f"  {label_path} not found — skipping {split}")
            continue

        if args.out is not None:
            out_path = Path(args.out)
            out_path.parent.mkdir(parents=True, exist_ok=True)
        else:
            out_path = odir / f"{split}.pt"
        if out_path.exists() and not args.overwrite:
            logger.info(f"  Cache exists: {out_path} — skip (use --overwrite to rebuild)")
            continue

        ds = PolicyDataset(
            manifests     = [label_path],
            split         = split,
            debug         = args.debug,
            debug_samples = args.debug_samples,
            n_frames      = args.n_frames,
            sampling      = args.sampling,
            source_filter = args.source_filter,
        )
        if len(ds) == 0:
            logger.warning(f"  {split}: dataset empty after filtering — skipping")
            continue
        loader = DataLoader(
            ds,
            batch_size  = args.batch_size,
            shuffle     = False,
            num_workers = args.num_workers,
            collate_fn  = policy_collate_fn,
            pin_memory  = True,
        )

        chunk_dir = None
        if args.chunk_size > 0:
            chunk_dir = out_path.parent / (out_path.stem + ".chunks")
        cache = build_cache(
            model, loader, split,
            args.cache_mode, sms, img_tok_id, args.n_frames,
            chunk_dir=chunk_dir,
            chunk_size=args.chunk_size,
            expected_n=len(ds),
        )

        # Preserve sample IDs / labels in meta for downstream alignment
        ids = [s.get("video_id") for s in ds.samples]
        labels = [int(s.get("action_label", -1)) for s in ds.samples]

        meta = {
            "schema_version": SCHEMA_VERSION,
            "cache_mode":     args.cache_mode,
            "hidden_dim":     model.hidden_dim,
            "n_frames":       args.n_frames,
            "sampling":       args.sampling,
            "source_filter":  args.source_filter,
            "n_samples":      int(next(iter(cache.values())).shape[0]),
            "spatial_merge_size": sms,
            "image_token_id": int(img_tok_id),
            "sft_checkpoint": str(args.sft_checkpoint),
            "label_path":     str(label_path),
            "ids":            ids,
            "action_labels":  labels,
        }
        cache_to_save = {k: v for k, v in cache.items() if k != "__meta__"}
        cache_to_save["meta"] = meta

        tmp_path = out_path.with_suffix(out_path.suffix + ".tmp")
        torch.save(cache_to_save, tmp_path)
        tmp_path.rename(out_path)
        logger.info(f"  Saved → {out_path}")
        with open(out_path.with_suffix(".meta.json"), "w") as f:
            meta_slim = {k: v for k, v in meta.items()
                         if k not in ("ids", "action_labels")}
            meta_slim["n_ids"] = len(ids)
            json.dump(meta_slim, f, indent=2)

        if chunk_dir is not None and chunk_dir.exists() and not args.keep_chunks:
            import shutil
            shutil.rmtree(chunk_dir)
            logger.info(f"  Removed chunk dir {chunk_dir}")

    logger.info("\nbelief_cache_v2 complete.")


if __name__ == "__main__":
    main()