onf-c5/scripts/cache_policy_actions.py

#!/usr/bin/env python3
"""Cache a frozen StableVLA policy's per-frame actions over a demo corpus.

Replays every raw frame of every demo HDF5 through the policy server under the deploy chunk
contract and writes a_pi_raw.npz, whose rows align 1:1 with NodeTable.q_raw.

Stage 2 needs the policy's OWN action, not the demo action: the learned blend weight's gradient is
proportional to (a_track - a_policy), so approximating a_policy by the demo action makes the two
branches coincide on clean windows and alpha never receives gradient in the normal operating regime.

CLI:  python scripts/cache_policy_actions.py <suite> [--dry-run] [--limit-tasks N] [--limit-demos N]
"""

from __future__ import annotations

import argparse
import hashlib
import json
import logging
import os
import shutil
import sys
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Final, Sequence

import h5py
import numpy as np
import yaml

_REPO_ROOT: Final[Path] = Path(__file__).resolve().parents[1]
_HARNESS_DIR: Final[Path] = _REPO_ROOT / "evals" / "libero_plus" / "harness"
for _path in (_REPO_ROOT / "src", _HARNESS_DIR):
    if str(_path) not in sys.path:
        sys.path.insert(0, str(_path))

from onf.config import default_paths  # noqa: E402
from onf.graph.build.from_demos import (  # noqa: E402
    MIN_DEMO_FRAMES,
    DemoLoader,
    GraphBuilder,
    LoadLimits,
    NodeBuilder,
)
from onf.graph.core import schema  # noqa: E402
from onf.graph.core.edges import EdgeSet  # noqa: E402
from onf.graph.core.geometry import graph_hash  # noqa: E402
from onf.graph.core.nodes import NodeTable  # noqa: E402

log = logging.getLogger("cache_policy_actions")

# Artifact names written into the output directory.
ACTIONS_NPZ: Final[str] = "a_pi_raw.npz"
MANIFEST_JSON: Final[str] = "a_pi_raw.manifest.json"
SHARD_DIRNAME: Final[str] = "a_pi_raw_shards"

# The one place configs/suites.yaml is read for the per-suite policy checkpoint.
SUITES_YAML: Final[Path] = _REPO_ROOT / "configs" / "suites.yaml"

# Policy input size, matching evals/libero_plus/harness/eval_libero.py RESIZE_SIZE.
RESIZE_SIZE: Final[list[int]] = [224, 224]

ACTION_DIM: Final[int] = 7
GRIP_COL: Final[int] = 6

# Rows in a stubbed chunk. Only the leading action_chunk_size of them are ever indexed, so this
# just has to be at least as long as whatever chunk the client is configured with.
STUB_CHUNK_ROWS: Final[int] = 64


# ==================================================================================
# Provenance
# ==================================================================================
@dataclass(frozen=True, slots=True)
class CheckpointId:
    """Content fingerprint of a policy checkpoint directory.

    Attributes:
        path: Resolved checkpoint directory.
        stats_sha256: sha256 of dataset_statistics.json, or "" if absent.
        weights_sha256: sha256 over the sorted (name, size) list of weight files.
    """

    path: Path
    stats_sha256: str
    weights_sha256: str

    @classmethod
    def probe(cls, ckpt_dir: Path) -> CheckpointId:
        """Fingerprint a checkpoint directory without reading the weights.

        Args:
            ckpt_dir: Checkpoint directory.

        Returns:
            The identifier.
        """
        stats = ckpt_dir / "dataset_statistics.json"
        stats_sha = (
            hashlib.sha256(stats.read_bytes()).hexdigest() if stats.exists() else ""
        )
        listing = sorted(
            (p.name, p.stat().st_size)
            for p in ckpt_dir.glob("*")
            if p.is_file() and p.suffix in (".safetensors", ".pt", ".bin")
        )
        digest = hashlib.sha256(json.dumps(listing).encode()).hexdigest()
        return cls(path=ckpt_dir, stats_sha256=stats_sha, weights_sha256=digest)

    def to_dict(self) -> dict[str, Any]:
        """JSON-serialisable form.

        Returns:
            Dict with the path flattened to str.
        """
        return {
            "path": os.fspath(self.path),
            "stats_sha256": self.stats_sha256,
            "weights_sha256": self.weights_sha256,
        }


