feat(robot,eval): unified robot stack + Shared runtime for concurrent rollouts

docs/v6/cookbooks/robot-benchmark.mdx

@@ -29,17 +29,17 @@ endpoint = RobotEndpoint(LiberoSimBridge(use_delta=True))  # drive the bridge th
 @env.initialize
 async def _up():
     await endpoint.start()
-    env.add_capability(await endpoint.capability(contract=CONTRACT))
+    env.add_capability(await endpoint.capability())
 
 @env.shutdown
 async def _down():
     await endpoint.stop()
 
 @env.template(id="libero_spatial")
 async def libero_spatial(libero_task_id: int, init_state_id: int = 0):
-    prompt = await endpoint.reset(task_suite="libero_spatial",
-                                  task_id=libero_task_id, init_state_id=init_state_id)
-    yield {"prompt": prompt}
+    ep = await endpoint.reset(task_suite="libero_spatial",
+                              task_id=libero_task_id, init_state_id=init_state_id)
+    yield {"prompt": ep["prompt"]}
     yield await endpoint.result()
 ```
 
@@ -67,8 +67,6 @@ from hud.eval import DockerRuntime, Task, Taskset
 CHECKPOINT = "lerobot/pi05_libero_finetuned"
 
 class PI05Agent(RobotAgent):
-    max_steps = 400
-
     def __init__(self):
         device = "cuda" if torch.cuda.is_available() else "cpu"
         policy = PI05Policy.from_pretrained(CHECKPOINT).to(device).eval()

docs/v6/reference/runtime.mdx

@@ -24,6 +24,7 @@ await taskset.run(agent, runtime=LocalRuntime("env.py"))   # serve env.py locall
 | `ModalRuntime("my-env")` | A fresh [Modal](https://modal.com/) sandbox per rollout | Cloud scale, no infra to manage |
 | `DaytonaRuntime("my-env")` | A fresh [Daytona](https://www.daytona.io/) sandbox per rollout | Cloud scale on Daytona |
 | `Runtime("tcp://host:8765")` | A substrate you already started | Attaching to a long-lived container or sandbox you own |
+| `Shared(provider, width=N)` | One substrate from `provider`, shared by up to N concurrent connections | Multi-agent evals - agents that share, compete, or coordinate over one live environment |
 | `HUDRuntime()` | A HUD-hosted env, leased by name and tunneled | Local agent loop against a deployed env |
 | `HostedRuntime()` | The whole rollout on a HUD-leased box | Agent and env run together off your machine |
 
@@ -182,6 +183,23 @@ Runtime(url, params=..., config=...)
 - **`url`** - control-channel address of an already-running substrate (e.g. `tcp://host:8765`).
 - **`params`** - connection-time data a transport may need (auth token, sandbox id).
 
+### `Shared`
+
+```python
+Shared(provider, *, width)
+```
+
+- **`provider`** - the `Provider` to provision once, e.g. `DockerRuntime("my-env")`.
+- **`width`** - how many concurrent connections that one substrate accepts.
+
+The first rollout to acquire it provisions through `provider`; every later rollout is just handed
+that same address instead of provisioning its own, so up to `width` rollouts end up connected to
+one live substrate - agents that need to share, compete, or coordinate in the same environment
+(a vectorized sim's `num_envs` slots, see [robots](/v6/advanced/robots#vectorized-envs-and-evals),
+are one case). Pair it with `taskset.run(..., group=width, max_concurrent=width)` so exactly
+`width` rollouts connect at once. `Runtime(url)` needs no wrapping since every caller already
+dials the same address.
+
 ## Run on your own infra
 
 A **runtime is just a function**: given a task, start a container somewhere and yield its

hud/__init__.py

@@ -60,6 +60,7 @@
     "instrument",
 ]
 
+
 try:
     from .version import __version__
 except ImportError:

hud/agents/robot/__init__.py

