Reference inference stack + validated live_features.py + UPDATE_GUIDE for target-machine merge
121869d verified 2 months ago
33.4 kB
"""Inference for trained ACT checkpoints.
Four modes:
* ``offline`` — sanity-check a checkpoint against ground-truth actions in an
  HDF5 episode. For each frame t we predict the chunk a[t : t+chunk_size] and
  measure per-step L1 against the recorded action.
* ``online-demo`` — simulate real-time input by streaming frames from an
  episode through the same ``ACTInference`` class you'd use in a robot loop.
* ``online-server`` — bind a TCP socket and serve ``step`` / ``predict_chunk``
  / ``reset`` over the wire. Robot stack acts as the client and sends obs
  dicts; the server responds with action chunks. Length-prefixed pickle
  framing. Intended for trusted LANs only (pickle is not safe across trust
  boundaries).
* ``online-client-demo`` — connect to a running ``online-server`` and replay
  an episode through it. Serves as the loopback sanity check for the socket
  pipeline.
Loading is fully self-contained: the checkpoint stores ``config`` (the
ACTConfig fields) and ``model`` (state_dict including normalization buffers),
so a single ``--checkpoint`` flag is enough — no need to repeat ``--state-dim``
etc.
"""

from __future__ import annotations

import argparse
import csv
import glob
import os
import pickle
import socket
import struct
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Iterable

import h5py
import numpy as np
import torch

from act_policy import ACTConfig, ACTPolicy
from data.dataset import _to_3chw


# --------------------------------------------------------------------------- #
# Inference wrapper                                                           #
# --------------------------------------------------------------------------- #


def _to_torch(x: Any, device: torch.device) -> torch.Tensor:
    if isinstance(x, torch.Tensor):
        return x.to(device)
    return torch.as_tensor(np.ascontiguousarray(x), device=device)