def suite_checkpoint(suite: str) -> Path:
    """Resolve a suite's StableVLA checkpoint from configs/suites.yaml.

    Args:
        suite: Suite short-name.

    Returns:
        The checkpoint directory.

    Raises:
        KeyError: The suite is not declared in configs/suites.yaml.
    """
    with open(SUITES_YAML) as handle:
        suites = (yaml.safe_load(handle) or {}).get("suites") or {}
    if suite not in suites:
        raise KeyError(f"suite {suite!r} not in {SUITES_YAML}; have {sorted(suites)}")
    return _REPO_ROOT / str(suites[suite]["stablevla_ckpt"])


# ==================================================================================
# Demo index -- the ordering NodeTable.q_raw is built in
# ==================================================================================
@dataclass(frozen=True, slots=True)
class DemoSource:
    """Locates one demo inside the corpus, in global demo-id order.

    Attributes:
        demo_id: Global demo id, i.e. the index into NodeTable.raw_ptr.
        hdf5_path: Source file.
        group_key: Demo group name inside data/.
        task_id: Sort position of the source file.
        task_name: Filename-derived task name, as from_demos.DemoLoader reports it.
        instruction: LIBERO language instruction fed to the policy.
        n_frames: Raw frame count T.
    """

    demo_id: int
    hdf5_path: Path
    group_key: str
    task_id: int
    task_name: str
    instruction: str
    n_frames: int

    @property
    def label(self) -> str:
        """Human-readable id used in logs and shard bookkeeping."""
        return f"{self.task_name}/{self.group_key}"


class DemoIndexer(DemoLoader):
    """DemoLoader that also reports WHICH (file, group) each demo came from.

    The base class discovers files, orders demo groups and drops too-short demos; q_raw is the
    concatenation of the surviving demos' joint_states in exactly that order. Subclassing rather
    than re-walking the directory is what keeps this script's row order identical to the graph's.
    """

    def index(self) -> tuple[DemoSource, ...]:
        """Enumerate every demo the corresponding graph build would have kept.

        Returns:
            Demo sources in global demo-id order.
        """
        sources: list[DemoSource] = []
        for task_id, path in enumerate(self.files):
            task_name = self.task_names[task_id]
            with h5py.File(path, "r") as handle:
                data = handle["data"]
                instruction = _instruction(data, task_name)
                for key in self._sorted_demo_keys(data):
                    n_frames = int(data[key]["obs"]["joint_states"].shape[0])
                    if n_frames < MIN_DEMO_FRAMES:
                        continue
                    sources.append(
                        DemoSource(
                            demo_id=len(sources),
                            hdf5_path=path,
                            group_key=key,
                            task_id=task_id,
                            task_name=task_name,
                            instruction=instruction,
                            n_frames=n_frames,
                        )
                    )
        return tuple(sources)


def _instruction(data: h5py.Group, task_name: str) -> str:
    """Read a file's LIBERO language instruction and tie it to the from_demos task name.

    The HDF5 problem_info attribute carries the same string LIBERO's benchmark serves as
    task.language, which is what the harness sends the policy. Cross-checking it against the
    filename-derived task name is what catches a mis-filed or renamed demo file.

    Args:
        data: The data group of an open HDF5 file.
        task_name: Task name reported by DemoLoader for the same file.

    Returns:
        The instruction string.

    Raises:
        ValueError: problem_info is missing, or its instruction is not part of the task name.
    """
    raw = data.attrs.get("problem_info")
    if raw is None:
        raise ValueError(f"{task_name}: HDF5 data group has no problem_info attribute")
    text = raw.decode() if isinstance(raw, bytes) else str(raw)
    instruction = str(json.loads(text)["language_instruction"]).strip()
    if instruction.replace(" ", "_").lower() not in task_name.lower():
        raise ValueError(
            f"{task_name}: problem_info instruction {instruction!r} does not match the filename"
        )
    return instruction