@@ -1,41 +1,39 @@
-"""Robot agent harness: drive a ``robot`` capability with a policy.
-
-The harness splits a policy rollout into three seams, each replaceable on its own:
-
-- :class:`~hud.agents.robot.agent.RobotAgent` — the loop: connect to the env's
-  ``robot`` capability, observe, act, stop.
-- :class:`~hud.agents.robot.model.Model` — *how to run* the policy (preprocess →
-  forward → postprocess). :class:`~hud.agents.robot.model.LeRobotModel` ships the
-  LeRobot checkpoint convention.
-- :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.
-
-This subpackage needs the ``robot`` extra (``pip install 'hud[robot]'``) for
-``numpy`` + ``msgpack``; importing :mod:`hud.agents` alone never pulls them in.
+"""Agent-side robot harness: drive a ``robot`` env with a VLA policy.
+
+- :class:`~.agent.RobotAgent` — the harness: connects to the ``robot``
+  capability (claiming a slot token from ``run.started``), reads the contract,
+  drives one scalar connection with an open-loop chunk queue. Subclass and
+  set ``self.model`` + ``self.adapter``.
+- :class:`~.model.Model` / :class:`~.model.LeRobotModel` /
+  :class:`~.model.RemoteModel` — the policy and its inference mechanics.
+- :class:`~.adapter.Adapter` / :class:`~.adapter.LeRobotAdapter` /
+  :class:`~.adapter.OpenPIAdapter` — env <-> policy space translation.
+- :class:`~.batching.BatchedAgent` — many concurrent single-env rollouts
+  sharing one batched model.
+- :class:`~.dataset.DatasetWriter` — opt-in LeRobot v3 dataset recording
+  (``agent.save = True``).
+
+This subpackage needs the ``robot`` extra (``pip install 'hud[robot]'``).
 """
 
 from __future__ import annotations
 
 from .adapter import Adapter, LeRobotAdapter, OpenPIAdapter
 from .agent import ROBOT_PROTOCOL, RobotAgent
 from .batching import BatchedAgent, BatchedModel
-from .model import LeRobotModel, Model
+from .dataset import DatasetWriter
+from .model import LeRobotModel, Model, RemoteModel
 
 __all__ = [
     "ROBOT_PROTOCOL",
     "Adapter",
     "BatchedAgent",
     "BatchedModel",
+    "DatasetWriter",
     "LeRobotAdapter",
     "LeRobotModel",
     "Model",
     "OpenPIAdapter",
+    "RemoteModel",
     "RobotAgent",
 ]

hud/agents/robot/adapter.py

@@ -1,56 +1,59 @@
 """Translate observations and actions between env and policy spaces.
 
-The loop calls ``bind``, ``reset``, ``adapt_observation``, and ``adapt_action``.
+The loop calls ``bind``, ``reset``, ``adapt_observation``, and the action hooks.
 Use :class:`LeRobotAdapter` for LeRobot models; subclass for custom wiring;
 ``adapter=None`` for pass-through.
 """
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 import numpy as np
+from numpy.typing import NDArray
 
-if TYPE_CHECKING:
-    from ._types import ActionArray
+#: A policy-emitted action / chunk array (the robot stack's shared alias).
+ActionArray = NDArray[np.floating[Any]]
 
-# ─── the abstraction ──────────────────────────────────────────────────────────
+#: Image types the bundled adapters treat as cameras (vs the state vector).
+IMAGE_TYPES = ("rgb", "bgr", "gray", "depth")
 
 
 class Adapter:
     """Translate between an env's observation/action spaces and a policy's.
 
-    Driven by :class:`~hud.agents.robot.agent.RobotAgent`: :meth:`bind` once after
-    connect, :meth:`reset` once per episode, then :meth:`adapt_observation` /
-    :meth:`adapt_action` each step. Construct with the policy's image-slot names;
-    everything env-side is learned in :meth:`bind`.
+    Driven by :class:`~hud.agents.robot.agent.RobotAgent`: :meth:`bind` once
+    after connect, :meth:`reset` once per episode, then :meth:`adapt_observation`
+    per inference and the action hooks per chunk/step. Construct with the
+    policy's image-slot names; everything env-side is learned in :meth:`bind`.
     """
 
-    def __init__(self, *, model_image_keys: list[str] | None = None) -> None:
+    def __init__(
+        self,
+        *,
+        model_image_keys: list[str] | None = None,
+        chunk_size: int | None = None,
+    ) -> None:
         #: The policy's ordered image-slot names (model side; known at load time).
         self.model_image_keys: list[str] = list(model_image_keys or [])
-        #: The env's selected action feature (set in :meth:`bind`).
+        #: Open-loop horizon: how many predicted actions to execute before re-querying.
+        #: ``None`` = full model chunk. Applied by :class:`~.agent.RobotAgent`, not here —
+        #: so subclass ``adapt_chunk`` overrides stay free of truncation duty.
+        self.chunk_size: int | None = chunk_size
+        #: The env's action feature and observation layout (set in :meth:`bind`).
         self.action_space: dict[str, Any] = {}
-        #: The env's image / state observation keys (set in :meth:`bind`).
         self.image_keys: list[str] = []
         self.state_key: str | None = None
 
     def bind(self, action_space: dict[str, Any], observation_space: dict[str, Any]) -> None:
-        """Learn the env's layout from the contract (``client.spaces()``).
-
-        Splits observation features into image keys vs the single state key and stores
-        the action feature. Override to derive extra env-side parameters.
-        """
-        # TODO CLEAN
+        """Learn the env's layout from the contract (``client.spaces()``): image
+        features become the camera keys (in contract order), the first non-image
+        feature is the state. Override to derive extra env-side parameters."""
         self.action_space = action_space or {}
-        image_types = ("rgb", "bgr", "gray", "depth")
-        self.image_keys = []
-        self.state_key = None
-        for name, feature in observation_space.items():
-            if feature.get("type") in image_types:
-                self.image_keys.append(name)
-            elif self.state_key is None:
-                self.state_key = name
+        self.image_keys = [n for n, f in observation_space.items() if f.get("type") in IMAGE_TYPES]
+        self.state_key = next(
+            (n for n, f in observation_space.items() if f.get("type") not in IMAGE_TYPES), None
+        )
 
     def reset(self) -> None:
         """Override only if the adapter is stateful across steps within an episode."""
@@ -59,47 +62,57 @@ def adapt_observation(self, obs: dict[str, Any], prompt: str) -> Any:
         """Translate an env observation + task prompt into the policy's input."""
         raise NotImplementedError
 