class ACTInference:
    """Stateless-by-default wrapper around a trained ACTPolicy checkpoint.
    Use this from a robot control loop:
        infer = ACTInference.from_checkpoint("checkpoints/apple/act_final.pt")
        infer.reset()                       # at episode start
        while running:
            obs = collect_obs_from_hardware()
            action = infer.step(obs)        # (action_dim,) numpy
            send_to_robot(action)
    ``obs`` is a dict containing keys listed in ``self.image_keys`` (each value
    is HxW or HxWx3 or 3xHxW, uint8 or float in [0, 1]), and:
      - ``"state"``   — (state_dim,) qpos  if state_dim > 0
      - ``"tactile"`` — (tactile_dim,) wrench  if tactile_dim > 0
    Anything else in the dict is ignored.
    ``step()`` re-queries the model every ``n_action_steps`` frames; in between,
    it pops actions from an internal queue, matching the training-time chunk
    semantics. Use ``predict_chunk()`` for the full chunk.
    """

    def __init__(self, policy: ACTPolicy, device: torch.device):
        self.policy = policy.to(device).eval()
        self.device = device
        self.config: ACTConfig = policy.config
        self.image_keys = list(self.config.image_keys)
        self.image_shape = tuple(self.config.image_shape)  # (C, H, W)
        self.state_dim = self.config.state_dim
        self.tactile_dim = self.config.tactile_dim
        self.action_dim = self.config.action_dim
        self.chunk_size = self.config.chunk_size
        # Optional per-frame index masks (None = use all dims). The robot loop
        # may pass full-dim qpos; _build_batch then slices to match the model.
        self.state_indices: list[int] | None = getattr(self.config, "state_indices", None)
        self.action_indices: list[int] | None = getattr(self.config, "action_indices", None)
        # Per-frame state dim (before history concat) — used to detect whether
        # an incoming "state" tensor has been pre-sliced by the caller or not.
        sh = max(1, getattr(self.config, "state_history", 1))
        self.state_dim_per_frame = self.state_dim // sh
        # In-graph tactile-fusion branch. When set, the obs dict is expected to
        # contain "tactile_frames" + "tactile_bg" instead of "tactile"; the
        # encoder produces the 960-D vector inside the policy.
        self.tactile_image_keys: list[str] = list(getattr(self.config, "tactile_image_keys", []) or [])
        self.uses_tactile_branch: bool = bool(self.tactile_image_keys)

    @classmethod
    def from_checkpoint(cls, path: str | Path, device: str | torch.device = "cuda") -> "ACTInference":
        device_t = torch.device(device if torch.cuda.is_available() or str(device) != "cuda" else "cuda")
        ckpt = torch.load(str(path), map_location=device_t, weights_only=False)
        # Strip any keys not recognised by the current ACTConfig (forward-compat).
        cfg_fields = {f.name for f in ACTConfig.__dataclass_fields__.values()}  # type: ignore[attr-defined]
        cfg_dict = {k: v for k, v in ckpt["config"].items() if k in cfg_fields}
        policy = ACTPolicy(ACTConfig(**cfg_dict))
        policy.load_state_dict(ckpt["model"])
        return cls(policy, device_t)

    # ------------------------------------------------------------------- API

    def reset(self) -> None:
        """Clear the action queue. Call at the start of each new episode."""
        self.policy.reset()

    @torch.no_grad()
    def predict_chunk(self, obs: dict[str, Any]) -> np.ndarray:
        """Predict a full action chunk. Returns (chunk_size, action_dim) numpy."""
        batch = self._build_batch(obs)
        chunk = self.policy.predict_action_chunk(batch)  # (1, T, A)
        return chunk[0].detach().cpu().numpy()

    @torch.no_grad()
    def step(self, obs: dict[str, Any]) -> np.ndarray:
        """Return one (action_dim,) action; refills internally every n_action_steps."""
        batch = self._build_batch(obs)
        action = self.policy.select_action(batch)  # (1, A)
        return action[0].detach().cpu().numpy()

    # ------------------------------------------------------------------- impl

    def _build_batch(self, obs: dict[str, Any]) -> dict[str, torch.Tensor]:
        """Convert a single-frame obs dict into a (B=1) batch in the model's expected dtype/shape."""
        target_hw = (self.image_shape[1], self.image_shape[2])
        batch: dict[str, torch.Tensor] = {}

        for key in self.image_keys:
            if key not in obs:
                raise KeyError(f"obs is missing image key {key!r} (config expects {self.image_keys})")
            img = obs[key]
            if isinstance(img, torch.Tensor):
                img_np = img.detach().cpu().numpy()
            else:
                img_np = np.asarray(img)
            # Accept (H,W), (H,W,3), or (3,H,W). _to_3chw handles all three and resizes.
            img_t = _to_3chw(img_np, target_hw)  # uint8 (3, H, W)
            if img_t.dtype != torch.uint8:
                # _to_3chw preserves dtype; if caller gave float, leave as-is and let
                # ACTPolicy._prepare_batch skip the /255 path. Convert to uint8 only
                # when we got uint8 in.
                pass
            batch[key] = img_t.unsqueeze(0).to(self.device)

        if self.state_dim > 0:
            if "state" not in obs:
                raise KeyError("obs is missing 'state' (state_dim > 0)")
            state = _to_torch(obs["state"], self.device).float()
            # If state_indices is set, the caller may pass either the pre-sliced
            # state (shape == state_dim) or the full unsliced state (history
            # frames × full_per_frame_dim). Slice on the fly in the latter case
            # so robot loops can stay agnostic of the index mask.
            if self.state_indices is not None and state.shape[-1] != self.state_dim:
                hist = max(1, getattr(self.config, "state_history", 1))
                full_per_frame = state.shape[-1] // hist
                if state.shape[-1] != full_per_frame * hist:
                    raise ValueError(
                        f"state dim {state.shape[-1]} not divisible by state_history {hist}"
                    )
                # (..., hist*full) → (..., hist, full) → slice → (..., hist*sliced)
                state = state.view(*state.shape[:-1], hist, full_per_frame)
                state = state[..., self.state_indices]
                state = state.reshape(*state.shape[:-2], hist * self.state_dim_per_frame)
            if state.shape[-1] != self.state_dim:
                raise ValueError(f"state has dim {state.shape[-1]}, expected {self.state_dim}")
            batch["state"] = state.unsqueeze(0)

        if self.tactile_dim > 0:
            if self.uses_tactile_branch:
                # In-graph variant: obs must carry the raw image pairs + bg per
                # finger; the policy's branch produces the tactile vector.
                if "tactile_frames" not in obs or "tactile_bg" not in obs:
                    raise KeyError(
                        "obs is missing 'tactile_frames' / 'tactile_bg' "
                        f"(config.tactile_image_keys={self.tactile_image_keys})"
                    )
                frames = _to_torch(obs["tactile_frames"], self.device)
                bg = _to_torch(obs["tactile_bg"], self.device)
                if frames.dtype != torch.uint8:
                    frames = frames.to(torch.uint8)
                if bg.dtype != torch.uint8:
                    bg = bg.to(torch.uint8)
                batch["tactile_frames"] = frames.unsqueeze(0)
                batch["tactile_bg"] = bg.unsqueeze(0)
            else:
                if "tactile" not in obs:
                    raise KeyError("obs is missing 'tactile' (tactile_dim > 0)")
                tactile = _to_torch(obs["tactile"], self.device).float()
                if tactile.shape[-1] != self.tactile_dim:
                    raise ValueError(f"tactile has dim {tactile.shape[-1]}, expected {self.tactile_dim}")
                batch["tactile"] = tactile.unsqueeze(0)

        return batch