# ==================================================================================
# Observations
# ==================================================================================
@dataclass(frozen=True, slots=True)
class DemoFrames:
    """Every observation of one demo, in the form the harness hands the policy.

    Attributes:
        agentview: [T, H, W, 3] u8, already rotated 180 degrees.
        wrist: [T, H, W, 3] u8, already rotated 180 degrees.
        state: [T, 8] f64 concat(ee_pos, ee_ori, gripper_states).
        joints: [T, D] f32 joint_states, for the q_raw alignment check.
    """

    agentview: np.ndarray
    wrist: np.ndarray
    state: np.ndarray
    joints: np.ndarray

    def __len__(self) -> int:
        return int(self.joints.shape[0])

    @classmethod
    def read(cls, source: DemoSource) -> DemoFrames:
        """Load one demo's observations.

        The 180-degree rotation matches eval_libero.py: the stored RGB is in the same
        OpenGL bottom-up convention the live env returns, so it needs the same fix. ee_ori is
        already robosuite's unnormalised axis-angle, i.e. what _quat2axisangle produces.

        Args:
            source: Demo to read.

        Returns:
            The observations.

        Raises:
            ValueError: The group's frame count disagrees with the index.
        """
        with h5py.File(source.hdf5_path, "r") as handle:
            obs = handle["data"][source.group_key]["obs"]
            agentview = _rotate_180(np.asarray(obs["agentview_rgb"]))
            wrist = _rotate_180(np.asarray(obs["eye_in_hand_rgb"]))
            state = np.concatenate(
                [
                    np.asarray(obs["ee_pos"], dtype=np.float64),
                    np.asarray(obs["ee_ori"], dtype=np.float64),
                    np.asarray(obs["gripper_states"], dtype=np.float64),
                ],
                axis=1,
            )
            joints = np.asarray(obs["joint_states"], dtype=np.float32)
        if joints.shape[0] != source.n_frames:
            raise ValueError(
                f"{source.label}: read {joints.shape[0]} frames, index says {source.n_frames}"
            )
        return cls(agentview=agentview, wrist=wrist, state=state, joints=joints)


def _rotate_180(frames: np.ndarray) -> np.ndarray:
    """Rotate a stack of images 180 degrees, per-frame contiguous.

    Args:
        frames: [T, H, W, 3] u8 stack.

    Returns:
        [T, H, W, 3] u8 rotated stack.
    """
    return np.ascontiguousarray(frames[:, ::-1, ::-1])


# ==================================================================================
# Alignment against the graph's q_raw
# ==================================================================================
@dataclass(frozen=True, slots=True)
class Reference:
    """The node table this cache must align to.

    Attributes:
        table: Node table supplying raw_ptr and q_raw.
        origin: Where it came from -- an on-disk artifact path, or "rebuilt".
        graph_hash: Fingerprint of the built graph, or "" when g_edges.npz is absent.
    """

    table: NodeTable
    origin: str
    graph_hash: str


class AlignmentGuard:
    """Refuses to write a cache whose rows would not line up with q_raw.

    Two independent checks: per-demo frame counts against raw_ptr (structure), and per-frame
    joint configurations against q_raw (identity). The second is the one that catches a demo
    ordering that happens to preserve every length.
    """

    def __init__(self, reference: Reference, sources: Sequence[DemoSource]) -> None:
        """Bind a reference table and the demo order to check.

        Args:
            reference: Node table to align to.
            sources: Demo sources in global demo-id order.
        """
        self._table = reference.table
        self._origin = reference.origin
        self._sources = tuple(sources)
        self._offsets = np.concatenate(
            [[0], np.cumsum([s.n_frames for s in self._sources])]
        ).astype(np.int64)

    @property
    def offsets(self) -> np.ndarray:
        """[n_demos + 1] i64 CSR boundaries this run will write at."""
        return self._offsets

    @property
    def n_frames(self) -> int:
        """Total raw frames across every indexed demo."""
        return int(self._offsets[-1])

    def check_counts(self) -> None:
        """Compare per-demo frame counts against the reference raw_ptr.

        Raises:
            ValueError: The demo count, the total frame count, or any per-demo count differs.
        """
        expected = self._table.raw_ptr.astype(np.int64)
        if expected.shape != self._offsets.shape:
            raise ValueError(
                f"demo count mismatch vs {self._origin}: index has {len(self._sources)} demos, "
                f"raw_ptr has {expected.shape[0] - 1}"
            )
        if not np.array_equal(expected, self._offsets):
            bad = int(np.flatnonzero(expected != self._offsets)[0])
            raise ValueError(
                f"frame-count mismatch vs {self._origin} at demo {max(0, bad - 1)} "
                f"({self._sources[max(0, bad - 1)].label}): raw_ptr={expected[bad]} "
                f"index={self._offsets[bad]}"
            )

    def check_frames(self, demo_id: int, joints: np.ndarray) -> None:
        """Compare one demo's joint configurations against the reference q_raw slice.

        Args:
            demo_id: Global demo id.
            joints: [T, D] f32 joint_states as read from the HDF5.

        Raises:
            ValueError: Any row differs from q_raw.
        """
        lo, hi = int(self._offsets[demo_id]), int(self._offsets[demo_id + 1])
        expected = self._table.q_raw[lo:hi]
        if not np.array_equal(expected, joints):
            raise ValueError(
                f"{self._sources[demo_id].label}: q_raw[{lo}:{hi}] does not match this demo's "
                f"joint_states -- the demo ordering does not reproduce the graph's"
            )


