fact/data.py

"""Episode-level data access for FACT training and evaluation.

Mirrors the sampling contract of ``the T-Rex single-episode ablation protocol``
(``ChunkBuilder``) and of ``fact.data_utils``: a 20 Hz logical
control grid built from timestamps, 16-step delta-base EEF62 action chunks,
one clean 10 Hz observation frame per autoregressive block, and four 5 Hz
tactile refresh slots per block.  Future states are the 16 EEF62 states that
follow each block anchor (offsets 1..16 at 20 Hz), expressed delta-base
relative to the block reference state and normalized with the action
statistics.

Caches (built by ``scripts/build_cache_dataset.py``):

* ``<cache_dir>/latents/episode_XXXXXX.npy``  [N,48,10,20] float16 Wan VAE
  latents of every frame (memory-mapped);
* ``<cache_dir>/text_cache.npz``  umT5-XXL token embeddings per task
  (``tokens`` [n_tasks,T,4096] float16, ``mask`` [n_tasks,T] bool,
  ``task_index`` [n_tasks]).
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq
import torch
from torch.utils.data import Dataset

from .data_utils import (
    eef62_delta_base,
    nearest_timestamp_indices,
    uniform_target_times,
)
from .force import (
    ACTION_HORIZON,
    FORCE_HISTORY_FRAMES,
    FORCE_OFFSETS,
)
from .runtime_stats import TrexRuntimeStatistics

ACTION_RATE_HZ = 20.0
TACTILE_RATE_HZ = 5.0
VIDEO_RATE_HZ = 10.0
AR_BLOCKS = 4
VIDEO_FRAMES_PER_BLOCK = 8
PHYSICAL_DIM = 62
PADDED_DIM = 64


def _column(table, name: str, dtype=np.float32) -> np.ndarray:
    values = table.column(name).to_numpy(zero_copy_only=False)
    values = np.asarray(values)
    while values.dtype == object:
        values = np.stack([np.stack(row) for row in values])
    return values.astype(dtype)


def _sample(timestamps, anchor, offsets, rate):
    return nearest_timestamp_indices(timestamps, uniform_target_times(anchor, offsets, rate))


def normalize_q99(values: np.ndarray, q01: np.ndarray, q99: np.ndarray) -> np.ndarray:
    scale = q99 - q01
    out = np.clip(2.0 * (values - q01) / np.where(scale == 0, 1.0, scale) - 1.0, -1.0, 1.0)
    return np.where(scale == 0, values, out).astype(np.float32)


def pad64(values: np.ndarray) -> np.ndarray:
    if values.shape[-1] == PADDED_DIM:
        return values.astype(np.float32)
    pad = [(0, 0)] * (values.ndim - 1) + [(0, PADDED_DIM - values.shape[-1])]
    return np.pad(values.astype(np.float32), pad)


def episode_parquet_path(root: Path, episode_index: int) -> Path:
    return root / "data" / f"chunk-{episode_index // 1000:03d}" / f"episode_{episode_index:06d}.parquet"


def latent_cache_path(cache_dir: Path, episode_index: int) -> Path:
    return cache_dir / "latents" / f"episode_{episode_index:06d}.npy"


def load_episode_tasks(root: Path) -> dict[int, tuple[int, str]]:
    """episode_index -> (task_index, task string) from meta/{episodes,tasks}.jsonl."""

    tasks: dict[str, int] = {}
    with (root / "meta" / "tasks.jsonl").open() as handle:
        for line in handle:
            if line.strip():
                rec = json.loads(line)
                tasks[rec["task"]] = int(rec["task_index"])
    out: dict[int, tuple[int, str]] = {}
    with (root / "meta" / "episodes.jsonl").open() as handle:
        for line in handle:
            if not line.strip():
                continue
            rec = json.loads(line)
            task = rec["tasks"][0]
            out[int(rec["episode_index"])] = (tasks.get(task, -1), task)
    return out


class TextCache:
    """Per-task umT5 token embeddings (built offline)."""

    def __init__(self, path: str | Path) -> None:
        data = np.load(path)
        self.tokens = data["tokens"]  # [n, T, 4096] float16
        self.mask = data["mask"].astype(bool)  # [n, T]
        self.index = {int(t): i for i, t in enumerate(data["task_index"])}
        self.max_tokens = self.tokens.shape[1]

    def get(self, task_index: int) -> tuple[np.ndarray, np.ndarray]:
        row = self.index[task_index]
        return self.tokens[row], self.mask[row]


class FactEpisode:
    """One LeRobot-v2 T-Rex episode plus its cached VAE latents."""

    def __init__(
        self,
        dataset_root: str | Path,
        episode_index: int = 0,
        latent_cache: str | Path | None = None,
        *,
        stats: TrexRuntimeStatistics | None = None,
        text_cache: TextCache | None = None,
        task: tuple[int, str] | None = None,
    ) -> None:
        self.root = Path(dataset_root)
        self.episode_index = episode_index
        # Read only the needed columns (the parquet also carries 250-point
        # tracks) and keep plain arrays, so thousands of episodes fit in RAM.
        table = pq.read_table(
            episode_parquet_path(self.root, episode_index),
            columns=["timestamp", "observation.state_eef62", "action.eef62_absolute", "observation.tactile_force"],
        )
        self.stats = stats if stats is not None else TrexRuntimeStatistics.from_dataset(self.root)
        self.timestamps = _column(table, "timestamp", np.float64).reshape(-1)
        self.state = _column(table, "observation.state_eef62")
        self.action_abs = _column(table, "action.eef62_absolute")
        self.force = _column(table, "observation.tactile_force").reshape(-1, 10, 6)
        del table
        self.length = len(self.timestamps)
        if task is None:
            task = load_episode_tasks(self.root).get(episode_index, (-1, ""))
        self.task_index, self.prompt = task
        self.text_cache = text_cache
        # Latents are opened lazily per access (a persistent memmap per episode
        # would exhaust the open-file limit with thousands of episodes).
        self.latent_path: Path | None = None
        self.latents: np.ndarray | None = None
        if latent_cache is not None:
            path = Path(latent_cache)
            if path.suffix == ".npy":
                self.latent_path = path
                frames = np.load(path, mmap_mode="r").shape[0]
            else:
                self.latents = np.load(path)["latents"]
                frames = self.latents.shape[0]
            if frames != self.length:
                raise ValueError(f"latent cache {path} has {frames} frames, parquet has {self.length}")

    @property
    def has_latents(self) -> bool:
        return self.latents is not None or self.latent_path is not None

    def load_latents(self, frame_indices: np.ndarray) -> np.ndarray:
        if self.latents is not None:
            return np.asarray(self.latents[frame_indices]).astype(np.float32)
        assert self.latent_path is not None
        mm = np.load(self.latent_path, mmap_mode="r")  # memmap; released when `mm` goes out of scope
        out = np.array(mm[frame_indices], dtype=np.float32)
        del mm
        return out

    # ------------------------------------------------------------------
    def split_time(self, fraction: float) -> float:
        """Absolute timestamp at ``fraction`` of the episode duration."""

        return float(self.timestamps[0] + fraction * (self.timestamps[-1] - self.timestamps[0]))

    def valid_anchor_times(
        self, blocks: int, *, t_min: float | None = None, t_max_end: float | None = None
    ) -> list[float]:
        """20 Hz anchors whose K-block span fits in [t_min, t_max_end]."""

        future = max(
            blocks * ACTION_HORIZON / ACTION_RATE_HZ,
            blocks * VIDEO_FRAMES_PER_BLOCK / VIDEO_RATE_HZ,
        )
        end = self.timestamps[-1] if t_max_end is None else min(t_max_end, self.timestamps[-1])
        grid = np.arange(self.timestamps[0], self.timestamps[-1] + 1e-9, 1.0 / ACTION_RATE_HZ)
        return [
            float(t)
            for t in grid
            if t + future <= end and (t_min is None or t >= t_min - 1e-9)
        ]

    def _norm_force(self, selection) -> np.ndarray:
        values = self.stats.normalize_force(self.force[selection.indices])
        values[selection.padding_mask] = 0.0
        return values

    def normalize_action(self, delta: np.ndarray) -> np.ndarray:
        return normalize_q99(delta, self.stats.action_q01, self.stats.action_q99)

    def build(self, anchor: float, *, blocks: int) -> dict[str, np.ndarray]:
        ts = self.timestamps
        block_anchors = [anchor + b * ACTION_HORIZON / ACTION_RATE_HZ for b in range(blocks)]
        action_sel = _sample(ts, anchor, range(blocks * ACTION_HORIZON), ACTION_RATE_HZ)
        state_sel = _sample(ts, anchor, range(0, blocks * ACTION_HORIZON, ACTION_HORIZON), ACTION_RATE_HZ)
        reference = self.state[state_sel.indices]  # [K,62]
        absolute = self.action_abs[action_sel.indices].reshape(blocks, ACTION_HORIZON, 62)
        delta = np.stack([eef62_delta_base(reference[b], absolute[b]) for b in range(blocks)])

        # future states: offsets 1..16 after each block anchor, delta-base to the block reference
        fstate_sel = [
            _sample(ts, b, range(1, ACTION_HORIZON + 1), ACTION_RATE_HZ) for b in block_anchors
        ]
        future_abs = np.stack([self.state[s.indices] for s in fstate_sel])  # [K,16,62]
        future_delta = np.stack(
            [eef62_delta_base(reference[b], future_abs[b]) for b in range(blocks)]
        )

        force_sels = [
            [
                _sample(ts, b + off / ACTION_RATE_HZ, range(-(FORCE_HISTORY_FRAMES - 1), 1), TACTILE_RATE_HZ)
                for off in FORCE_OFFSETS
            ]
            for b in block_anchors
        ]
        force_history = np.stack(
            [[self._norm_force(sel) for sel in block_sel] for block_sel in force_sels]
        )  # [K,4,16,10,6]
        force_valid = np.stack(
            [[~sel.padding_mask for sel in block_sel] for block_sel in force_sels]
        )  # [K,4,16]

        obs_indices = np.array(
            [int(_sample(ts, b, range(1), VIDEO_RATE_HZ).indices[0]) for b in block_anchors]
        )
        sample: dict[str, np.ndarray] = {
            "state": pad64(self.stats.normalize_state(reference)),
            "action": pad64(self.normalize_action(delta)),
            "future_state": pad64(self.normalize_action(future_delta)),
            "current_force": force_history[:, :, -1].astype(np.float32),
            "force_history": force_history.astype(np.float32),
            "force_history_valid": force_valid,
            "obs_frame_index": obs_indices,
            "action_delta_raw": delta.astype(np.float32),
            "reference_state": reference.astype(np.float32),
        }
        if self.has_latents:
            sample["obs_latents"] = self.load_latents(obs_indices)
        if self.text_cache is not None:
            tokens, mask = self.text_cache.get(self.task_index)
            sample["text_tokens"] = tokens.astype(np.float32)
            sample["text_mask"] = mask
        return sample


class FactAnchorDataset(Dataset):
    """Every valid causal K-block anchor of one episode."""

    def __init__(self, episode: FactEpisode, blocks: int = AR_BLOCKS,
                 train_fraction: float | None = None) -> None:
        self.episode = episode
        self.blocks = blocks
        t_max_end = None if train_fraction is None else episode.split_time(train_fraction)
        self.anchors = episode.valid_anchor_times(blocks, t_max_end=t_max_end)

    def __len__(self) -> int:
        return len(self.anchors)

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        sample = self.episode.build(self.anchors[index], blocks=self.blocks)
        return {
            key: torch.as_tensor(value)
            for key, value in sample.items()
            if key not in ("action_delta_raw", "reference_state")
        }


class FactMultiEpisodeDataset(Dataset):
    """Every valid causal K-block anchor over many episodes of one dataset."""

    def __init__(
        self,
        dataset_root: str | Path,
        cache_dir: str | Path,
        *,
        episodes: list[int] | None = None,
        blocks: int = AR_BLOCKS,
        use_text: bool = True,
    ) -> None:
        self.root = Path(dataset_root)
        self.cache_dir = Path(cache_dir)
        self.blocks = blocks
        stats = TrexRuntimeStatistics.from_dataset(self.root)
        tasks = load_episode_tasks(self.root)
        text_cache = TextCache(self.cache_dir / "text_cache.npz") if use_text else None
        if episodes is None:
            episodes = sorted(tasks)
        self.episodes: list[FactEpisode] = []
        self.index: list[tuple[int, float]] = []
        missing = []
        for ep in episodes:
            path = latent_cache_path(self.cache_dir, ep)
            if not path.exists():
                missing.append(ep)
                continue
            episode = FactEpisode(
                self.root, ep, path, stats=stats, text_cache=text_cache, task=tasks.get(ep)
            )
            anchors = episode.valid_anchor_times(blocks)
            if not anchors:
                continue
            slot = len(self.episodes)
            self.episodes.append(episode)
            self.index.extend((slot, anchor) for anchor in anchors)
        if missing:
            raise FileNotFoundError(
                f"{len(missing)} episodes have no latent cache under {self.cache_dir}/latents "
                f"(first: {missing[:5]}); run scripts/build_cache_dataset.py"
            )
        if not self.index:
            raise ValueError("no valid anchors")

    def __len__(self) -> int:
        return len(self.index)

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        slot, anchor = self.index[index]
        sample = self.episodes[slot].build(anchor, blocks=self.blocks)
        return {
            key: torch.as_tensor(value)
            for key, value in sample.items()
            if key not in ("action_delta_raw", "reference_state")
        }


def collate(samples: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    return {key: torch.stack([s[key] for s in samples]) for key in samples[0]}


__all__ = [
    "ACTION_RATE_HZ",
    "AR_BLOCKS",
    "FactAnchorDataset",
    "FactEpisode",
    "FactMultiEpisodeDataset",
    "TextCache",
    "collate",
    "episode_parquet_path",
    "latent_cache_path",
    "load_episode_tasks",
    "normalize_q99",
    "pad64",
]