hud-python 0.6.4__tar.gz → 0.6.5__tar.gz

{hud_python-0.6.4 → hud_python-0.6.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hud-python
-Version: 0.6.4
+Version: 0.6.5
 Summary: SDK for the HUD platform.
 Project-URL: Homepage, https://github.com/hud-evals/hud-python
 Project-URL: Bug Tracker, https://github.com/hud-evals/hud-python/issues
@@ -70,6 +70,7 @@ Requires-Dist: ruff<0.15.0,>=0.11.8; extra == 'dev'
 Provides-Extra: modal
 Requires-Dist: modal>=1.0; extra == 'modal'
 Provides-Extra: robot
+Requires-Dist: av>=12; extra == 'robot'
 Requires-Dist: numpy>=1.24; extra == 'robot'
 Requires-Dist: openpi-client>=0.1.2; extra == 'robot'
 Provides-Extra: train

{hud_python-0.6.4 → hud_python-0.6.5}/hud/agents/robot/__init__.py

@@ -10,6 +10,9 @@ The harness splits a policy rollout into three seams, each replaceable on its ow
 - :class:`~hud.agents.robot.adapter.Adapter` — translate between the env's
   observation/action spaces (from the contract) and the policy's.
 
+Wrap an agent in :class:`~hud.agents.robot.batching.BatchedAgent` to run many rollouts
+concurrently off one batched GPU forward (``max_concurrent`` rollouts, shared model).
+
 Per-tick platform tracing is emitted by the loop itself: each step records an
 :class:`~hud.agents.types.ObservationStep`, and each re-inference an
 :class:`~hud.agents.types.InferenceStep`, so runs stream live into the HUD trace viewer.
@@ -20,16 +23,19 @@ This subpackage needs the ``robot`` extra (``pip install 'hud-python[robot]'``)
 
 from __future__ import annotations
 
-from .adapter import Adapter, LeRobotAdapter
+from .adapter import Adapter, LeRobotAdapter, OpenPIAdapter
 from .agent import ROBOT_PROTOCOL, RobotAgent
-from .model import LeRobotModel, Model, lerobot_infer
+from .batching import BatchedAgent, BatchedModel
+from .model import LeRobotModel, Model
 
 __all__ = [
     "ROBOT_PROTOCOL",
     "Adapter",
+    "BatchedAgent",
+    "BatchedModel",
     "LeRobotAdapter",
     "LeRobotModel",
     "Model",
+    "OpenPIAdapter",
     "RobotAgent",
-    "lerobot_infer",
 ]

{hud_python-0.6.4 → hud_python-0.6.5}/hud/agents/robot/adapter.py

@@ -89,7 +89,17 @@ class LeRobotAdapter(Adapter):
         return action
 
 
+class OpenPIAdapter(Adapter):
+    """unwraps obs['data'] to OpenPI wire keys, attaches prompt; actions are passthrough"""
+
+    def adapt_observation(self, obs: dict[str, Any], prompt: str) -> dict[str, Any]:
+        out = dict(obs["data"])
+        out.setdefault("prompt", prompt)
+        return out
+
+
 __all__ = [
     "Adapter",
     "LeRobotAdapter",
+    "OpenPIAdapter",
 ]

{hud_python-0.6.4 → hud_python-0.6.5}/hud/agents/robot/agent.py

@@ -5,8 +5,8 @@ Subclass :class:`RobotAgent`, set ``self.model`` and ``self.adapter`` in
 
 The base calls the adapter and model at the right moments::
 
-    setup_robot      -> adapter.bind(spaces)                          # once after connect
-    on_episode_start -> model.reset(); adapter.reset()                # once per episode
+    setup_robot      -> adapter.bind(spaces)       # once after connect
+    on_episode_start -> adapter.reset()            # per episode; model is stateless
     select_action    -> adapt_observation -> model.ainfer -> pop chunk -> adapt_action
 
 ``model.ainfer`` always returns a ``[T, A]`` chunk; :meth:`RobotAgent.select_action`
@@ -24,9 +24,10 @@ from typing import TYPE_CHECKING, Any, ClassVar
 import numpy as np
 
 from hud.agents.base import Agent
-from hud.agents.types import InferenceStep, ObservationStep
 from hud.capabilities.robot import RobotClient
 
+from .record import Recorder
+
 if TYPE_CHECKING:
     from hud.eval.run import Run
 
@@ -57,6 +58,9 @@ class RobotAgent(Agent):
     robot_protocol: ClassVar[str] = ROBOT_PROTOCOL
     #: How often (in steps) to print a step-progress line. 0 = off.
     log_every: ClassVar[int] = 20
+    #: Opt-in: also save a LeRobot v3 dataset of every (obs, action) pair to disk
+    #: (the ``--save`` flag). Telemetry streams regardless; see :mod:`.record`.
+    save: bool = False
 
     #: Runs the policy (preprocess → forward → postprocess). Subclasses set this.
     model: Model | None = None
@@ -70,9 +74,11 @@ class RobotAgent(Agent):
     _env_obs_space: dict[str, Any]
     #: Unexecuted tail of the current policy chunk; popped one action per step.
     _active_chunk: deque[ActionArray]
-    #: The live run + control-tick index, so ``select_action`` can record its own InferenceStep.
-    _run: Run
+    #: Control-tick index, incremented per executed action.
     _tick: int
+    #: Records all telemetry (observation/inference steps + video) and, when ``save``, a
+    #: LeRobot dataset. Agent-lifetime (the dataset spans every episode); created lazily.
+    _recorder: Recorder | None = None
 
     def setup_robot(self, client: RobotClient) -> None:
         """Discover the env's action/observation layout and bind the adapter to it."""
@@ -81,16 +87,19 @@ class RobotAgent(Agent):
             self.adapter.bind(self._env_action_space, self._env_obs_space)
 
     def on_episode_start(self, run: Run, client: RobotClient, *, prompt: str) -> None:
-        """Store the prompt and reset the model and adapter before the act loop.
+        """Store the prompt and reset per-episode state before the act loop.
 
-        Override (calling ``super()`` first) only for extra per-episode setup.
+        The model is stateless (per-episode state lives here, not on the shared model), so
+        only the adapter is reset. Override (calling ``super()`` first) for extra setup.
         """
         self._prompt = prompt
         self._active_chunk = deque()
-        self._run = run
         self._tick = 0
-        if self.model is not None:
-            self.model.reset()
+        # One recorder for the agent's life so its LeRobot dataset spans every episode;
+        # begin() opens this episode (fresh video stream, prompt) and takes the run it records onto.
+        if self._recorder is None:
+            self._recorder = Recorder(client, save=self.save)
+        self._recorder.begin(run, prompt)
         if self.adapter is not None:
             self.adapter.reset()
 
@@ -110,9 +119,8 @@ class RobotAgent(Agent):
             )
             chunk = np.atleast_2d(await self.model.ainfer(batch))  # [T, A]
             self._active_chunk = deque(chunk)
-            self._run.record(
-                InferenceStep(tick=self._tick, chunk=chunk.tolist(), chunk_length=len(chunk))
-            )
+            assert self._recorder is not None  # set in on_episode_start
+            self._recorder.record_inference(chunk, tick=self._tick)
         self._tick += 1
         raw = self._active_chunk.popleft()
         return raw if self.adapter is None else self.adapter.adapt_action(raw, obs)
@@ -131,15 +139,17 @@ class RobotAgent(Agent):
             self.on_episode_start(run, client, prompt=prompt)
             print(f"[agent] episode started: {prompt!r} (max_steps={step_limit})", flush=True)
 
+            assert self._recorder is not None  # set in on_episode_start above
             for step in range(step_limit):
                 obs = await client.get_observation()
-                run.record(ObservationStep.from_obs(obs, tick=step, obs_space=self._env_obs_space))
+                self._recorder.record_observation(obs, tick=step)
 
                 if self.should_stop(obs, step=step, max_steps=step_limit):
                     print(f"[agent] env reported terminated at step {step}", flush=True)
                     break
 
                 action = await self.select_action(obs)
+                self._recorder.record_action(action)
                 await client.send_action(action)
 
                 if self.log_every and step % self.log_every == 0:
@@ -151,6 +161,8 @@ class RobotAgent(Agent):
             run.trace.status = "completed"
             run.trace.content = "done"
         finally:
+            if self._recorder is not None:
+                self._recorder.end()  # flush video tails + commit the LeRobot episode
             await client.close()
 
 

hud_python-0.6.5/hud/agents/robot/batching.py

@@ -0,0 +1,130 @@
+"""Batched inference for concurrent robot rollouts.
+
+- BatchedModel: stacks concurrent ainfer calls into one infer
+- BatchedAgent: gives each rollout its own state, shares one batched model
+"""
+
+from __future__ import annotations
+
+import asyncio
+import copy
+import importlib
+from typing import TYPE_CHECKING, Any
+
+from hud.agents.base import Agent
+
+from .model import Model
+
+if TYPE_CHECKING:
+    from hud.eval.run import Run
+
+    from ._types import ActionArray
+    from .agent import RobotAgent
+
+
+class BatchedModel(Model):
+    """Coalesce concurrent ``ainfer`` calls into one stacked ``inner.infer``.
+
+    A lazily-started worker drains up to ``batch_size`` queued calls (or waits up to
+    ``max_wait_s`` for stragglers — which avoids stalling when fewer rollouts are live,
+    e.g. the tail of a suite), stacks them into one ``[N, ...]`` batch, runs a single
+    forward, and scatters the ``[N, T, A]`` rows back to each caller.
+
+    ``inner`` must be an in-process, stateless model whose :meth:`~Model.infer` runs the
+    whole ``[N, ...]`` batch in one forward (e.g. :class:`~hud.agents.robot.model.LeRobotModel`).
+    :class:`~hud.agents.robot.model.RemoteModel` is **not** supported: it does one WebSocket
+    request per env and the OpenPI server protocol has no batched-request shape, so a stacked
+    batch would be mis-sent as a single env. Run one agent per rollout against it instead.
+    """
+
+    def __init__(self, inner: Model, *, batch_size: int, max_wait_s: float = 0.05) -> None:
+        self.inner = inner
+        self.batch_size = int(batch_size)
+        self.max_wait_s = float(max_wait_s)
+        # Bound to the running loop on first ainfer (the harness owns the loop).
+        self._queue: asyncio.Queue[tuple[Any, asyncio.Future[ActionArray]]] | None = None
+        self._worker: asyncio.Task[None] | None = None
+
+    def infer(self, batch: Any) -> ActionArray:
+        return self.inner.infer(batch)
+
+    async def ainfer(self, batch: Any) -> ActionArray:
+        loop = asyncio.get_running_loop()
+        if self._worker is None:
+            self._queue = asyncio.Queue()
+            self._worker = loop.create_task(self._batch_loop())
+        assert self._queue is not None
+        fut: asyncio.Future[ActionArray] = loop.create_future()
+        await self._queue.put((batch, fut))
+        return await fut
+
+    async def _batch_loop(self) -> None:
+        assert self._queue is not None
+        loop = asyncio.get_running_loop()
+        while True:
+            items = [await self._queue.get()]  # block for the first caller
+            deadline = loop.time() + self.max_wait_s
+            while len(items) < self.batch_size:
+                timeout = deadline - loop.time()
+                if timeout <= 0:
+                    break
+                try:
+                    items.append(await asyncio.wait_for(self._queue.get(), timeout))
+                except TimeoutError:
+                    break
+            samples = [b for b, _ in items]
+            try:
+                torch: Any = importlib.import_module("torch")
+
+                # Collate N raw observations into one [N, ...] batch: stack tensor
+                # fields on a new leading dim, gather scalars/strings into a list.
+                stacked: dict[str, Any] = {
+                    k: torch.stack([s[k] for s in samples])
+                    if torch.is_tensor(samples[0][k])
+                    else [s[k] for s in samples]
+                    for k in samples[0]
+                }
+                arr = await asyncio.to_thread(self.inner.infer, stacked)  # [N, T, A]
+                for (_, fut), chunk in zip(items, arr, strict=True):
+                    if not fut.done():
+                        fut.set_result(chunk)
+            except Exception as exc:  # isolate: a bad batch fails only its own callers
+                for _, fut in items:
+                    if not fut.done():
+                        fut.set_exception(exc)
+
+
+class BatchedAgent(Agent):
+    """Drive many rollouts concurrently against one shared, batched model.
+
+    Per run: a shallow clone of ``agent`` (its own episode state) sharing a per-run
+    adapter copy and the single :class:`BatchedModel`, so concurrent ``ainfer`` calls
+    coalesce into one forward. Relies on the agent keeping per-run state out of
+    ``__init__`` (assigned in ``on_episode_start``) so the clones stay isolated, and on
+    the model being stateless (no per-episode ``reset``) since it is shared across clones.
+
+    Requires an in-process batchable model; :class:`~hud.agents.robot.model.RemoteModel`
+    is not supported (the OpenPI server protocol has no batched-request shape).
+
+    Takes ownership of ``agent``: it swaps ``agent.model`` for a :class:`BatchedModel` wrapper
+    in place (so the wrapper is shared by every per-run clone). The passed-in instance is
+    therefore permanently batched — hand :class:`BatchedAgent` a dedicated agent and don't
+    also use that same instance for direct, unbatched :class:`RobotAgent` rollouts.
+    """
+
+    def __init__(self, agent: RobotAgent, *, batch_size: int, max_wait_s: float = 0.05) -> None:
+        if agent.model is None:
+            raise RuntimeError("BatchedAgent needs agent.model set")
+        self._template = agent
+        # Wrap once, in place: the passed-in agent is now permanently batched (see class doc).
+        # Every per-run clone shares this batcher by reference.
+        agent.model = BatchedModel(agent.model, batch_size=batch_size, max_wait_s=max_wait_s)
+
+    async def __call__(self, run: Run, **kwargs: Any) -> None:
+        worker = copy.copy(self._template)  # fresh __dict__; shares the batched model
+        if worker.adapter is not None:  # defensive: a stateful custom adapter must be per-run
+            worker.adapter = copy.copy(worker.adapter)
+        await worker(run, **kwargs)
+
+
+__all__ = ["BatchedAgent", "BatchedModel"]

hud_python-0.6.5/hud/agents/robot/model.py

@@ -0,0 +1,127 @@
+"""The ``Model``: wraps a policy and owns its inference mechanics.
+
+A ``Model`` knows *how to run* a policy (preprocess → forward → postprocess); the
+harness only awaits ``model.ainfer(batch)``. Use :class:`LeRobotModel` for stock
+LeRobot checkpoints; subclass :class:`Model` and implement ``infer`` otherwise.
+
+:meth:`Model.infer` is batch-shaped (one batch dict in, an ``[N, T, A]`` chunk out) and
+stateless across calls, so one model can be shared and batched across concurrent rollouts
+(see :mod:`hud.agents.robot.batching`); per-episode state belongs on the agent.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import importlib
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from ._types import ActionArray
+
+
+class Model:
+    """Owns a policy and its inference mechanics.
+
+    Stateless by contract: the agent owns all per-episode state (the open-loop chunk), so a
+    single model can be shared and batched across concurrent rollouts. There is deliberately
+    no ``reset`` hook — anything that resets per episode belongs on the agent, not here.
+    Driven by :class:`~hud.agents.robot.agent.RobotAgent`, which awaits :meth:`ainfer`.
+    """
+
+    def infer(self, batch: Any) -> ActionArray:
+        """Run the policy on an ``[N, ...]`` batch, return an ``[N, T, A]`` chunk.
+
+        Implementations MUST keep the leading batch dim ``N`` (even for ``N == 1``):
+        :meth:`ainfer` indexes ``[0]`` and :class:`~hud.agents.robot.batching.BatchedModel`
+        scatters rows along it, so a squeezed ``[T, A]`` silently breaks both.
+        """
+        raise NotImplementedError
+
+    async def ainfer(self, batch: Any) -> ActionArray:
+        """Awaited single-rollout entry: run :meth:`infer` in a thread, return its single
+        ``[T, A]`` row. Indexing ``[0]`` assumes :meth:`infer` honors the ``[N, T, A]`` contract.
+        """
+        return (await asyncio.to_thread(self.infer, batch))[0]
+
+
+class LeRobotModel(Model):
+    """LeRobot policy with pre/post-processors: ``preprocess`` → ``predict_action_chunk`` →
+    ``postprocess``. ``preprocess`` adds the batch dim for an unbatched sample and is a no-op
+    for an already-stacked one, so :meth:`infer` handles both single and batched inputs.
+
+    Stateless: ``predict_action_chunk`` is a pure forward and the agent owns the open-loop
+    chunk, so LeRobot's internal action queue is never consumed here — hence no ``reset``.
+    """
+
+    def __init__(self, policy: Any, preprocess: Any, postprocess: Any) -> None:
+        self.policy = policy
+        self.preprocess = preprocess
+        self.postprocess = postprocess
+        #: Flipped to False after the first forward; used to print the one-time
+        #: CUDA/flow-matching warmup message.
+        self._first_inference = True
+
+    def infer(self, batch: Any) -> ActionArray:
+        """run batch dict (N dim) → [N, T, A] chunk"""
+        torch: Any = importlib.import_module("torch")
+        if self._first_inference:
+            print(
+                "[agent] first inference — flow-matching/CUDA warmup; this may take a while",
+                flush=True,
+            )
+        with torch.no_grad():
+            chunk = self.postprocess(self.policy.predict_action_chunk(self.preprocess(batch)))
+        if self._first_inference:
+            print("[agent] first inference done — inference is now fast", flush=True)
+            self._first_inference = False
+        arr = chunk.float().cpu().numpy()
+        assert arr.ndim == 3, (
+            f"expected [N, T, A] chunk, got {arr.shape}"
+        )  # LeRobot keeps the N dim
+        return arr
+
+
+class RemoteModel(Model):
+    """Weightless client to an OpenPI-WebSocket policy server: ships the adapter's request
+    dict, returns the server's chunk. All pre/post-processing lives in the adapter + server.
+
+    Not batchable: each :meth:`infer` is one WebSocket request for one env and always adds a
+    single leading batch dim, and the OpenPI server protocol currently has no batched-request
+    shape. Do not wrap in :class:`~hud.agents.robot.batching.BatchedModel` — use one
+    :class:`~hud.agents.robot.agent.RobotAgent` per concurrent rollout instead.
+    """
+
+    def __init__(
+        self, host: str = "localhost", port: int = 8000, *, response_key: str = "actions"
+    ) -> None:
+        self.host = host
+        self.port = port
+        #: Server chunk key — "actions" (stock OpenPI) or "action" (Cosmos).
+        self.response_key = response_key
+        self._client: Any = None
+
+    def connect(self) -> None:
+        """Open the websocket (idempotent); blocks until the server is up."""
+        if self._client is None:
+            mod: Any = importlib.import_module("openpi_client.websocket_client_policy")
+
+            print(
+                f"[agent] connecting to openpi server ws://{self.host}:{self.port} — on hold...",
+                flush=True,
+            )
+            self._client = mod.WebsocketClientPolicy(self.host, self.port)
+
+    def infer(self, batch: Any) -> ActionArray:
+        """Ship one request dict → the server's ``[T, A]`` chunk, returned as ``[1, T, A]``."""
+        self.connect()  # lazy connect on first call (blocks until the server is up)
+        chunk = np.asarray(self._client.infer(batch)[self.response_key], dtype=np.float32)
+        return chunk[None]  # add the leading N=1 batch dim
+
+
+__all__ = [
+    "LeRobotModel",
+    "Model",
+    "RemoteModel",
+]

hud_python-0.6.5/hud/agents/robot/record.py

@@ -0,0 +1,230 @@
+"""Per-episode recording for robot rollouts — telemetry, plus an optional LeRobot dataset.
+
+The agent loop hands every tick to one :class:`Recorder`. It always streams the telemetry
+the HUD viewer needs (an :class:`~hud.agents.types.ObservationStep` of numeric state +
+per-camera H.264 video); when ``save`` is on it *also* appends each
+``(observation, executed action)`` pair to a LeRobot v3 dataset for offline
+training/finetuning.
+
+Saving is opt-in (the agent's ``save`` flag — the ``--save`` runner flag), so the heavy
+LeRobot/PyAV imports stay deferred until a dataset is actually built. One dataset spans the
+whole run (every episode the shared agent drives appends to it) and is finalized at process
+exit, optionally pushed to the HF Hub. Destination + push come from the environment:
+
+- ``RECORD_DIR``  — dataset root (default ``./data`` from where the rollout launched)
+- ``HF_REPO``     — HF namespace to also push to (needs ``HF_TOKEN``)
+- ``HF_PRIVATE``  — push the dataset private
+"""
+
+from __future__ import annotations
+
+import atexit
+import importlib.util
+import logging
+import os
+import time
+import uuid
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+
+from hud.agents.types import InferenceStep, ObservationStep
+from hud.telemetry.context import get_current_trace_id
+
+from .video import VideoStreamer
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+    from hud.capabilities.robot import RobotClient
+    from hud.eval.run import Run
+
+logger = logging.getLogger(__name__)
+
+
+def _lerobot_features(contract: dict[str, Any]) -> tuple[dict[str, dict[str, Any]], dict[str, str]]:
+    """Map a robot contract to LeRobot ``features`` + a wire-key -> LeRobot-key map.
+
+    Image obs -> ``observation.images.<leaf>`` (video); the lone vector obs ->
+    ``observation.state`` (else ``observation.<leaf>``); the action -> ``action``. String
+    obs are dropped (LeRobot carries the prompt as its per-frame ``task``).
+    """
+    feats = contract.get("features", {})
+    vectors = [
+        n
+        for n, f in feats.items()
+        if f.get("role") == "observation" and f.get("dtype") not in ("image", "string")
+    ]
+    single_state = len(vectors) == 1
+
+    features: dict[str, dict[str, Any]] = {}
+    key_map: dict[str, str] = {}
+    for name, f in feats.items():
+        role, dtype, shape = f.get("role"), f.get("dtype"), tuple(f.get("shape") or ())
+        leaf = name.split("/")[-1]  # contract keys are slash-paths; LeRobot wants the leaf
+        if role == "observation" and dtype != "string":
+            if dtype == "image":
+                key, dtype = f"observation.images.{leaf}", "video"
+            elif leaf == "state" or single_state:
+                key = "observation.state"
+            else:
+                key = f"observation.{leaf}"
+            features[key] = {"dtype": dtype, "shape": shape, "names": _feature_names(f, leaf)}
+            key_map[name] = key
+        elif role == "action":
+            features["action"] = {"dtype": dtype, "shape": shape, "names": _feature_names(f, "act")}
+    return features, key_map
+
+
+def _feature_names(feature: dict[str, Any], base: str) -> list[str]:
+    """Contract per-element labels, else positional defaults sized to the (rank-1) shape."""
+    if names := feature.get("names"):
+        return list(names)
+    if feature.get("dtype") == "image":
+        return ["height", "width", "channel"]
+    return [f"{base}_{i}" for i in range(int((feature.get("shape") or [1])[0]))]
+
+
+class Recorder:
+    """Records one agent's rollouts: always telemetry, optionally a LeRobot dataset.
+
+    The agent owns a single instance for its lifetime and routes *all* recording through
+    it: :meth:`begin`/:meth:`end` bracket each episode, :meth:`record_observation` /
+    :meth:`record_inference` / :meth:`record_action` feed each tick (the first two write
+    telemetry steps onto the run passed to :meth:`begin`; the last completes a LeRobot
+    frame), and :meth:`save` (also an ``atexit`` hook) finalizes the cross-episode dataset.
+    With ``save=False`` only the telemetry path runs and the LeRobot deps are never imported.
+    """
+
+    def __init__(self, client: RobotClient, *, save: bool = False) -> None:
+        self._obs_space = client.spaces()[1]
+        self._fps = client.get_control_rate()
+        self._contract = client.contract
+        # Telemetry is always on; saving also needs lerobot installed.
+        if save and importlib.util.find_spec("lerobot") is None:
+            logger.warning(
+                "save=True but lerobot is not installed; streaming telemetry only "
+                "(pip install 'lerobot[dataset]')"
+            )
+            save = False
+        self._save = save
+        self._features: dict[str, dict[str, Any]] = {}
+        self._key_map: dict[str, str] = {}
+        if save:
+            self._features, self._key_map = _lerobot_features(self._contract)
+
+        self._video: VideoStreamer | None = None  # per-episode
+        self._run: Run | None = None
+        self._task = ""
+        self._pending: dict[str, Any] | None = None  # last obs awaiting its action
+        # LeRobot dataset spans every episode; created lazily on the first frame.
+        self._ds: Any | None = None
+        self._root: Path | None = None
+        self._repo_id = ""
+        if save:
+            atexit.register(self.save)  # finalize even on an abrupt exit (parquet footer)
+
+    # ── episode lifecycle (called from the agent harness) ─────────────────────
+    def begin(self, run: Run, prompt: str) -> None:
+        """Open an episode: fresh per-camera video stream + the task prompt."""
+        self._run = run
+        self._task = prompt
+        self._pending = None
+        self._video = VideoStreamer(fps=self._fps, trace_id=get_current_trace_id())
+
+    def record_observation(self, obs: dict[str, Any], *, tick: int) -> None:
+        """One observation: numeric-state span + per-camera video (always streamed)."""
+        assert self._run is not None and self._video is not None  # set in begin()
+        self._run.record(ObservationStep.from_obs(obs, tick=tick, obs_space=self._obs_space))
+        self._video.record(obs)
+        self._pending = obs.get("data")  # paired with the action in record_action()
+
+    def record_inference(self, chunk: NDArray[Any], *, tick: int) -> None:
+        """One re-inference: the freshly inferred ``[T, A]`` action chunk, onto the run."""
+        assert self._run is not None  # set in begin()
+        self._run.record(InferenceStep(tick=tick, chunk=chunk.tolist(), chunk_length=len(chunk)))
+
+    def record_action(self, action: NDArray[Any]) -> None:
+        """The executed (env-space) action: completes the pending LeRobot frame."""
+        if self._save and self._pending is not None:
+            self._add_frame(self._pending, action)
+        self._pending = None
+
+    def end(self) -> None:
+        """Close the episode: flush video tails; commit the LeRobot episode (if any frames)."""
+        if self._video is not None:
+            self._video.finalize()
+        if self._ds is not None and self._ds.has_pending_frames():
+            self._ds.save_episode()
+
+    def save(self) -> None:
+        """Finalize the dataset (writes the parquet footer) + optionally push to the Hub.
+
+        Idempotent; registered with ``atexit`` so the dataset stays loadable even if the
+        process exits without an explicit call.
+        """
+        if not self._save or self._ds is None:
+            return
+        self._save = False  # idempotent across the explicit call + the atexit hook
+        self._ds.finalize()
+        print(f"[agent] saved LeRobot dataset -> {self._root}", flush=True)
+        if not os.environ.get("HF_REPO"):
+            return
+        private = os.environ.get("HF_PRIVATE", "0") not in ("0", "", "false", "False")
+        try:  # best-effort: the on-disk dataset is the source of truth
+            self._ds.push_to_hub(private=private)
+            print(f"[agent] pushed -> https://huggingface.co/datasets/{self._repo_id}", flush=True)
+        except Exception as exc:
+            logger.exception("HF push failed for %s", self._repo_id)
+            print(f"[agent] WARNING: HF push failed: {exc!r} (dataset still on disk)", flush=True)
+
+    # ── LeRobot writing ───────────────────────────────────────────────────────
+    def _add_frame(self, data: dict[str, Any], action: NDArray[Any]) -> None:
+        ds = self._ensure_dataset()
+        row: dict[str, Any] = {}
+        for wire, key in self._key_map.items():
+            value = data.get(wire)
+            if value is None:
+                logger.warning("obs missing contract feature %r; skipping frame", wire)
+                return
+            ft = self._features[key]
+            row[key] = (
+                np.ascontiguousarray(value, dtype=np.uint8)  # bridge images are uint8 HWC
+                if ft["dtype"] in ("video", "image")
+                else np.asarray(value, dtype=ft["dtype"]).reshape(ft["shape"])
+            )
+        act_ft = self._features["action"]
+        row["action"] = np.asarray(action, dtype=act_ft["dtype"]).reshape(act_ft["shape"])
+        row["task"] = self._task
+        ds.add_frame(row)
+
+    def _ensure_dataset(self) -> Any:
+        if self._ds is not None:
+            return self._ds
+        lerobot_dataset: Any = importlib.import_module("lerobot.datasets.lerobot_dataset")
+
+        name = self._contract.get("robot_type") or "robot"
+        stamp = time.strftime("%Y%m%d_%H%M%S")
+        # Unique per recorder so concurrent (batched) rollouts never share a root;
+        # tie it to the trace id when there is one so a shard maps back to its trace.
+        tag = (get_current_trace_id() or uuid.uuid4().hex)[:8]
+        # Default under ./data (relative to where the rollout was launched), created if absent.
+        record_dir = Path(os.environ.get("RECORD_DIR", "data"))
+        record_dir.mkdir(parents=True, exist_ok=True)
+        self._root = record_dir / f"{name}_{stamp}_{tag}"
+        self._repo_id = f"{os.environ.get('HF_REPO') or 'hud'}/{name}_{stamp}_{tag}"
+        # LeRobotDataset.create requires a fresh root; images encode to per-episode video.
+        self._ds = lerobot_dataset.LeRobotDataset.create(
+            repo_id=self._repo_id,
+            fps=self._fps,
+            features=self._features,
+            root=self._root,
+            robot_type=self._contract.get("robot_type"),
+            use_videos=True,
+        )
+        print(f"[agent] recording LeRobot dataset -> {self._root}", flush=True)
+        return self._ds
+
+
+__all__ = ["Recorder"]