def load_reference(
    suite: str,
    nodes_npz: Path | None,
    loader: DemoIndexer,
    coarsen: int | None,
    *,
    rebuild: bool = False,
) -> Reference:
    """Resolve the node table to align against.

    An existing g_nodes.npz is authoritative. When there is none, or when caps make the corpus a
    strict subset of the built graph, the table is rebuilt in memory from the same DemoLoader that
    produced the index -- which still exercises NodeTable.from_demos' ordering.

    Args:
        suite: Suite short-name.
        nodes_npz: Explicit artifact path, or None to resolve one for the suite.
        loader: Loader whose caps define the corpus.
        coarsen: Coarsening factor for a rebuild. None uses the schema default.
        rebuild: Skip the artifact and always rebuild.

    Returns:
        The reference.
    """
    path = nodes_npz or (default_paths().graph(suite) / schema.NODES_NPZ)
    if path.exists() and not rebuild:
        table = NodeTable.load(os.fspath(path))
        edges_npz = path.parent / schema.EDGES_NPZ
        fingerprint = ""
        if edges_npz.exists():
            fingerprint = graph_hash(table, EdgeSet.load(os.fspath(edges_npz), strict=False))
        return Reference(table=table, origin=os.fspath(path), graph_hash=fingerprint)

    log.warning("not using %s -- rebuilding the reference table in memory", path)
    table = NodeBuilder(coarsen=coarsen or schema.COARSEN).build(loader.load())
    return Reference(table=table, origin="rebuilt", graph_hash="")


# ==================================================================================
# Policy replay
# ==================================================================================
@dataclass(frozen=True, slots=True)
class DemoActions:
    """One demo's cached policy output.

    Attributes:
        commanded: [T, 7] f32 action as eval_libero would have sent it, gripper binarized.
        grip: [T] f32 pre-binarization gripper channel.
    """

    commanded: np.ndarray
    grip: np.ndarray

    def __len__(self) -> int:
        return int(self.commanded.shape[0])


class StubPolicyClient:
    """Drop-in for WebsocketClientPolicy that answers offline, for --dry-run.

    Every returned chunk row equals the episode step it will be executed at, so a correctly
    indexed cache reproduces 0, 1, 2, ... within each demo -- an end-to-end check of the chunk
    contract and the demo ordering that needs no GPU and no server.
    """

    def __init__(self, host: str = "", port: int | None = None) -> None:
        """Accept and ignore the websocket client's constructor signature.

        Args:
            host: Ignored.
            port: Ignored.
        """
        self.host = host
        self.port = port

    @staticmethod
    def get_server_metadata() -> dict[str, Any]:
        """Metadata standing in for the server's.

        Returns:
            A dict marking the payload as stubbed.
        """
        return {"env": "stub", "model": "stub"}

    @staticmethod
    def infer(payload: dict[str, Any]) -> dict[str, Any]:
        """Return a deterministic chunk keyed to the requested episode step.

        Args:
            payload: The vla_input M1Inference.step builds.

        Returns:
            A response in the server's {"data": {"actions": [B, chunk, D]}} shape.
        """
        step = int(payload["episode_step"][0])
        chunk = np.arange(step, step + STUB_CHUNK_ROWS, dtype=np.float32)
        return {"data": {"actions": chunk[None, :, None].repeat(ACTION_DIM, axis=2)}}