# --------------------------------------------------------------------------- #
# Episode iterator (used by both offline and online-demo modes)               #
# --------------------------------------------------------------------------- #


# Tactile keys aren't recoverable from the checkpoint config (only the dim is stored),
# so the user must pass them at inference time. Default mirrors the training default.
DEFAULT_TACTILE_KEYS = ["wrench_1"]


@dataclass
class EpisodeFrame:
    t: int
    obs: dict[str, np.ndarray]
    gt_action: np.ndarray  # (action_dim,) ground-truth action at this t


def _tactile_path(name: str) -> str:
    """Map a tactile key like 'wrench_1' or 'sitr_0' to '/wrench/1' or '/sitr/0'.
    Bare leaf names ('1') default to '/wrench/1'. Mirrors EpisodicHDF5Dataset.
    """
    if "_" in name:
        group, leaf = name.split("_", 1)
        return f"/{group}/{leaf}"
    return f"/wrench/{name}"


def iter_episode_frames(
    path: str | Path,
    image_keys: Iterable[str],
    tactile_keys: Iterable[str] | None,
    has_state: bool,
    state_history: int = 1,
    tactile_history: int = 1,
    state_indices: list[int] | None = None,
    action_indices: list[int] | None = None,
    tactile_image_keys: Iterable[str] | None = None,
    tactile_image_group: str = "tactile_imgs",
    tactile_bg_group: str = "tactile_bg_224",
    tactile_image_pair: int = 2,
) -> Iterable[EpisodeFrame]:
    """Yield one EpisodeFrame per timestep. Loads the file lazily inside a `with`.
    When ``state_history > 1`` or ``tactile_history > 1``, the corresponding obs
    entry is the concatenation of the last N frames (oldest first), front-padded
    with the first frame for t < N-1 — matching ``EpisodicHDF5Dataset``.
    When ``state_indices`` / ``action_indices`` is set, qpos and the emitted
    ``gt_action`` are sliced to those dims (matches what the model was trained on).
    """
    image_keys = list(image_keys)
    tactile_keys = list(tactile_keys) if tactile_keys else []
    tactile_image_keys = list(tactile_image_keys) if tactile_image_keys else []
    sh = max(1, state_history)
    th = max(1, tactile_history)
    n_pair = max(1, tactile_image_pair)
    with h5py.File(path, "r") as f:
        T = f["/action"].shape[0]
        actions = f["/action"][:]
        if action_indices is not None:
            actions = actions[:, action_indices]
        qpos_all = f["/qpos"][:].astype(np.float32) if has_state else None  # (T, state_dim)
        if qpos_all is not None and state_indices is not None:
            qpos_all = qpos_all[:, state_indices]
        # Pre-read full per-key tactile arrays only if we need them (small for wrench, larger for sitr).
        tactile_arrays = {k: f[_tactile_path(k)][:].astype(np.float32) for k in tactile_keys}
        # Pre-read per-finger image frames + bg for the in-graph branch (one read per episode).
        if tactile_image_keys:
            tactile_img_arrays = {
                k: np.asarray(f[f"/{tactile_image_group}/{k}"])     # (T, H, W) uint8
                for k in tactile_image_keys
            }
            tactile_bg_arrays = {
                k: np.asarray(f[f"/{tactile_bg_group}/{k}"])        # (H, W) uint8
                for k in tactile_image_keys
            }
        for t in range(T):
            obs: dict[str, np.ndarray] = {}
            for key in image_keys:
                obs[key] = f[f"/images/{key}"][t]
            if has_state:
                ts = [max(0, t - sh + 1 + k) for k in range(sh)]
                obs["state"] = qpos_all[ts].reshape(-1)
            if tactile_keys:
                ts = [max(0, t - th + 1 + k) for k in range(th)]
                per_t = []
                for ti in ts:
                    parts = [tactile_arrays[k][ti] for k in tactile_keys]
                    per_t.append(np.concatenate(parts, axis=-1))
                obs["tactile"] = np.concatenate(per_t, axis=-1)
            if tactile_image_keys:
                ts = [max(0, t - n_pair + 1 + k) for k in range(n_pair)]
                frames_per_finger = np.stack(
                    [np.stack([tactile_img_arrays[k][ti] for ti in ts], axis=0) for k in tactile_image_keys],
                    axis=0,
                )                                                          # (n_f, n_pair, H, W) uint8
                bg_per_finger = np.stack(
                    [tactile_bg_arrays[k] for k in tactile_image_keys], axis=0
                )                                                          # (n_f, H, W) uint8
                obs["tactile_frames"] = frames_per_finger
                obs["tactile_bg"] = bg_per_finger
            yield EpisodeFrame(t=t, obs=obs, gt_action=actions[t].astype(np.float32))


