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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hud-python
-Version: 0.6.3
+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.5/cookbooks/connect4-selfplay/README.md

@@ -0,0 +1,57 @@
+# Connect Four self-play
+
+Symmetric self-play RL on a 6×7 Connect Four board. Draws are rare (you need a
+full 42-cell board with no four-in-a-row), so the win/loss reward signal
+persists as the policy improves and the GRPO advantage stays non-zero.
+
+## How it works
+
+- One agent ("outer") plays a full game against an inner model on the **same
+  slug** — true self-play. `seed % 2` decides who drops first, for symmetric
+  first-move coverage.
+- Each game trains **both sides at once**: the outer agent's `Run` (reward from
+  its perspective) plus a hand-built `TrajectoryPayload` for the inner model
+  with the flipped reward (`1 - outer_reward`).
+- `group_size=2` pairs each game's two trajectories so the GRPO advantage is
+  `reward - 0.5` per game.
+- `loss_fn="ppo"` clips the importance-sampling ratio, so a single lucky game
+  can't blow up the update.
+
+The training loop uses the public API directly — `forward_backward` accepts
+`Run` and `TrajectoryPayload` mixed, so no private helpers are needed.
+
+## Setup
+
+```bash
+hud models fork Qwen/Qwen3.5-4B --name c4-selfplay   # prints a slug like c4-selfplay-<id>
+```
+
+Put your `HUD_API_KEY` in a `.env` here (or the environment).
+
+## Run
+
+Local sanity check (one game, cheap external model as the outer agent):
+
+```bash
+hud eval env.py claude --model claude-haiku-4-5
+```
+
+Train:
+
+```bash
+python train.py --model c4-selfplay-<id> --steps 20 --group 4 --lr 1e-5
+```
+
+## Tuning notes
+
+- **Memory scales with `tasks × group`.** Each task×rollout is a fresh `env.py`
+  subprocess. With 8 tasks and `--group 4` that's 32 concurrent games. Connect
+  Four games can run up to 42 plies, so they cost more tokens and time per game —
+  start at `--group 4` and raise only if you have RAM headroom.
+- **Watch the server-side metrics.** The loop prints local win/draw/loss counts
+  each step and the last few checkpoints' `mean_reward` / `reward_std` via
+  `trainer.checkpoints()` at the end. A healthy run keeps non-trivial
+  `reward_std` (within-group spread); if it collapses, the policy has saturated.
+- **Reset on changes.** If you edit the reward or the board, roll the head back
+  to a clean checkpoint (`hud models head <slug> --set <id>`) or fork fresh —
+  don't keep training a policy shaped by the old objective.

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

@@ -8,7 +8,12 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any, cast
 
 from hud.types import AgentType
-from hud.utils.gateway import build_gateway_client, list_gateway_models
+from hud.utils.gateway import (
+    build_gateway_client,
+    gateway_model_aliases,
+    list_gateway_models,
+    normalize_gateway_model_id,
+)
 
 if TYPE_CHECKING:
     from typing import TypeAlias
@@ -27,6 +32,8 @@ def create_agent(model: str, **kwargs: Any) -> GatewayAgent:
 
     For direct API access with provider API keys, instantiate the agent classes directly.
     """
+    requested_model = model
+    model = normalize_gateway_model_id(model)
     agent_type = next((candidate for candidate in AgentType if candidate.value == model), None)
     if agent_type is not None:
         model_id = model
@@ -73,7 +80,8 @@ def create_agent(model: str, **kwargs: Any) -> GatewayAgent:
                 for n in (gm.id, gm.name, gm.model_name)
                 if isinstance(n, str)
             ]
-            near = difflib.get_close_matches(model, known, n=3, cutoff=0.5)
+            known.extend(gateway_model_aliases())
+            near = difflib.get_close_matches(requested_model, known, n=3, cutoff=0.5)
             hint = (
                 f" Did you mean: {', '.join(near)}?"
                 if near
@@ -84,7 +92,7 @@ def create_agent(model: str, **kwargs: Any) -> GatewayAgent:
                 if gateway_models
                 else "the HUD gateway registry (empty — is HUD_API_KEY set?)"
             )
-            raise ValueError(f"Model {model!r} not found in {source}.{hint}")
+            raise ValueError(f"Model {requested_model!r} not found in {source}.{hint}")
 
     kwargs.setdefault("model", model_id)
     kwargs.setdefault("model_client", build_gateway_client(provider_name))

{hud_python-0.6.3 → hud_python-0.6.5}/hud/agents/openai_compatible/agent.py

@@ -193,16 +193,27 @@ class OpenAIChatAgent(ToolAgent[ChatCompletionMessageParam, OpenAIChatConfig]):
         sample: Sample | None = None
         if return_token_ids:
             prompt_token_ids = getattr(choice, "prompt_token_ids", None)
+            # Multimodal prompt (text + image chunks): the only prompt representation
+            # that survives image inputs; flat prompt_token_ids is null in that case.
+            prompt_chunks = getattr(choice, "prompt_chunks", None)
             token_ids = getattr(choice, "token_ids", None)
-            if prompt_token_ids is not None and token_ids is not None:
-                chat_state.continuation_token_ids = list(prompt_token_ids) + list(token_ids)
-                chat_state.continuation_message_count = len(messages)
+            has_prompt = prompt_token_ids is not None or prompt_chunks is not None
+            if token_ids is not None and has_prompt:
                 content_lp = choice.logprobs.content if choice.logprobs else None
                 sample = Sample(
-                    prompt_token_ids=list(prompt_token_ids),
+                    prompt_token_ids=list(prompt_token_ids) if prompt_token_ids is not None else [],
+                    prompt_chunks=list(prompt_chunks) if prompt_chunks is not None else None,
                     output_token_ids=list(token_ids),
                     output_logprobs=[tok.logprob for tok in content_lp] if content_lp else [],
                 )
+                # KV-cache continuation only applies to flat text prompts; clear any
+                # stale state when the gateway returns chunks-only (multimodal turn).
+                if prompt_token_ids is not None:
+                    chat_state.continuation_token_ids = list(prompt_token_ids) + list(token_ids)
+                    chat_state.continuation_message_count = len(messages)
+                else:
+                    chat_state.continuation_token_ids = None
+                    chat_state.continuation_message_count = None
 
         tool_calls: list[MCPToolCall] = []
         for tc in function_calls:

{hud_python-0.6.3 → 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.3 → 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.3 → 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",
+]