class PolicyReplayer:
    """Drives M1Inference over HDF5 frames under the deploy chunk contract."""

    def __init__(
        self,
        checkpoint: Path,
        *,
        host: str,
        port: int,
        worker_id: int = 0,
        stub: bool = False,
    ) -> None:
        """Connect to a policy server (or stub it) and build the inference client.

        Args:
            checkpoint: Checkpoint dir, read for action normalization stats.
            host: Policy server host.
            port: Policy server port.
            worker_id: Stable client id, forwarded verbatim like the harness's shard index.
            stub: Replace the websocket client with StubPolicyClient.
        """
        import model2libero_interface as interface

        _configure_logging()
        if stub:
            interface.WebsocketClientPolicy = StubPolicyClient
        self._model = interface.M1Inference(
            policy_ckpt_path=str(checkpoint),
            host=host,
            port=port,
            image_size=RESIZE_SIZE,
            worker_id=worker_id,
        )

    @property
    def chunk_size(self) -> int:
        """Action-chunk length the client caches and replays open-loop."""
        return int(self._model.action_chunk_size)

    @property
    def server_metadata(self) -> dict[str, Any]:
        """Whatever the server announced on connect."""
        return dict(self._model.client.get_server_metadata())

    def replay(self, source: DemoSource, frames: DemoFrames) -> DemoActions:
        """Step the policy over every frame of one demo.

        Args:
            source: The demo being replayed, for its instruction.
            frames: That demo's observations.

        Returns:
            The demo's cached actions.
        """
        self._model.reset(task_description=source.instruction)
        commanded = np.empty((len(frames), ACTION_DIM), dtype=np.float32)
        grip = np.empty(len(frames), dtype=np.float32)
        for step in range(len(frames)):
            response = self._model.step(
                images=[frames.agentview[step], frames.wrist[step]],
                task_description=source.instruction,
                state=frames.state[step][None],
                step=step,
            )
            raw = response["raw_action"]
            open_gripper = float(np.asarray(raw["open_gripper"]).reshape(-1)[0])
            commanded[step, :3] = np.asarray(raw["world_vector"], dtype=np.float32).reshape(-1)
            commanded[step, 3:6] = np.asarray(raw["rotation_delta"], dtype=np.float32).reshape(-1)
            # eval_libero.py's _binarize_gripper_open, inlined: importing it would pull in the
            # whole LIBERO sim stack for one comparison.
            commanded[step, GRIP_COL] = 1.0 - 2.0 * (open_gripper > 0.5)
            grip[step] = open_gripper
        return DemoActions(commanded=commanded, grip=grip)


# ==================================================================================
# Resumable per-file shards
# ==================================================================================
class ShardStore:
    """Per-task-file checkpoints, rewritten after every demo so a session death costs one demo."""

    def __init__(self, root: Path) -> None:
        """Bind a shard directory.

        Args:
            root: Directory holding the shard files.
        """
        self._root = root
        self._root.mkdir(parents=True, exist_ok=True)

    def path(self, task_id: int, task_name: str) -> Path:
        """Shard file for one task.

        Args:
            task_id: Sort position of the source HDF5 file.
            task_name: Task name, for a readable filename.

        Returns:
            The shard path.
        """
        return self._root / f"{task_id:03d}_{task_name}.npz"

    def load(self, task_id: int, task_name: str, keys: Sequence[str]) -> list[DemoActions]:
        """Read the demos already cached for one task.

        A shard whose demo keys are not a prefix of the expected order is stale and is ignored,
        so a resumed run can never splice work done under a different demo ordering.

        Args:
            task_id: Sort position of the source file.
            task_name: Task name.
            keys: Expected demo keys for this task, in order.

        Returns:
            Cached actions for the leading demos, possibly empty.
        """
        path = self.path(task_id, task_name)
        if not path.exists():
            return []
        with np.load(os.fspath(path), allow_pickle=False) as archive:
            done = [str(k) for k in archive["demo_keys"]]
            counts = archive["counts"].astype(np.int64)
            actions = archive["actions"].astype(np.float32)
            grip = archive["grip"].astype(np.float32)
        if list(keys[: len(done)]) != done:
            log.warning("%s: stale shard (demo order changed) -- recomputing", path.name)
            return []
        bounds = np.concatenate([[0], np.cumsum(counts)]).astype(np.int64)
        return [
            DemoActions(commanded=actions[lo:hi], grip=grip[lo:hi])
            for lo, hi in zip(bounds[:-1], bounds[1:])
        ]

    def save(
        self, task_id: int, task_name: str, keys: Sequence[str], done: Sequence[DemoActions]
    ) -> None:
        """Write the demos cached so far for one task, atomically.

        Args:
            task_id: Sort position of the source file.
            task_name: Task name.
            keys: Demo keys covered, in order.
            done: Cached actions, one per key.
        """
        target = self.path(task_id, task_name)
        # np.savez appends .npz to any name that lacks it, so the temp name must already end in it.
        tmp = target.with_name(target.name + ".tmp.npz")
        np.savez(
            os.fspath(tmp),
            demo_keys=np.array(list(keys)),
            counts=np.array([len(d) for d in done], dtype=np.int64),
            actions=np.concatenate([d.commanded for d in done], axis=0),
            grip=np.concatenate([d.grip for d in done], axis=0),
        )
        os.replace(tmp, target)

    def clear(self) -> None:
        """Delete the shard directory once the final artifact is on disk."""
        shutil.rmtree(self._root, ignore_errors=True)