# --------------------------------------------------------------------------- #
# Offline evaluation                                                          #
# --------------------------------------------------------------------------- #


def offline_evaluate(
    infer: ACTInference,
    episode_paths: list[str],
    tactile_keys: list[str] | None,
    max_frames_per_ep: int | None,
    out_csv: str | None,
) -> dict[str, float]:
    """Frame-by-frame: per-frame predict_chunk vs ground-truth action[t : t+chunk_size]."""
    chunk_size = infer.chunk_size
    has_state = infer.state_dim > 0
    rows: list[dict[str, float]] = []
    per_episode_l1: list[float] = []

    for ep_path in episode_paths:
        ep_l1: list[float] = []
        # Materialize one episode in memory so we can index into actions[t : t+chunk_size]
        # without re-opening the file inside the loop.
        with h5py.File(ep_path, "r") as f:
            actions = f["/action"][:].astype(np.float32)  # (T, full_A)
            if infer.action_indices is not None:
                actions = actions[:, infer.action_indices]   # slice GT to model's output dims
            T = actions.shape[0]

        infer.reset()
        for frame in iter_episode_frames(
            ep_path, infer.image_keys, tactile_keys, has_state,
            state_history=getattr(infer.config, "state_history", 1),
            tactile_history=getattr(infer.config, "tactile_history", 1),
            state_indices=infer.state_indices,
            action_indices=infer.action_indices,
            tactile_image_keys=infer.tactile_image_keys if infer.uses_tactile_branch else None,
            tactile_image_group=getattr(infer.config, "_tactile_image_group", "tactile_imgs"),
            tactile_bg_group=getattr(infer.config, "_tactile_bg_group", "tactile_bg_224"),
            tactile_image_pair=getattr(infer.config, "tactile_image_pair", 2),
        ):
            if max_frames_per_ep is not None and frame.t >= max_frames_per_ep:
                break
            chunk = infer.predict_chunk(frame.obs)  # (chunk_size, A)
            end = min(frame.t + chunk_size, T)
            n_valid = end - frame.t
            gt = actions[frame.t:end]                 # (n_valid, A)
            pred = chunk[:n_valid]                    # (n_valid, A)
            per_step_l1 = np.abs(pred - gt).mean(axis=-1)  # (n_valid,)
            l1 = float(per_step_l1.mean())
            ep_l1.append(l1)
            rows.append({
                "episode": os.path.basename(ep_path),
                "t": frame.t,
                "n_valid": n_valid,
                "l1_mean": l1,
                "l1_step0": float(per_step_l1[0]),
                "l1_last": float(per_step_l1[-1]),
            })

        if ep_l1:
            ep_mean = float(np.mean(ep_l1))
            per_episode_l1.append(ep_mean)
            print(f"[offline] {os.path.basename(ep_path)}: frames={len(ep_l1)} mean_l1={ep_mean:.4f}")

    if out_csv:
        Path(out_csv).parent.mkdir(parents=True, exist_ok=True)
        with open(out_csv, "w", newline="") as fh:
            w = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
            w.writeheader()
            w.writerows(rows)
        print(f"[offline] wrote per-frame CSV: {out_csv}")

    overall = {"per_episode_mean_l1": float(np.mean(per_episode_l1)) if per_episode_l1 else float("nan"),
               "n_episodes": len(per_episode_l1),
               "n_frames": len(rows)}
    print(f"[offline] overall: episodes={overall['n_episodes']} frames={overall['n_frames']} "
          f"mean_l1={overall['per_episode_mean_l1']:.4f}")
    return overall