+    def adapt_chunk(self, chunk: ActionArray, obs: dict[str, Any]) -> ActionArray:
+        """Translate a freshly-inferred ``[T, A]`` chunk to env space, given the
+        (per-slot) query-time observation it was inferred from (default identity).
+
+        Called once per slot at inference time, so a chunk expressed relative to
+        the query state — e.g. joint *deltas* to be added to the query-time
+        joints — converts in one shot.
+        """
+        return chunk
+
     def adapt_action(self, action: ActionArray, obs: dict[str, Any]) -> ActionArray:
-        """Translate a policy action into the env's action space (default identity)."""
+        """Per-step execution-time hook on the popped action (default identity)."""
         return action
 
 
 class LeRobotAdapter(Adapter):
-    """Vanilla LeRobot adapter for a standard image/state env.
+    """Vanilla LeRobot adapter for a standard image/state env, single or batched.
 
-    Maps env cameras onto the model's image slots in order, converts HWC ``uint8`` to
-    CHW ``float`` in ``[0, 1]``, and passes state + prompt through. Actions are identity
-    (postprocess already returns env-space actions); subclass for resize/pad/reshaping.
+    Maps env cameras onto the model's image slots in order and converts HWC
+    ``uint8`` to CHW ``float`` in ``[0, 1]``; state and prompt pass through.
+    A batched observation (state ``[N, S]``) maps in one go — cameras become
+    ``[N, C, H, W]`` and the shared task is repeated to ``N``. Actions are
+    identity (postprocess already returns env-space actions).
     """
 
     def adapt_observation(self, obs: dict[str, Any], prompt: str) -> dict[str, Any]:
         import torch  # pyright: ignore[reportMissingImports]
 
-        torch_mod: Any = torch
+        torch_mod: Any = torch  # torch ships no stubs; keep strict mode quiet
         data = obs["data"]
+        state = np.asarray(data[self.state_key], dtype=np.float32)
+        batched = state.ndim > 1  # [N, S] vs [S]
         batch: dict[str, Any] = {
-            "observation.state": torch_mod.from_numpy(data[self.state_key].astype(np.float32)),
-            "task": prompt,
+            "observation.state": torch_mod.from_numpy(state),
+            "task": [prompt] * len(state) if batched else prompt,
         }
         for model_key, env_key in zip(self.model_image_keys, self.image_keys, strict=False):
-            batch[model_key] = torch_mod.from_numpy(data[env_key]).permute(2, 0, 1).float() / 255.0
+            img = torch_mod.from_numpy(np.asarray(data[env_key]))
+            perm = (0, 3, 1, 2) if batched else (2, 0, 1)
+            batch[model_key] = img.permute(*perm).float() / 255.0
         return batch
 
-    def adapt_action(self, action: ActionArray, obs: dict[str, Any]) -> ActionArray:
-        return action
-
 
 class OpenPIAdapter(Adapter):
-    """unwraps obs['data'] to OpenPI wire keys, attaches prompt; actions are passthrough"""
+    """Unwraps ``obs['data']`` to OpenPI wire keys and attaches the prompt;
+    actions are pass-through."""
 
     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",
-]
+__all__ = ["ActionArray", "Adapter", "LeRobotAdapter", "OpenPIAdapter"]