# ==================================================================================
# Progress
# ==================================================================================
class Progress:
    """Frames-per-second and ETA over the frames this process actually computes."""

    def __init__(self, total_frames: int) -> None:
        """Start the clock.

        Args:
            total_frames: Frames still to compute.
        """
        self._total = max(1, int(total_frames))
        self._done = 0
        self._start = time.monotonic()

    def advance(self, frames: int) -> str:
        """Record progress and format a one-line status.

        Args:
            frames: Frames just completed.

        Returns:
            A "done/total (pct) fps eta" status string.
        """
        self._done += int(frames)
        elapsed = max(1e-6, time.monotonic() - self._start)
        fps = self._done / elapsed
        eta_s = (self._total - self._done) / max(1e-6, fps)
        return (
            f"{self._done}/{self._total} ({100.0 * self._done / self._total:.1f}%) "
            f"{fps:.1f} frame/s eta {eta_s / 60.0:.1f} min"
        )


# ==================================================================================
# Orchestrator
# ==================================================================================
class ActionCacheBuilder:
    """Runs index -> align -> replay -> write for one suite."""

    def __init__(
        self,
        suite: str,
        *,
        hdf5_dir: Path,
        out_dir: Path,
        checkpoint: Path,
        reference: Reference,
        sources: Sequence[DemoSource],
        replayer: PolicyReplayer,
        host: str,
        port: int,
        dry_run: bool,
    ) -> None:
        """Configure a build.

        Args:
            suite: Suite short-name.
            hdf5_dir: Resolved demo directory.
            out_dir: Where the artifact and manifest are written.
            checkpoint: Policy checkpoint directory.
            reference: Node table to align against.
            sources: Demo sources in global demo-id order.
            replayer: Configured policy driver.
            host: Policy server host, recorded in the manifest.
            port: Policy server port, recorded in the manifest.
            dry_run: Whether the policy client is stubbed.
        """
        self._suite = suite
        self._hdf5_dir = hdf5_dir
        self._out_dir = out_dir
        self._checkpoint = checkpoint
        self._reference = reference
        self._sources = tuple(sources)
        self._replayer = replayer
        self._host = host
        self._port = port
        self._dry_run = dry_run
        self._guard = AlignmentGuard(reference, self._sources)
        self._shards = ShardStore(out_dir / SHARD_DIRNAME)

    def run(self) -> Path:
        """Build the cache, resuming from any shards already on disk.

        Returns:
            The written a_pi_raw.npz path.
        """
        self._guard.check_counts()
        log.info(
            "aligned: %d demos, %d frames, reference=%s",
            len(self._sources),
            self._guard.n_frames,
            self._reference.origin,
        )

        by_task = self._group_by_task()
        cached = {
            task_id: self._shards.load(task_id, group[0].task_name, [s.group_key for s in group])
            for task_id, group in by_task.items()
        }
        remaining = sum(
            sum(s.n_frames for s in group[len(cached[task_id]):])
            for task_id, group in by_task.items()
        )
        log.info("resuming: %d frames already cached", self._guard.n_frames - remaining)

        progress = Progress(remaining)
        for task_id, group in by_task.items():
            done = cached[task_id]
            keys = [s.group_key for s in group]
            for source in group[len(done):]:
                done.append(self._replay_one(source))
                self._shards.save(task_id, source.task_name, keys[: len(done)], done)
                log.info("%s  %s", source.label, progress.advance(source.n_frames))

        return self._write(by_task, cached)

    def _group_by_task(self) -> dict[int, list[DemoSource]]:
        """Bucket the demo index by source file, preserving global order.

        Returns:
            Mapping of task_id to its demos.
        """
        grouped: dict[int, list[DemoSource]] = {}
        for source in self._sources:
            grouped.setdefault(source.task_id, []).append(source)
        return grouped

    def _replay_one(self, source: DemoSource) -> DemoActions:
        """Read, verify and replay one demo.

        Args:
            source: The demo to run.

        Returns:
            Its cached actions.
        """
        frames = DemoFrames.read(source)
        self._guard.check_frames(source.demo_id, frames.joints)
        return self._replayer.replay(source, frames)

    def _write(
        self, by_task: dict[int, list[DemoSource]], cached: dict[int, list[DemoActions]]
    ) -> Path:
        """Concatenate every shard in global demo order and write the artifact plus manifest.

        Args:
            by_task: Demo index bucketed by source file.
            cached: Per-task cached actions, complete.

        Returns:
            The written a_pi_raw.npz path.

        Raises:
            ValueError: The assembled array's length disagrees with the reference.
        """
        order = [d for task_id in sorted(by_task) for d in cached[task_id]]
        actions = np.concatenate([d.commanded for d in order], axis=0).astype(np.float32)
        grip = np.concatenate([d.grip for d in order], axis=0).astype(np.float32)
        if actions.shape[0] != self._guard.n_frames:
            raise ValueError(
                f"assembled {actions.shape[0]} rows, expected {self._guard.n_frames}"
            )
        if self._dry_run:
            self._check_stub_alignment(actions)

        self._out_dir.mkdir(parents=True, exist_ok=True)
        target = self._out_dir / ACTIONS_NPZ
        np.savez_compressed(
            os.fspath(target),
            a_pi_raw=actions,
            a_pi_grip=grip,
            raw_ptr=self._guard.offsets.astype(np.int32),
        )
        (self._out_dir / MANIFEST_JSON).write_text(json.dumps(self._manifest(), indent=2) + "\n")
        self._shards.clear()
        log.info("wrote %s  a_pi_raw=%s", target, actions.shape)
        return target

    def _check_stub_alignment(self, actions: np.ndarray) -> None:
        """Verify the stub's step-numbered chunks came back in per-demo frame order.

        Args:
            actions: [N, 7] assembled cache.

        Raises:
            ValueError: Any row does not carry its own within-demo frame index.
        """
        expected = np.concatenate([np.arange(s.n_frames) for s in self._sources])
        if not np.array_equal(actions[:, 0].astype(np.int64), expected):
            bad = int(np.flatnonzero(actions[:, 0].astype(np.int64) != expected)[0])
            raise ValueError(
                f"chunk/order check failed at row {bad}: got {actions[bad, 0]}, "
                f"expected {expected[bad]}"
            )
        log.info("dry-run: chunk indexing and demo ordering verified over %d rows", len(expected))

    def _manifest(self) -> dict[str, Any]:
        """Everything needed to decide whether this cache may be reused.

        Returns:
            The JSON-serialisable manifest.
        """
        return {
            "created_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
            "suite": self._suite,
            "hdf5_dir": os.fspath(self._hdf5_dir),
            "dry_run": self._dry_run,
            "policy": {
                "checkpoint": CheckpointId.probe(self._checkpoint).to_dict(),
                "server_url": f"ws://{self._host}:{self._port}",
                "server_metadata": self._replayer.server_metadata,
                "chunk_size": self._replayer.chunk_size,
            },
            "preprocessing": {
                "image_rotation": "[:, ::-1, ::-1]",
                "resize": RESIZE_SIZE,
                "state": "concat(ee_pos[3], ee_ori[3], gripper_states[2])",
                "source_image_size": 128,
            },
            "reference": {
                "origin": self._reference.origin,
                "graph_hash": self._reference.graph_hash,
                "n_raw": int(self._reference.table.n_raw),
            },
            "n_demos": len(self._sources),
            "n_frames": self._guard.n_frames,
            "arrays": {
                "a_pi_raw": "[N, 7] f32 commanded action; gripper binarized to +-1 like eval_libero",
                "a_pi_grip": "[N] f32 pre-binarization gripper channel",
                "raw_ptr": "[n_demos + 1] i32 CSR boundaries, equal to NodeTable.raw_ptr",
            },
            "demos": [
                {
                    "task_id": s.task_id,
                    "task_name": s.task_name,
                    "group": s.group_key,
                    "instruction": s.instruction,
                    "n_frames": s.n_frames,
                }
                for s in self._sources
            ],
        }