# --------------------------------------------------------------------------- #
# Online demo (simulated real-time stream)                                    #
# --------------------------------------------------------------------------- #


def online_demo(
    infer: ACTInference,
    episode_path: str,
    tactile_keys: list[str] | None,
    target_hz: float,
    max_frames: int | None,
) -> None:
    """Pretend the episode HDF5 is a live sensor stream.
    The same `ACTInference` instance is used as in offline mode — only the
    obs source changes. Replace ``iter_episode_frames`` with your hardware
    callback to wire this into a real robot.
    """
    period = 1.0 / target_hz if target_hz > 0 else 0.0
    has_state = infer.state_dim > 0
    infer.reset()

    print(f"[online] streaming from {os.path.basename(episode_path)} at {target_hz:.1f} Hz, "
          f"chunk_size={infer.chunk_size} n_action_steps={infer.config.n_action_steps}")
    times_ms: list[float] = []
    for frame in iter_episode_frames(
        episode_path, infer.image_keys, tactile_keys, has_state,
        state_history=getattr(infer.config, "state_history", 1),
        tactile_history=getattr(infer.config, "tactile_history", 1),
        state_indices=infer.state_indices,
        action_indices=infer.action_indices,
        tactile_image_keys=infer.tactile_image_keys if infer.uses_tactile_branch else None,
        tactile_image_pair=getattr(infer.config, "tactile_image_pair", 2),
    ):
        if max_frames is not None and frame.t >= max_frames:
            break
        loop_start = time.perf_counter()
        action = infer.step(frame.obs)            # (action_dim,)
        infer_ms = (time.perf_counter() - loop_start) * 1000
        times_ms.append(infer_ms)

        if frame.t % 10 == 0:
            l1 = float(np.abs(action - frame.gt_action).mean())
            print(f"  t={frame.t:>4d}  action[:4]={action[:4].round(3).tolist()}  "
                  f"gt[:4]={frame.gt_action[:4].round(3).tolist()}  l1={l1:.3f}  "
                  f"infer={infer_ms:.1f}ms")

        # Pace the loop to the requested rate (skip if inference itself was slower).
        sleep_left = period - (time.perf_counter() - loop_start)
        if sleep_left > 0:
            time.sleep(sleep_left)

    if times_ms:
        arr = np.asarray(times_ms)
        print(f"[online] inference latency: mean={arr.mean():.2f}ms p50={np.median(arr):.2f}ms "
              f"p95={np.percentile(arr, 95):.2f}ms max={arr.max():.2f}ms")