# ==================================================================================
# CLI
# ==================================================================================
def _parse_args(argv: Sequence[str]) -> argparse.Namespace:
    """Parse command-line arguments.

    Args:
        argv: Argument list, excluding the program name.

    Returns:
        Parsed namespace.
    """
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("suite", choices=("object", "spatial", "goal", "long"))
    parser.add_argument("--out-dir", default=None, help="default: the suite's graph artifact dir")
    parser.add_argument("--nodes-npz", default=None, help="reference g_nodes.npz")
    parser.add_argument("--ckpt", default=None, help="default: configs/suites.yaml")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=10093)
    parser.add_argument("--worker-id", type=int, default=0)
    parser.add_argument("--coarsen", type=int, default=None)
    parser.add_argument("--limit-tasks", type=int, default=None)
    parser.add_argument("--limit-demos", type=int, default=None, help="per task file")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="stub the policy server and assert the chunk/order contract",
    )
    return parser.parse_args(list(argv))


def _configure_logging() -> None:
    """Give this module its own stdout handler, detached from the root logger.

    Idempotent, and called again after the policy client is imported: that import reaches a
    logging.dictConfig whose default disable_existing_loggers switches every logger built before it
    off, silently swallowing the per-demo progress a multi-hour resumable job is steered by.
    """
    handler = logging.StreamHandler(sys.stdout)
    handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
    log.handlers[:] = [handler]
    log.setLevel(logging.INFO)
    log.propagate = False
    log.disabled = False


def main(argv: Sequence[str] | None = None) -> int:
    """CLI entry point.

    Args:
        argv: Argument list, or None to read sys.argv.

    Returns:
        Process exit code.
    """
    args = _parse_args(sys.argv[1:] if argv is None else argv)
    _configure_logging()

    limits = LoadLimits(tasks=args.limit_tasks, demos=args.limit_demos)
    hdf5_dir = GraphBuilder(args.suite).hdf5_dir
    loader = DemoIndexer(hdf5_dir, limits=limits)
    sources = loader.index()
    log.info("%s: %d demos under %s", args.suite, len(sources), hdf5_dir)

    # Caps make the corpus a strict subset of the built graph, so the on-disk artifact cannot be
    # the reference -- rebuild one from the same capped loader instead.
    capped = limits.tasks is not None or limits.demos is not None
    reference = load_reference(
        args.suite,
        Path(args.nodes_npz) if args.nodes_npz else None,
        loader,
        args.coarsen,
        rebuild=capped,
    )

    checkpoint = Path(args.ckpt) if args.ckpt else suite_checkpoint(args.suite)
    out_dir = Path(args.out_dir) if args.out_dir else default_paths().graph(args.suite)
    replayer = PolicyReplayer(
        checkpoint,
        host=args.host,
        port=args.port,
        worker_id=args.worker_id,
        stub=args.dry_run,
    )
    ActionCacheBuilder(
        args.suite,
        hdf5_dir=hdf5_dir,
        out_dir=out_dir,
        checkpoint=checkpoint,
        reference=reference,
        sources=sources,
        replayer=replayer,
        host=args.host,
        port=args.port,
        dry_run=args.dry_run,
    ).run()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())