# --------------------------------------------------------------------------- #
# Socket server / client                                                      #
# --------------------------------------------------------------------------- #
#
# Wire protocol: length-prefixed pickle.
#
#   request                            response
#   -----------------------------------------------------------------------
#   {"cmd": "step", "obs": {...}}      {"ok": True, "action": np.ndarray(A,), "infer_ms": float}
#   {"cmd": "predict_chunk", "obs": …} {"ok": True, "chunk":  np.ndarray(T,A)}
#   {"cmd": "reset"}                   {"ok": True}
#   {"cmd": "close"}                   {"ok": True}
#   <anything else>                    {"ok": False, "error": "..."}
#
# Pickle is fast and supports numpy arrays out of the box but is NOT safe to
# expose across a trust boundary — only use this on the robot LAN.


def _recv_exact(sock: socket.socket, n: int) -> bytes:
    buf = bytearray()
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            raise ConnectionError("peer closed mid-message")
        buf.extend(chunk)
    return bytes(buf)


def _recv_msg(sock: socket.socket) -> Any:
    (length,) = struct.unpack(">I", _recv_exact(sock, 4))
    return pickle.loads(_recv_exact(sock, length))


def _send_msg(sock: socket.socket, obj: Any) -> None:
    payload = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
    sock.sendall(struct.pack(">I", len(payload)) + payload)


class ACTSocketClient:
    """Thin client to call a remote ``online-server``.
    Use this from the robot stack:
        client = ACTSocketClient(host="192.168.1.42", port=7654)
        client.reset()                         # at episode start
        while running:
            action = client.step(obs)          # obs same dict as ACTInference.step
            send_to_robot(action)
        client.close()
    """

    def __init__(self, host: str = "127.0.0.1", port: int = 7654, timeout: float | None = 10.0):
        self.sock = socket.create_connection((host, port), timeout=timeout)
        # Once connected, allow indefinite blocking on per-message recv (inference can
        # take longer than the connect timeout); callers can set their own timeout.
        self.sock.settimeout(None)

    def _round_trip(self, msg: dict) -> dict:
        _send_msg(self.sock, msg)
        reply = _recv_msg(self.sock)
        if not reply.get("ok"):
            raise RuntimeError(f"server error: {reply.get('error')}")
        return reply

    def reset(self) -> None:
        self._round_trip({"cmd": "reset"})

    def step(self, obs: dict[str, Any]) -> np.ndarray:
        return self._round_trip({"cmd": "step", "obs": obs})["action"]

    def predict_chunk(self, obs: dict[str, Any]) -> np.ndarray:
        return self._round_trip({"cmd": "predict_chunk", "obs": obs})["chunk"]

    def close(self) -> None:
        try:
            _send_msg(self.sock, {"cmd": "close"})
            try:
                _recv_msg(self.sock)
            except (ConnectionError, EOFError):
                pass
        finally:
            self.sock.close()


def online_server(
    infer: ACTInference,
    host: str = "0.0.0.0",
    port: int = 7654,
) -> None:
    """Block forever accepting one client at a time and serving inference calls.
    The action queue is reset on every new client connection. Press Ctrl-C
    to shut down.
    """
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind((host, port))
    srv.listen(1)
    print(f"[server] listening on {host}:{port}  (image_keys={infer.image_keys} "
          f"state_dim={infer.state_dim} tactile_dim={infer.tactile_dim} "
          f"action_dim={infer.action_dim} chunk_size={infer.chunk_size})")

    try:
        while True:
            conn, addr = srv.accept()
            print(f"[server] client connected: {addr}")
            infer.reset()
            n_step = 0
            try:
                while True:
                    msg = _recv_msg(conn)
                    cmd = msg.get("cmd")
                    if cmd == "step":
                        t0 = time.perf_counter()
                        action = infer.step(msg["obs"])
                        infer_ms = (time.perf_counter() - t0) * 1000
                        _send_msg(conn, {"ok": True, "action": action, "infer_ms": infer_ms})
                        n_step += 1
                    elif cmd == "predict_chunk":
                        chunk = infer.predict_chunk(msg["obs"])
                        _send_msg(conn, {"ok": True, "chunk": chunk})
                    elif cmd == "reset":
                        infer.reset()
                        _send_msg(conn, {"ok": True})
                    elif cmd == "close":
                        _send_msg(conn, {"ok": True})
                        break
                    else:
                        _send_msg(conn, {"ok": False, "error": f"unknown cmd {cmd!r}"})
            except (ConnectionError, EOFError) as e:
                print(f"[server] client disconnected: {e}")
            except Exception as e:                              # noqa: BLE001
                # Best-effort error report back to the client, then drop the connection.
                print(f"[server] error while handling client: {type(e).__name__}: {e}")
                try:
                    _send_msg(conn, {"ok": False, "error": f"{type(e).__name__}: {e}"})
                except Exception:                                # noqa: BLE001
                    pass
            finally:
                conn.close()
                print(f"[server] connection closed after {n_step} step(s)")
    except KeyboardInterrupt:
        print("\n[server] interrupted, shutting down")
    finally:
        srv.close()


def online_client_demo(
    host: str,
    port: int,
    episode_path: str,
    image_keys: list[str],
    tactile_keys: list[str] | None,
    has_state: bool,
    target_hz: float,
    max_frames: int | None,
) -> None:
    """Connect to ``online-server`` and stream an episode through it as a sanity check."""
    period = 1.0 / target_hz if target_hz > 0 else 0.0
    client = ACTSocketClient(host=host, port=port)
    print(f"[client] connected to {host}:{port}")
    client.reset()

    rtt_ms: list[float] = []
    try:
        for frame in iter_episode_frames(episode_path, image_keys, tactile_keys, has_state):
            if max_frames is not None and frame.t >= max_frames:
                break
            loop_start = time.perf_counter()
            action = client.step(frame.obs)
            rtt_ms.append((time.perf_counter() - loop_start) * 1000)
            if frame.t % 10 == 0:
                l1 = float(np.abs(action - frame.gt_action).mean())
                print(f"  t={frame.t:>4d}  l1={l1:.3f}  rtt={rtt_ms[-1]:.1f}ms")
            sleep_left = period - (time.perf_counter() - loop_start)
            if sleep_left > 0:
                time.sleep(sleep_left)
    finally:
        client.close()

    if rtt_ms:
        arr = np.asarray(rtt_ms)
        print(f"[client] round-trip latency: mean={arr.mean():.2f}ms p50={np.median(arr):.2f}ms "
              f"p95={np.percentile(arr, 95):.2f}ms max={arr.max():.2f}ms")


# --------------------------------------------------------------------------- #
# CLI                                                                         #
# --------------------------------------------------------------------------- #


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser()
    ap.add_argument("--checkpoint", type=str, default=None,
                    help="Required for offline / online-demo / online-server; ignored for online-client-demo.")
    ap.add_argument("--mode", choices=["offline", "online-demo", "online-server", "online-client-demo"],
                    default="offline")
    ap.add_argument("--data", type=str, default=None,
                    help="Episode dir (offline) or single .hdf5 path (online-demo / online-client-demo).")
    ap.add_argument("--episode-glob", type=str, default="episode_*.hdf5")
    ap.add_argument("--tactile-keys", nargs="*", default=None,
                    help=f"Override tactile wrench keys. Default {DEFAULT_TACTILE_KEYS}; only the "
                         "config's tactile_dim is stored in the checkpoint, so the keys must match "
                         "what the model was trained with.")
    ap.add_argument("--max-frames", type=int, default=None,
                    help="Cap frames per episode (offline) or total (online-demo / online-client-demo).")
    ap.add_argument("--out-csv", type=str, default=None, help="Offline only.")
    ap.add_argument("--hz", type=float, default=30.0, help="Demo loop rate (online-demo / online-client-demo).")
    ap.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")

    # Socket-mode flags.
    ap.add_argument("--host", type=str, default="0.0.0.0",
                    help="Bind addr for online-server; connect addr for online-client-demo.")
    ap.add_argument("--port", type=int, default=7654)
    # online-client-demo needs to know which obs keys to assemble locally before sending.
    ap.add_argument("--client-image-keys", nargs="*", default=None,
                    help="(online-client-demo) image keys to read from the HDF5; must match the server's checkpoint.")
    ap.add_argument("--client-state-dim", type=int, default=22,
                    help="(online-client-demo) whether to include qpos in obs (0 disables).")
    return ap.parse_args()


def main() -> None:
    args = parse_args()

    # online-client-demo is the only mode that does NOT load a checkpoint.
    if args.mode == "online-client-demo":
        if args.data is None:
            raise SystemExit("--data is required in online-client-demo (path to one episode .hdf5)")
        if args.client_image_keys is None:
            raise SystemExit("--client-image-keys is required in online-client-demo")
        ep = args.data if args.data.endswith(".hdf5") else \
            sorted(glob.glob(os.path.join(args.data, args.episode_glob)))[0]
        tactile_keys = args.tactile_keys  # may be None
        online_client_demo(
            host=args.host,
            port=args.port,
            episode_path=ep,
            image_keys=list(args.client_image_keys),
            tactile_keys=tactile_keys,
            has_state=args.client_state_dim > 0,
            target_hz=args.hz,
            max_frames=args.max_frames,
        )
        return

    if args.checkpoint is None:
        raise SystemExit(f"--checkpoint is required for mode {args.mode!r}")

    infer = ACTInference.from_checkpoint(args.checkpoint, device=args.device)
    print(f"[load] {args.checkpoint}  image_keys={infer.image_keys} state_dim={infer.state_dim} "
          f"tactile_dim={infer.tactile_dim} action_dim={infer.action_dim} chunk_size={infer.chunk_size}")

    tactile_keys = args.tactile_keys
    if infer.tactile_dim > 0 and tactile_keys is None:
        tactile_keys = list(DEFAULT_TACTILE_KEYS)

    if args.mode == "offline":
        if args.data is None:
            raise SystemExit("--data is required in offline mode (directory of episode_*.hdf5)")
        episode_paths = sorted(glob.glob(os.path.join(args.data, args.episode_glob))) \
            if os.path.isdir(args.data) else [args.data]
        if not episode_paths:
            raise SystemExit(f"no episodes matched {args.data}/{args.episode_glob}")
        offline_evaluate(
            infer,
            episode_paths=episode_paths,
            tactile_keys=tactile_keys,
            max_frames_per_ep=args.max_frames,
            out_csv=args.out_csv,
        )
    elif args.mode == "online-demo":
        if args.data is None:
            raise SystemExit("--data is required in online-demo (path to one episode .hdf5)")
        ep = args.data if args.data.endswith(".hdf5") else \
            sorted(glob.glob(os.path.join(args.data, args.episode_glob)))[0]
        online_demo(
            infer,
            episode_path=ep,
            tactile_keys=tactile_keys,
            target_hz=args.hz,
            max_frames=args.max_frames,
        )
    elif args.mode == "online-server":
        online_server(infer, host=args.host, port=args.port)


if __name__ == "__main__":
    main()