From 3966c3186ea43d60097fec8264f515da81a8145e Mon Sep 17 00:00:00 2001 From: lukass16 Date: Wed, 1 Jul 2026 21:51:30 +0000 Subject: [PATCH 01/43] wip: save robot agent/recorder/bridge state before vectorized taskset work --- hud/__init__.py | 3 + hud/agents/robot/__init__.py | 8 +- hud/agents/robot/adapter.py | 33 ++++ hud/agents/robot/agent.py | 6 +- hud/agents/robot/record.py | 304 +++++++++++++++++++++++++++++- hud/agents/robot/vec_agent.py | 123 ++++++++++++ hud/capabilities/robot.py | 4 +- hud/environment/robot/__init__.py | 4 +- hud/environment/robot/bridge.py | 154 +++++++++++++-- 9 files changed, 611 insertions(+), 28 deletions(-) create mode 100644 hud/agents/robot/vec_agent.py diff --git a/hud/__init__.py b/hud/__init__.py index 79c520916..a1854a682 100644 --- a/hud/__init__.py +++ b/hud/__init__.py @@ -8,6 +8,7 @@ # Apply patches to third-party libraries early, before other imports from . import patches as _patches # noqa: F401 from ._legacy import install as _install_v5_compat +from .agents.robot.record import Recorder, VecRecorder from .clients import connect from .environment import Environment from .eval import ( @@ -44,6 +45,7 @@ "HostedRuntime", "Job", "LocalRuntime", + "Recorder", "Run", "Runtime", "RuntimeConfig", @@ -56,6 +58,7 @@ "Taskset", "Trace", "TrainingClient", + "VecRecorder", "connect", "instrument", ] diff --git a/hud/agents/robot/__init__.py b/hud/agents/robot/__init__.py index cba9db77b..58ad8f89f 100644 --- a/hud/agents/robot/__init__.py +++ b/hud/agents/robot/__init__.py @@ -23,10 +23,11 @@ from __future__ import annotations -from .adapter import Adapter, LeRobotAdapter, OpenPIAdapter +from .adapter import Adapter, LeRobotAdapter, OpenPIAdapter, VecLeRobotAdapter from .agent import ROBOT_PROTOCOL, RobotAgent from .batching import BatchedAgent, BatchedModel -from .model import LeRobotModel, Model +from .model import LeRobotModel, Model, RemoteModel +from .vec_agent import VecRobotAgent __all__ = [ "ROBOT_PROTOCOL", @@ -37,5 +38,8 @@ "LeRobotModel", "Model", "OpenPIAdapter", + "RemoteModel", "RobotAgent", + "VecLeRobotAdapter", + "VecRobotAgent", ] diff --git a/hud/agents/robot/adapter.py b/hud/agents/robot/adapter.py index 08c5fca72..d5451f777 100644 --- a/hud/agents/robot/adapter.py +++ b/hud/agents/robot/adapter.py @@ -63,6 +63,16 @@ def adapt_action(self, action: ActionArray, obs: dict[str, Any]) -> ActionArray: """Translate a policy action into the env's action space (default identity).""" return action + def adapt_chunk(self, chunk: ActionArray, obs: dict[str, Any]) -> ActionArray: + """Translate a freshly-inferred ``[T, A]`` chunk to env space, given the query-time + observation it was inferred from (default identity). + + The vectorized harness calls this once per slot at inference time (not per step), so a + chunk expressed relative to the query state — e.g. DROID joint *deltas* that must be + added to the query-time joints for absolute targets — can be converted in one shot. + """ + return chunk + class LeRobotAdapter(Adapter): """Vanilla LeRobot adapter for a standard image/state env. @@ -89,6 +99,28 @@ def adapt_action(self, action: ActionArray, obs: dict[str, Any]) -> ActionArray: return action +class VecLeRobotAdapter(LeRobotAdapter): + """Batched :class:`LeRobotAdapter` for a vectorized env (:class:`~hud.agents.robot.vec_agent.VecRobotAgent`). + + Same wiring, but the obs arrays carry a leading ``N`` and the whole batch maps in one go: + state stays ``[N, S]``, each camera ``[N, H, W, C]`` uint8 becomes ``[N, C, H, W]`` float in + ``[0, 1]``, and the shared task is repeated to ``N`` (one prompt per env in the batch). + """ + + def adapt_observation(self, obs: dict[str, Any], prompt: str) -> dict[str, Any]: + import torch # pyright: ignore[reportMissingImports] + + data = obs["data"] + n = len(np.asarray(data[self.state_key])) + batch: dict[str, Any] = { + "observation.state": torch.from_numpy(np.asarray(data[self.state_key], dtype=np.float32)), + "task": [prompt] * n, + } + for model_key, env_key in zip(self.model_image_keys, self.image_keys, strict=False): + batch[model_key] = torch.from_numpy(np.asarray(data[env_key])).permute(0, 3, 1, 2).float() / 255.0 + return batch + + class OpenPIAdapter(Adapter): """unwraps obs['data'] to OpenPI wire keys, attaches prompt; actions are passthrough""" @@ -102,4 +134,5 @@ def adapt_observation(self, obs: dict[str, Any], prompt: str) -> dict[str, Any]: "Adapter", "LeRobotAdapter", "OpenPIAdapter", + "VecLeRobotAdapter", ] diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py index 0b05310f4..fb4a11a8f 100644 --- a/hud/agents/robot/agent.py +++ b/hud/agents/robot/agent.py @@ -26,7 +26,7 @@ from hud.agents.base import Agent from hud.capabilities.robot import RobotClient -from .record import Recorder +from .record import EpisodeRecorder if TYPE_CHECKING: from hud.eval.run import Run @@ -78,7 +78,7 @@ class RobotAgent(Agent): _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 + _recorder: EpisodeRecorder | None = None def setup_robot(self, client: RobotClient) -> None: """Discover the env's action/observation layout and bind the adapter to it.""" @@ -98,7 +98,7 @@ def on_episode_start(self, run: Run, client: RobotClient, *, prompt: str) -> Non # 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 = EpisodeRecorder(client, save=self.save) self._recorder.begin(run, prompt) if self.adapter is not None: self.adapter.reset() diff --git a/hud/agents/robot/record.py b/hud/agents/robot/record.py index 2a053710a..ece3c117f 100644 --- a/hud/agents/robot/record.py +++ b/hud/agents/robot/record.py @@ -1,10 +1,18 @@ -"""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. +"""Recording for robot/sim rollouts: telemetry streaming, plus an optional LeRobot dataset. + +Three recorders, smallest to largest, all streaming the *same* robot telemetry the HUD +viewer plays (numeric :class:`~hud.agents.types.ObservationStep` state + per-camera H.264 +video) via the existing exporter: + +- :class:`Recorder` — one trace. Emits spans with an *explicit* ``trace_id`` (no rollout + contextvar, no ``Run``), so it works from a bare synchronous loop or any thread. The + primitive the vectorized recorder is built from. +- :class:`VecRecorder` — a whole vectorized env as one **Job** of per-episode traces. One + :class:`Recorder` per recorded env slot; on each ``done[i]`` (auto-reset) it closes that + slot's trace and opens a fresh one, so every episode becomes its own fully-saved trace. +- :class:`EpisodeRecorder` — the agent-loop recorder: records onto a :class:`Run` (so steps + land on ``run.trace`` for grading/training) and, when ``save`` is on, *also* appends each + ``(observation, executed action)`` pair to a LeRobot v3 dataset for offline training. 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 @@ -14,6 +22,10 @@ - ``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 + +Everything HUD-specific (job/trace lifecycle, trace-id assignment, span building, video +encoding, upload) lives here; a benchmark just does ``VecRecorder(...)`` -> ``record`` -> +``close``. """ from __future__ import annotations @@ -29,12 +41,15 @@ import numpy as np -from hud.agents.types import InferenceStep, ObservationStep +from hud.agents.types import InferenceStep, ObservationStep, StateFeature from hud.telemetry.context import get_current_trace_id +from hud.utils.platform import PlatformClient from .video import VideoStreamer if TYPE_CHECKING: + from typing import Self + from numpy.typing import NDArray from hud.capabilities.robot import RobotClient @@ -43,6 +58,34 @@ logger = logging.getLogger(__name__) +def _reporting_enabled() -> bool: + from hud.settings import settings + + return bool(settings.telemetry_enabled and settings.api_key) + + +def _report_sync(path: str, payload: dict[str, Any]) -> None: + """Best-effort sync POST to the platform (job/trace lifecycle). No-op without an API key. + + Mirrors the async ``hud.eval.job`` reporting so a sync vec-env loop never needs an event + loop, and never fails the run when the platform is unreachable. + """ + if not _reporting_enabled(): + return + body = {k: v for k, v in payload.items() if v is not None} + try: + PlatformClient.from_settings().post(path, json=body) + except Exception as exc: # reporting is fire-and-forget + logger.warning("platform report %s failed: %s", path, exc) + + +def _to_numpy(x: Any) -> NDArray[Any]: + """Coerce a torch tensor / array / scalar to a numpy array (no torch dependency).""" + if hasattr(x, "detach"): # torch.Tensor + x = x.detach().cpu().numpy() + return np.asarray(x) + + 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. @@ -86,7 +129,7 @@ def _feature_names(feature: dict[str, Any], base: str) -> list[str]: return [f"{base}_{i}" for i in range(int((feature.get("shape") or [1])[0]))] -class Recorder: +class EpisodeRecorder: """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 @@ -227,4 +270,245 @@ def _ensure_dataset(self) -> Any: return self._ds -__all__ = ["Recorder"] +# ── lightweight telemetry recorders (vec-env / bare-loop, no Run, no LeRobot) ──────────── + + +def _observation_step(obs: dict[str, Any], *, tick: int) -> ObservationStep: + """Build an :class:`ObservationStep` (numeric state) from a flat per-slot obs dict. + + Camera frames travel as H.264 video, so array obs with rank >= 2 are skipped; flat + numeric vectors become unlabelled :class:`StateFeature`s. Robust to scalars (unlike + :meth:`ObservationStep.from_obs`, which expects a robot-contract ``obs['data']``). + """ + state: dict[str, StateFeature] = {} + for name, val in obs.items(): + arr = _to_numpy(val) + if arr.ndim >= 2: + continue + flat = np.atleast_1d(arr).astype(float).ravel() + if 0 < flat.size <= 512: + state[name] = StateFeature(values=flat.tolist()) + return ObservationStep(tick=tick, state=state) + + +class Recorder: + """Streams one trace's telemetry to the platform: state, per-camera video, actions. + + Standalone: emits spans with an explicit ``trace_id`` rather than relying on the rollout + contextvar, so it runs inside a plain (or vectorized) loop. Reuses the robot step schema + and :class:`~hud.agents.robot.video.VideoStreamer`, so the existing exporter and trace + viewer ingest it unchanged. + """ + + def __init__( + self, + *, + trace_id: str, + job_id: str | None = None, + group_id: str | None = None, + fps: int = 10, + model: str | None = None, + metadata: dict[str, Any] | None = None, + enter: bool = True, + ) -> None: + self.trace_id = trace_id + self.job_id = job_id + self._fps = fps + self._tick = 0 + self._reward = 0.0 + self._metadata = metadata or {} + self._video: VideoStreamer | None = None # lazy, one per trace + self._closed = False + if enter: + _report_sync( + f"/trace/{trace_id}/enter", + {"job_id": job_id, "group_id": group_id, "model": model}, + ) + + def add_reward(self, reward: float) -> None: + """Accumulate per-step reward into this trace's total (reported on close).""" + self._reward += float(reward) + + def record( + self, + *, + obs: dict[str, Any] | None = None, + frames: dict[str, Any] | None = None, + action: Any | None = None, + ) -> None: + """Record one tick: numeric-state span, per-camera video, and the executed action. + + ``obs`` is this slot's observation dict (flat numeric vectors); ``frames`` maps a + camera name to its ``HxWxC`` frame; ``action`` is the executed action vector. + """ + if obs is not None: + _observation_step(obs, tick=self._tick).emit(trace_id=self.trace_id) + if frames: + if self._video is None: + self._video = VideoStreamer(fps=self._fps, trace_id=self.trace_id) + self._video.record({"data": frames}) + if action is not None: + chunk = _to_numpy(action).astype(float).ravel().tolist() + InferenceStep(tick=self._tick, chunk=[chunk], chunk_length=1).emit( + trace_id=self.trace_id + ) + self._tick += 1 + + def close( + self, + *, + status: str = "completed", + reward: float | None = None, + error: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> None: + """Flush video tails and report the trace exit (status / reward / metadata).""" + if self._closed: + return + self._closed = True + if self._video is not None: + self._video.finalize() + meta = {**self._metadata, **(metadata or {})} + _report_sync( + f"/trace/{self.trace_id}/exit", + { + "status": status, + "reward": self._reward if reward is None else reward, + "error": error, + "metadata": meta or None, + }, + ) + + +class VecRecorder: + """Records a vectorized env as one Job of per-episode traces. + + Construct it once with the batch size, call :meth:`record` after every ``env.step`` with + the batched tensors, and :meth:`close` at the end. A configurable subset of slots + (``record_indices``) gets rich traces (state + video); each ``done[i]`` closes that slot's + current trace and opens a new one, so a slot that plays many episodes produces many + traces — all under one Job at :attr:`job_url`. + """ + + def __init__( + self, + name: str, + num_envs: int, + *, + record_indices: list[int] | None = None, + fps: int = 10, + seed: int | None = None, + group_id: str | None = None, + model: str | None = None, + job_id: str | None = None, + ) -> None: + self.name = name + self.num_envs = num_envs + self.fps = fps + self.seed = seed + self.group_id = group_id + self.model = model + self.job_id = job_id or uuid.uuid4().hex + if record_indices is None: + record_indices = list(range(min(num_envs, 4))) # a few representative slots + self.record_indices = [i for i in record_indices if 0 <= i < num_envs] + + self._episode = dict.fromkeys(self.record_indices, 0) + self._rec: dict[int, Recorder | None] = dict.fromkeys(self.record_indices) + _report_sync(f"/trace/job/{self.job_id}/enter", {"name": name, "group": 1}) + logger.info("hud vec job: %s", self.job_url) + + @property + def job_url(self) -> str: + from hud.settings import settings + + return f"{settings.hud_web_url}/jobs/{self.job_id}" + + def _new_trace_id(self, i: int) -> str: + # Deterministic (reproducible, idempotent re-uploads) when a seed is given. + if self.seed is not None: + key = f"hud.vec:{self.job_id}:{i}:{self._episode[i]}:{self.seed}" + return uuid.uuid5(uuid.NAMESPACE_URL, key).hex + return uuid.uuid4().hex + + def _open(self, i: int) -> Recorder: + rec = Recorder( + trace_id=self._new_trace_id(i), + job_id=self.job_id, + group_id=self.group_id, + fps=self.fps, + model=self.model, + metadata={"env_index": i, "episode_index": self._episode[i], "seed": self.seed}, + ) + self._rec[i] = rec + return rec + + def record( + self, + *, + obs: Any | None = None, + frames: dict[str, Any] | None = None, + action: Any | None = None, + reward: Any | None = None, + done: Any | None = None, + info: dict[str, Any] | None = None, # reserved for per-env metrics + ) -> None: + """Record one batched step. Pass the tensors straight from ``env.step``. + + On ``done[i]`` the slot's current episode is closed (final reward attributed to it) + and a fresh trace is opened for the next episode; the post-reset observation returned + on the same step is skipped so frames never bleed across the episode boundary. + """ + done_np = _to_numpy(done) if done is not None else None + reward_np = _to_numpy(reward) if reward is not None else None + + for i in self.record_indices: + rec = self._rec[i] or self._open(i) + if reward_np is not None: + rec.add_reward(float(reward_np[i])) + if done_np is not None and bool(done_np[i]): + rec.close() # closing episode keeps its own env/episode metadata + self._episode[i] += 1 + self._rec[i] = None + continue # the returned obs belongs to the next episode + rec.record( + obs=_slice_obs(obs, i), + frames=_slice_frames(frames, i), + action=None if action is None else _to_numpy(action)[i], + ) + + def close(self) -> None: + """Close any open traces and flush all telemetry to the platform.""" + for i in self.record_indices: + rec = self._rec[i] + if rec is not None: + rec.close() + self._rec[i] = None + from hud.telemetry.exporter import flush + + flush() + + def __enter__(self) -> Self: + return self + + def __exit__(self, *exc: object) -> None: + self.close() + + +def _slice_obs(obs: Any | None, i: int) -> dict[str, Any] | None: + """One slot's observation as a name->array dict (accepts a dict or a bare batched array).""" + if obs is None: + return None + if isinstance(obs, dict): + return {k: _to_numpy(v)[i] for k, v in obs.items()} + return {"obs": _to_numpy(obs)[i]} + + +def _slice_frames(frames: dict[str, Any] | None, i: int) -> dict[str, Any] | None: + """One slot's camera frames as a name->``HxWxC`` array dict.""" + if not frames: + return None + return {name: _to_numpy(arr)[i] for name, arr in frames.items()} + + +__all__ = ["EpisodeRecorder", "Recorder", "VecRecorder"] diff --git a/hud/agents/robot/vec_agent.py b/hud/agents/robot/vec_agent.py new file mode 100644 index 000000000..3e866c2b6 --- /dev/null +++ b/hud/agents/robot/vec_agent.py @@ -0,0 +1,123 @@ +"""``VecRobotAgent`` — drive a whole vectorized robot env over one connection. + +The agent-side counterpart to :class:`~hud.environment.robot.isaac_bridge.IsaacBridge`: one +batched forward per tick. Receive an ``[N, ...]`` observation, run the model once to an +``[N, T, A]`` chunk, execute it open-loop per slot, and send one ``[N, A]`` action. The N +parallel episodes stream to the platform as one **Job** of per-episode traces via +:class:`~hud.agents.robot.record.VecRecorder` (slots split into fresh traces on each reset). + +Set ``self.model`` (and ``self.adapter``) exactly as for :class:`~hud.agents.robot.agent.RobotAgent`; +the model's ``infer`` is already ``[N, ...] -> [N, T, A]``, so :class:`~hud.agents.robot.model.LeRobotModel` +works unchanged. Pair with :class:`~hud.agents.robot.adapter.VecLeRobotAdapter` for the batched obs. +""" + +from __future__ import annotations + +import asyncio +from collections import deque +from typing import TYPE_CHECKING, Any, ClassVar + +import numpy as np + +from hud.capabilities.robot import RobotClient + +from .record import VecRecorder + +if TYPE_CHECKING: + from hud.capabilities.base import Capability + + from .adapter import Adapter + from .model import Model + + +class VecRobotAgent: + """Drive an N-env robot batch as one Job of per-episode traces. + + **Subclass contract:** set ``self.model`` (a :class:`~hud.agents.robot.model.Model`) and + ``self.adapter`` (a :class:`~hud.agents.robot.adapter.Adapter`) in ``__init__``. + """ + + model: Model | None = None + adapter: Adapter | None = None + #: Max control ticks before the run is cut off (the env auto-resets episodes within this). + max_steps: ClassVar[int] = 520 + #: Step-progress print frequency. 0 = off. + log_every: ClassVar[int] = 20 + + async def run( + self, + cap: Capability, + prompt: str, + *, + name: str, + num_record: int = 4, + seed: int | None = None, + group_id: str | None = None, + model_name: str | None = None, + ) -> str: + """Connect, drive the batch to ``max_steps`` (or all-terminated), return the Job URL.""" + if self.model is None: + raise RuntimeError(f"{type(self).__name__} must set self.model in __init__") + client = await RobotClient.connect(cap) + try: + action_space, obs_space = client.spaces() + if self.adapter is not None: + self.adapter.bind(action_space, obs_space) + return await self._drive( + client, prompt, name=name, num_record=num_record, seed=seed, + group_id=group_id, model_name=model_name, + ) + finally: + await client.close() + + async def _drive( + self, client: RobotClient, prompt: str, *, name: str, num_record: int, + seed: int | None, group_id: str | None, model_name: str | None, + ) -> str: + adapter = self.adapter + state_key = adapter.state_key if adapter else None + image_keys = adapter.image_keys if adapter else [] + + obs = await client.get_observation() + # Batch size from the leading dim of any obs array (state if present, else the first). + probe = obs["data"][state_key] if state_key else next(iter(obs["data"].values())) + n = int(np.asarray(probe).shape[0]) + rec = VecRecorder( + name, num_envs=n, record_indices=list(range(min(num_record, n))), + fps=client.get_control_rate(), seed=seed, group_id=group_id, model=model_name, + ) + print(f"[vec-agent] job: {rec.job_url} (N={n})", flush=True) + + chunks: list[deque[np.ndarray]] = [deque() for _ in range(n)] + for step in range(self.max_steps): + done = np.asarray(obs["terminated"], dtype=bool).reshape(-1) + for i in np.nonzero(done)[0]: # a freshly-reset slot re-infers for its new episode + chunks[i].clear() + if step and done.all(): + break + + if any(not c for c in chunks): # refill spent slots with a fresh batched chunk + batch = adapter.adapt_observation(obs, prompt) if adapter else obs + chunk = np.asarray(await asyncio.to_thread(self.model.infer, batch)) # [N, T, A] + for i, c in enumerate(chunks): + if not c: + rows = chunk[i] + if adapter is not None: # e.g. DROID delta -> absolute against slot i's query obs + rows = adapter.adapt_chunk(rows, {"data": {k: v[i] for k, v in obs["data"].items()}}) + c.extend(rows) + action = np.stack([chunks[i].popleft() for i in range(n)]) + + state = {state_key: obs["data"][state_key]} if state_key else None + frames = {k: obs["data"][k] for k in image_keys} if image_keys else None + rec.record(obs=state, frames=frames, action=action, done=done) + await client.send_action(action) + + if self.log_every and step % self.log_every == 0: + print(f"[vec-agent] step {step}/{self.max_steps} live={int((~done).sum())}/{n}", flush=True) + obs = await client.get_observation() + + rec.close() + return rec.job_url + + +__all__ = ["VecRobotAgent"] diff --git a/hud/capabilities/robot.py b/hud/capabilities/robot.py index 25af6a337..5a0c88960 100644 --- a/hud/capabilities/robot.py +++ b/hud/capabilities/robot.py @@ -95,7 +95,9 @@ async def get_observation(self) -> dict[str, Any]: msg = await self._queue.get() if "error" in msg: raise RuntimeError(f"robot env error:\n{msg['error']}") - terminated = bool(msg.pop("terminated", False)) + # Passthrough (no bool() coercion): a scalar bool for single-env bridges, an [N] + # mask for vectorized ones (RobotClient serves both; the wire decides the shape). + terminated = msg.pop("terminated", False) meta = msg.pop("meta", None) out: dict[str, Any] = {"data": msg, "terminated": terminated} if meta is not None: diff --git a/hud/environment/robot/__init__.py b/hud/environment/robot/__init__.py index 538e0c784..77c7b4fc2 100644 --- a/hud/environment/robot/__init__.py +++ b/hud/environment/robot/__init__.py @@ -15,15 +15,17 @@ from __future__ import annotations -from .bridge import RobotBridge +from .bridge import IsaacBridge, RobotBridge, VecRobotBridge from .endpoint import RobotEndpoint from .sim_runner import InlineSimRunner, MainThreadSimRunner, SimRunner, ThreadSimRunner __all__ = [ "InlineSimRunner", + "IsaacBridge", "MainThreadSimRunner", "RobotBridge", "RobotEndpoint", "SimRunner", "ThreadSimRunner", + "VecRobotBridge", ] diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py index 2fc5303c2..4bef06576 100644 --- a/hud/environment/robot/bridge.py +++ b/hud/environment/robot/bridge.py @@ -1,20 +1,26 @@ -"""Env-side ``robot`` bridge: the base class users subclass to wrap their sim. +"""Env-side ``robot`` bridges: the base classes users subclass to wrap their sim. The *server* side of the ``robot`` protocol (agent-side client: -:class:`~hud.capabilities.robot.RobotClient`); both share the wire codec defined -there. Subclass :class:`RobotBridge` and implement ``step`` / ``get_observation`` to -serve a sim over WebSocket — it steps the sim once per received action. +:class:`~hud.capabilities.robot.RobotClient`); both share the wire codec defined there. +Three layers, smallest to largest: -An injected :class:`~.sim_runner.SimRunner` owns *which thread runs the -(thread-affine) sim*, so subclasses stay thread-naive. +- :class:`RobotBridge` — single-agent, one sim step per received action. +- :class:`VecRobotBridge` — the same loop but every frame carries batched ``[N, ...]`` arrays + (``terminated`` an ``[N]`` mask, action ``[N, A]``), for a vectorized env served in lockstep. +- :class:`IsaacBridge` — a :class:`VecRobotBridge` that owns an Isaac Lab env: main-thread sim + threading, rebuild-on-instance-change, and a :meth:`~IsaacBridge.serve_forever` pump loop. + +An injected :class:`~.sim_runner.SimRunner` owns *which thread runs the (thread-affine) sim*, +so subclasses stay thread-naive. """ from __future__ import annotations import contextlib from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any +from typing import Any +import numpy as np import websockets import websockets.exceptions @@ -22,10 +28,11 @@ # ends of the protocol stay in lockstep (env -> capabilities is the correct direction). from hud.capabilities.robot import _packb, _unpackb -from .sim_runner import InlineSimRunner, SimRunner +from .sim_runner import InlineSimRunner, MainThreadSimRunner, SimRunner -if TYPE_CHECKING: - import numpy as np +# Task args that vary *within* one built Isaac env (cheap reset), vs. env-defining args that +# need a full rebuild (Isaac bakes assets at construction, so a new family/size is a new env). +EPISODIC_KEYS = ("seed", "background_id", "instruction_id") class RobotBridge(ABC): @@ -123,6 +130,10 @@ def url(self) -> str: return f"ws://{self._host}:{self._port}" async def start(self) -> None: + # Idempotent: a long-lived bridge serves sequential agents, so re-``start`` (e.g. a + # second run against the same sim) is a no-op rather than an EADDRINUSE rebind. + if self._server is not None: + return self._server = await websockets.serve( self._handle_client, self._host, self._port, max_size=None, reuse_address=True ) @@ -173,4 +184,125 @@ async def _send_observation(self) -> None: await self._client.send(_packb(msg)) -__all__ = ["RobotBridge"] +class VecRobotBridge(RobotBridge): + """A :class:`RobotBridge` that serves a whole *vectorized* env in lockstep. + + Same single-agent WebSocket loop, but every frame carries batched ``[N, ...]`` arrays + and ``terminated`` is an ``[N]`` per-env mask (not a scalar); the agent sends back one + ``[N, A]`` action. Subclasses implement the batched :meth:`step` / :meth:`get_observation` + (which returns ``(data{name: [N, ...]}, terminated[N])``). The wire codec already carries + arrays of any rank, so the only single-env assumption to drop is the scalar ``terminated``. + """ + + async def _send_observation(self) -> None: + if self._client is None: + return + out = await self._sim_runner.call(self.get_observation) + if out is None: + return + data, terminated = out + msg = {**data, "terminated": np.asarray(terminated, dtype=bool)} # [N] mask, not a scalar + with contextlib.suppress(websockets.exceptions.ConnectionClosed): + await self._client.send(_packb(msg)) + + +class IsaacBridge(VecRobotBridge): + """Serve an Isaac Lab env's whole vectorized batch over ``robot``. + + The reusable base for any Isaac Lab / Omniverse benchmark: one process owns one + ``num_envs`` env and serves the whole batch in lockstep (``[N, ...]`` obs in, ``[N, A]`` + action out). It encodes the Isaac gotchas once — Isaac/Omniverse pins the simulator to the + **main thread** and ``env.reset()`` nests a ``run_until_complete`` for USD loading, so every + sim touch is routed through a :class:`~.sim_runner.MainThreadSimRunner` and executed by + :meth:`serve_forever`'s pump loop *outside* any asyncio task. + + **Subclass contract:** implement :meth:`make_env` and :meth:`observe`. Optionally override + :meth:`prompt` (defaults to the env's ``instruction``), :meth:`instance_key` (which task + args trigger a rebuild), and :meth:`result` (defaults to the env's ``extras["metrics"]``). + """ + + def __init__(self, *, sim_runner: SimRunner | None = None, **kwargs: Any) -> None: + super().__init__(sim_runner=sim_runner or MainThreadSimRunner(), **kwargs) + self.env: Any = None + self.base: Any = None # env.unwrapped + self._instance: Any = None # current env-defining key; a mismatch forces a rebuild + self._done: np.ndarray | None = None + + # ── subclass hooks ──────────────────────────────────────────────────────── + @abstractmethod + def make_env(self, **task_args: Any) -> Any: + """Build (and return) the gym env for the resolved task. Called on the sim thread.""" + + @abstractmethod + def observe(self) -> dict[str, np.ndarray]: + """The current batched observation as ``{contract_key: [N, ...] ndarray}``.""" + + def prompt(self) -> str: + return self.base.instruction + + def instance_key(self, task_args: dict[str, Any]) -> Any: + """The env-defining subset of the task args; a change here rebuilds the env.""" + return tuple(sorted((k, v) for k, v in task_args.items() if k not in EPISODIC_KEYS)) + + # ── bridge protocol (all sim touches run on the main thread) ─────────────── + async def reset(self, **task_args: Any) -> str: + return await self._sim_runner.call(self._sync_reset, task_args) + + def _sync_reset(self, task_args: dict[str, Any]) -> str: + key = self.instance_key(task_args) + if self.env is None or key != self._instance: + if self.env is not None: + self.env.close() + self.env = self.make_env(**task_args) + self.base = self.env.unwrapped + self._instance = key + self.env.reset(seed=task_args.get("seed")) + self._done = np.zeros(self.base.num_envs, dtype=bool) + return self.prompt() + + def step(self, action: np.ndarray) -> None: + import torch + + # np.array (copy) not asarray: the wire-decoded buffer is read-only, which torch warns on. + act = torch.as_tensor(np.array(action, dtype=np.float32), device=self.base.device) + _, _, terminated, truncated, _ = self.env.step(act) + self._done = (terminated | truncated).detach().cpu().numpy().astype(bool) + + def get_observation(self) -> tuple[dict[str, np.ndarray], np.ndarray] | None: + if self.base is None: + return None + return self.observe(), self._done + + def result(self, **extra: Any) -> dict[str, Any]: + """Episode-batch score from the env's ``extras["metrics"]`` (set pre-reset by the env).""" + m = dict(self.base.extras.get("metrics", {})) if self.base is not None else {} + return { + "score": float(m.get("force_penalized_score", m.get("success_rate", 0.0))), + "success": float(m.get("success_rate", 0.0)), + **m, + **extra, + } + + # ── serving: own the Kit main thread, drain sim touches between frames ────── + def serve_forever(self, simulation_app: Any, *, host: str = "0.0.0.0", port: int = 9100) -> None: + """Serve the control endpoint + robot WebSocket on the Kit main thread, blocking for the + process lifetime. Each pass services socket IO once, drains queued sim touches (task-free, + on main), then pumps Kit — the one loop Isaac and asyncio can share. + """ + import asyncio + + from .endpoint import RobotEndpoint + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(RobotEndpoint(self).serve(host, port)) + drain = getattr(self._sim_runner, "drain", None) + while simulation_app.is_running(): + loop.run_until_complete(asyncio.sleep(0)) # advance control + robot WS tasks once + if drain is not None: + drain() # execute queued sim touches on main, outside any task + simulation_app.update() # pump Kit (renders cameras, steps physics queue) + loop.run_until_complete(self.stop()) + + +__all__ = ["EPISODIC_KEYS", "IsaacBridge", "RobotBridge", "VecRobotBridge"] From ebd1cbe522271e3d02c6e010729152f3f833d8cb Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 20:56:57 +0000 Subject: [PATCH 02/43] refactor(robot): serve every sim with one process shape (SimThread) The sim always owns the process main thread; serving (control channel + robot WebSocket) runs on a background loop thread, with every sim touch queued back to main. Cheap CPU envs block on the queue; Isaac/Omniverse gets its Kit pump from the idle hook - no per-sim threading choices. serve_blocking is the single entry hud serve and the module runner use, replacing the serve_pumped/asyncio.run branching. --- hud/cli/serve.py | 5 +- hud/environment/robot/sim_thread.py | 148 ++++++++++++++++++++++++++++ hud/environment/server.py | 22 ++++- 3 files changed, 169 insertions(+), 6 deletions(-) create mode 100644 hud/environment/robot/sim_thread.py diff --git a/hud/cli/serve.py b/hud/cli/serve.py index c12385486..7719ceaf2 100644 --- a/hud/cli/serve.py +++ b/hud/cli/serve.py @@ -7,7 +7,6 @@ from __future__ import annotations -import asyncio from pathlib import Path from typing import Any @@ -60,10 +59,10 @@ def _serve_environment(env: Any, host: str, port: int) -> None: highlight=False, ) hud_console.hint("Press Ctrl+C to stop.") - from hud.environment.server import serve + from hud.environment.server import serve_blocking try: - asyncio.run(serve(env, host, port)) + serve_blocking(env, host, port) except KeyboardInterrupt: hud_console.info("Stopped.") diff --git a/hud/environment/robot/sim_thread.py b/hud/environment/robot/sim_thread.py new file mode 100644 index 000000000..67b2bcc46 --- /dev/null +++ b/hud/environment/robot/sim_thread.py @@ -0,0 +1,148 @@ +"""One way to run any simulation: the sim owns the process main thread. + +A simulator is usually *thread-affine* (every touch must run on the thread that +created its GL/device context) and some — Isaac/Omniverse — must own the process +main thread outright: Kit drives its own main-thread loop and ``env.reset()`` +nests ``run_until_complete``, which cannot run inside an asyncio task. + +So HUD serves every sim with one process shape: serving (control channel + +robot WebSocket) runs on a background loop thread, and every sim touch is +queued to the main thread via a :class:`SimThread`. :func:`run_with_sim` is +that shape — cheap CPU envs just block on the queue; when Kit is loaded, the +idle hook pumps it between touches. + +Before :meth:`SimThread.run` starts (tests, in-loop use) calls execute inline +on the caller — the degenerate single-thread case. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import queue +import signal +import sys +import threading +from concurrent.futures import Future +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Callable, Coroutine + + +class SimThread: + """Executes sim touches on the thread that owns the simulator. + + ``call`` is awaitable from any event loop and routes the touch to the + owning thread; :meth:`run` is the blocking loop that thread drives. + Re-entrant calls (a sim touch calling back in) and calls made before + :meth:`run` starts execute inline. + """ + + _shared: SimThread | None = None + + def __init__(self) -> None: + self._q: queue.Queue[tuple[Callable[[], Any], Future]] = queue.Queue() + self._ident: int | None = None + + @classmethod + def shared(cls) -> SimThread: + """The process-wide instance (one main thread, one sim thread).""" + if cls._shared is None: + cls._shared = cls() + return cls._shared + + async def call(self, fn: Callable[..., Any], *args: Any) -> Any: + """Run ``fn(*args)`` on the sim thread, awaited on the caller's loop.""" + if self._ident is None or threading.get_ident() == self._ident: + return fn(*args) # not serving yet, or already on the sim thread + fut: Future = Future() + self._q.put((lambda: fn(*args), fut)) + return await asyncio.wrap_future(fut) + + def bind(self) -> None: + """Claim the calling thread as the sim thread (before serving starts, + so no touch can slip through inline on the wrong thread).""" + self._ident = threading.get_ident() + + def run(self, until: Callable[[], bool]) -> None: + """Own the calling thread: execute queued touches until ``until()``. + + Each pass drains the queue, then runs the idle hook — a Kit update when + Omniverse is loaded (it needs continuous pumping), else nothing (the + queue ``get`` timeout is the wait). + """ + self.bind() + while not until(): + kit = sys.modules.get("omni.kit.app") + try: + item = self._q.get(timeout=0.002 if kit else 0.05) + except queue.Empty: + pass + else: + self._execute(*item) + with contextlib.suppress(queue.Empty): # drain the backlog + while True: + self._execute(*self._q.get_nowait()) + if kit: + kit.get_app().update() + + @staticmethod + def _execute(fn: Callable[[], Any], fut: Future) -> None: + if not fut.set_running_or_notify_cancel(): + return + try: + fut.set_result(fn()) + except BaseException as exc: # propagate to the awaiting caller + fut.set_exception(exc) + + +def run_with_sim( + serve: Callable[[], Coroutine[Any, Any, Any]], + *, + sim: SimThread | None = None, +) -> None: + """THE process shape for serving a sim, blocking for the serve's lifetime. + + ``await serve()`` runs on a background loop thread while the shared + :class:`SimThread` owns the calling (main) thread. SIGTERM and Ctrl-C + cancel the serve coroutine; the sim keeps draining through its teardown + (``env.stop()`` touches the sim too), then this returns. + """ + sim = sim or SimThread.shared() + sim.bind() # claim main before serving starts, so no touch runs inline elsewhere + loop = asyncio.new_event_loop() + done = threading.Event() + task_box: dict[str, asyncio.Task[Any]] = {} + + async def _main() -> None: + task_box["task"] = asyncio.current_task() # type: ignore[assignment] + with contextlib.suppress(asyncio.CancelledError): + await serve() + + def _thread() -> None: + asyncio.set_event_loop(loop) + try: + loop.run_until_complete(_main()) + finally: + done.set() + + def _cancel(*_: object) -> None: + task = task_box.get("task") + if task is not None: + loop.call_soon_threadsafe(task.cancel) + + with contextlib.suppress(ValueError): # signals are main-thread-only; tests may not be + signal.signal(signal.SIGTERM, _cancel) + + thread = threading.Thread(target=_thread, name="hud-serve", daemon=True) + thread.start() + try: + sim.run(until=done.is_set) + except KeyboardInterrupt: + _cancel() + sim.run(until=done.is_set) # keep draining through teardown + thread.join(timeout=30) + + +__all__ = ["SimThread", "run_with_sim"] diff --git a/hud/environment/server.py b/hud/environment/server.py index f4e264c22..c968a681f 100644 --- a/hud/environment/server.py +++ b/hud/environment/server.py @@ -418,6 +418,24 @@ async def _serve_until_terminated(env: Environment, host: str, port: int) -> Non await serve(env, host, port) +def serve_blocking(env: Environment, host: str, port: int) -> None: + """Serve *env*, blocking until terminated — the one entry every CLI path uses. + + An env with attached sims (``env.gym(...)``) runs the sim-main shape: the + sim owns this (main) thread and serving runs on a background loop thread, + with every sim touch queued back here (see + :mod:`hud.environment.robot.sim_thread`). Anything else serves on a plain + ``asyncio.run``. + """ + if not getattr(env, "_sims", None): + asyncio.run(_serve_until_terminated(env, host, port)) + return + # Submodule import: only sim-serving needs it (keeps non-robot envs free of the extra). + from hud.environment.robot.sim_thread import run_with_sim + + run_with_sim(lambda: serve(env, host, port)) + + def main() -> None: from hud.environment import load_environment @@ -429,9 +447,7 @@ def main() -> None: ) parser.add_argument("--port", type=int, default=0, help="Port to bind (0 = ephemeral).") args = parser.parse_args() - asyncio.run( - _serve_until_terminated(load_environment(args.path, name=args.env), args.host, args.port) - ) + serve_blocking(load_environment(args.path, name=args.env), args.host, args.port) if __name__ == "__main__": From 032143590440fa368bbcdd267b193d69687ab3e2 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 20:57:11 +0000 Subject: [PATCH 03/43] refactor(telemetry): extract shared robot recording into hud.telemetry.robot Recording is shared infrastructure (agent harness, hud.wrap, gym bridges), not agent code. TraceRecorder (one trace, span emission only - lifecycle stays with the caller) and JobRecorder (a vectorized env as one Job of per-episode traces) replace Recorder/VecRecorder/EpisodeRecorder and their dual-mode flags; video streaming moves alongside. The LeRobot dataset half of record.py becomes agents/robot/dataset.py (DatasetWriter). InferenceStep gains contract-derived per-dim action names so the viewer can label plots. --- hud/__init__.py | 4 +- hud/agents/robot/dataset.py | 169 ++++++++ hud/agents/robot/record.py | 514 ----------------------- hud/agents/types.py | 3 + hud/telemetry/robot/__init__.py | 12 + hud/telemetry/robot/recorder.py | 312 ++++++++++++++ hud/{agents => telemetry}/robot/video.py | 0 7 files changed, 497 insertions(+), 517 deletions(-) create mode 100644 hud/agents/robot/dataset.py delete mode 100644 hud/agents/robot/record.py create mode 100644 hud/telemetry/robot/__init__.py create mode 100644 hud/telemetry/robot/recorder.py rename hud/{agents => telemetry}/robot/video.py (100%) diff --git a/hud/__init__.py b/hud/__init__.py index a1854a682..fddc41fa8 100644 --- a/hud/__init__.py +++ b/hud/__init__.py @@ -8,7 +8,6 @@ # Apply patches to third-party libraries early, before other imports from . import patches as _patches # noqa: F401 from ._legacy import install as _install_v5_compat -from .agents.robot.record import Recorder, VecRecorder from .clients import connect from .environment import Environment from .eval import ( @@ -45,7 +44,6 @@ "HostedRuntime", "Job", "LocalRuntime", - "Recorder", "Run", "Runtime", "RuntimeConfig", @@ -58,11 +56,11 @@ "Taskset", "Trace", "TrainingClient", - "VecRecorder", "connect", "instrument", ] + try: from .version import __version__ except ImportError: diff --git a/hud/agents/robot/dataset.py b/hud/agents/robot/dataset.py new file mode 100644 index 000000000..7da02da59 --- /dev/null +++ b/hud/agents/robot/dataset.py @@ -0,0 +1,169 @@ +"""Opt-in LeRobot v3 dataset writing for robot rollouts. + +One :class:`DatasetWriter` spans an agent's lifetime: every episode it drives +appends ``(observation, executed action)`` frames, and the dataset is finalized +at process exit (or an explicit :meth:`finalize`), optionally pushed to the HF +Hub. The contract drives the schema with no extra wiring. 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 +import logging +import os +import time +import uuid +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np + +from hud.telemetry.context import get_current_trace_id + +if TYPE_CHECKING: + from numpy.typing import NDArray + +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.`` (video); the lone vector obs -> + ``observation.state`` (else ``observation.``); 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": _names(f, leaf)} + key_map[name] = key + elif role == "action": + features["action"] = {"dtype": dtype, "shape": shape, "names": _names(f, "act")} + return features, key_map + + +def _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 DatasetWriter: + """Appends robot frames to a LeRobot v3 dataset; a no-op shell when lerobot + is missing (warned once) so telemetry-only runs never break.""" + + def __init__(self, contract: dict[str, Any], *, fps: int) -> None: + self._contract = contract + self._fps = fps + self._features, self._key_map = _lerobot_features(contract) + self._ds: Any | None = None # created lazily on the first frame + self._root: Path | None = None + self._repo_id = "" + self._enabled = importlib.util.find_spec("lerobot") is not None + if not self._enabled: + logger.warning( + "save=True but lerobot is not installed; streaming telemetry only " + "(pip install 'lerobot[dataset]')" + ) + else: + atexit.register(self.finalize) # keep the dataset loadable on abrupt exits + + def add(self, data: dict[str, Any], action: NDArray[Any], *, task: str) -> None: + """One frame: the wire observation dict + the executed env-space action.""" + if not self._enabled: + return + 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"] = task + ds.add_frame(row) + + def end_episode(self) -> None: + """Commit the episode's pending frames.""" + if self._ds is not None and self._ds.has_pending_frames(): + self._ds.save_episode() + + def finalize(self) -> None: + """Write the parquet footer + optionally push to the Hub. Idempotent.""" + if self._ds is None: + return + ds, self._ds = self._ds, None + 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 + 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) + + 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 writer so concurrent rollouts never share a root; tied 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] + 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__ = ["DatasetWriter"] diff --git a/hud/agents/robot/record.py b/hud/agents/robot/record.py deleted file mode 100644 index ece3c117f..000000000 --- a/hud/agents/robot/record.py +++ /dev/null @@ -1,514 +0,0 @@ -"""Recording for robot/sim rollouts: telemetry streaming, plus an optional LeRobot dataset. - -Three recorders, smallest to largest, all streaming the *same* robot telemetry the HUD -viewer plays (numeric :class:`~hud.agents.types.ObservationStep` state + per-camera H.264 -video) via the existing exporter: - -- :class:`Recorder` — one trace. Emits spans with an *explicit* ``trace_id`` (no rollout - contextvar, no ``Run``), so it works from a bare synchronous loop or any thread. The - primitive the vectorized recorder is built from. -- :class:`VecRecorder` — a whole vectorized env as one **Job** of per-episode traces. One - :class:`Recorder` per recorded env slot; on each ``done[i]`` (auto-reset) it closes that - slot's trace and opens a fresh one, so every episode becomes its own fully-saved trace. -- :class:`EpisodeRecorder` — the agent-loop recorder: records onto a :class:`Run` (so steps - land on ``run.trace`` for grading/training) and, when ``save`` is on, *also* appends each - ``(observation, executed action)`` pair to a LeRobot v3 dataset for offline training. - -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 - -Everything HUD-specific (job/trace lifecycle, trace-id assignment, span building, video -encoding, upload) lives here; a benchmark just does ``VecRecorder(...)`` -> ``record`` -> -``close``. -""" - -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, StateFeature -from hud.telemetry.context import get_current_trace_id -from hud.utils.platform import PlatformClient - -from .video import VideoStreamer - -if TYPE_CHECKING: - from typing import Self - - from numpy.typing import NDArray - - from hud.capabilities.robot import RobotClient - from hud.eval.run import Run - -logger = logging.getLogger(__name__) - - -def _reporting_enabled() -> bool: - from hud.settings import settings - - return bool(settings.telemetry_enabled and settings.api_key) - - -def _report_sync(path: str, payload: dict[str, Any]) -> None: - """Best-effort sync POST to the platform (job/trace lifecycle). No-op without an API key. - - Mirrors the async ``hud.eval.job`` reporting so a sync vec-env loop never needs an event - loop, and never fails the run when the platform is unreachable. - """ - if not _reporting_enabled(): - return - body = {k: v for k, v in payload.items() if v is not None} - try: - PlatformClient.from_settings().post(path, json=body) - except Exception as exc: # reporting is fire-and-forget - logger.warning("platform report %s failed: %s", path, exc) - - -def _to_numpy(x: Any) -> NDArray[Any]: - """Coerce a torch tensor / array / scalar to a numpy array (no torch dependency).""" - if hasattr(x, "detach"): # torch.Tensor - x = x.detach().cpu().numpy() - return np.asarray(x) - - -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.`` (video); the lone vector obs -> - ``observation.state`` (else ``observation.``); 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 EpisodeRecorder: - """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 - - -# ── lightweight telemetry recorders (vec-env / bare-loop, no Run, no LeRobot) ──────────── - - -def _observation_step(obs: dict[str, Any], *, tick: int) -> ObservationStep: - """Build an :class:`ObservationStep` (numeric state) from a flat per-slot obs dict. - - Camera frames travel as H.264 video, so array obs with rank >= 2 are skipped; flat - numeric vectors become unlabelled :class:`StateFeature`s. Robust to scalars (unlike - :meth:`ObservationStep.from_obs`, which expects a robot-contract ``obs['data']``). - """ - state: dict[str, StateFeature] = {} - for name, val in obs.items(): - arr = _to_numpy(val) - if arr.ndim >= 2: - continue - flat = np.atleast_1d(arr).astype(float).ravel() - if 0 < flat.size <= 512: - state[name] = StateFeature(values=flat.tolist()) - return ObservationStep(tick=tick, state=state) - - -class Recorder: - """Streams one trace's telemetry to the platform: state, per-camera video, actions. - - Standalone: emits spans with an explicit ``trace_id`` rather than relying on the rollout - contextvar, so it runs inside a plain (or vectorized) loop. Reuses the robot step schema - and :class:`~hud.agents.robot.video.VideoStreamer`, so the existing exporter and trace - viewer ingest it unchanged. - """ - - def __init__( - self, - *, - trace_id: str, - job_id: str | None = None, - group_id: str | None = None, - fps: int = 10, - model: str | None = None, - metadata: dict[str, Any] | None = None, - enter: bool = True, - ) -> None: - self.trace_id = trace_id - self.job_id = job_id - self._fps = fps - self._tick = 0 - self._reward = 0.0 - self._metadata = metadata or {} - self._video: VideoStreamer | None = None # lazy, one per trace - self._closed = False - if enter: - _report_sync( - f"/trace/{trace_id}/enter", - {"job_id": job_id, "group_id": group_id, "model": model}, - ) - - def add_reward(self, reward: float) -> None: - """Accumulate per-step reward into this trace's total (reported on close).""" - self._reward += float(reward) - - def record( - self, - *, - obs: dict[str, Any] | None = None, - frames: dict[str, Any] | None = None, - action: Any | None = None, - ) -> None: - """Record one tick: numeric-state span, per-camera video, and the executed action. - - ``obs`` is this slot's observation dict (flat numeric vectors); ``frames`` maps a - camera name to its ``HxWxC`` frame; ``action`` is the executed action vector. - """ - if obs is not None: - _observation_step(obs, tick=self._tick).emit(trace_id=self.trace_id) - if frames: - if self._video is None: - self._video = VideoStreamer(fps=self._fps, trace_id=self.trace_id) - self._video.record({"data": frames}) - if action is not None: - chunk = _to_numpy(action).astype(float).ravel().tolist() - InferenceStep(tick=self._tick, chunk=[chunk], chunk_length=1).emit( - trace_id=self.trace_id - ) - self._tick += 1 - - def close( - self, - *, - status: str = "completed", - reward: float | None = None, - error: str | None = None, - metadata: dict[str, Any] | None = None, - ) -> None: - """Flush video tails and report the trace exit (status / reward / metadata).""" - if self._closed: - return - self._closed = True - if self._video is not None: - self._video.finalize() - meta = {**self._metadata, **(metadata or {})} - _report_sync( - f"/trace/{self.trace_id}/exit", - { - "status": status, - "reward": self._reward if reward is None else reward, - "error": error, - "metadata": meta or None, - }, - ) - - -class VecRecorder: - """Records a vectorized env as one Job of per-episode traces. - - Construct it once with the batch size, call :meth:`record` after every ``env.step`` with - the batched tensors, and :meth:`close` at the end. A configurable subset of slots - (``record_indices``) gets rich traces (state + video); each ``done[i]`` closes that slot's - current trace and opens a new one, so a slot that plays many episodes produces many - traces — all under one Job at :attr:`job_url`. - """ - - def __init__( - self, - name: str, - num_envs: int, - *, - record_indices: list[int] | None = None, - fps: int = 10, - seed: int | None = None, - group_id: str | None = None, - model: str | None = None, - job_id: str | None = None, - ) -> None: - self.name = name - self.num_envs = num_envs - self.fps = fps - self.seed = seed - self.group_id = group_id - self.model = model - self.job_id = job_id or uuid.uuid4().hex - if record_indices is None: - record_indices = list(range(min(num_envs, 4))) # a few representative slots - self.record_indices = [i for i in record_indices if 0 <= i < num_envs] - - self._episode = dict.fromkeys(self.record_indices, 0) - self._rec: dict[int, Recorder | None] = dict.fromkeys(self.record_indices) - _report_sync(f"/trace/job/{self.job_id}/enter", {"name": name, "group": 1}) - logger.info("hud vec job: %s", self.job_url) - - @property - def job_url(self) -> str: - from hud.settings import settings - - return f"{settings.hud_web_url}/jobs/{self.job_id}" - - def _new_trace_id(self, i: int) -> str: - # Deterministic (reproducible, idempotent re-uploads) when a seed is given. - if self.seed is not None: - key = f"hud.vec:{self.job_id}:{i}:{self._episode[i]}:{self.seed}" - return uuid.uuid5(uuid.NAMESPACE_URL, key).hex - return uuid.uuid4().hex - - def _open(self, i: int) -> Recorder: - rec = Recorder( - trace_id=self._new_trace_id(i), - job_id=self.job_id, - group_id=self.group_id, - fps=self.fps, - model=self.model, - metadata={"env_index": i, "episode_index": self._episode[i], "seed": self.seed}, - ) - self._rec[i] = rec - return rec - - def record( - self, - *, - obs: Any | None = None, - frames: dict[str, Any] | None = None, - action: Any | None = None, - reward: Any | None = None, - done: Any | None = None, - info: dict[str, Any] | None = None, # reserved for per-env metrics - ) -> None: - """Record one batched step. Pass the tensors straight from ``env.step``. - - On ``done[i]`` the slot's current episode is closed (final reward attributed to it) - and a fresh trace is opened for the next episode; the post-reset observation returned - on the same step is skipped so frames never bleed across the episode boundary. - """ - done_np = _to_numpy(done) if done is not None else None - reward_np = _to_numpy(reward) if reward is not None else None - - for i in self.record_indices: - rec = self._rec[i] or self._open(i) - if reward_np is not None: - rec.add_reward(float(reward_np[i])) - if done_np is not None and bool(done_np[i]): - rec.close() # closing episode keeps its own env/episode metadata - self._episode[i] += 1 - self._rec[i] = None - continue # the returned obs belongs to the next episode - rec.record( - obs=_slice_obs(obs, i), - frames=_slice_frames(frames, i), - action=None if action is None else _to_numpy(action)[i], - ) - - def close(self) -> None: - """Close any open traces and flush all telemetry to the platform.""" - for i in self.record_indices: - rec = self._rec[i] - if rec is not None: - rec.close() - self._rec[i] = None - from hud.telemetry.exporter import flush - - flush() - - def __enter__(self) -> Self: - return self - - def __exit__(self, *exc: object) -> None: - self.close() - - -def _slice_obs(obs: Any | None, i: int) -> dict[str, Any] | None: - """One slot's observation as a name->array dict (accepts a dict or a bare batched array).""" - if obs is None: - return None - if isinstance(obs, dict): - return {k: _to_numpy(v)[i] for k, v in obs.items()} - return {"obs": _to_numpy(obs)[i]} - - -def _slice_frames(frames: dict[str, Any] | None, i: int) -> dict[str, Any] | None: - """One slot's camera frames as a name->``HxWxC`` array dict.""" - if not frames: - return None - return {name: _to_numpy(arr)[i] for name, arr in frames.items()} - - -__all__ = ["EpisodeRecorder", "Recorder", "VecRecorder"] diff --git a/hud/agents/types.py b/hud/agents/types.py index 8bafff7e9..632215649 100644 --- a/hud/agents/types.py +++ b/hud/agents/types.py @@ -422,6 +422,9 @@ class InferenceStep(Step): # post model inference (a single action is a length-1 chunk) chunk: list[list[float]] = Field(default_factory=list[list[float]]) chunk_length: int = 1 + #: Per action-dim labels from the env contract (e.g. ["eef.dx", ..., "gripper"]), + #: so the viewer can label action plots and match them to state features. + names: list[str] = Field(default_factory=list[str]) class VideoSegmentStep(Step): diff --git a/hud/telemetry/robot/__init__.py b/hud/telemetry/robot/__init__.py new file mode 100644 index 000000000..5f4288902 --- /dev/null +++ b/hud/telemetry/robot/__init__.py @@ -0,0 +1,12 @@ +"""Shared robot telemetry: trace/job recorders + per-camera H.264 video streaming. + +Imported by both sides of the robot stack (agent harness, ``hud.wrap``, gym +bridges); needs the ``robot`` extra (numpy, PyAV). +""" + +from __future__ import annotations + +from .recorder import JobRecorder, TraceRecorder, to_numpy +from .video import SegmentEncoder, VideoStreamer + +__all__ = ["JobRecorder", "SegmentEncoder", "TraceRecorder", "VideoStreamer", "to_numpy"] diff --git a/hud/telemetry/robot/recorder.py b/hud/telemetry/robot/recorder.py new file mode 100644 index 000000000..24ca66fb1 --- /dev/null +++ b/hud/telemetry/robot/recorder.py @@ -0,0 +1,312 @@ +"""Robot trace recording: numeric state, per-camera H.264 video, action chunks. + +Shared by both sides of the robot stack (the agent harness, ``hud.wrap``, the +gym bridges), so it lives under telemetry rather than either side: + +- :class:`TraceRecorder` — one trace. Emits spans with an explicit ``trace_id`` + (no rollout contextvar needed, works from any thread), or through a ``Run`` + so steps also land on ``run.trace`` for training. Span emission only — + trace lifecycle (enter/exit) is the caller's job. +- :class:`JobRecorder` — a whole vectorized env as one **Job** of per-episode + traces: one :class:`TraceRecorder` per recorded slot; ``done[i]`` closes that + slot's trace (reporting its exit) and opens a fresh one for the next episode. +""" + +from __future__ import annotations + +import logging +import uuid +from typing import TYPE_CHECKING, Any + +import mcp.types as mcp_types +import numpy as np + +from hud.agents.types import InferenceStep, ObservationStep, StateFeature +from hud.telemetry.context import get_current_trace_id +from hud.types import Step +from hud.utils.platform import PlatformClient + +from .video import VideoStreamer + +if TYPE_CHECKING: + from typing import Self + + from numpy.typing import NDArray + + from hud.eval.run import Run + +logger = logging.getLogger(__name__) + + +def to_numpy(x: Any) -> NDArray[Any]: + """Coerce a torch tensor / array / scalar to a numpy array (no torch dependency).""" + if hasattr(x, "detach"): # torch.Tensor + x = x.detach().cpu().numpy() + return np.asarray(x) + + +def _report_sync(path: str, payload: dict[str, Any]) -> None: + """Best-effort sync POST to the platform (job/trace lifecycle). No-op without an + API key; never fails the run when the platform is unreachable.""" + from hud.settings import settings + + if not (settings.telemetry_enabled and settings.api_key): + return + body = {k: v for k, v in payload.items() if v is not None} + try: + PlatformClient.from_settings().post(path, json=body) + except Exception as exc: # reporting is fire-and-forget + logger.warning("platform report %s failed: %s", path, exc) + + +class TraceRecorder: + """Streams one trace's robot telemetry; lifecycle stays with the caller. + + Construct with ``trace_id`` (spans emitted explicitly — any thread, no + ambient context) or ``run`` (steps recorded through it, so they also land + on ``run.trace``). ``obs_space`` labels observations from the env contract; + ``state_names``/``action_names`` are the contract-free fallback labels. + """ + + def __init__( + self, + *, + trace_id: str | None = None, + run: Run | None = None, + fps: int = 10, + prompt: str | None = None, + obs_space: dict[str, Any] | None = None, + action_names: list[str] | None = None, + state_names: dict[str, list[str]] | None = None, + ) -> None: + assert trace_id or run, "TraceRecorder needs a trace_id or a run" + self._run = run + # Video encodes on a background thread, so it always needs an explicit id; + # a live run's id is the ambient rollout context. + self.trace_id = trace_id or (run.trace_id if run else None) or get_current_trace_id() + self._fps = fps + self._obs_space = obs_space + self._action_names = action_names or [] + self._state_names = state_names or {} + self.reward = 0.0 + self._video: VideoStreamer | None = None # lazy, one per trace + # The task instruction as an opening user step (shows on the timeline). + # A Run records its own prompt step, so only the bare-trace path emits one. + if prompt and run is None: + self._emit( + Step( + source="user", + messages=[ + mcp_types.PromptMessage( + role="user", content=mcp_types.TextContent(type="text", text=prompt) + ) + ], + ) + ) + + def _emit(self, step: Step) -> None: + if self._run is not None: + self._run.record(step) + else: + step.emit(trace_id=self.trace_id) + + def record_observation(self, data: dict[str, Any], *, tick: int) -> None: + """One tick's observation: numeric-state span + per-camera video. + + ``data`` maps feature names to per-env arrays; camera frames (``HxWxC``) + go to the video path, flat numeric vectors become labelled state. + """ + if self._obs_space is not None: # contract-aware labeling/slicing + step = ObservationStep.from_obs({"data": data}, tick=tick, obs_space=self._obs_space) + else: + state: dict[str, StateFeature] = {} + for name, val in data.items(): + arr = to_numpy(val) + if arr.ndim >= 2: + continue # camera frames travel as video + flat = np.atleast_1d(arr).astype(float).ravel() + if 0 < flat.size <= 512: + labels = self._state_names.get(name, []) + state[name] = StateFeature( + names=labels if len(labels) == flat.size else [], values=flat.tolist() + ) + step = ObservationStep(tick=tick, state=state) + self._emit(step) + if self._video is None: + self._video = VideoStreamer(fps=self._fps, trace_id=self.trace_id) + self._video.record({"data": data}) + + def record_inference(self, chunk: Any, *, tick: int) -> None: + """One inference: the freshly inferred ``[T, A]`` chunk (an executed + single action is a length-1 chunk).""" + rows = np.atleast_2d(to_numpy(chunk).astype(float)) + self._emit( + InferenceStep( + tick=tick, chunk=rows.tolist(), chunk_length=len(rows), names=self._action_names + ) + ) + + def add_reward(self, reward: float) -> None: + """Accumulate per-step reward into this trace's total.""" + self.reward += float(reward) + + def close(self) -> None: + """Flush video tails (idempotent). Trace exit reporting is the caller's job.""" + if self._video is not None: + self._video.finalize() + self._video = None + + +class JobRecorder: + """Records a vectorized env as one Job of per-episode traces. + + Construct once with the batch size, call :meth:`record` after every + ``env.step`` with the batched tensors, :meth:`close` at the end. A subset of + slots (``record_indices``) gets rich traces; each ``done[i]`` closes that + slot's trace (reporting reward — env-reported ``success`` outranks + accumulated shaped reward) and opens a fresh one, all under one Job. + """ + + def __init__( + self, + name: str, + num_envs: int, + *, + record_indices: list[int] | None = None, + fps: int = 10, + seed: int | None = None, + group_id: str | None = None, + model: str | None = None, + job_id: str | None = None, + prompt: str | None = None, + action_names: list[str] | None = None, + state_names: dict[str, list[str]] | None = None, + ) -> None: + self.name = name + self.num_envs = num_envs + self.fps = fps + self.seed = seed + self.group_id = group_id + self.model = model + self.job_id = job_id or uuid.uuid4().hex + self._prompt = prompt + self._action_names = action_names + self._state_names = state_names + #: Extra metadata merged into each newly opened trace (e.g. the episode's + #: task parametrization). Mutable: set it before the episode's traces open. + self.extra_metadata: dict[str, Any] = {} + if record_indices is None: + record_indices = list(range(min(num_envs, 4))) # a few representative slots + self.record_indices = [i for i in record_indices if 0 <= i < num_envs] + self._episode = dict.fromkeys(self.record_indices, 0) + self._tick = dict.fromkeys(self.record_indices, 0) # per-slot, per-episode tick + self._rec: dict[int, TraceRecorder | None] = dict.fromkeys(self.record_indices) + self._meta: dict[int, dict[str, Any]] = {} + _report_sync(f"/trace/job/{self.job_id}/enter", {"name": name, "group": 1}) + logger.info("hud vec job: %s", self.job_url) + + @property + def job_url(self) -> str: + from hud.settings import settings + + return f"{settings.hud_web_url}/jobs/{self.job_id}" + + def _open(self, i: int) -> TraceRecorder: + # Deterministic trace ids (reproducible, idempotent re-uploads) when seeded. + if self.seed is not None: + key = f"hud.vec:{self.job_id}:{i}:{self._episode[i]}:{self.seed}" + trace_id = uuid.uuid5(uuid.NAMESPACE_URL, key).hex + else: + trace_id = uuid.uuid4().hex + self._meta[i] = { + "env_index": i, + "episode_index": self._episode[i], + "seed": self.seed, + **self.extra_metadata, + } + _report_sync( + f"/trace/{trace_id}/enter", + {"job_id": self.job_id, "group_id": self.group_id, "model": self.model}, + ) + rec = TraceRecorder( + trace_id=trace_id, + fps=self.fps, + prompt=self._prompt, + action_names=self._action_names, + state_names=self._state_names, + ) + self._rec[i] = rec + return rec + + def _close(self, i: int, rec: TraceRecorder, *, success: bool | None = None) -> None: + rec.close() + reward = rec.reward if success is None else float(success) + meta = {**self._meta.pop(i, {}), **({"success": success} if success is not None else {})} + _report_sync( + f"/trace/{rec.trace_id}/exit", + {"status": "completed", "reward": reward, "metadata": meta or None}, + ) + self._rec[i] = None + self._episode[i] += 1 # the next obs opens a fresh trace + self._tick[i] = 0 + + def record( + self, + *, + obs: Any | None = None, + frames: dict[str, Any] | None = None, + action: Any | None = None, + reward: Any | None = None, + done: Any | None = None, + success: Any | None = None, + ) -> None: + """Record one batched step. Pass the tensors straight from ``env.step``. + + On ``done[i]`` the slot's episode is closed (final reward attributed) and + the post-reset observation returned on the same step is skipped, so frames + never bleed across the episode boundary. + """ + done_np = to_numpy(done) if done is not None else None + reward_np = to_numpy(reward) if reward is not None else None + success_np = to_numpy(success) if success is not None else None + + for i in self.record_indices: + rec = self._rec[i] or self._open(i) + if reward_np is not None: + rec.add_reward(float(reward_np[i])) + if done_np is not None and bool(done_np[i]): + self._close(i, rec, success=bool(success_np[i]) if success_np is not None else None) + continue # the returned obs belongs to the next episode + data: dict[str, Any] = {} + if obs is not None: + sliced = obs if isinstance(obs, dict) else {"obs": obs} + data.update({k: to_numpy(v)[i] for k, v in sliced.items()}) + data.update({k: to_numpy(v)[i] for k, v in (frames or {}).items()}) + if data: + rec.record_observation(data, tick=self._tick[i]) + if action is not None: + rec.record_inference(to_numpy(action)[i], tick=self._tick[i]) + self._tick[i] += 1 + + def close_slots(self) -> None: + """Close every open per-slot trace (an explicit mid-run reset ends its episodes).""" + for i in self.record_indices: + rec = self._rec[i] + if rec is not None: + self._close(i, rec) + + def close(self) -> None: + """Close open traces and flush all telemetry to the platform.""" + self.close_slots() + from hud.telemetry.exporter import flush + + flush() + + def __enter__(self) -> Self: + return self + + def __exit__(self, *exc: object) -> None: + self.close() + + +__all__ = ["JobRecorder", "TraceRecorder", "to_numpy"] diff --git a/hud/agents/robot/video.py b/hud/telemetry/robot/video.py similarity index 100% rename from hud/agents/robot/video.py rename to hud/telemetry/robot/video.py From ed4cb29f96b70606c99172a33a14434fdd4a8587 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 20:57:25 +0000 Subject: [PATCH 04/43] refactor(robot): batched-first RobotBridge with slots grading; add GymBridge One bridge serves num_envs slots in lockstep; a plain single env is a batch of one speaking the scalar wire framing on the same code path, so VecRobotBridge and IsaacBridge collapse into the base. GymBridge is the generic bridge over any gym-style factory (arg partitioning by factory signature, rebuild-on-instance-change, success probing, torch/int action handling). Grading becomes result() -> {score, success, slots: [...]} (one dict per slot), replacing the result_batch plumbing through the endpoint's JSON-RPC. SimRunner strategies are gone - bridges route every sim touch through the shared SimThread - and the endpoint gains serve_blocking() for split-process sims (replacing serve_forever). --- hud/environment/robot/bridge.py | 374 ++++++++++++++++------------ hud/environment/robot/endpoint.py | 21 +- hud/environment/robot/introspect.py | 148 +++++++++++ hud/environment/robot/sim_runner.py | 111 --------- 4 files changed, 375 insertions(+), 279 deletions(-) create mode 100644 hud/environment/robot/introspect.py delete mode 100644 hud/environment/robot/sim_runner.py diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py index 4bef06576..d16dfd27b 100644 --- a/hud/environment/robot/bridge.py +++ b/hud/environment/robot/bridge.py @@ -1,22 +1,25 @@ -"""Env-side ``robot`` bridges: the base classes users subclass to wrap their sim. +"""Env-side ``robot`` bridges: the base class and every shipped variant. The *server* side of the ``robot`` protocol (agent-side client: -:class:`~hud.capabilities.robot.RobotClient`); both share the wire codec defined there. -Three layers, smallest to largest: - -- :class:`RobotBridge` — single-agent, one sim step per received action. -- :class:`VecRobotBridge` — the same loop but every frame carries batched ``[N, ...]`` arrays - (``terminated`` an ``[N]`` mask, action ``[N, A]``), for a vectorized env served in lockstep. -- :class:`IsaacBridge` — a :class:`VecRobotBridge` that owns an Isaac Lab env: main-thread sim - threading, rebuild-on-instance-change, and a :meth:`~IsaacBridge.serve_forever` pump loop. - -An injected :class:`~.sim_runner.SimRunner` owns *which thread runs the (thread-affine) sim*, -so subclasses stay thread-naive. +:class:`~hud.capabilities.robot.RobotClient`); both share the wire codec defined +there. Bridges are **batched-first**: one bridge serves ``num_envs`` slots in +lockstep (``[N, ...]`` obs frames, ``[N, A]`` actions, an ``[N]`` ``terminated`` +mask). ``num_envs == 1`` — a plain single-env sim — speaks the scalar framing +(per-env arrays, scalar ``terminated``) on the same code path. + +- :class:`RobotBridge` — subclass with your sim: ``reset`` / ``step`` / + ``get_observation``. The base owns the WebSocket serve loop. +- :class:`GymBridge` — the generic bridge over any gym-style env factory; + users get one via ``env.gym(make_env)`` and never subclass it. + +Every sim touch routes through the process :class:`~.sim_thread.SimThread`, so +bridges stay thread-naive (see :mod:`~.sim_thread` for the one process shape). """ from __future__ import annotations import contextlib +import inspect from abc import ABC, abstractmethod from typing import Any @@ -27,39 +30,31 @@ # The openpi/0 wire codec is defined alongside the agent-side client; reuse it so both # ends of the protocol stay in lockstep (env -> capabilities is the correct direction). from hud.capabilities.robot import _packb, _unpackb +from hud.telemetry.robot import to_numpy -from .sim_runner import InlineSimRunner, MainThreadSimRunner, SimRunner - -# Task args that vary *within* one built Isaac env (cheap reset), vs. env-defining args that -# need a full rebuild (Isaac bakes assets at construction, so a new family/size is a new env). -EPISODIC_KEYS = ("seed", "background_id", "instruction_id") +from .introspect import probe_success, split_observation +from .sim_thread import SimThread class RobotBridge(ABC): """Serves ``robot`` over WebSocket; subclass and implement the env hooks. - **Subclass contract:** implement :meth:`step`, :meth:`get_observation`, and - :meth:`reset`. The base owns the WebSocket serve loop; subclasses own the sim. + **Subclass contract:** implement :meth:`reset`, :meth:`step`, and + :meth:`get_observation`. The base owns the WebSocket serve loop; subclasses + own the sim and set ``num_envs`` (default 1). - :meth:`reset` initialises the sim for a new episode and returns the task - prompt. The base resets scoring state and pushes the first frame for you. - - :meth:`step` advances the sim by one action. - - :meth:`get_observation` returns ``(data, terminated)`` for the current state - or ``None`` if not ready. - - :meth:`result` returns the episode score dict. The default implementation - covers the common binary-success case; override for richer scoring (e.g. - fractional subtask progress or realtime stats). Concrete bridges must set - ``self.success``, ``self.total_reward``, and ``self.terminated`` during - :meth:`step` for the default to work. + prompt. The base resets scoring state and pushes the first frame. + - :meth:`step` advances the sim by one action (``[A]``, or ``[N, A]`` batched). + - :meth:`get_observation` returns ``(data, terminated)`` — per-env arrays + + scalar bool for ``num_envs == 1``, ``[N, ...]`` arrays + ``[N]`` mask + otherwise — or ``None`` if not ready. + - :meth:`result_slots` returns one score dict per slot. The default covers + the common single-env binary-success case from ``self.success`` / + ``self.total_reward``; batched bridges override it. """ - def __init__( - self, - *, - host: str = "127.0.0.1", - port: int = 0, - sim_runner: SimRunner | None = None, - ) -> None: + def __init__(self, *, host: str = "127.0.0.1", port: int = 0) -> None: # Loopback + ephemeral by default; the concrete address is published in the # manifest post-``start()`` and tunneled, so no env manages bridge ports. self._host = host @@ -68,10 +63,11 @@ def __init__( self._server: Any = None # Connect-time metadata frame (sent first on each connection); subclasses may set it. self.metadata: dict[str, Any] = {} - # Which thread runs the (thread-affine) sim. Default InlineSimRunner (loop - # thread); inject a ThreadSimRunner (or custom) when render-heavy or thread-bound. - self._sim_runner: SimRunner = sim_runner or InlineSimRunner() - # Episode scoring read by ``result()``; subclasses update in ``reset``/``step``. + # Every sim touch runs on the process sim thread (see sim_thread.py). + self._sim = SimThread.shared() + self.num_envs: int = 1 + # Episode scoring read by ``result()``; single-env subclasses update these + # in ``reset``/``step`` (batched bridges override result_slots instead). self.task_description: str = "" self.total_reward: float = 0.0 self.success: bool = False @@ -89,32 +85,38 @@ async def _reset(self, **kwargs: Any) -> str: @abstractmethod async def reset(self, **kwargs: Any) -> str: - """Reset the sim for a new episode; return the task prompt. - - Take whatever task kwargs you need (e.g. ``task_id``, ``seed``). The base - resets scoring + sends the first obs — just reset your sim and return the prompt. - """ + """Reset the sim for a new episode; return the task prompt.""" @abstractmethod def step(self, action: np.ndarray) -> None: """Advance the sim by one action.""" @abstractmethod - def get_observation(self) -> tuple[dict[str, np.ndarray], bool] | None: + def get_observation(self) -> tuple[dict[str, np.ndarray], Any] | None: """Return ``(data, terminated)`` for the current state, or ``None`` if not ready.""" + def result_slots(self) -> list[dict[str, Any]]: + """One score dict per slot. Default: single-env binary success.""" + return [ + { + "score": 1.0 if self.success else 0.0, + "success": bool(self.success), + "total_reward": float(self.total_reward), + } + ] + def result(self) -> dict[str, Any]: - """Return the episode score dict after the episode ends. + """The episode grade: per-slot dicts under ``"slots"``, means at the top. - Default: binary success score + total reward. Override when the bridge - tracks richer scoring (fractional subtask progress, realtime stats, …). - The returned dict is forwarded to the harness, so include any fields the - downstream consumers expect. + The grouped eval path grades each trace from ``slots[i]``; single-result + consumers read the aggregate ``score``/``success`` unchanged. """ + slots = self.result_slots() return { - "score": 1.0 if self.success else 0.0, - "success": bool(self.success), - "total_reward": float(self.total_reward), + "score": float(np.mean([s["score"] for s in slots])), + "success": float(np.mean([float(s["success"]) for s in slots])), + "total_reward": float(np.mean([s.get("total_reward", 0.0) for s in slots])), + "slots": slots, } @property @@ -122,8 +124,7 @@ def url(self) -> str: """The bridge's concrete ``ws://`` address — publish this in the manifest. With an ephemeral port (the default) the address only exists once - :meth:`start` has bound the socket, so publish from an - ``@env.initialize`` hook *after* ``await bridge.start()``. + :meth:`start` has bound the socket, so publish after ``await bridge.start()``. """ if self._port == 0: raise RuntimeError("bridge bound to an ephemeral port; call start() before reading url") @@ -155,7 +156,7 @@ async def _handle_client(self, ws: Any) -> None: await self._send_observation() # current obs on connect (if ready) async for raw in ws: action = _unpackb(raw)["actions"] # codec already returns an ndarray - await self._sim_runner.call(self.step, action) # on the sim thread + await self._sim.call(self.step, action) # on the sim thread await self._send_observation() except websockets.exceptions.ConnectionClosed: pass @@ -171,138 +172,179 @@ async def _handle_client(self, ws: Any) -> None: self._client = None async def _send_observation(self) -> None: - """Send the current observation to the connected agent (if any).""" - if self._client is None: - return - out = await self._sim_runner.call(self.get_observation) - if out is None: - return - data, terminated = out - # openpi-style flat obs dict: array fields at the top level, terminated alongside. - msg = {**data, "terminated": bool(terminated)} - with contextlib.suppress(websockets.exceptions.ConnectionClosed): - await self._client.send(_packb(msg)) - - -class VecRobotBridge(RobotBridge): - """A :class:`RobotBridge` that serves a whole *vectorized* env in lockstep. + """Send the current observation to the connected agent (if any). - Same single-agent WebSocket loop, but every frame carries batched ``[N, ...]`` arrays - and ``terminated`` is an ``[N]`` per-env mask (not a scalar); the agent sends back one - ``[N, A]`` action. Subclasses implement the batched :meth:`step` / :meth:`get_observation` - (which returns ``(data{name: [N, ...]}, terminated[N])``). The wire codec already carries - arrays of any rank, so the only single-env assumption to drop is the scalar ``terminated``. - """ - - async def _send_observation(self) -> None: + Framing follows ``num_envs``: scalar ``terminated`` for a single env, an + ``[N]`` mask for a batch — the one place the two wire shapes meet. + """ if self._client is None: return - out = await self._sim_runner.call(self.get_observation) + out = await self._sim.call(self.get_observation) if out is None: return data, terminated = out - msg = {**data, "terminated": np.asarray(terminated, dtype=bool)} # [N] mask, not a scalar + done = np.asarray(terminated, dtype=bool) + msg = {**data, "terminated": bool(done.ravel()[0]) if self.num_envs == 1 else done} with contextlib.suppress(websockets.exceptions.ConnectionClosed): await self._client.send(_packb(msg)) -class IsaacBridge(VecRobotBridge): - """Serve an Isaac Lab env's whole vectorized batch over ``robot``. - - The reusable base for any Isaac Lab / Omniverse benchmark: one process owns one - ``num_envs`` env and serves the whole batch in lockstep (``[N, ...]`` obs in, ``[N, A]`` - action out). It encodes the Isaac gotchas once — Isaac/Omniverse pins the simulator to the - **main thread** and ``env.reset()`` nests a ``run_until_complete`` for USD loading, so every - sim touch is routed through a :class:`~.sim_runner.MainThreadSimRunner` and executed by - :meth:`serve_forever`'s pump loop *outside* any asyncio task. +class GymBridge(RobotBridge): + """Serve any gym-style env factory over the ``robot`` protocol, generically. - **Subclass contract:** implement :meth:`make_env` and :meth:`observe`. Optionally override - :meth:`prompt` (defaults to the env's ``instruction``), :meth:`instance_key` (which task - args trigger a rebuild), and :meth:`result` (defaults to the env's ``extras["metrics"]``). + Task args are partitioned by the factory's signature: args the factory + accepts define the env (a change rebuilds it — so ``num_envs`` in the + factory signature is the vectorization declaration); everything else is + episodic and flows to ``env.reset(seed=..., options=...)``. """ - def __init__(self, *, sim_runner: SimRunner | None = None, **kwargs: Any) -> None: - super().__init__(sim_runner=sim_runner or MainThreadSimRunner(), **kwargs) + def __init__(self, factory: Any, **kwargs: Any) -> None: + super().__init__(**kwargs) + self._factory = factory + self._factory_params = { + n + for n, p in inspect.signature(factory).parameters.items() + if p.kind in (p.POSITIONAL_OR_KEYWORD, p.KEYWORD_ONLY) + } self.env: Any = None - self.base: Any = None # env.unwrapped - self._instance: Any = None # current env-defining key; a mismatch forces a rebuild - self._done: np.ndarray | None = None - - # ── subclass hooks ──────────────────────────────────────────────────────── - @abstractmethod - def make_env(self, **task_args: Any) -> Any: - """Build (and return) the gym env for the resolved task. Called on the sim thread.""" - - @abstractmethod - def observe(self) -> dict[str, np.ndarray]: - """The current batched observation as ``{contract_key: [N, ...] ndarray}``.""" + self.batched = False # env carries a leading [N] dim (any env exposing num_envs) + self._obs: Any = None # latest observation (reset or step) + self._instance: Any = None # current env-defining args; a mismatch rebuilds + self._is_torch = False + # Per-slot episode scoring, sticky until the next reset. + self._done: np.ndarray = np.zeros(1, dtype=bool) + self._success: np.ndarray = np.zeros(1, dtype=bool) + self._acc_reward: np.ndarray = np.zeros(1) + self._seen_success = False # any env-reported success signal this episode + + # ── env lifecycle (all sim touches on the sim thread) ─────────────────────── + + async def ensure_env(self, **task_args: Any) -> None: + """Build the env (factory defaults unless task args say otherwise). Idempotent.""" + if self.env is None: + await self._sim.call(self._sync_reset, task_args) - def prompt(self) -> str: - return self.base.instruction - - def instance_key(self, task_args: dict[str, Any]) -> Any: - """The env-defining subset of the task args; a change here rebuilds the env.""" - return tuple(sorted((k, v) for k, v in task_args.items() if k not in EPISODIC_KEYS)) - - # ── bridge protocol (all sim touches run on the main thread) ─────────────── async def reset(self, **task_args: Any) -> str: - return await self._sim_runner.call(self._sync_reset, task_args) + return await self._sim.call(self._sync_reset, task_args) + + async def stop(self) -> None: + await super().stop() + if self.env is not None: + await self._sim.call(self.env.close) + self.env = None def _sync_reset(self, task_args: dict[str, Any]) -> str: - key = self.instance_key(task_args) + build = {k: v for k, v in task_args.items() if k in self._factory_params} + episodic = {k: v for k, v in task_args.items() if k not in self._factory_params} + key = tuple(sorted(build.items())) if self.env is None or key != self._instance: if self.env is not None: self.env.close() - self.env = self.make_env(**task_args) - self.base = self.env.unwrapped + self.env = self._factory(**build) self._instance = key - self.env.reset(seed=task_args.get("seed")) - self._done = np.zeros(self.base.num_envs, dtype=bool) - return self.prompt() + base = getattr(self.env, "unwrapped", self.env) + n = getattr(self.env, "num_envs", getattr(base, "num_envs", None)) + self.batched = n is not None + self.num_envs = int(n or 1) + seed = episodic.pop("seed", None) + obs, _ = self.env.reset(seed=seed, options=episodic or None) + self._obs = obs # the first frame an agent sees on connect/reset + self._is_torch = "torch" in type(_first_leaf(obs)).__module__ + self._done = np.zeros(self.num_envs, dtype=bool) + self._success = np.zeros(self.num_envs, dtype=bool) + self._acc_reward = np.zeros(self.num_envs) + self._seen_success = False + return self._prompt(task_args) + + def _prompt(self, task_args: dict[str, Any]) -> str: + base = getattr(self.env, "unwrapped", self.env) + for attr in ("task_description", "instruction"): + text = getattr(getattr(base, "cfg", None), attr, None) or getattr(base, attr, None) + if isinstance(text, str) and text: + return text + return ", ".join(f"{k}={v}" for k, v in sorted(task_args.items())) or "run the task" + + def sample_observation(self) -> tuple[dict[str, Any], dict[str, Any]]: + """One per-env ``(state, frames)`` sample for contract derivation (post-build).""" + state, frames = split_observation(self._obs, batched=self.batched) + if self.batched: + state = {k: to_numpy(v)[0] for k, v in state.items()} + frames = {k: to_numpy(v)[0] for k, v in frames.items()} + return state, frames + + # ── bridge protocol ────────────────────────────────────────────────────────── def step(self, action: np.ndarray) -> None: - import torch - - # np.array (copy) not asarray: the wire-decoded buffer is read-only, which torch warns on. - act = torch.as_tensor(np.array(action, dtype=np.float32), device=self.base.device) - _, _, terminated, truncated, _ = self.env.step(act) - self._done = (terminated | truncated).detach().cpu().numpy().astype(bool) - - def get_observation(self) -> tuple[dict[str, np.ndarray], np.ndarray] | None: - if self.base is None: + act: Any = np.array(action, dtype=np.float32) # wire buffer is read-only + if not self.batched: + act = act[0] if act.ndim > 1 else act # single plain env: drop the batch dim + elif act.ndim == 1: + act = act[None] # batched-of-one served scalar-framed: restore the [N] dim + # The wire carries floats; discrete/int action spaces need their dtype + shape back. + space = getattr(self.env, "action_space", None) + dtype = getattr(space, "dtype", None) + if dtype is not None and np.issubdtype(dtype, np.integer): + act = act.astype(dtype).reshape(getattr(space, "shape", act.shape) or ()) + if self._is_torch: + import torch + + base = getattr(self.env, "unwrapped", self.env) + act = torch.as_tensor(act, device=getattr(base, "device", None)) + obs, reward, terminated, truncated, info = self.env.step(act) + self._obs = obs + done = np.atleast_1d(to_numpy(terminated)).astype(bool) | np.atleast_1d( + to_numpy(truncated) + ).astype(bool) + self._acc_reward += np.atleast_1d(to_numpy(reward)) * ~self._done + newly = done & ~self._done + if newly.any(): + success = self._resolve_success(info) + if success is not None: + self._seen_success = True + self._success |= success & newly + self._done |= done + + def _resolve_success(self, info: Any) -> np.ndarray | None: + """Env-reported success at a done step: info keys, else Isaac's termination term.""" + found = probe_success(info, num_envs=self.num_envs) + if found is not None: + return found + base = getattr(self.env, "unwrapped", self.env) + manager = getattr(base, "termination_manager", None) + if manager is not None: + try: + return np.atleast_1d(to_numpy(manager.get_term("success"))).astype(bool) + except Exception: + return None + return None + + def get_observation(self) -> tuple[dict[str, np.ndarray], Any] | None: + if self.env is None or self._obs is None: return None - return self.observe(), self._done - - def result(self, **extra: Any) -> dict[str, Any]: - """Episode-batch score from the env's ``extras["metrics"]`` (set pre-reset by the env).""" - m = dict(self.base.extras.get("metrics", {})) if self.base is not None else {} - return { - "score": float(m.get("force_penalized_score", m.get("success_rate", 0.0))), - "success": float(m.get("success_rate", 0.0)), - **m, - **extra, - } - - # ── serving: own the Kit main thread, drain sim touches between frames ────── - def serve_forever(self, simulation_app: Any, *, host: str = "0.0.0.0", port: int = 9100) -> None: - """Serve the control endpoint + robot WebSocket on the Kit main thread, blocking for the - process lifetime. Each pass services socket IO once, drains queued sim touches (task-free, - on main), then pumps Kit — the one loop Isaac and asyncio can share. - """ - import asyncio - - from .endpoint import RobotEndpoint - - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - loop.run_until_complete(RobotEndpoint(self).serve(host, port)) - drain = getattr(self._sim_runner, "drain", None) - while simulation_app.is_running(): - loop.run_until_complete(asyncio.sleep(0)) # advance control + robot WS tasks once - if drain is not None: - drain() # execute queued sim touches on main, outside any task - simulation_app.update() # pump Kit (renders cameras, steps physics queue) - loop.run_until_complete(self.stop()) - - -__all__ = ["EPISODIC_KEYS", "IsaacBridge", "RobotBridge", "VecRobotBridge"] + state, frames = split_observation(self._obs, batched=self.batched) + data = {k: to_numpy(v) for k, v in {**state, **frames}.items()} + if self.batched and self.num_envs == 1: + data = {k: v[0] for k, v in data.items()} # batched-of-one: squeeze to scalar framing + return data, self._done if self.num_envs > 1 else bool(self._done[0]) + + def result_slots(self) -> list[dict[str, Any]]: + """Per-slot grades: env-reported success when available, else accumulated reward.""" + # The env's own success check outranks accumulated shaped reward. + scores = self._success if self._seen_success else self._acc_reward + return [ + { + "score": float(scores[i]), + "success": bool(self._success[i]), + "total_reward": float(self._acc_reward[i]), + } + for i in range(self.num_envs) + ] + + +def _first_leaf(obs: Any) -> Any: + while isinstance(obs, dict): + obs = next(iter(obs.values())) + return obs + + +__all__ = ["GymBridge", "RobotBridge"] diff --git a/hud/environment/robot/endpoint.py b/hud/environment/robot/endpoint.py index 28517920d..e20138425 100644 --- a/hud/environment/robot/endpoint.py +++ b/hud/environment/robot/endpoint.py @@ -119,8 +119,6 @@ async def result(self, **extra: Any) -> dict[str, Any]: ) return res - """ in your simulation program where bridge is started """ - # ── serving: expose a local bridge so a remote endpoint can drive it ── async def serve(self, host: str = "127.0.0.1", port: int = 9100) -> asyncio.AbstractServer: """Serve this (local) bridge's control surface over JSON-RPC. @@ -136,6 +134,25 @@ async def serve(self, host: str = "127.0.0.1", port: int = 9100) -> asyncio.Abst print(f"[env] control endpoint listening on {host}:{port}", flush=True) return server + def serve_blocking(self, host: str = "0.0.0.0", port: int = 9100) -> None: # noqa: S104 — split-process sims serve cross-host + """Serve this (local) bridge for the process's lifetime, with the sim owning + the main thread — the entry a split-process sim program calls last. + + Same shape as ``hud.environment.server``: the control endpoint runs on a + background loop thread; every sim touch drains here on main. + """ + from .sim_thread import run_with_sim + + async def _serve() -> None: + server = await self.serve(host, port) + try: + await asyncio.Event().wait() # until SIGTERM / Ctrl-C cancels + finally: + server.close() + await self._local_bridge().stop() + + run_with_sim(_serve) + async def _handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: with contextlib.suppress(ConnectionResetError, asyncio.IncompleteReadError): while (msg := await read_frame(reader)) is not None: diff --git a/hud/environment/robot/introspect.py b/hud/environment/robot/introspect.py new file mode 100644 index 000000000..337cd60a9 --- /dev/null +++ b/hud/environment/robot/introspect.py @@ -0,0 +1,148 @@ +"""Gym-env introspection: split observations, probe success, derive a minimal contract. + +Pure functions shared by :func:`hud.wrap` (in-process trace streaming), the +:class:`~.bridge.GymBridge`, and the :class:`~.gym.Gym` handle — no telemetry, +no wrapper state. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import numpy as np + +from hud.telemetry.robot import to_numpy + +#: Conventional info keys carrying env-reported success, probed in order. +SUCCESS_KEYS = ("success", "is_success", "task_success") + + +def flatten_observation(obs: Any, prefix: str = "") -> dict[str, Any]: + """Flatten nested dict observations to slash-keyed leaves (non-dicts -> ``{"obs": x}``).""" + if not isinstance(obs, dict): + return {prefix or "obs": obs} + flat: dict[str, Any] = {} + for k, v in obs.items(): + key = f"{prefix}/{k}" if prefix else str(k) + flat.update(flatten_observation(v, key) if isinstance(v, dict) else {key: v}) + return flat + + +def split_observation(obs: Any, *, batched: bool = False) -> tuple[dict[str, Any], dict[str, Any]]: + """Split an observation into ``(state, frames)``, both name -> array. + + A camera frame is a channel-last image (rank 3, or rank 4 when batched); state is + any flat numeric vector. Anything else is dropped. Batched arrays keep their + leading ``[N]`` dim — the recorder slices per slot. + """ + state: dict[str, Any] = {} + frames: dict[str, Any] = {} + img_rank, vec_rank = (4, 2) if batched else (3, 1) + for name, val in flatten_observation(obs).items(): + arr = to_numpy(val) + if arr.ndim == img_rank and arr.shape[-1] in (1, 3, 4): + frames[name] = arr + elif arr.ndim <= vec_rank and np.issubdtype(arr.dtype, np.number): + state[name] = arr if batched else np.atleast_1d(arr) + return state, frames + + +def probe_success(info: Any, *, num_envs: int = 1) -> np.ndarray | None: + """Per-env success bools from conventional info keys; ``None`` when the env reports none.""" + if not isinstance(info, dict): + return None + for key in SUCCESS_KEYS: + if info.get(key) is not None: + return np.broadcast_to(to_numpy(info[key]).astype(bool).ravel(), (num_envs,)) + return None + + +def detect_fps(env: Any) -> int: + """Control rate from env metadata (``render_fps``) or Isaac's ``step_dt``; default 10.""" + fps = (getattr(env, "metadata", None) or {}).get("render_fps") + if fps: + return round(fps) + dt = getattr(getattr(env, "unwrapped", env), "step_dt", None) + return round(1 / dt) if dt else 10 + + +def capture_task_params(kwargs: dict[str, Any]) -> dict[str, Any]: + """Reset parametrization as json-safe trace metadata — the full kwargs, never a label.""" + + def safe(v: Any) -> Any: + if v is None or isinstance(v, (str, int, float, bool)): + return v + if isinstance(v, dict): + return {str(k): safe(x) for k, x in v.items()} + if isinstance(v, (list, tuple)): + return [safe(x) for x in v] + return str(v) + + return {k: safe(v) for k, v in kwargs.items() if v is not None} + + +def derive_contract( + state: dict[str, Any], frames: dict[str, Any], action_dim: int, fps: int +) -> dict[str, Any]: + """Minimal contract from one (per-env) sample observation + the action size. + + Just enough to structure the spaces and label plots — cameras as rgb observations, + state vectors with positional names, one action feature. Users edit the written + ``contract.json`` to rename dimensions; nothing else is derived on purpose. + """ + features: dict[str, Any] = {} + for name in frames: + features[name] = {"role": "observation", "type": "rgb"} + for name, vec in state.items(): + leaf = name.split("/")[-1] + features[name] = { + "role": "observation", + "names": [f"{leaf}_{i}" for i in range(int(np.asarray(vec).size))], + } + features["action"] = { + "role": "action", + "names": [f"act_{i}" for i in range(action_dim)], + } + return {"control_rate": fps, "features": features} + + +def load_or_write_contract( + path: str | Path | None, + state: dict[str, Any], + frames: dict[str, Any], + action_dim: int, + fps: int, +) -> dict[str, Any]: + """The contract round-trip: load the user-edited file if present, else derive a + minimal contract from the sample observation and write it once for inspection.""" + file = Path(path) if path else None + if file is not None and file.exists(): + return json.loads(file.read_text()) + contract = derive_contract(state, frames, action_dim, fps) + if file is not None: + file.write_text(json.dumps(contract, indent=2) + "\n") + return contract + + +def action_dim_of(env: Any, *, batched: bool) -> int: + """Per-env action size from the env's (possibly batched) action space.""" + space = getattr(env, "single_action_space", None) or getattr(env, "action_space", None) + shape = tuple(getattr(space, "shape", None) or ()) + if batched and len(shape) > 1: + shape = shape[1:] # batched Isaac spaces carry the [N] dim + return int(np.prod(shape)) if shape else 1 + + +__all__ = [ + "SUCCESS_KEYS", + "action_dim_of", + "capture_task_params", + "derive_contract", + "detect_fps", + "flatten_observation", + "load_or_write_contract", + "probe_success", + "split_observation", +] diff --git a/hud/environment/robot/sim_runner.py b/hud/environment/robot/sim_runner.py deleted file mode 100644 index 74b278ab1..000000000 --- a/hud/environment/robot/sim_runner.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Sim execution strategies: *which thread* runs the (thread-affine) simulator. - -A sim (MuJoCo/EGL, Isaac, a hardware SDK) is usually thread-affine — every touch must -run on the thread that created it — but the bridge's asyncio loop can't be stalled by a -blocking step. A :class:`SimRunner` hides that choice behind one :meth:`~SimRunner.call` -verb: - -- :class:`InlineSimRunner` — runs on the loop thread. Default; for cheap/CPU sims + tests. -- :class:`ThreadSimRunner` — sim on a dedicated worker thread, loop kept free. For - heavy/blocking sims; used by the realtime bridges. -- :class:`MainThreadSimRunner` — sim on the main thread, for runtimes that own *both* the - main thread and the asyncio loop themselves (Isaac/Omniverse: ``omni.kit.async_engine`` - drives one main-thread loop, and ``env.reset()`` internally calls ``run_until_complete`` - for USD loading — which must not nest inside a running task). HUD's servers run on that - same loop; the owner's pump loop calls :meth:`~MainThreadSimRunner.drain` to run queued - sim touches on the main thread *outside* any task. -""" - -from __future__ import annotations - -import asyncio -import queue -import threading -from abc import ABC, abstractmethod -from concurrent.futures import Future, ThreadPoolExecutor -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from collections.abc import Callable - - -class SimRunner(ABC): - """Strategy for *which thread* runs the (thread-affine) sim; bridges route every - sim touch through :meth:`call`, so the choice is a one-line injection.""" - - @abstractmethod - async def call(self, fn: Callable[..., Any], *args: Any) -> Any: - """Run ``fn(*args)`` on the sim thread, awaited on the loop.""" - - def shutdown(self) -> None: # noqa: B027 # optional hook: default no-op, subclasses override if they own threads - """Release any owned thread(s). Idempotent.""" - - -class InlineSimRunner(SimRunner): - """Run sim work inline on the caller's (loop) thread. The default; for cheap/CPU - sims and tests.""" - - async def call(self, fn: Callable[..., Any], *args: Any) -> Any: - return fn(*args) - - -class ThreadSimRunner(SimRunner): - """Sim on a dedicated worker thread: the GL/device context binds to the worker, - leaving the loop free during a blocking step. Used by the realtime bridges.""" - - def __init__(self, *, thread_name_prefix: str = "sim") -> None: - self._worker_ident: int | None = None - # max_workers=1 -> the worker spawns lazily on first submit; its initializer - # records the ident so re-entrant calls (already on the sim thread) run inline. - self._executor = ThreadPoolExecutor( - max_workers=1, thread_name_prefix=thread_name_prefix, initializer=self._record_ident - ) - - def _record_ident(self) -> None: - self._worker_ident = threading.get_ident() - - async def call(self, fn: Callable[..., Any], *args: Any) -> Any: - if threading.get_ident() == self._worker_ident: # avoid self-dispatch deadlock - return fn(*args) - loop = asyncio.get_running_loop() - return await loop.run_in_executor(self._executor, lambda: fn(*args)) - - def shutdown(self) -> None: - self._executor.shutdown(wait=False) - - -class MainThreadSimRunner(SimRunner): - """Sim on the main thread, for runtimes that own both it and the loop (Isaac/Omniverse: - Kit drives one main-thread loop and ``env.reset()`` nests ``run_until_complete``, which - can't run inside a task). A handler's :meth:`call` queues each sim touch; the owner's - pump loop :meth:`drain`\\ s it between ticks so it runs on main, *outside* any task.""" - - def __init__(self) -> None: - self._q: queue.Queue[tuple[Callable[[], Any], Future]] = queue.Queue() - # loop/tasks/drain share one thread, so thread id can't tell "in a task" from "in - # drain" — this flag can: a task queues, a sim touch re-entering call() runs inline. - self._draining = False - - async def call(self, fn: Callable[..., Any], *args: Any) -> Any: - if self._draining: - return fn(*args) # re-entrant from a sim touch — already task-free - fut: Future = Future() - self._q.put((lambda: fn(*args), fut)) - return await asyncio.wrap_future(fut) - - def drain(self) -> None: - """Run all queued sim touches on main, task-free. Call between the owner's loop ticks.""" - self._draining = True - try: - while not self._q.empty(): - fn, fut = self._q.get_nowait() - if fut.set_running_or_notify_cancel(): - try: - fut.set_result(fn()) - except BaseException as exc: # propagate to the awaiting caller - fut.set_exception(exc) - finally: - self._draining = False - - -__all__ = ["InlineSimRunner", "MainThreadSimRunner", "SimRunner", "ThreadSimRunner"] From 37826e4664b446382ef84ae50542c5fa89595195 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 20:57:36 +0000 Subject: [PATCH 05/43] feat(robot): declarative gym envs (env.gym) and hud.wrap trace streaming env.gym(make_env) turns any gym-style factory into a served robot sim: contract derived from a sample observation (round-tripped through an editable contract.json), capability minted, lifecycle wired to the env's hooks, episodes driven via the returned handle (sim.reset / sim.result). A factory accepting num_envs is the vectorization declaration. hud.wrap is the loop-owning counterpart: wrap an env you drive yourself and every episode streams to the platform as a trace under one job (lazy top-level attr so core hud stays free of the robot extra). --- hud/__init__.py | 9 ++ hud/environment/env.py | 29 ++++ hud/environment/robot/__init__.py | 41 +++--- hud/environment/robot/gym.py | 124 +++++++++++++++++ hud/environment/robot/wrap.py | 222 ++++++++++++++++++++++++++++++ 5 files changed, 407 insertions(+), 18 deletions(-) create mode 100644 hud/environment/robot/gym.py create mode 100644 hud/environment/robot/wrap.py diff --git a/hud/__init__.py b/hud/__init__.py index fddc41fa8..6ebf8a276 100644 --- a/hud/__init__.py +++ b/hud/__init__.py @@ -61,6 +61,15 @@ ] +def __getattr__(name: str) -> object: + # Lazy: hud.wrap pulls in the robot extra (numpy, websockets) only when used. + if name == "wrap": + from .environment.robot import wrap + + return wrap + raise AttributeError(f"module 'hud' has no attribute {name!r}") + + try: from .version import __version__ except ImportError: diff --git a/hud/environment/env.py b/hud/environment/env.py index b3072e78c..25aa5933c 100644 --- a/hud/environment/env.py +++ b/hud/environment/env.py @@ -167,6 +167,9 @@ def __init__( # stands up). Run once by the serving substrate around its lifetime. self._on_start: list[Callable[[], Awaitable[None]]] = [] self._on_stop: list[Callable[[], Awaitable[None]]] = [] + #: Sims attached via :meth:`gym`; when present, the server runs the + #: sim-main process shape (sim on the main thread, serving beside it). + self._sims: list[Any] = [] self._init_legacy() # ─── task registration ─────────────────────────────────────────── @@ -305,6 +308,32 @@ async def _down() -> None: return ws + def gym(self, factory: Any, *, name: str = "robot", **kwargs: Any) -> Any: + """Attach a gym-style sim serving ``name`` over the ``robot`` protocol. + + ``factory`` is any callable returning a gym-style env (the same ``make_env`` + an EnvHub repo exposes). Registers the start → publish → stop lifecycle on + this env's hooks; nothing is built until the env serves. Extra kwargs go to + :class:`~hud.environment.robot.Gym` (``fps=``, ``contract=``, ...). + Returns the handle templates drive episodes through (``sim.reset`` / + ``sim.result``). + """ + from hud.environment.robot import Gym + + sim = Gym(factory, **kwargs) + self._sims.append(sim) # the server serves sim-main when any are attached + + @self.initialize + async def _up() -> None: + await sim.start() + self.add_capability(sim.capability(name)) + + @self.shutdown + async def _down() -> None: + await sim.stop() + + return sim + # ─── substrate-run daemon lifecycle ────────────────────────────────── async def start(self) -> None: diff --git a/hud/environment/robot/__init__.py b/hud/environment/robot/__init__.py index 77c7b4fc2..4a61fdc32 100644 --- a/hud/environment/robot/__init__.py +++ b/hud/environment/robot/__init__.py @@ -1,31 +1,36 @@ -"""Env-side robot runtime: the ``robot`` bridge + its building blocks. +"""Env-side robot runtime: bridges, the control endpoint, and gym integration. -This package holds everything an *environment* needs to own a simulator and serve it to -an agent over the ``robot`` WebSocket protocol: +Everything an *environment* needs to own a simulator and serve it to an agent +over the ``robot`` WebSocket protocol: -- :class:`~hud.environment.robot.bridge.RobotBridge` — the server-side (synchronous) - bridge: one sim step per received action. -- :class:`~hud.environment.robot.sim_runner.SimRunner` (``Inline`` / ``Thread`` / - ``MainThread``) — the strategy for *which thread* runs the thread-affine simulator. +- :class:`~.bridge.RobotBridge` — the batched-first bridge base (``num_envs`` + slots in lockstep; a plain single env is a batch of one). +- :class:`~.bridge.GymBridge` / :class:`~.gym.Gym` — the generic gym-factory + path (``env.gym(make_env)``): contract derivation, capability, episode control. +- :class:`~.endpoint.RobotEndpoint` — the control handle, local or remote. +- :func:`hud.wrap` (:mod:`~.wrap`) — one-line trace streaming for any gym env. +- :class:`~.sim_thread.SimThread` — the one process shape: the sim owns the + main thread, serving runs on a background loop thread. -The agent-side counterpart, :class:`~hud.capabilities.robot.RobotClient`, lives under -:mod:`hud.capabilities` (it is a capability *client*, dialed by the agent); these two ends -share the ``robot`` wire codec defined there. +The agent-side counterpart, :class:`~hud.capabilities.robot.RobotClient`, lives +under :mod:`hud.capabilities`; both ends share the wire codec defined there. """ from __future__ import annotations -from .bridge import IsaacBridge, RobotBridge, VecRobotBridge +from .bridge import GymBridge, RobotBridge from .endpoint import RobotEndpoint -from .sim_runner import InlineSimRunner, MainThreadSimRunner, SimRunner, ThreadSimRunner +from .gym import Gym +from .sim_thread import SimThread, run_with_sim +from .wrap import TracedEnv, wrap __all__ = [ - "InlineSimRunner", - "IsaacBridge", - "MainThreadSimRunner", + "Gym", + "GymBridge", "RobotBridge", "RobotEndpoint", - "SimRunner", - "ThreadSimRunner", - "VecRobotBridge", + "SimThread", + "TracedEnv", + "run_with_sim", + "wrap", ] diff --git a/hud/environment/robot/gym.py b/hud/environment/robot/gym.py new file mode 100644 index 000000000..5f8e53dcc --- /dev/null +++ b/hud/environment/robot/gym.py @@ -0,0 +1,124 @@ +"""``Gym`` — the declarative sim handle over a gym-style env factory. + +The robotics analog of :class:`~hud.environment.workspace.Workspace`: +construction is pure data, ``start()`` materializes everything (env build, +contract derivation, the ``robot`` WebSocket, a JSON-RPC control endpoint for +split-process setups), and ``capability()`` mints the wire capability. +``env.gym(make_env)`` wires the lifecycle. + +Templates drive episodes through the handle:: + + sim = env.gym(make_env) + + + @env.template(id="pawn_lift") + async def pawn_lift(task: str = "solo_pawn_lift", seed: int = 0, num_envs: int = 1): + yield {"prompt": await sim.reset(task=task, seed=seed, num_envs=num_envs)} + yield await sim.result() +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from .bridge import GymBridge +from .endpoint import RobotEndpoint +from .introspect import action_dim_of, detect_fps, load_or_write_contract + +if TYPE_CHECKING: + import asyncio + + from hud.capabilities import Capability + +logger = logging.getLogger(__name__) + + +class Gym: + """One gym-style sim as a HUD building block: contract, capability, episode control. + + Nothing is built at import — the factory runs at :meth:`start` (serve time). + All control goes through the bridge's public surface via a local + :class:`~.endpoint.RobotEndpoint`. + """ + + def __init__( + self, + factory: Any, + *, + fps: int | None = None, + contract: str | Path | None = "contract.json", + host: str = "127.0.0.1", + port: int = 0, + control_port: int = 9100, + ) -> None: + self._bridge = GymBridge(factory, host=host, port=port) + self._endpoint = RobotEndpoint(self._bridge) + self._fps = fps + self._contract_path = contract + self._contract: dict[str, Any] = {} + self._control_host = host + self._control_port = control_port + self._control_server: asyncio.AbstractServer | None = None + + # ── lifecycle (driven by env.gym()'s hooks, or directly) ──────────────────── + + async def start(self) -> None: + """Build the env (factory defaults), derive/load the contract, bring up the wire. + + Also serves the bridge's JSON-RPC control endpoint so a split-process + env can dial this process directly. + """ + await self._bridge.ensure_env() + state, frames = self._bridge.sample_observation() + existed = self._contract_path is not None and Path(self._contract_path).exists() + self._contract = load_or_write_contract( + self._contract_path, + state, + frames, + action_dim_of(self._bridge.env, batched=self._bridge.batched), + self._fps or detect_fps(self._bridge.env), + ) + if not existed and self._contract_path is not None: + logger.info("gym: wrote %s (edit names to relabel plots)", self._contract_path) + await self._bridge.start() + if self._control_server is None: + self._control_server = await self._endpoint.serve( + self._control_host, self._control_port + ) + + async def stop(self) -> None: + if self._control_server is not None: + self._control_server.close() + self._control_server = None + await self._bridge.stop() # also closes the built env + + def capability(self, name: str = "robot") -> Capability: + """The concrete ``robot`` capability — mirrors ``Workspace.capability()``.""" + from hud.capabilities import Capability + + return Capability.robot(name=name, url=self._bridge.url, contract=self._contract) + + # ── the template surface ───────────────────────────────────────────────────── + + async def reset(self, **task_args: Any) -> str: + """Start an episode (rebuilding the env if an env-defining arg changed); + returns the task prompt.""" + return await self._endpoint.reset(**task_args) + + async def result(self) -> dict[str, Any]: + """The episode grade: per-slot dicts under ``"slots"``, means at the top.""" + return await self._endpoint.result() + + @property + def env(self) -> Any: + """The live env — privileged sim access for custom grading in templates.""" + return self._bridge.env + + @property + def contract(self) -> dict[str, Any]: + return self._contract + + +__all__ = ["Gym"] diff --git a/hud/environment/robot/wrap.py b/hud/environment/robot/wrap.py new file mode 100644 index 000000000..20042f280 --- /dev/null +++ b/hud/environment/robot/wrap.py @@ -0,0 +1,222 @@ +"""``hud.wrap`` — one-line trace streaming for gym-style envs. + +Wrap any ``gym.Env``, ``gym.vector.VectorEnv``, or batched-tensor Isaac env and keep +using it exactly as before; every episode streams to the platform as a trace (numeric +state, per-camera H.264 video, actions, reward/success) under one job. The user — or +``lerobot-eval`` — keeps owning the loop; the wrapper only observes ``reset``/``step``. + + env = hud.wrap(make_env(...), job="chess-eval") + +On first reset a minimal ``contract.json`` is written next to your script describing how +the observation/action spaces were interpreted; edit its ``names`` to relabel the +platform's plots. An existing file is loaded instead, so edits stick. +""" + +from __future__ import annotations + +import atexit +import logging +from pathlib import Path +from typing import Any, Self + +import numpy as np + +from hud.telemetry.robot import JobRecorder, to_numpy + +from .introspect import ( + action_dim_of, + capture_task_params, + detect_fps, + load_or_write_contract, + probe_success, + split_observation, +) + +logger = logging.getLogger(__name__) + + +def wrap( + env: Any, + *, + job: str | None = None, + job_id: str | None = None, + task: str | None = None, + fps: int | None = None, + contract: str | Path | None = "contract.json", + record_indices: list[int] | None = None, +) -> TracedEnv: + """Stream a gym-style env's episodes to the platform; returns the env, traced. + + - ``job`` — job name on the platform (default: the env's spec id / class name). + - ``job_id`` — share one job across several wrapped envs (multi-task suites). + - ``task`` — optional instruction/label shown on each trace's timeline. + - ``fps`` — control rate override (default: detected from the env). + - ``contract`` — path for the derived contract round-trip; ``None`` disables it. + - ``record_indices`` — which env slots get rich traces (default: first 4). + """ + return TracedEnv( + env, + job=job, + job_id=job_id, + task=task, + fps=fps, + contract=contract, + record_indices=record_indices, + ) + + +class TracedEnv: + """The observing wrapper: same ``reset``/``step``/``close`` surface, plus telemetry. + + Plain envs are recorded as a batch of one; vectorized/batched envs fan out into one + trace per episode per slot (``done[i]`` closes slot ``i``'s trace and opens the next). + """ + + def __init__( + self, + env: Any, + *, + job: str | None, + job_id: str | None, + task: str | None, + fps: int | None, + contract: str | Path | None, + record_indices: list[int] | None, + ) -> None: + self.env = env + self._batched = _is_batched(env) + self._n = int(_num_envs(env) or 1) + self._fps = fps or detect_fps(env) + self._job = job or getattr(getattr(env, "spec", None), "id", None) or type(env).__name__ + self._job_id = job_id + self._task = task + self._contract_path = Path(contract) if contract else None + self._record_indices = record_indices + self._rec: JobRecorder | None = None + self._closed = False + atexit.register(self.close) # flush traces even without an explicit close() + + def __getattr__(self, name: str) -> Any: + return getattr(self.env, name) + + # ── the observed surface ────────────────────────────────────────────────── + + def reset(self, **kwargs: Any) -> Any: + result = self.env.reset(**kwargs) + obs = result[0] if isinstance(result, tuple) else result + if self._rec is None: + self._rec = self._start(obs) + else: + self._rec.close_slots() # an explicit mid-run reset ends open episodes + # The episode's full parametrization (reset kwargs + options), never a label. + params = capture_task_params( + {k: v for k, v in kwargs.items() if k != "options"} | (kwargs.get("options") or {}) + ) + self._rec.extra_metadata = {"task_params": params} if params else {} + return result + + def step(self, action: Any) -> Any: + result = self.env.step(action) + obs, reward, terminated, truncated = result[:4] + info = result[4] if len(result) > 4 else {} + if self._rec is not None: + state, frames = split_observation(obs, batched=self._batched) + done = np.atleast_1d(to_numpy(terminated)).astype(bool) | np.atleast_1d( + to_numpy(truncated) + ).astype(bool) + if not self._batched: # record a plain env as a batch of one + state = {k: v[None] for k, v in state.items()} + frames = {k: v[None] for k, v in frames.items()} + action = to_numpy(action)[None] + reward = np.atleast_1d(to_numpy(reward)) + self._rec.record( + obs=state or None, + frames=frames or None, + action=action, + reward=reward, + done=done, + success=probe_success(info, num_envs=self._n), + ) + return result + + def close(self) -> None: + if self._closed: + return + self._closed = True + if self._rec is not None: + self._rec.close() + self.env.close() + + # ── setup ──────────────────────────────────────────────────────────────── + + def _start(self, obs: Any) -> JobRecorder: + """First reset: derive/load the contract, then open the job's recorder.""" + state, frames = split_observation(obs, batched=self._batched) + sample = {k: to_numpy(v)[0] if self._batched else v for k, v in state.items()} + contract = self._contract(sample, frames) + feats = contract.get("features", {}) + rec = JobRecorder( + self._job, + self._n, + record_indices=self._record_indices, + fps=int(contract.get("control_rate") or self._fps), + job_id=self._job_id, + prompt=self._task, + action_names=next( + (f.get("names") for f in feats.values() if f.get("role") == "action"), None + ), + state_names={ + k: f["names"] + for k, f in feats.items() + if f.get("role") == "observation" and f.get("names") + }, + ) + return rec + + def _contract(self, state: dict[str, Any], frames: dict[str, Any]) -> dict[str, Any]: + """Load the user-edited contract if present; else derive one and write it once.""" + existed = self._contract_path is not None and self._contract_path.exists() + contract = load_or_write_contract( + self._contract_path, + state, + frames, + action_dim_of(self.env, batched=self._batched), + self._fps, + ) + if not existed and self._contract_path is not None: + logger.info("hud.wrap: wrote %s (edit names to relabel plots)", self._contract_path) + return contract + + def __enter__(self) -> Self: + return self + + def __exit__(self, *exc: object) -> None: + self.close() + + +def _num_envs(env: Any) -> int | None: + """``num_envs`` from the env or its unwrapped core (gymnasium 1.x wrappers + no longer forward attributes, so an Isaac env inside ``gym.make`` hides it).""" + n = getattr(env, "num_envs", None) + if n is None: + n = getattr(getattr(env, "unwrapped", env), "num_envs", None) + return None if n is None else int(n) + + +def _is_batched(env: Any) -> bool: + """Vectorized (gym VectorEnv) or batched-tensor (Isaac) envs carry a leading [N] dim. + + Any env exposing ``num_envs`` is batched — Isaac keeps batch semantics even at + ``num_envs == 1``; plain ``gym.Env``s don't have the attribute at all. + """ + if _num_envs(env) is not None: + return True + try: + import gymnasium as gym + + return isinstance(getattr(env, "unwrapped", env), gym.vector.VectorEnv) + except ImportError: + return False + + +__all__ = ["TracedEnv", "wrap"] From 0f84f9c3f54985ab49c87deb0a51685594a0fabc Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 20:57:47 +0000 Subject: [PATCH 06/43] refactor(agents): unify RobotAgent to drive N>=1 env slots One open-loop chunk queue drives single and vectorized envs alike: the wire framing (scalar terminated vs [N] mask) sets the batch size, spent slots refill from one batched forward, adapt_chunk converts per slot at inference time. __call__(run) is the generic rollout contract (a group of one, still coalescing through ainfer so BatchedAgent works unchanged); drive(runs, client) is the grouped-eval entry recording spans per slot. VecRobotAgent and VecLeRobotAdapter are deleted - LeRobotAdapter maps batched and unbatched observations with the same wiring. --- hud/agents/robot/__init__.py | 43 +++--- hud/agents/robot/_types.py | 12 -- hud/agents/robot/adapter.py | 116 ++++++-------- hud/agents/robot/agent.py | 276 +++++++++++++++++++--------------- hud/agents/robot/batching.py | 11 +- hud/agents/robot/model.py | 2 +- hud/agents/robot/vec_agent.py | 123 --------------- 7 files changed, 222 insertions(+), 361 deletions(-) delete mode 100644 hud/agents/robot/_types.py delete mode 100644 hud/agents/robot/vec_agent.py diff --git a/hud/agents/robot/__init__.py b/hud/agents/robot/__init__.py index 58ad8f89f..57ccfa7b0 100644 --- a/hud/agents/robot/__init__.py +++ b/hud/agents/robot/__init__.py @@ -1,45 +1,38 @@ -"""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, reads the contract, drives N >= 1 env slots with one 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, VecLeRobotAdapter +from .adapter import Adapter, LeRobotAdapter, OpenPIAdapter from .agent import ROBOT_PROTOCOL, RobotAgent from .batching import BatchedAgent, BatchedModel +from .dataset import DatasetWriter from .model import LeRobotModel, Model, RemoteModel -from .vec_agent import VecRobotAgent __all__ = [ "ROBOT_PROTOCOL", "Adapter", "BatchedAgent", "BatchedModel", + "DatasetWriter", "LeRobotAdapter", "LeRobotModel", "Model", "OpenPIAdapter", "RemoteModel", "RobotAgent", - "VecLeRobotAdapter", - "VecRobotAgent", ] diff --git a/hud/agents/robot/_types.py b/hud/agents/robot/_types.py deleted file mode 100644 index a55208e07..000000000 --- a/hud/agents/robot/_types.py +++ /dev/null @@ -1,12 +0,0 @@ -"""Shared robot-agent typing helpers.""" - -from __future__ import annotations - -from typing import Any - -import numpy as np -from numpy.typing import NDArray - -ActionArray = NDArray[np.floating[Any]] - -__all__ = ["ActionArray"] diff --git a/hud/agents/robot/adapter.py b/hud/agents/robot/adapter.py index d5451f777..51b01021f 100644 --- a/hud/agents/robot/adapter.py +++ b/hud/agents/robot/adapter.py @@ -1,56 +1,50 @@ """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: #: 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`). + #: 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,70 +53,51 @@ 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_action(self, action: ActionArray, obs: dict[str, Any]) -> ActionArray: - """Translate a policy action into the env's action space (default identity).""" - return action - def adapt_chunk(self, chunk: ActionArray, obs: dict[str, Any]) -> ActionArray: - """Translate a freshly-inferred ``[T, A]`` chunk to env space, given the query-time - observation it was inferred from (default identity). + """Translate a freshly-inferred ``[T, A]`` chunk to env space, given the + (per-slot) query-time observation it was inferred from (default identity). - The vectorized harness calls this once per slot at inference time (not per step), so a - chunk expressed relative to the query state — e.g. DROID joint *deltas* that must be - added to the query-time joints for absolute targets — can be converted in one shot. + 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 - -class LeRobotAdapter(Adapter): - """Vanilla LeRobot adapter for a standard image/state env. - - 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. - """ - - def adapt_observation(self, obs: dict[str, Any], prompt: str) -> dict[str, Any]: - import torch # pyright: ignore[reportMissingImports] - - torch_mod: Any = torch - data = obs["data"] - batch: dict[str, Any] = { - "observation.state": torch_mod.from_numpy(data[self.state_key].astype(np.float32)), - "task": 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 - return batch - def adapt_action(self, action: ActionArray, obs: dict[str, Any]) -> ActionArray: + """Per-step execution-time hook on the popped action (default identity).""" return action -class VecLeRobotAdapter(LeRobotAdapter): - """Batched :class:`LeRobotAdapter` for a vectorized env (:class:`~hud.agents.robot.vec_agent.VecRobotAgent`). +class LeRobotAdapter(Adapter): + """Vanilla LeRobot adapter for a standard image/state env, single or batched. - Same wiring, but the obs arrays carry a leading ``N`` and the whole batch maps in one go: - state stays ``[N, S]``, each camera ``[N, H, W, C]`` uint8 becomes ``[N, C, H, W]`` float in - ``[0, 1]``, and the shared task is repeated to ``N`` (one prompt per env in the batch). + 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] + import torch data = obs["data"] - n = len(np.asarray(data[self.state_key])) + 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.from_numpy(np.asarray(data[self.state_key], dtype=np.float32)), - "task": [prompt] * n, + "observation.state": torch.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.from_numpy(np.asarray(data[env_key])).permute(0, 3, 1, 2).float() / 255.0 + img = torch.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 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"]) @@ -130,9 +105,4 @@ def adapt_observation(self, obs: dict[str, Any], prompt: str) -> dict[str, Any]: return out -__all__ = [ - "Adapter", - "LeRobotAdapter", - "OpenPIAdapter", - "VecLeRobotAdapter", -] +__all__ = ["ActionArray", "Adapter", "LeRobotAdapter", "OpenPIAdapter"] diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py index fb4a11a8f..9ce3dafd5 100644 --- a/hud/agents/robot/agent.py +++ b/hud/agents/robot/agent.py @@ -1,23 +1,25 @@ -"""Base v6 agent for any env that exposes a ``robot`` capability. +"""``RobotAgent`` — the one robot harness, driving N >= 1 env slots. -Subclass :class:`RobotAgent`, set ``self.model`` and ``self.adapter`` in -``__init__``, and the base owns the rest. +Subclass, set ``self.model`` and ``self.adapter`` in ``__init__``, and the base +owns the rest: connect to the ``robot`` capability, read the contract, run the +open-loop chunk queue until the env terminates. A plain single env is a batch +of one — the same loop drives a vectorized env's whole ``[N, ...]`` batch over +one connection (one batched forward per refill). -The base calls the adapter and model at the right moments:: +Two entries, one loop: - 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 +- ``__call__(run)`` — the generic rollout contract (one run, one trace). +- ``drive(runs, client)`` — the grouped-eval entry + (:func:`hud.eval.run.rollout_group`): N runs sharing one env instance, spans + recorded per slot onto each run's trace. -``model.ainfer`` always returns a ``[T, A]`` chunk; :meth:`RobotAgent.select_action` -executes it open-loop, re-inferring only once the active chunk is spent. - -Most policies use :class:`~hud.agents.robot.adapter.LeRobotAdapter`; a policy whose -spaces match the env natively can set ``adapter = None`` (raw pass-through). +Most policies use :class:`~.adapter.LeRobotAdapter`; a policy whose spaces +match the env natively can set ``adapter = None`` (raw pass-through). """ from __future__ import annotations +import asyncio from collections import deque from typing import TYPE_CHECKING, Any, ClassVar @@ -25,13 +27,12 @@ from hud.agents.base import Agent from hud.capabilities.robot import RobotClient - -from .record import EpisodeRecorder +from hud.telemetry.robot import TraceRecorder if TYPE_CHECKING: + from hud.clients.client import HudClient from hud.eval.run import Run - from ._types import ActionArray from .adapter import Adapter from .model import Model @@ -39,131 +40,164 @@ class RobotAgent(Agent): - """Drive a ``robot`` side-channel for one :class:`~hud.client.Run`. + """Drive a ``robot`` env — single or vectorized — with one open-loop chunk queue. **Subclass contract:** in ``__init__`` set ``self.model`` (a - :class:`~hud.agents.robot.model.Model`) and ``self.adapter`` (an - :class:`~hud.agents.robot.adapter.Adapter`, or ``None`` for raw pass-through). - - **Override if needed:** - - - :attr:`robot_protocol` — class attr if not ``openpi/0`` - - :meth:`on_episode_start` — mostly internal; override (with ``super()``) to - add per-episode setup (e.g. reading the env contract). - - :meth:`should_stop` — custom early-exit condition beyond ``obs["terminated"]`` - - :meth:`select_action` — only for a wholly different inference path - - :attr:`log_every` — class-level print frequency (0 = off) + :class:`~.model.Model`) and ``self.adapter`` (an :class:`~.adapter.Adapter`, + or ``None`` for raw pass-through). ``model.infer`` is batch-shaped + (``[N, ...] -> [N, T, A]``), so the same subclass drives both shapes. """ robot_protocol: ClassVar[str] = ROBOT_PROTOCOL + #: Max control ticks before the episode is cut off. + max_steps: ClassVar[int] = 520 #: 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`. + #: Opt-in: also save a LeRobot v3 dataset of every (obs, action) pair + #: (single-env runs only). Telemetry streams regardless; see :mod:`.dataset`. save: bool = False - #: Runs the policy (preprocess → forward → postprocess). Subclasses set this. + #: Runs the policy (preprocess -> forward -> postprocess). Subclasses set this. model: Model | None = None #: Translates env<->policy spaces. Subclasses set this; ``None`` = raw pass-through. adapter: Adapter | None = None - _prompt: str = "" - #: The env's action / observation contract features (from ``client.spaces()``), - #: named ``_env_*`` to mark them as env-side values (not the policy's spaces). - _env_action_space: dict[str, Any] - _env_obs_space: dict[str, Any] - #: Unexecuted tail of the current policy chunk; popped one action per step. - _active_chunk: deque[ActionArray] - #: 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: EpisodeRecorder | None = None - - def setup_robot(self, client: RobotClient) -> None: - """Discover the env's action/observation layout and bind the adapter to it.""" - self._env_action_space, self._env_obs_space = client.spaces() - if self.adapter is not None: - 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 per-episode state before the act loop. - - 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._tick = 0 - # 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 = EpisodeRecorder(client, save=self.save) - self._recorder.begin(run, prompt) - if self.adapter is not None: - self.adapter.reset() - - def should_stop(self, obs: dict[str, Any], *, step: int, max_steps: int) -> bool: - """Return True to break out of the step loop (before ``select_action``).""" - return bool(obs.get("terminated")) - - async def select_action(self, obs: dict[str, Any]) -> ActionArray: - """Pop the next action, re-inferring a ``[T, A]`` chunk once the active one is - spent, then adapt it to env space. Override only for a different inference path. + async def __call__(self, run: Run, *, max_steps: int | None = None) -> None: + """The generic rollout contract: one run, one trace.""" + await self.drive([run], run.client, max_steps=max_steps) + run.trace.status = "completed" + run.trace.content = "done" + + async def drive( + self, runs: list[Run], client: HudClient, *, max_steps: int | None = None + ) -> None: + """Drive every env slot to termination, recording onto ``runs[i]``'s trace. + + ``len(runs)`` must match the env's batch size (the wire framing tells us: + scalar ``terminated`` = 1, an ``[N]`` mask = N). """ if self.model is None: raise RuntimeError(f"{type(self).__name__} must set self.model in __init__") - if not self._active_chunk: - batch = ( - obs if self.adapter is None else self.adapter.adapt_observation(obs, self._prompt) - ) - chunk = np.atleast_2d(await self.model.ainfer(batch)) # [T, A] - self._active_chunk = deque(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) + prompt = runs[0].prompt + if not isinstance(prompt, str): + raise TypeError(f"run.prompt must be a str, got {type(prompt).__name__}: {prompt!r}") - async def __call__(self, run: Run, *, max_steps: int | None = None) -> None: - step_limit = max_steps if max_steps is not None else int(getattr(self, "max_steps", 520)) - cap = run.client.binding(self.robot_protocol) - client = await RobotClient.connect(cap) + robot = await RobotClient.connect(client.binding(self.robot_protocol)) try: - self.setup_robot(client) - prompt = run.prompt - if not isinstance(prompt, str): - raise TypeError( - f"run.prompt must be a str, got {type(prompt).__name__}: {prompt!r}" - ) - 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() - 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: - preview = np.array2string(action, precision=3, suppress_small=True) - print(f"[agent] step {step}/{step_limit} action={preview}", flush=True) - else: - print(f"[agent] reached max_steps={step_limit}", flush=True) - - run.trace.status = "completed" - run.trace.content = "done" + _, obs_space = robot.spaces() + if self.adapter is not None: + self.adapter.bind(*robot.spaces()) + self.adapter.reset() + + obs = await robot.get_observation() + single = np.ndim(obs["terminated"]) == 0 # wire framing: scalar vs [N] mask + n = 1 if single else int(np.asarray(obs["terminated"]).shape[0]) + if len(runs) != n: + raise ValueError(f"got {len(runs)} runs for an env batch of {n}") + + fps = robot.get_control_rate() + recorders = [ + # A live run (single path) records through it so steps land on + # run.trace for training; grouped receipts emit by trace id. + TraceRecorder(run=r, fps=fps, obs_space=obs_space) + if r._client is not None + else TraceRecorder(trace_id=r.trace_id, fps=fps, obs_space=obs_space) + for r in runs + ] + writer = None + if self.save: + if n == 1: + from .dataset import DatasetWriter + + writer = DatasetWriter(robot.contract, fps=fps) + else: + print("[agent] save=True is single-env only; streaming telemetry", flush=True) + + print(f"[agent] episode started: {prompt!r} (n={n})", flush=True) + await self._loop( + robot, + obs, + prompt, + recorders, + writer, + single=single, + max_steps=max_steps or self.max_steps, + ) + for rec in recorders: + rec.close() + if writer is not None: + writer.end_episode() finally: - if self._recorder is not None: - self._recorder.end() # flush video tails + commit the LeRobot episode - await client.close() + await robot.close() + + async def _loop( + self, + robot: RobotClient, + obs: dict[str, Any], + prompt: str, + recorders: list[TraceRecorder], + writer: Any, + *, + single: bool, + max_steps: int, + ) -> None: + """One batched forward per refill; execute chunks open-loop per slot.""" + adapter = self.adapter + n = len(recorders) + chunks: list[deque[np.ndarray]] = [deque() for _ in range(n)] + ever_done = np.zeros(n, dtype=bool) + + for step in range(max_steps): + done = np.atleast_1d(np.asarray(obs["terminated"], dtype=bool)).reshape(-1) + for i in np.nonzero(done)[0]: # a reset slot re-infers for its new episode + chunks[i].clear() + ever_done |= done + if step and ever_done.all(): + print(f"[agent] all slots terminated at step {step}", flush=True) + break + + # Batched view of the observation: single framing lifts to a batch of one. + data = obs["data"] if not single else {k: v[None] for k, v in obs["data"].items()} + for i, rec in enumerate(recorders): + if not ever_done[i]: + rec.record_observation({k: v[i] for k, v in data.items()}, tick=step) + + if any(not c for c in chunks): # refill spent slots with a fresh forward + # The adapter sees the wire framing (unbatched when single), so a + # BatchedModel can still stack samples across concurrent rollouts. + batch = adapter.adapt_observation(obs, prompt) if adapter else obs + if n == 1: + # ainfer is the coalescing point for cross-rollout batching + # (BatchedModel), so the single slot goes through it. + chunk = np.atleast_2d(await self.model.ainfer(batch))[None] # [1, T, A] + else: + chunk = np.asarray(await asyncio.to_thread(self.model.infer, batch)) + for i, c in enumerate(chunks): + if not c: + rows = chunk[i] + if adapter is not None: # e.g. deltas -> absolute vs the query obs + slot = {"data": {k: v[i] for k, v in data.items()}} + rows = adapter.adapt_chunk(rows, slot) + c.extend(rows) + recorders[i].record_inference(rows, tick=step) + + raw = [chunks[i].popleft() for i in range(n)] + if adapter is not None: # per-step execution-time hook (default identity) + raw = [ + adapter.adapt_action(a, {"data": {k: v[i] for k, v in data.items()}}) + for i, a in enumerate(raw) + ] + action = raw[0] if single else np.stack(raw) + if writer is not None: + writer.add(obs["data"], np.asarray(raw[0]), task=prompt) + await robot.send_action(action) + + if self.log_every and step % self.log_every == 0: + live = int((~ever_done).sum()) + print(f"[agent] step {step}/{max_steps} live={live}/{n}", flush=True) + obs = await robot.get_observation() + else: + print(f"[agent] reached max_steps={max_steps}", flush=True) __all__ = ["ROBOT_PROTOCOL", "RobotAgent"] diff --git a/hud/agents/robot/batching.py b/hud/agents/robot/batching.py index a24594488..9613a4149 100644 --- a/hud/agents/robot/batching.py +++ b/hud/agents/robot/batching.py @@ -18,7 +18,7 @@ if TYPE_CHECKING: from hud.eval.run import Run - from ._types import ActionArray + from .adapter import ActionArray from .agent import RobotAgent @@ -97,11 +97,10 @@ async def _batch_loop(self) -> None: 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. + Per run: a shallow clone of ``agent`` sharing a per-run adapter copy and the + single :class:`BatchedModel`, so concurrent ``ainfer`` calls coalesce into one + forward. The adapter copy keeps per-env bindings isolated; the model is + stateless by contract, so sharing it across clones is safe. Requires an in-process batchable model; :class:`~hud.agents.robot.model.RemoteModel` is not supported (the OpenPI server protocol has no batched-request shape). diff --git a/hud/agents/robot/model.py b/hud/agents/robot/model.py index 3429e4a7f..9504c489e 100644 --- a/hud/agents/robot/model.py +++ b/hud/agents/robot/model.py @@ -18,7 +18,7 @@ import numpy as np if TYPE_CHECKING: - from ._types import ActionArray + from .adapter import ActionArray class Model: diff --git a/hud/agents/robot/vec_agent.py b/hud/agents/robot/vec_agent.py deleted file mode 100644 index 3e866c2b6..000000000 --- a/hud/agents/robot/vec_agent.py +++ /dev/null @@ -1,123 +0,0 @@ -"""``VecRobotAgent`` — drive a whole vectorized robot env over one connection. - -The agent-side counterpart to :class:`~hud.environment.robot.isaac_bridge.IsaacBridge`: one -batched forward per tick. Receive an ``[N, ...]`` observation, run the model once to an -``[N, T, A]`` chunk, execute it open-loop per slot, and send one ``[N, A]`` action. The N -parallel episodes stream to the platform as one **Job** of per-episode traces via -:class:`~hud.agents.robot.record.VecRecorder` (slots split into fresh traces on each reset). - -Set ``self.model`` (and ``self.adapter``) exactly as for :class:`~hud.agents.robot.agent.RobotAgent`; -the model's ``infer`` is already ``[N, ...] -> [N, T, A]``, so :class:`~hud.agents.robot.model.LeRobotModel` -works unchanged. Pair with :class:`~hud.agents.robot.adapter.VecLeRobotAdapter` for the batched obs. -""" - -from __future__ import annotations - -import asyncio -from collections import deque -from typing import TYPE_CHECKING, Any, ClassVar - -import numpy as np - -from hud.capabilities.robot import RobotClient - -from .record import VecRecorder - -if TYPE_CHECKING: - from hud.capabilities.base import Capability - - from .adapter import Adapter - from .model import Model - - -class VecRobotAgent: - """Drive an N-env robot batch as one Job of per-episode traces. - - **Subclass contract:** set ``self.model`` (a :class:`~hud.agents.robot.model.Model`) and - ``self.adapter`` (a :class:`~hud.agents.robot.adapter.Adapter`) in ``__init__``. - """ - - model: Model | None = None - adapter: Adapter | None = None - #: Max control ticks before the run is cut off (the env auto-resets episodes within this). - max_steps: ClassVar[int] = 520 - #: Step-progress print frequency. 0 = off. - log_every: ClassVar[int] = 20 - - async def run( - self, - cap: Capability, - prompt: str, - *, - name: str, - num_record: int = 4, - seed: int | None = None, - group_id: str | None = None, - model_name: str | None = None, - ) -> str: - """Connect, drive the batch to ``max_steps`` (or all-terminated), return the Job URL.""" - if self.model is None: - raise RuntimeError(f"{type(self).__name__} must set self.model in __init__") - client = await RobotClient.connect(cap) - try: - action_space, obs_space = client.spaces() - if self.adapter is not None: - self.adapter.bind(action_space, obs_space) - return await self._drive( - client, prompt, name=name, num_record=num_record, seed=seed, - group_id=group_id, model_name=model_name, - ) - finally: - await client.close() - - async def _drive( - self, client: RobotClient, prompt: str, *, name: str, num_record: int, - seed: int | None, group_id: str | None, model_name: str | None, - ) -> str: - adapter = self.adapter - state_key = adapter.state_key if adapter else None - image_keys = adapter.image_keys if adapter else [] - - obs = await client.get_observation() - # Batch size from the leading dim of any obs array (state if present, else the first). - probe = obs["data"][state_key] if state_key else next(iter(obs["data"].values())) - n = int(np.asarray(probe).shape[0]) - rec = VecRecorder( - name, num_envs=n, record_indices=list(range(min(num_record, n))), - fps=client.get_control_rate(), seed=seed, group_id=group_id, model=model_name, - ) - print(f"[vec-agent] job: {rec.job_url} (N={n})", flush=True) - - chunks: list[deque[np.ndarray]] = [deque() for _ in range(n)] - for step in range(self.max_steps): - done = np.asarray(obs["terminated"], dtype=bool).reshape(-1) - for i in np.nonzero(done)[0]: # a freshly-reset slot re-infers for its new episode - chunks[i].clear() - if step and done.all(): - break - - if any(not c for c in chunks): # refill spent slots with a fresh batched chunk - batch = adapter.adapt_observation(obs, prompt) if adapter else obs - chunk = np.asarray(await asyncio.to_thread(self.model.infer, batch)) # [N, T, A] - for i, c in enumerate(chunks): - if not c: - rows = chunk[i] - if adapter is not None: # e.g. DROID delta -> absolute against slot i's query obs - rows = adapter.adapt_chunk(rows, {"data": {k: v[i] for k, v in obs["data"].items()}}) - c.extend(rows) - action = np.stack([chunks[i].popleft() for i in range(n)]) - - state = {state_key: obs["data"][state_key]} if state_key else None - frames = {k: obs["data"][k] for k in image_keys} if image_keys else None - rec.record(obs=state, frames=frames, action=action, done=done) - await client.send_action(action) - - if self.log_every and step % self.log_every == 0: - print(f"[vec-agent] step {step}/{self.max_steps} live={int((~done).sum())}/{n}", flush=True) - obs = await client.get_observation() - - rec.close() - return rec.job_url - - -__all__ = ["VecRobotAgent"] From 8b27cae90061c70e04bd7972ee3e9ea073decce7 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 20:58:04 +0000 Subject: [PATCH 07/43] feat(eval): grouped rollouts via rollout_group and Taskset.run(num_envs=) One vectorized env instance becomes num_envs graded runs sharing a group_id, behind the same Taskset.run surface and with no robot concepts in hud/eval: the atom goes through the normal env client and template (num_envs injected into task args; an env that doesn't accept it fails loudly), the contract rides the capability manifest, the agent's drive(runs, client) entry drives the slots, and run i grades from the template result's slots[i]. Replaces vec_rollout and the vectorized=/ contract= kwargs; num_envs composes with group= (k instances x N traces, group ids per instance) and requires a self-managed runtime. --- hud/eval/__init__.py | 3 +- hud/eval/run.py | 94 +++++++++++++++++++++++++++++++++++++++++++- hud/eval/taskset.py | 66 ++++++++++++++++++++++--------- 3 files changed, 143 insertions(+), 20 deletions(-) diff --git a/hud/eval/__init__.py b/hud/eval/__init__.py index 2824645c7..3ab966020 100644 --- a/hud/eval/__init__.py +++ b/hud/eval/__init__.py @@ -32,7 +32,7 @@ from .chat import Chat from .job import Job -from .run import Grade, Run, rollout +from .run import Grade, Run, rollout, rollout_group from .runtime import ( DaytonaRuntime, DockerRuntime, @@ -75,4 +75,5 @@ "Taskset", "Trace", "rollout", + "rollout_group", ] diff --git a/hud/eval/run.py b/hud/eval/run.py index 5949786bd..b99f397f7 100644 --- a/hud/eval/run.py +++ b/hud/eval/run.py @@ -415,4 +415,96 @@ async def _drive() -> None: return run -__all__ = ["Grade", "Run", "rollout"] +async def rollout_group( + task: Task, + agent: Any, + *, + runtime: Provider, + num_envs: int, + job_id: str | None = None, + group_id: str | None = None, + rollout_timeout: float | None = None, +) -> list[Run]: + """Drive one env instance's ``num_envs`` slots to graded runs (one group). + + The grouped counterpart of :func:`rollout`, domain-agnostic: one runtime, + one task start (``num_envs`` is injected into the task args — an env whose + template doesn't accept it fails loudly), then the agent drives every slot + over one connection via its ``drive(runs, client)`` entry, and the grade's + ``"slots"`` list (one dict per slot, the soft convention grouped envs + return) grades each run. The runs are *receipts*: the agent records spans + onto the trace ids minted here; this atom owns trace lifecycle and grading. + + A future scheduler may pack *different* compatible tasks into one + instance's slots; the atom's task -> runs shape already permits it. + """ + if not hasattr(agent, "drive"): + raise TypeError( + f"{type(agent).__name__} cannot drive a grouped rollout " + "(needs a drive(runs, client) entry, e.g. hud.agents.robot.RobotAgent)" + ) + if job_id is None: # a lone grouped rollout is a job of one instance + job_id = uuid.uuid4().hex + await job_enter(job_id, name=task.id, group=1) + group_id = group_id or uuid.uuid4().hex + + runs = [Run(None, task.id, task.args) for _ in range(num_envs)] + for run in runs: + run.trace.trace_id = uuid.uuid4().hex + run.job_id = job_id + run.group_id = group_id + run.slug = task.slug or task.default_slug() + + loop = asyncio.get_running_loop() + deadline = None if rollout_timeout is None else loop.time() + rollout_timeout + + async def _bounded(awaitable: Any) -> Any: + # One shared wall-clock deadline across provision, start, and the agent + # loop (see rollout._bounded for why a read-timeout is not enough). + if deadline is None: + return await awaitable + return await asyncio.wait_for(awaitable, max(deadline - loop.time(), 0.0)) + + _phase = "provisioning" + try: + async with contextlib.AsyncExitStack() as stack: + addr = cast("Runtime", await _bounded(stack.enter_async_context(runtime(task)))) + _phase = "starting task" + client = cast("HudClient", await _bounded(stack.enter_async_context(connect(addr)))) + args = {**task.args, "num_envs": num_envs} + prompt = (await _bounded(client.start_task(task.id, args))).get("prompt") + for run in runs: + run.prompt = prompt + run._runtime = addr.url + await trace_enter(run.trace_id, job_id=job_id, group_id=group_id, model=None) + with set_trace_context(run.trace_id): # opening user step per trace + run.record(Step(source="user", messages=run.prompt_messages)) + _phase = "agent loop" + await _bounded(agent.drive(runs, client)) + _phase = "grading" + evaluation = await client.grade({"answer": None}) + slots = evaluation.get("slots") or [] + if len(slots) != num_envs: + raise ValueError( + f"grade returned {len(slots)} slots for {num_envs} envs " + "(grouped envs must return one 'slots' entry per slot)" + ) + for run, slot in zip(runs, slots, strict=True): + run.grade = Grade.from_dict(slot) + run.trace.status = "completed" + except Exception as exc: # isolate: one bad instance never kills the batch + detail = ( + f"timed out after {rollout_timeout:.0f}s" + if isinstance(exc, TimeoutError) and rollout_timeout + else str(exc) + ) + logger.warning("grouped rollout failed (%s): %s", _phase, detail) + for run in runs: + run.trace.status = "error" + run.record(Step(source="system", error=f"[{_phase}] {detail}")) + for run in runs: + await trace_exit(run) + return runs + + +__all__ = ["Grade", "Run", "rollout", "rollout_group"] diff --git a/hud/eval/taskset.py b/hud/eval/taskset.py index 2ef0d223a..9c9bcd8fe 100644 --- a/hud/eval/taskset.py +++ b/hud/eval/taskset.py @@ -22,7 +22,7 @@ from hud.utils.platform import PlatformClient from .job import Job, job_enter -from .run import rollout +from .run import rollout, rollout_group from .runtime import HostedRuntime, HUDRuntime, LocalRuntime, _declared_env, _declared_names from .sync import fetch_taskset_tasks, resolve_taskset_id @@ -245,6 +245,7 @@ async def run( *, runtime: Provider | HostedRuntime | None = None, group: int | None = None, + num_envs: int | None = None, max_concurrent: int | None = None, job: Job | None = None, rollout_timeout: float | None = None, @@ -268,6 +269,14 @@ async def run( one id. Returned ``job.runs`` preserves expansion order (task-major, then group). + ``num_envs`` selects grouped execution (:func:`~hud.eval.run.rollout_group`): + each task instance runs one *vectorized* env whose N slots (each its own + seeded perturbation) become N graded traces sharing a group_id. Distinct + from ``group`` — statistical repeats as separate instances — and they + compose: ``group=3, num_envs=5`` is 3 instances x 5 traces per task. + ``max_concurrent`` always counts instances. Requires a self-managed + runtime and a ``drive``-capable agent (e.g. ``RobotAgent``). + ``rollout_timeout`` is a hard per-rollout wall-clock cap (seconds) for the local (Provider) path: a rollout that exceeds it is cancelled and recorded as a failed/errored run so one wedged rollout (e.g. a stuck sampling @@ -279,14 +288,19 @@ async def run( raise ValueError("group must be >= 1") if max_concurrent is not None and max_concurrent < 1: raise ValueError("max_concurrent must be >= 1") + if num_envs is not None and num_envs < 1: + raise ValueError("num_envs must be >= 1") - # Tasks are pure rows, shared across rollouts; the ``group`` repeats of - # one task share a group_id (the GRPO group). + # Tasks are pure rows, shared across rollouts. The ``group`` repeats of one + # task share a group_id (the GRPO group) — except under ``num_envs``, where + # each instance's N slot-traces are their own group. expanded: list[tuple[Task, str]] = [] task_list = list(self) for task in task_list: group_id = uuid.uuid4().hex - expanded.extend((task, group_id) for _ in range(group)) + expanded.extend( + (task, uuid.uuid4().hex if num_envs else group_id) for _ in range(group) + ) if job is None: job = Job( @@ -307,35 +321,51 @@ async def run( # an error naming the forms to pass. # An empty taskset schedules nothing, so it needs no placement. placement = runtime if runtime is not None or not task_list else self._resolve_placement() + if num_envs is not None and isinstance(placement, HostedRuntime): + raise ValueError("num_envs (grouped rollouts) requires a self-managed runtime") sem = asyncio.Semaphore(max_concurrent) if max_concurrent else None - async def _run(task: Task, group_id: str) -> Run: + async def _run(task: Task, group_id: str) -> list[Run]: assert placement is not None # only reached when tasks were expanded if isinstance(placement, HostedRuntime): - return await placement.run(task, agent, job_id=job_id, group_id=group_id) - return await rollout( - task, - agent, - runtime=placement, - job_id=job_id, - group_id=group_id, - rollout_timeout=rollout_timeout, - ) - - async def _one(task: Task, group_id: str) -> Run: + return [await placement.run(task, agent, job_id=job_id, group_id=group_id)] + if num_envs is not None: # grouped: one instance -> num_envs graded runs + return await rollout_group( + task, + agent, + runtime=placement, + num_envs=num_envs, + job_id=job_id, + group_id=group_id, + rollout_timeout=rollout_timeout, + ) + return [ + await rollout( + task, + agent, + runtime=placement, + job_id=job_id, + group_id=group_id, + rollout_timeout=rollout_timeout, + ) + ] + + async def _one(task: Task, group_id: str) -> list[Run]: if sem is None: return await _run(task, group_id) async with sem: return await _run(task, group_id) logger.info( - "running %d rollouts (%d tasks x %d group)%s", + "running %d rollouts (%d tasks x %d group)%s%s", len(expanded), len(task_list), group, + f" x {num_envs} envs" if num_envs else "", f", max_concurrent={max_concurrent}" if max_concurrent else "", ) - job.runs.extend(await asyncio.gather(*(_one(t, gid) for t, gid in expanded))) + waves = await asyncio.gather(*(_one(t, gid) for t, gid in expanded)) + job.runs.extend(run for wave in waves for run in wave) # Drain telemetry before returning. The exporter uploads in parallel and # flush is completion-based (waits for in-flight uploads, not a fixed # sleep), so the timeout is only a safety cap for a wedged network. From 9e222ff080bc606dd51d5c00f7a4bccefbaaea02 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 20:58:04 +0000 Subject: [PATCH 08/43] docs(robots): update for the unified robot stack env.gym-first environment side, the one sim-serving process shape (SimRunner section replaced), slots-based grading, serve_blocking for split-process sims, a new vectorized-envs-and-grouped-evals section, and a refreshed API summary. --- docs/v6/advanced/robots.mdx | 119 ++++++++++++++++++++++++------------ 1 file changed, 81 insertions(+), 38 deletions(-) diff --git a/docs/v6/advanced/robots.mdx b/docs/v6/advanced/robots.mdx index 40d78d03e..249db2b3a 100644 --- a/docs/v6/advanced/robots.mdx +++ b/docs/v6/advanced/robots.mdx @@ -91,7 +91,25 @@ The agent wires observations to policy inputs purely from the manifest; there is ## Environment side -You implement one class - the **bridge**. +For a gym-style sim you implement nothing: `env.gym(make_env)` takes any callable returning a +gym-style env, derives the [contract](#contract) from a sample observation, serves the sim over the +``robot`` WebSocket, and returns the handle templates drive episodes through: + +```python env.py +from hud import Environment + +env = Environment(name="my-sim") +sim = env.gym(make_env) # make_env: any callable returning a gym-style env + +@env.template() +async def pick_and_place(task: str = "default", seed: int = 0): + yield {"prompt": await sim.reset(task=task, seed=seed)} + yield await sim.result() +``` + +Task args are partitioned by the factory's signature: args the factory accepts define the env (a +change rebuilds it), everything else is episodic and flows to `env.reset(seed=..., options=...)`. +For a sim that isn't gym-shaped, subclass the **bridge** instead: ```python from hud.environment.robot import RobotBridge @@ -182,9 +200,13 @@ async def _down(): @env.template() async def pick_and_place(task_id: str, seed: int = 0): prompt = yield {"prompt": await endpoint.reset(task_id=task_id, seed=seed)} - yield await endpoint.result() # {"score", "success", "total_reward"} + yield await endpoint.result() # {"score", "success", "total_reward", "slots": [...]} ``` +The result dict carries a `"slots"` list with one score dict per env slot (a single env is a batch +of one). Aggregates sit at the top; the [grouped eval](#vectorized-envs-and-grouped-evals) path +grades each trace from its slot. + ## Agent side The harness lives in `hud.agents.robot`. @@ -208,7 +230,7 @@ is a complete agent in a few lines. The two bundled seams *are* the LeRobot conv pre/post-processors, so the checkpoint behaves exactly as it does upstream. Pass an `Ensembler` to reduce overlapping action chunks to one action per step. - **`LeRobotAdapter(model_image_keys=...)`** maps the env's cameras and state onto the policy's - inputs from the [contract](#the-contract) - HWC `uint8` → CHW float, state and prompt passed + inputs from the [contract](#contract) - HWC `uint8` → CHW float, state and prompt passed through. ```python @@ -291,6 +313,33 @@ has no batched-request shape); run one agent per rollout against it instead. `Ba ownership of the agent you pass (it swaps in the batched model in place), so give it a dedicated instance rather than one you also use for unbatched runs. +## Vectorized envs and grouped evals + +A GPU sim like Isaac runs many env copies in one instance, so batching at the model isn't enough - +the *environment itself* is vectorized. Three roles stay decoupled: + +- **The environment declares what it is.** A sim factory accepting `num_envs` is the declaration; + the template passes it through (`await sim.reset(task=task, seed=seed, num_envs=num_envs)`). Any + vectorized env still runs as a plain single env when `num_envs` is 1. +- **The agent declares what it can drive.** A stock `RobotAgent` drives N >= 1 slots over one + connection with one batched forward per refill - nothing extra to implement. +- **The run selects how.** `num_envs` on `Taskset.run` turns each task into one vectorized + instance whose N slots (each its own seeded perturbation) become N graded traces sharing a + group: + +```python +job = await taskset.run( + MyAgent(), + runtime=Runtime("tcp://127.0.0.1:9100"), + num_envs=5, # slots per instance -> 5 graded traces per task + max_concurrent=2, # instances in flight (each has its own runtime) +) +``` + +Distinct from `group` - statistical repeats as separate instances - and they compose: +`group=3, num_envs=5` is 3 instances x 5 traces per task. Grading is per slot: the engine reads +the template result's `"slots"` list and grades trace `i` from `slots[i]`. + ## Contract Embodiments and policies disagree on cameras, state layout, action semantics, and control rate, so @@ -361,21 +410,18 @@ spec - the closed symbol sets and known traps - lives outside the SDK alongside -## Sim threading +## How sims run The loop is lockstep - the bridge steps the sim once per received action. A simulator is usually -**thread-affine** (every touch must run on the thread that created its GL/device context), but the -bridge's asyncio loop can't be stalled by a blocking step. **`SimRunner`** is the one-line injection -that decides *which thread* runs the sim; the bridge routes every sim touch through it: - -- **`InlineSimRunner`** - runs on the event-loop thread. The default; for cheap/CPU sims and tests. -- **`ThreadSimRunner`** - sim on a dedicated worker thread, leaving the loop free during a blocking - step. For render-heavy or thread-bound sims. -- **`MainThreadSimRunner`** - sim on the main thread, for runtimes that own *both* the main thread - and the loop (Isaac/Omniverse); the owner's pump loop drains queued sim touches between ticks. +**thread-affine** (every touch must run on the thread that created its GL/device context), and some +runtimes - Isaac/Omniverse - must own the process main thread outright. HUD therefore serves every +sim with one process shape: **the sim owns the main thread**, serving (the control channel and the +robot WebSocket) runs on a background loop thread, and the framework queues every sim touch back to +the main thread. A cheap CPU sim blocks on that queue; when Omniverse is loaded, the main thread +also pumps Kit between touches. -Pass one to the bridge (`RobotBridge(sim_runner=ThreadSimRunner())`), or subclass `SimRunner` for an -exotic topology. +There is nothing to configure. `hud serve` and `python -m hud.environment.server` detect an +attached sim and pick the shape; any fix to the loop applies to every sim the same way. ## Recording datasets @@ -388,8 +434,9 @@ it already produces, so it runs in *your* process - not the environment containe sims (e.g. Isaac/RoboLab) whose dependency stack conflicts with `lerobot`; only your machine needs `pip install 'lerobot[dataset]'`. -One dataset spans the whole run - every episode the shared agent drives appends to it - and is -finalized at process exit. Destination and Hub push come from the environment: +Each rollout writes its own dataset (tagged with its trace id, so a shard maps back to its +trace), finalized at process exit. Saving is single-env; a vectorized run streams telemetry only. +Destination and Hub push come from the environment: | Env var | Effect | |---------|--------| @@ -397,7 +444,7 @@ finalized at process exit. Destination and Hub push come from the environment: | `HF_REPO` | Also push the finalized dataset to this HF namespace (needs `HF_TOKEN`) | | `HF_PRIVATE` | Push the dataset private | -The [contract](#the-contract) drives the schema with no extra wiring: image features become +The [contract](#contract) drives the schema with no extra wiring: image features become `observation.images.` (encoded to per-episode video), the lone state vector becomes `observation.state`, the action becomes `action`, and the task prompt rides along as each frame's `task`. @@ -405,18 +452,17 @@ The [contract](#the-contract) drives the schema with no extra wiring: image feat ## Running a sim in another process -Some simulators must **own the process main thread** - most notably **Isaac Sim / Omniverse**, where -Kit drives its own main-thread event loop and `env.reset()` loads USD through a nested -`run_until_complete`. That can't run inside `hud serve`, which already owns the asyncio loop. The fix -is to move the sim into its own process and keep the env code essentially unchanged. +Serving in one process is the default (see [how sims run](#how-sims-run)), but the sim can also +live in its own process - on another machine, or built against a dependency stack the env +container can't carry. The env code stays essentially unchanged. `RobotEndpoint` is built for exactly this: the same control surface (`start` / `reset` / `result` / `stop`) works whether the bridge is local or remote. - **Env process** - publish a *remote* handle with `RobotEndpoint.remote(host, port)`. It dials the sim process and forwards every control call over JSON-RPC. -- **Sim process** - wrap the real bridge and expose it with `RobotEndpoint(bridge).serve(host, port)`, - using a [`MainThreadSimRunner`](#sim-threading) so every sim touch runs on the main thread. +- **Sim process** - wrap the real bridge and expose it with + `RobotEndpoint(bridge).serve_blocking(host, port)`, the same sim-owns-main shape `hud serve` uses. The two planes split cleanly, which is why the agent never knows the sim is remote: @@ -451,19 +497,14 @@ async def pick_and_place(task_id: str, seed: int = 0): yield await endpoint.result() ``` -**Sim process** - your Isaac program builds the bridge and serves its control surface, then runs for -the process's lifetime: +**Sim process** - your program builds the bridge and serves its control surface, blocking for the +process's lifetime with the sim owning the main thread: ```python sim_main.py -import asyncio -from hud.environment.robot import RobotEndpoint, MainThreadSimRunner - -async def main(): - bridge = MySimBridge(sim_runner=MainThreadSimRunner()) # sim touches run on main - server = await RobotEndpoint(bridge).serve("127.0.0.1", 9100) - await server.wait_closed() +from hud.environment.robot import RobotEndpoint -asyncio.run(main()) # launched on the main thread the sim owns +bridge = MySimBridge() +RobotEndpoint(bridge).serve_blocking("127.0.0.1", 9100) ``` Bring the two up together - the env's `connect()` retries until the sim is listening. Everything @@ -474,15 +515,17 @@ downstream (`hud eval`, tasksets, the agent) is unchanged; only *where the bridg | Symbol | Where | Role | |--------|-------|------| +| `env.gym(make_env)` / `Gym` | `hud.environment.robot` | Declarative sim: contract, capability, and serving derived from a gym factory | +| `hud.wrap(env)` | `hud` | One-line trace streaming for any gym-style env you drive yourself | +| `RobotBridge` / `GymBridge` | `hud.environment.robot` | Env-side serve loop (batched-first); subclass `RobotBridge` for a custom sim | +| `RobotEndpoint` | `hud.environment.robot` | Episode bookkeeping + results (local or `.remote()`); `serve_blocking()` for a split process | | `RobotEndpoint.capability(contract=...)` | `hud.environment.robot` | Build the `openpi/0` capability after `start()` | | `Capability.robot(name, url, contract)` | `hud.capabilities` | Lower-level constructor (usually via `endpoint.capability`) | | `RobotClient` | `hud.capabilities.robot` | Agent-side wire client (`spaces`, `get_observation`, `send_action`, `send_chunk`) | -| `RobotBridge` | `hud.environment.robot` | Env-side serve loop; subclass with your sim | -| `RobotEndpoint` | `hud.environment.robot` | Episode bookkeeping + results (local or `.remote()`) | -| `SimRunner` (`Inline`/`Thread`/`MainThread`) | `hud.environment.robot` | Which thread runs the sim | -| `RobotAgent` | `hud.agents.robot` | The episode-loop harness | +| `RobotAgent` | `hud.agents.robot` | The harness: one open-loop chunk queue driving N >= 1 env slots | | `Model` / `LeRobotModel`, `Adapter` / `LeRobotAdapter` | `hud.agents.robot` | Policy + space-translation seams | | `BatchedAgent` / `BatchedModel` | `hud.agents.robot.batching` | Many concurrent rollouts against one shared, batched model | +| `Taskset.run(..., num_envs=N)` | `hud.eval` | Grouped eval: one vectorized instance per task, N graded traces | ## See also From 47264de427320260d60017766e2540dcd30ef284 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 23:03:22 +0000 Subject: [PATCH 09/43] feat(agents): size BatchedModel to the live concurrency by default batch_size is now an optional cap: with no value the worker stacks every ainfer call queued within the max_wait_s window, so the forward is as wide as whatever is in flight and max_concurrent never needs manual pairing. A set cap still flushes the window early once reached. --- hud/agents/robot/batching.py | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/hud/agents/robot/batching.py b/hud/agents/robot/batching.py index 9613a4149..b300274bd 100644 --- a/hud/agents/robot/batching.py +++ b/hud/agents/robot/batching.py @@ -25,10 +25,13 @@ 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. + A lazily-started worker waits up to ``max_wait_s`` for callers to queue, stacks + them into one ``[N, ...]`` batch, runs a single forward, and scatters the + ``[N, T, A]`` rows back to each caller. With no ``batch_size`` the batch sizes + itself to whatever is in flight — the scheduler's ``max_concurrent`` already + bounds that, so the two never need manual pairing. Pass ``batch_size`` only to + cap the forward below the live concurrency (e.g. VRAM headroom); a set cap also + flushes the window early once reached, saving the tail of ``max_wait_s``. ``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`). @@ -37,9 +40,11 @@ class BatchedModel(Model): 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: + def __init__( + self, inner: Model, *, batch_size: int | None = None, max_wait_s: float = 0.05 + ) -> None: self.inner = inner - self.batch_size = int(batch_size) + self.batch_size = None if batch_size is None else 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 @@ -64,7 +69,7 @@ async def _batch_loop(self) -> None: 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: + while self.batch_size is None or len(items) < self.batch_size: timeout = deadline - loop.time() if timeout <= 0: break @@ -111,7 +116,9 @@ class BatchedAgent(Agent): 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: + def __init__( + self, agent: RobotAgent, *, batch_size: int | None = None, max_wait_s: float = 0.05 + ) -> None: if agent.model is None: raise RuntimeError("BatchedAgent needs agent.model set") self._template = agent From 4d05aae010b74c97fdce4a4c6a115bcb13b62344 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 23:03:22 +0000 Subject: [PATCH 10/43] refactor(eval): rename rollout_group to vec_rollout "group" already means statistical repeats (group=k); the vectorized atom deserves its own word. Docstrings and error messages now say "vectorized" throughout and the atom's docstring is tightened. Also narrows the minted trace id before trace_enter (pyright). --- hud/environment/robot/bridge.py | 2 +- hud/eval/__init__.py | 4 ++-- hud/eval/run.py | 40 ++++++++++++++++----------------- hud/eval/taskset.py | 10 ++++----- 4 files changed, 28 insertions(+), 28 deletions(-) diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py index d16dfd27b..187fe9247 100644 --- a/hud/environment/robot/bridge.py +++ b/hud/environment/robot/bridge.py @@ -108,7 +108,7 @@ def result_slots(self) -> list[dict[str, Any]]: def result(self) -> dict[str, Any]: """The episode grade: per-slot dicts under ``"slots"``, means at the top. - The grouped eval path grades each trace from ``slots[i]``; single-result + The vectorized eval path grades each trace from ``slots[i]``; single-result consumers read the aggregate ``score``/``success`` unchanged. """ slots = self.result_slots() diff --git a/hud/eval/__init__.py b/hud/eval/__init__.py index 3ab966020..e27896f80 100644 --- a/hud/eval/__init__.py +++ b/hud/eval/__init__.py @@ -32,7 +32,7 @@ from .chat import Chat from .job import Job -from .run import Grade, Run, rollout, rollout_group +from .run import Grade, Run, rollout, vec_rollout from .runtime import ( DaytonaRuntime, DockerRuntime, @@ -75,5 +75,5 @@ "Taskset", "Trace", "rollout", - "rollout_group", + "vec_rollout", ] diff --git a/hud/eval/run.py b/hud/eval/run.py index b99f397f7..41592cd9c 100644 --- a/hud/eval/run.py +++ b/hud/eval/run.py @@ -415,7 +415,7 @@ async def _drive() -> None: return run -async def rollout_group( +async def vec_rollout( task: Task, agent: Any, *, @@ -425,25 +425,23 @@ async def rollout_group( group_id: str | None = None, rollout_timeout: float | None = None, ) -> list[Run]: - """Drive one env instance's ``num_envs`` slots to graded runs (one group). - - The grouped counterpart of :func:`rollout`, domain-agnostic: one runtime, - one task start (``num_envs`` is injected into the task args — an env whose - template doesn't accept it fails loudly), then the agent drives every slot - over one connection via its ``drive(runs, client)`` entry, and the grade's - ``"slots"`` list (one dict per slot, the soft convention grouped envs - return) grades each run. The runs are *receipts*: the agent records spans - onto the trace ids minted here; this atom owns trace lifecycle and grading. - - A future scheduler may pack *different* compatible tasks into one - instance's slots; the atom's task -> runs shape already permits it. + """Drive one vectorized env instance to ``num_envs`` graded runs. + + The vectorized counterpart of :func:`rollout`, still domain-agnostic: one + runtime, one task start (``num_envs`` injected into the task args; an env + whose template doesn't accept it fails loudly), the agent's + ``drive(runs, client)`` entry drives every slot over one connection, and + the grade's ``"slots"`` list grades run ``i``. The runs are *receipts* — + this atom owns trace lifecycle and grading. Its task -> runs shape also + permits a future scheduler packing different compatible tasks into one + instance's slots. """ if not hasattr(agent, "drive"): raise TypeError( - f"{type(agent).__name__} cannot drive a grouped rollout " + f"{type(agent).__name__} cannot drive a vectorized rollout " "(needs a drive(runs, client) entry, e.g. hud.agents.robot.RobotAgent)" ) - if job_id is None: # a lone grouped rollout is a job of one instance + if job_id is None: # a lone vectorized rollout is a job of one instance job_id = uuid.uuid4().hex await job_enter(job_id, name=task.id, group=1) group_id = group_id or uuid.uuid4().hex @@ -476,8 +474,10 @@ async def _bounded(awaitable: Any) -> Any: for run in runs: run.prompt = prompt run._runtime = addr.url - await trace_enter(run.trace_id, job_id=job_id, group_id=group_id, model=None) - with set_trace_context(run.trace_id): # opening user step per trace + tid = run.trace.trace_id + assert tid is not None # minted above + await trace_enter(tid, job_id=job_id, group_id=group_id, model=None) + with set_trace_context(tid): # opening user step per trace run.record(Step(source="user", messages=run.prompt_messages)) _phase = "agent loop" await _bounded(agent.drive(runs, client)) @@ -487,7 +487,7 @@ async def _bounded(awaitable: Any) -> Any: if len(slots) != num_envs: raise ValueError( f"grade returned {len(slots)} slots for {num_envs} envs " - "(grouped envs must return one 'slots' entry per slot)" + "(vectorized envs must return one 'slots' entry per slot)" ) for run, slot in zip(runs, slots, strict=True): run.grade = Grade.from_dict(slot) @@ -498,7 +498,7 @@ async def _bounded(awaitable: Any) -> Any: if isinstance(exc, TimeoutError) and rollout_timeout else str(exc) ) - logger.warning("grouped rollout failed (%s): %s", _phase, detail) + logger.warning("vectorized rollout failed (%s): %s", _phase, detail) for run in runs: run.trace.status = "error" run.record(Step(source="system", error=f"[{_phase}] {detail}")) @@ -507,4 +507,4 @@ async def _bounded(awaitable: Any) -> Any: return runs -__all__ = ["Grade", "Run", "rollout", "rollout_group"] +__all__ = ["Grade", "Run", "rollout", "vec_rollout"] diff --git a/hud/eval/taskset.py b/hud/eval/taskset.py index 9c9bcd8fe..5251ef783 100644 --- a/hud/eval/taskset.py +++ b/hud/eval/taskset.py @@ -22,7 +22,7 @@ from hud.utils.platform import PlatformClient from .job import Job, job_enter -from .run import rollout, rollout_group +from .run import rollout, vec_rollout from .runtime import HostedRuntime, HUDRuntime, LocalRuntime, _declared_env, _declared_names from .sync import fetch_taskset_tasks, resolve_taskset_id @@ -269,7 +269,7 @@ async def run( one id. Returned ``job.runs`` preserves expansion order (task-major, then group). - ``num_envs`` selects grouped execution (:func:`~hud.eval.run.rollout_group`): + ``num_envs`` selects vectorized execution (:func:`~hud.eval.run.vec_rollout`): each task instance runs one *vectorized* env whose N slots (each its own seeded perturbation) become N graded traces sharing a group_id. Distinct from ``group`` — statistical repeats as separate instances — and they @@ -322,15 +322,15 @@ async def run( # An empty taskset schedules nothing, so it needs no placement. placement = runtime if runtime is not None or not task_list else self._resolve_placement() if num_envs is not None and isinstance(placement, HostedRuntime): - raise ValueError("num_envs (grouped rollouts) requires a self-managed runtime") + raise ValueError("num_envs (vectorized rollouts) requires a self-managed runtime") sem = asyncio.Semaphore(max_concurrent) if max_concurrent else None async def _run(task: Task, group_id: str) -> list[Run]: assert placement is not None # only reached when tasks were expanded if isinstance(placement, HostedRuntime): return [await placement.run(task, agent, job_id=job_id, group_id=group_id)] - if num_envs is not None: # grouped: one instance -> num_envs graded runs - return await rollout_group( + if num_envs is not None: # vectorized: one instance -> num_envs graded runs + return await vec_rollout( task, agent, runtime=placement, From d4f7efa4b9e8bb00258bbf66ab29e42d7eb9a539 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 23:03:40 +0000 Subject: [PATCH 11/43] fix(agents): satisfy strict pyright in the robot agent stack hud/agents is strict-typed in CI: restore the stubless-torch Any alias in LeRobotAdapter, type the chunk queues as deque[ActionArray], narrow the optional model to a local before the loop, and suppress the deliberate private _client check when picking recorder mode. --- hud/agents/robot/adapter.py | 7 ++++--- hud/agents/robot/agent.py | 25 ++++++++++++++----------- 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/hud/agents/robot/adapter.py b/hud/agents/robot/adapter.py index 51b01021f..8020ffca6 100644 --- a/hud/agents/robot/adapter.py +++ b/hud/agents/robot/adapter.py @@ -79,17 +79,18 @@ class LeRobotAdapter(Adapter): """ def adapt_observation(self, obs: dict[str, Any], prompt: str) -> dict[str, Any]: - import torch + import torch # pyright: ignore[reportMissingImports] + 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.from_numpy(state), + "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): - img = torch.from_numpy(np.asarray(data[env_key])) + 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 diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py index 9ce3dafd5..7fa7d7fa1 100644 --- a/hud/agents/robot/agent.py +++ b/hud/agents/robot/agent.py @@ -9,8 +9,8 @@ Two entries, one loop: - ``__call__(run)`` — the generic rollout contract (one run, one trace). -- ``drive(runs, client)`` — the grouped-eval entry - (:func:`hud.eval.run.rollout_group`): N runs sharing one env instance, spans +- ``drive(runs, client)`` — the vectorized-eval entry + (:func:`hud.eval.run.vec_rollout`): N runs sharing one env instance, spans recorded per slot onto each run's trace. Most policies use :class:`~.adapter.LeRobotAdapter`; a policy whose spaces @@ -33,7 +33,7 @@ from hud.clients.client import HudClient from hud.eval.run import Run - from .adapter import Adapter + from .adapter import ActionArray, Adapter from .model import Model ROBOT_PROTOCOL = "openpi/0" @@ -98,9 +98,9 @@ async def drive( fps = robot.get_control_rate() recorders = [ # A live run (single path) records through it so steps land on - # run.trace for training; grouped receipts emit by trace id. + # run.trace for training; vectorized receipts emit by trace id. TraceRecorder(run=r, fps=fps, obs_space=obs_space) - if r._client is not None + if r._client is not None # pyright: ignore[reportPrivateUsage] else TraceRecorder(trace_id=r.trace_id, fps=fps, obs_space=obs_space) for r in runs ] @@ -142,15 +142,18 @@ async def _loop( max_steps: int, ) -> None: """One batched forward per refill; execute chunks open-loop per slot.""" + model = self.model + assert model is not None # checked in drive() adapter = self.adapter n = len(recorders) - chunks: list[deque[np.ndarray]] = [deque() for _ in range(n)] + chunks: list[deque[ActionArray]] = [deque() for _ in range(n)] ever_done = np.zeros(n, dtype=bool) for step in range(max_steps): done = np.atleast_1d(np.asarray(obs["terminated"], dtype=bool)).reshape(-1) - for i in np.nonzero(done)[0]: # a reset slot re-infers for its new episode - chunks[i].clear() + for c, d in zip(chunks, done, strict=True): + if d: # a reset slot re-infers for its new episode + c.clear() ever_done |= done if step and ever_done.all(): print(f"[agent] all slots terminated at step {step}", flush=True) @@ -169,9 +172,9 @@ async def _loop( if n == 1: # ainfer is the coalescing point for cross-rollout batching # (BatchedModel), so the single slot goes through it. - chunk = np.atleast_2d(await self.model.ainfer(batch))[None] # [1, T, A] + chunk = np.atleast_2d(await model.ainfer(batch))[None] # [1, T, A] else: - chunk = np.asarray(await asyncio.to_thread(self.model.infer, batch)) + chunk = np.asarray(await asyncio.to_thread(model.infer, batch)) for i, c in enumerate(chunks): if not c: rows = chunk[i] @@ -181,7 +184,7 @@ async def _loop( c.extend(rows) recorders[i].record_inference(rows, tick=step) - raw = [chunks[i].popleft() for i in range(n)] + raw: list[ActionArray] = [chunks[i].popleft() for i in range(n)] if adapter is not None: # per-step execution-time hook (default identity) raw = [ adapter.adapt_action(a, {"data": {k: v[i] for k, v in data.items()}}) From 2fe6bfd5314cedb0f44bb5219cf68f8cc8c4c01f Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 23:03:40 +0000 Subject: [PATCH 12/43] fix(agents): make DatasetWriter work with derived contracts Cameras in derived gym contracts are tagged type: rgb (no dtype: image), so frames mapped to non-video LeRobot columns. Image detection now accepts both conventions; missing dtypes default to float32 and missing shapes are filled from the first real frame before the dataset is created. Also imports importlib.util explicitly (pyright). --- hud/agents/robot/dataset.py | 30 ++++++++++++++++++++++++------ 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/hud/agents/robot/dataset.py b/hud/agents/robot/dataset.py index 7da02da59..1076e2ad1 100644 --- a/hud/agents/robot/dataset.py +++ b/hud/agents/robot/dataset.py @@ -14,7 +14,7 @@ from __future__ import annotations import atexit -import importlib +import importlib.util import logging import os import time @@ -43,7 +43,7 @@ def _lerobot_features(contract: dict[str, Any]) -> tuple[dict[str, dict[str, Any vectors = [ n for n, f in feats.items() - if f.get("role") == "observation" and f.get("dtype") not in ("image", "string") + if f.get("role") == "observation" and not _is_image(f) and f.get("dtype") != "string" ] single_state = len(vectors) == 1 @@ -53,24 +53,36 @@ def _lerobot_features(contract: dict[str, Any]) -> tuple[dict[str, dict[str, Any 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": + if _is_image(f): 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": _names(f, leaf)} + # Derived contracts omit dtype/shape; default the dtype, and leave a + # missing shape empty for add() to fill from the first real frame. + features[key] = {"dtype": dtype or "float32", "shape": shape, "names": _names(f, leaf)} key_map[name] = key elif role == "action": - features["action"] = {"dtype": dtype, "shape": shape, "names": _names(f, "act")} + features["action"] = { + "dtype": dtype or "float32", + "shape": shape, + "names": _names(f, "act"), + } return features, key_map +def _is_image(feature: dict[str, Any]) -> bool: + """A camera feature: authored contracts say ``dtype: image``, derived ones tag + the (load-bearing) image ``type`` — accept both.""" + return feature.get("dtype") == "image" or feature.get("type") in ("rgb", "bgr", "gray", "depth") + + def _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": + if _is_image(feature): return ["height", "width", "channel"] return [f"{base}_{i}" for i in range(int((feature.get("shape") or [1])[0]))] @@ -99,6 +111,12 @@ def add(self, data: dict[str, Any], action: NDArray[Any], *, task: str) -> None: """One frame: the wire observation dict + the executed env-space action.""" if not self._enabled: return + if self._ds is None: # derived contracts carry no shapes; fill from the first frame + for wire, key in self._key_map.items(): + if not self._features[key]["shape"] and wire in data: + self._features[key]["shape"] = tuple(np.shape(data[wire])) + if not self._features["action"]["shape"]: + self._features["action"]["shape"] = tuple(np.shape(action)) ds = self._ensure_dataset() row: dict[str, Any] = {} for wire, key in self._key_map.items(): From 923e546703e5c9dbfabeea04a760fd08c4e8f46d Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 23:03:40 +0000 Subject: [PATCH 13/43] fix(agents): stop recording inference for finished env slots Finished slots keep acting so the [N, A] action frame stays complete, but their chunk refills no longer emit InferenceSteps - matching the existing observation guard, so a finished slot's trace gains no orphan inference. --- hud/agents/robot/agent.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py index 7fa7d7fa1..8670a83a8 100644 --- a/hud/agents/robot/agent.py +++ b/hud/agents/robot/agent.py @@ -182,7 +182,8 @@ async def _loop( slot = {"data": {k: v[i] for k, v in data.items()}} rows = adapter.adapt_chunk(rows, slot) c.extend(rows) - recorders[i].record_inference(rows, tick=step) + if not ever_done[i]: # finished slots still act, but stop recording + recorders[i].record_inference(rows, tick=step) raw: list[ActionArray] = [chunks[i].popleft() for i in range(n)] if adapter is not None: # per-step execution-time hook (default identity) From 7710838af4cd976fabe6eba300abba9d0a247fab Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 6 Jul 2026 23:03:50 +0000 Subject: [PATCH 14/43] docs(robots): run parameters table; batching and vec_rollout updates Adds a Taskset.run parameter reference (runtime / group / num_envs / max_concurrent / rollout_timeout / job) under the vectorized-evals section, updates the batching example for the self-sizing BatchedModel, and aligns wording with the vec_rollout rename. --- docs/v6/advanced/robots.mdx | 32 +++++++++++++++++++++++++------- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/docs/v6/advanced/robots.mdx b/docs/v6/advanced/robots.mdx index 249db2b3a..1f6ec5766 100644 --- a/docs/v6/advanced/robots.mdx +++ b/docs/v6/advanced/robots.mdx @@ -204,7 +204,7 @@ async def pick_and_place(task_id: str, seed: int = 0): ``` The result dict carries a `"slots"` list with one score dict per env slot (a single env is a batch -of one). Aggregates sit at the top; the [grouped eval](#vectorized-envs-and-grouped-evals) path +of one). Aggregates sit at the top; the [vectorized eval](#vectorized-envs-and-evals) path grades each trace from its slot. ## Agent side @@ -297,15 +297,16 @@ into a single stacked forward. ```python from hud.agents.robot.batching import BatchedAgent -agent = BatchedAgent(MyRobotAgent(), batch_size=8) # set batch_size to your max_concurrent +agent = BatchedAgent(MyRobotAgent()) # sizes itself to the live concurrency job = await taskset.run(agent, runtime=DockerRuntime("hud-libero-env"), max_concurrent=8) ``` -Each tick, up to `batch_size` concurrent observations are stacked into one `[N, ...]` batch, run -through a single forward, and scattered back to each rollout; a short `max_wait_s` (default 50 ms) -flushes early when fewer rollouts are live, so the tail of a suite never stalls. Set `batch_size` to -match `max_concurrent` so a full batch can form. +Each tick, the concurrent observations queued within a short window (`max_wait_s`, default 50 ms) +are stacked into one `[N, ...]` batch, run through a single forward, and scattered back to each +rollout - so the batch is as wide as whatever is in flight, and the tail of a suite never stalls. +`max_concurrent` already bounds the width; pass `batch_size=` only to cap the forward below it +(e.g. VRAM headroom). `BatchedAgent` needs an **in-process, stateless** model whose `infer` runs the whole `[N, ...]` batch in one forward - `LeRobotModel` qualifies. A remote OpenPI policy server is not supported (its protocol @@ -313,7 +314,7 @@ has no batched-request shape); run one agent per rollout against it instead. `Ba ownership of the agent you pass (it swaps in the batched model in place), so give it a dedicated instance rather than one you also use for unbatched runs. -## Vectorized envs and grouped evals +## Vectorized envs and evals A GPU sim like Isaac runs many env copies in one instance, so batching at the model isn't enough - the *environment itself* is vectorized. Three roles stay decoupled: @@ -340,6 +341,23 @@ Distinct from `group` - statistical repeats as separate instances - and they com `group=3, num_envs=5` is 3 instances x 5 traces per task. Grading is per slot: the engine reads the template result's `"slots"` list and grades trace `i` from `slots[i]`. +### Run parameters + +Everything about *how* an eval executes lives on `Taskset.run` (the future `hud eval` CLI maps +onto the same parameters): + +| Parameter | What it determines | +|-----------|--------------------| +| `runtime` | *Where* each instance runs: a provider (`Runtime(url)`, `LocalRuntime`, `DockerRuntime`, ...) acquired once per instance, or `HostedRuntime` to delegate rollouts to the platform. Unset: local source if the tasks share one, else the HUD tunnel. | +| `group` | Statistical repeats: each task runs as `group` *independent instances* (own runtime, reset, connection). Without `num_envs` the repeats share a `group_id`. | +| `num_envs` | Slots per instance: each task runs as *one vectorized instance* whose N seeded perturbations become N graded traces sharing a `group_id`. Composes with `group` (instances x slots). Needs a `drive`-capable agent and a self-managed runtime. | +| `max_concurrent` | How many *instances* are in flight at once - it never counts slots. | +| `rollout_timeout` | Wall-clock cap (seconds) per instance on the local path; a breach errors that instance's run(s) without stalling the batch. | +| `job` | An open `Job` to accumulate into, so one platform job spans several `run(...)` calls (e.g. a training arc). Unset: a fresh job per call. | + +Runs produced per task = `group` x `num_envs` (each defaulting to 1). Everything about *what* runs +stays on the task rows (task args) and the agent - the run parameters never change task content. + ## Contract Embodiments and policies disagree on cameras, state layout, action semantics, and control rate, so From a622ba5184d0986d87aec279826252d3f749debb Mon Sep 17 00:00:00 2001 From: Jaideep <67646710+jdchawla29@users.noreply.github.com> Date: Fri, 17 Jul 2026 18:28:20 -0700 Subject: [PATCH 15/43] feat(environment): serve concurrent control sessions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Key the control channel's suspended tasks by session id: each connection starts and grades its own task instead of the last start cancelling whatever any other connection had in flight. A connection drop parks the session's task; a later connection grades it by resuming the session (hello with its session_id) or, when exactly one session is parked, by a plain tasks.grade — so the split 'hud task start' / 'hud task grade' flow keeps working unchanged, and the ambiguous case errors loudly instead of guessing. --- hud/clients/client.py | 12 +- hud/environment/server.py | 211 +++++++++++++++---------- hud/environment/tests/test_sessions.py | 89 +++++++++++ 3 files changed, 224 insertions(+), 88 deletions(-) create mode 100644 hud/environment/tests/test_sessions.py diff --git a/hud/clients/client.py b/hud/clients/client.py index da46763b0..91fb6d024 100644 --- a/hud/clients/client.py +++ b/hud/clients/client.py @@ -130,9 +130,15 @@ async def close(self) -> None: # ─── handshake ──────────────────────────────────────────────────── - async def hello(self) -> Manifest: - """Send ``hello``; cache and return the parsed ``Manifest``.""" - result = await self._call("hello", {}) + async def hello(self, session_id: str | None = None) -> Manifest: + """Send ``hello``; cache and return the parsed ``Manifest``. + + ``session_id`` resumes that parked session on the env — its suspended + task, e.g. one a prior connection started — instead of minting a + fresh session. + """ + params: dict[str, Any] = {} if session_id is None else {"session_id": session_id} + result = await self._call("hello", params) env = result.get("env") or {} bindings = [ await self._reachable(Capability.from_manifest(b)) diff --git a/hud/environment/server.py b/hud/environment/server.py index f4e264c22..a414eabd9 100644 --- a/hud/environment/server.py +++ b/hud/environment/server.py @@ -200,33 +200,51 @@ async def _frames( class _ControlChannel: """Serving-time state for one bound control channel. - Owns what the declaration must not: runtime state — at most one suspended - task at a time, living on the channel itself (scoped to this server; two - servers for one env never share). ``start`` replaces it, ``grade`` - consumes it, ``cancel`` clears it, and a connection drop leaves it in - place — which is exactly the split start/grade flow (e.g. harbor's - verifier reconnecting to grade) with no parking handoff to manage. + Owns what the declaration must not: runtime state — one suspended task per + session, keyed by session id (scoped to this server; two servers for one + env never share). A session's ``start`` replaces its own runner, ``grade`` + consumes it, ``cancel`` clears it. A connection drop parks the session's + runner in place; a later connection grades it by resuming the session + (``hello`` with its id) or, when exactly one session is parked, by a plain + ``tasks.grade`` — the split start/grade flow (e.g. ``hud task grade`` or + harbor's verifier reconnecting to grade) with no parking handoff to manage. """ def __init__(self, env: Environment) -> None: self.env = env - self._runner: TaskRunner | None = None - - async def start(self, task_id: str, args: dict[str, Any]) -> dict[str, Any]: - await self.cancel() - self._runner = TaskRunner(self.env.tasks[task_id], args) - return await self._runner.start() - - async def grade(self, payload: dict[str, Any]) -> dict[str, Any]: - runner, self._runner = self._runner, None + #: Suspended tasks by session id. An entry whose session has no live + #: connection is parked, awaiting a resume (or adoption) to grade. + self._runners: dict[str, TaskRunner] = {} + self._live: set[str] = set() + + async def start(self, session_id: str, task_id: str, args: dict[str, Any]) -> dict[str, Any]: + await self.cancel(session_id) + runner = TaskRunner(self.env.tasks[task_id], args) + self._runners[session_id] = runner + return await runner.start() + + async def grade(self, session_id: str, payload: dict[str, Any]) -> dict[str, Any]: + runner = self._runners.pop(session_id, None) if runner is None: - raise _NoTaskInProgress("no task in progress") + runner = self._adopt_parked() return await runner.grade(payload) - async def cancel(self) -> None: - if self._runner is not None: - await self._runner.cancel() - self._runner = None + def _adopt_parked(self) -> TaskRunner: + """Claim the parked session iff unambiguous — the blind-reconnect grade path.""" + parked = sorted(sid for sid in self._runners if sid not in self._live) + if not parked: + raise _NoTaskInProgress("no task in progress") + if len(parked) > 1: + raise _NoTaskInProgress( + f"{len(parked)} parked sessions ({', '.join(parked)}); " + "resume one by sending hello with its session_id" + ) + return self._runners.pop(parked[0]) + + async def cancel(self, session_id: str) -> None: + runner = self._runners.pop(session_id, None) + if runner is not None: + await runner.cancel() async def session( self, @@ -237,6 +255,7 @@ async def session( """One control session: JSON-RPC dispatch for the connection's lifetime.""" env = self.env session_id = "sess-" + secrets.token_hex(4) + self._live.add(session_id) async def reply_to(msg_id: int | None, result: dict[str, Any]) -> None: if msg_id is not None: @@ -246,71 +265,93 @@ async def error_to(msg_id: int | None, code: int, message: str) -> None: if msg_id is not None: await send_frame(writer, error(msg_id, code, message)) - async for msg in _frames(first, reader): - method = msg.get("method", "") - params = msg.get("params") or {} - msg_id = msg.get("id") - - try: - if method == "hello": - # env.start() ran before serving, so hook-published - # capabilities (e.g. a workspace's ssh address) are - # already concrete here. - bindings = [c.to_manifest() for c in env.capabilities] - await reply_to( - msg_id, - { - "session_id": session_id, - "env": {"name": env.name, "version": env.version}, - "bindings": bindings, - }, - ) - - elif method == "tasks.list": - await reply_to( - msg_id, - {"tasks": [t.manifest_entry() for t in env.tasks.values()]}, - ) - - elif method == "tasks.start": - task_id = params.get("id") - if not isinstance(task_id, str): - await error_to(msg_id, -32602, "tasks.start: 'id' must be a string") - continue - args = params.get("args") or {} - if not isinstance(args, dict): - await error_to(msg_id, -32602, "tasks.start: 'args' must be an object") - continue - try: - prompt = await self.start(task_id, args) - except KeyError: - await error_to(msg_id, -32602, f"unknown task: {task_id!r}") - continue - await reply_to(msg_id, prompt) - - elif method == "tasks.grade": - try: - evaluation = await self.grade(params) - except _NoTaskInProgress: - await error_to(msg_id, -32600, "no task in progress") - continue - await reply_to(msg_id, evaluation) - - elif method == "tasks.cancel": - await self.cancel() - await reply_to(msg_id, {"cancelled": True}) - - elif method == "bye": - await self.cancel() - await reply_to(msg_id, {"goodbye": True}) - return - - else: - await error_to(msg_id, -32601, f"method not found: {method}") - - except Exception as exc: - LOGGER.exception("error handling %s", method) - await error_to(msg_id, -32000, str(exc)) + try: + async for msg in _frames(first, reader): + method = msg.get("method", "") + params = msg.get("params") or {} + msg_id = msg.get("id") + + try: + if method == "hello": + requested = params.get("session_id") + if requested is not None and requested != session_id: + if not isinstance(requested, str): + await error_to( + msg_id, -32602, "hello: 'session_id' must be a string" + ) + continue + if requested in self._live: + await error_to( + msg_id, -32600, f"session {requested!r} has a live connection" + ) + continue + if requested not in self._runners: + await error_to(msg_id, -32600, f"unknown session: {requested!r}") + continue + # Resume: this connection becomes the parked session. + self._live.discard(session_id) + session_id = requested + self._live.add(session_id) + # env.start() ran before serving, so hook-published + # capabilities (e.g. a workspace's ssh address) are + # already concrete here. + bindings = [c.to_manifest() for c in env.capabilities] + await reply_to( + msg_id, + { + "session_id": session_id, + "env": {"name": env.name, "version": env.version}, + "bindings": bindings, + }, + ) + + elif method == "tasks.list": + await reply_to( + msg_id, + {"tasks": [t.manifest_entry() for t in env.tasks.values()]}, + ) + + elif method == "tasks.start": + task_id = params.get("id") + if not isinstance(task_id, str): + await error_to(msg_id, -32602, "tasks.start: 'id' must be a string") + continue + args = params.get("args") or {} + if not isinstance(args, dict): + await error_to(msg_id, -32602, "tasks.start: 'args' must be an object") + continue + try: + prompt = await self.start(session_id, task_id, args) + except KeyError: + await error_to(msg_id, -32602, f"unknown task: {task_id!r}") + continue + await reply_to(msg_id, prompt) + + elif method == "tasks.grade": + try: + evaluation = await self.grade(session_id, params) + except _NoTaskInProgress as exc: + await error_to(msg_id, -32600, str(exc)) + continue + await reply_to(msg_id, evaluation) + + elif method == "tasks.cancel": + await self.cancel(session_id) + await reply_to(msg_id, {"cancelled": True}) + + elif method == "bye": + await self.cancel(session_id) + await reply_to(msg_id, {"goodbye": True}) + return + + else: + await error_to(msg_id, -32601, f"method not found: {method}") + + except Exception as exc: + LOGGER.exception("error handling %s", method) + await error_to(msg_id, -32000, str(exc)) + finally: + self._live.discard(session_id) async def _stream( diff --git a/hud/environment/tests/test_sessions.py b/hud/environment/tests/test_sessions.py new file mode 100644 index 000000000..1145b462e --- /dev/null +++ b/hud/environment/tests/test_sessions.py @@ -0,0 +1,89 @@ +"""Control-channel sessions: each connection drives its own task. + +Suspended tasks are keyed by session id: concurrent connections start and +grade independently, a dropped connection parks its task, and a later +connection grades a parked task by resuming the session (``hello`` with its +id) or — only when exactly one is parked — by a plain ``tasks.grade`` (the +split ``hud task start`` / ``hud task grade`` flow). +""" + +from __future__ import annotations + +import pytest + +from hud.clients import HudProtocolError, connect +from hud.environment import Environment +from hud.eval.runtime import _local + + +def _env() -> Environment: + env = Environment("sessions") + + @env.template() + async def echo(tag: str): + yield f"go {tag}" + yield {"score": 1.0, "tag": tag} + + return env + + +async def test_concurrent_sessions_grade_their_own_tasks() -> None: + async with _local(_env()) as runtime, connect(runtime) as a, connect(runtime) as b: + await a.start_task("echo", {"tag": "a"}) + await b.start_task("echo", {"tag": "b"}) # must not disturb a's task + assert (await a.grade({"answer": "x"}))["tag"] == "a" + assert (await b.grade({"answer": "x"}))["tag"] == "b" + + +async def test_restart_replaces_only_the_sessions_own_task() -> None: + async with _local(_env()) as runtime, connect(runtime) as a, connect(runtime) as b: + await b.start_task("echo", {"tag": "b"}) + await a.start_task("echo", {"tag": "first"}) + await a.start_task("echo", {"tag": "second"}) + assert (await a.grade({"answer": "x"}))["tag"] == "second" + assert (await b.grade({"answer": "x"}))["tag"] == "b" + + +async def test_disconnect_parks_the_task_for_a_later_connection() -> None: + async with _local(_env()) as runtime: + async with connect(runtime) as first: + await first.start_task("echo", {"tag": "parked"}) + async with connect(runtime) as later: + assert (await later.grade({"answer": "x"}))["tag"] == "parked" + + +async def test_grade_with_multiple_parked_sessions_errors_loudly() -> None: + async with _local(_env()) as runtime: + for tag in ("one", "two"): + async with connect(runtime) as client: + await client.start_task("echo", {"tag": tag}) + async with connect(runtime) as later: + with pytest.raises(HudProtocolError, match="parked sessions"): + await later.grade({"answer": "x"}) + + +async def test_hello_resumes_a_parked_session_by_id() -> None: + async with _local(_env()) as runtime: + ids: dict[str, str] = {} + for tag in ("one", "two"): + async with connect(runtime) as client: + assert client.manifest is not None + ids[tag] = client.manifest.session_id + await client.start_task("echo", {"tag": tag}) + async with connect(runtime) as later: + await later.hello(session_id=ids["two"]) + assert (await later.grade({"answer": "x"}))["tag"] == "two" + + +async def test_hello_with_an_unknown_session_id_errors() -> None: + async with _local(_env()) as runtime, connect(runtime) as client: + with pytest.raises(HudProtocolError, match="unknown session"): + await client.hello(session_id="sess-nope") + + +async def test_hello_cannot_resume_a_live_session() -> None: + async with _local(_env()) as runtime, connect(runtime) as a, connect(runtime) as b: + assert a.manifest is not None + await a.start_task("echo", {"tag": "a"}) + with pytest.raises(HudProtocolError, match="live connection"): + await b.hello(session_id=a.manifest.session_id) From 0faf8570a7121963c504eb915f37ab05ae8dce71 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Sat, 18 Jul 2026 18:24:59 +0000 Subject: [PATCH 16/43] refactor(agents): drop per-agent max_steps override on RobotAgent max_steps was configurable per-subclass but every override just repeated the same default; keep one constant (520) via the call-site argument instead of a class attribute agents had to remember to set. Co-authored-by: Cursor --- docs/v6/cookbooks/robot-benchmark.mdx | 2 -- hud/agents/robot/agent.py | 16 ++++------------ 2 files changed, 4 insertions(+), 14 deletions(-) diff --git a/docs/v6/cookbooks/robot-benchmark.mdx b/docs/v6/cookbooks/robot-benchmark.mdx index f9e6521e1..a598d2141 100644 --- a/docs/v6/cookbooks/robot-benchmark.mdx +++ b/docs/v6/cookbooks/robot-benchmark.mdx @@ -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() diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py index 8670a83a8..5d17b3e37 100644 --- a/hud/agents/robot/agent.py +++ b/hud/agents/robot/agent.py @@ -49,8 +49,6 @@ class RobotAgent(Agent): """ robot_protocol: ClassVar[str] = ROBOT_PROTOCOL - #: Max control ticks before the episode is cut off. - max_steps: ClassVar[int] = 520 #: 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 @@ -62,14 +60,14 @@ class RobotAgent(Agent): #: Translates env<->policy spaces. Subclasses set this; ``None`` = raw pass-through. adapter: Adapter | None = None - async def __call__(self, run: Run, *, max_steps: int | None = None) -> None: + async def __call__(self, run: Run, *, max_steps: int = 520) -> None: """The generic rollout contract: one run, one trace.""" await self.drive([run], run.client, max_steps=max_steps) run.trace.status = "completed" run.trace.content = "done" async def drive( - self, runs: list[Run], client: HudClient, *, max_steps: int | None = None + self, runs: list[Run], client: HudClient, *, max_steps: int = 520 ) -> None: """Drive every env slot to termination, recording onto ``runs[i]``'s trace. @@ -115,13 +113,7 @@ async def drive( print(f"[agent] episode started: {prompt!r} (n={n})", flush=True) await self._loop( - robot, - obs, - prompt, - recorders, - writer, - single=single, - max_steps=max_steps or self.max_steps, + robot, obs, prompt, recorders, writer, single=single, max_steps=max_steps ) for rec in recorders: rec.close() @@ -139,7 +131,7 @@ async def _loop( writer: Any, *, single: bool, - max_steps: int, + max_steps: int = 520, ) -> None: """One batched forward per refill; execute chunks open-loop per slot.""" model = self.model From 6e4a8111a25505e9258404eab069262941d46c4f Mon Sep 17 00:00:00 2001 From: lukass16 Date: Sat, 18 Jul 2026 18:25:16 +0000 Subject: [PATCH 17/43] refactor(robot): spawn every sim as its own fork-free process; declarative gym targets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The env server no longer forks itself around a sim; every sim runs as a separate process, always the same way. env.gym(...) spawns it (or a custom sim program calls serve_bridge directly), and the env server's serving path (server.serve_blocking) is a plain asyncio.run — no more sim-main special case. - bridge.py absorbs the sim program (formerly serve.py): serve_bridge, gym_command, and the CLI (`python -m hud.environment.robot.bridge`) now live next to RobotBridge/GymBridge, since both ends of the argv format belong together. - GymBridge / env.gym(target) take three declarative target forms: a module-level factory (unchanged), a gymnasium registry id, or a constructed registry env — reduced by gym_command to its EnvSpec JSON and closed, since only the declaration can cross the fork; the sim process rebuilds it via gym.make/gym.make_vec. - sim_thread.py -> sim.py: the sim-process shape (SimThread/run_with_sim) is unconditional infrastructure every sim uses, not a special case to fork around, so it's named and documented as the one shape, not an alternative. - introspect.py + wrap.py merge into gym.py: one module for the gym integration both GymBridge and hud.wrap share (observation splitting, contract derivation, TracedEnv). - endpoint.py (RobotEndpoint) and env.py (Environment.gym) updated for the merged module layout and broadened target type. Co-authored-by: Cursor --- hud/environment/env.py | 29 +- hud/environment/robot/__init__.py | 38 +- hud/environment/robot/bridge.py | 271 +++++++++++-- hud/environment/robot/endpoint.py | 288 ++++++-------- hud/environment/robot/gym.py | 372 +++++++++++++----- hud/environment/robot/introspect.py | 148 ------- .../robot/{sim_thread.py => sim.py} | 32 +- hud/environment/robot/wrap.py | 222 ----------- hud/environment/server.py | 16 +- 9 files changed, 704 insertions(+), 712 deletions(-) delete mode 100644 hud/environment/robot/introspect.py rename hud/environment/robot/{sim_thread.py => sim.py} (83%) delete mode 100644 hud/environment/robot/wrap.py diff --git a/hud/environment/env.py b/hud/environment/env.py index 25aa5933c..8c77a5bf6 100644 --- a/hud/environment/env.py +++ b/hud/environment/env.py @@ -167,9 +167,6 @@ def __init__( # stands up). Run once by the serving substrate around its lifetime. self._on_start: list[Callable[[], Awaitable[None]]] = [] self._on_stop: list[Callable[[], Awaitable[None]]] = [] - #: Sims attached via :meth:`gym`; when present, the server runs the - #: sim-main process shape (sim on the main thread, serving beside it). - self._sims: list[Any] = [] self._init_legacy() # ─── task registration ─────────────────────────────────────────── @@ -308,25 +305,29 @@ async def _down() -> None: return ws - def gym(self, factory: Any, *, name: str = "robot", **kwargs: Any) -> Any: + def gym(self, target: Any, *, name: str = "robot", **kwargs: Any) -> Any: """Attach a gym-style sim serving ``name`` over the ``robot`` protocol. - ``factory`` is any callable returning a gym-style env (the same ``make_env`` - an EnvHub repo exposes). Registers the start → publish → stop lifecycle on - this env's hooks; nothing is built until the env serves. Extra kwargs go to - :class:`~hud.environment.robot.Gym` (``fps=``, ``contract=``, ...). - Returns the handle templates drive episodes through (``sim.reset`` / - ``sim.result``). + ``target`` declares the env: a module-level factory (the same + ``make_env`` an EnvHub repo exposes), a gymnasium registry id + (``"CartPole-v1"``), or a constructed registry env — reduced here to + its spec and closed, since only the declaration crosses the fork. The + sim runs in its own process (spawned at serve time; see + :mod:`hud.environment.robot.bridge`); this registers the spawn → + publish → teardown lifecycle on this env's hooks — nothing runs until + the env serves. Extra kwargs go to the sim program (``fps=``, + ``contract=``). Returns the + :class:`~hud.environment.robot.RobotEndpoint` templates drive episodes + through (``sim.reset`` / ``sim.result``). """ - from hud.environment.robot import Gym + from hud.environment.robot import RobotEndpoint, gym_command - sim = Gym(factory, **kwargs) - self._sims.append(sim) # the server serves sim-main when any are attached + sim = RobotEndpoint.spawn(gym_command(target, **kwargs)) @self.initialize async def _up() -> None: await sim.start() - self.add_capability(sim.capability(name)) + self.add_capability(await sim.capability(name)) @self.shutdown async def _down() -> None: diff --git a/hud/environment/robot/__init__.py b/hud/environment/robot/__init__.py index 4a61fdc32..91ac461b1 100644 --- a/hud/environment/robot/__init__.py +++ b/hud/environment/robot/__init__.py @@ -1,16 +1,24 @@ -"""Env-side robot runtime: bridges, the control endpoint, and gym integration. +"""Env-side robot runtime: bridges, the sim program, and the control endpoint. -Everything an *environment* needs to own a simulator and serve it to an agent -over the ``robot`` WebSocket protocol: +A simulator always runs in its own process — the **sim program** +(``python -m hud.environment.robot.bridge``), where the sim owns the main +thread and a bridge serves the wire. The env server holds a +:class:`~.endpoint.RobotEndpoint` on it: - :class:`~.bridge.RobotBridge` — the batched-first bridge base (``num_envs`` - slots in lockstep; a plain single env is a batch of one). -- :class:`~.bridge.GymBridge` / :class:`~.gym.Gym` — the generic gym-factory - path (``env.gym(make_env)``): contract derivation, capability, episode control. -- :class:`~.endpoint.RobotEndpoint` — the control handle, local or remote. -- :func:`hud.wrap` (:mod:`~.wrap`) — one-line trace streaming for any gym env. -- :class:`~.sim_thread.SimThread` — the one process shape: the sim owns the - main thread, serving runs on a background loop thread. + slots in lockstep; a plain single env is a batch of one): the agent's + ``robot`` WebSocket plus the JSON-RPC control side channel. +- :class:`~.bridge.GymBridge` / :func:`~.bridge.gym_command` — the generic + gym path (``env.gym(...)`` — a factory, registry id, or declared env): + contract derivation, capability, episode control. +- :class:`~.endpoint.RobotEndpoint` — the env server's handle: spawn (owned) + or remote (attached), same control surface either way. +- :func:`~.bridge.serve_bridge` — the sim program's blocking entry (custom + bridges call it last). +- :func:`hud.wrap` (:mod:`~.gym`) — one-line trace streaming for any gym env + you drive yourself, plus the shared gym introspection. +- :class:`~.sim.SimThread` — the sim-process shape: the sim owns the main + thread, serving runs on a background loop thread. The agent-side counterpart, :class:`~hud.capabilities.robot.RobotClient`, lives under :mod:`hud.capabilities`; both ends share the wire codec defined there. @@ -18,19 +26,19 @@ from __future__ import annotations -from .bridge import GymBridge, RobotBridge +from .bridge import GymBridge, RobotBridge, gym_command, serve_bridge from .endpoint import RobotEndpoint -from .gym import Gym -from .sim_thread import SimThread, run_with_sim -from .wrap import TracedEnv, wrap +from .gym import TracedEnv, wrap +from .sim import SimThread, run_with_sim __all__ = [ - "Gym", "GymBridge", "RobotBridge", "RobotEndpoint", "SimThread", "TracedEnv", + "gym_command", "run_with_sim", + "serve_bridge", "wrap", ] diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py index 187fe9247..a4f7d8eca 100644 --- a/hud/environment/robot/bridge.py +++ b/hud/environment/robot/bridge.py @@ -1,26 +1,42 @@ -"""Env-side ``robot`` bridges: the base class and every shipped variant. +"""Sim-process ``robot`` bridges — and the sim program that serves them. The *server* side of the ``robot`` protocol (agent-side client: :class:`~hud.capabilities.robot.RobotClient`); both share the wire codec defined -there. Bridges are **batched-first**: one bridge serves ``num_envs`` slots in -lockstep (``[N, ...]`` obs frames, ``[N, A]`` actions, an ``[N]`` ``terminated`` -mask). ``num_envs == 1`` — a plain single-env sim — speaks the scalar framing +there. A bridge lives in the sim's own process and serves two channels: the +agent's obs/action **WebSocket** and the env server's JSON-RPC **control side +channel** (a :class:`~.endpoint.RobotEndpoint` drives episodes through it). +Bridges are **batched-first**: one bridge serves ``num_envs`` slots in lockstep +(``[N, ...]`` obs frames, ``[N, A]`` actions, an ``[N]`` ``terminated`` mask). +``num_envs == 1`` — a plain single-env sim — speaks the scalar framing (per-env arrays, scalar ``terminated``) on the same code path. - :class:`RobotBridge` — subclass with your sim: ``reset`` / ``step`` / - ``get_observation``. The base owns the WebSocket serve loop. -- :class:`GymBridge` — the generic bridge over any gym-style env factory; - users get one via ``env.gym(make_env)`` and never subclass it. - -Every sim touch routes through the process :class:`~.sim_thread.SimThread`, so -bridges stay thread-naive (see :mod:`~.sim_thread` for the one process shape). + ``get_observation``, and set ``self.contract``. The base owns both channels. +- :class:`GymBridge` — the generic bridge over any gym-style target (factory, + registry id, or a declared env's spec); users get one via ``env.gym(...)`` + and never subclass it. + +This module is also **the sim program** — the one process shape every sim runs +(the env server never hosts one, it spawns or attaches through an endpoint): + +- ``python -m hud.environment.robot.bridge `` — what ``env.gym()`` + spawns; :func:`gym_command` builds the argv, so both ends of the format live + here. +- :func:`serve_bridge` — the blocking entry a custom sim program calls last. + +Inside the sim process the sim owns the main thread and serving runs on a +background loop; every sim touch routes through the shared +:class:`~.sim.SimThread`, so bridges stay thread-naive (see :mod:`~.sim`). """ from __future__ import annotations +import asyncio import contextlib import inspect +import sys from abc import ABC, abstractmethod +from pathlib import Path from typing import Any import numpy as np @@ -30,18 +46,32 @@ # The openpi/0 wire codec is defined alongside the agent-side client; reuse it so both # ends of the protocol stay in lockstep (env -> capabilities is the correct direction). from hud.capabilities.robot import _packb, _unpackb +from hud.environment.utils import error, read_frame, reply, send_frame from hud.telemetry.robot import to_numpy -from .introspect import probe_success, split_observation -from .sim_thread import SimThread +from .gym import ( + action_dim_of, + detect_fps, + load_or_write_contract, + num_envs_of, + probe_success, + split_observation, +) +from .sim import SimThread, run_with_sim + +#: Line the sim program prints once its control channel is bound; the spawning +#: RobotEndpoint reads it from this process's stdout. +PORT_ANNOUNCEMENT = "HUD_SIM_PORT=" class RobotBridge(ABC): - """Serves ``robot`` over WebSocket; subclass and implement the env hooks. + """Serves a sim over ``robot`` WebSocket + a JSON-RPC control side channel. **Subclass contract:** implement :meth:`reset`, :meth:`step`, and - :meth:`get_observation`. The base owns the WebSocket serve loop; subclasses - own the sim and set ``num_envs`` (default 1). + :meth:`get_observation`, and set ``self.contract`` (the wire contract the + capability publishes) before serving. The base owns the WebSocket serve + loop and the control listener; subclasses own the sim and set ``num_envs`` + (default 1). - :meth:`reset` initialises the sim for a new episode and returns the task prompt. The base resets scoring state and pushes the first frame. @@ -63,7 +93,10 @@ def __init__(self, *, host: str = "127.0.0.1", port: int = 0) -> None: self._server: Any = None # Connect-time metadata frame (sent first on each connection); subclasses may set it. self.metadata: dict[str, Any] = {} - # Every sim touch runs on the process sim thread (see sim_thread.py). + #: The env's self-describing wire contract (features, control_rate); the + #: endpoint fetches it over the control channel to mint the capability. + self.contract: dict[str, Any] = {} + # Every sim touch runs on the process sim thread (see sim.py). self._sim = SimThread.shared() self.num_envs: int = 1 # Episode scoring read by ``result()``; single-env subclasses update these @@ -188,24 +221,77 @@ async def _send_observation(self) -> None: with contextlib.suppress(websockets.exceptions.ConnectionClosed): await self._client.send(_packb(msg)) + # ── the control side channel (driven by a RobotEndpoint) ──────────────────── + + async def serve_control(self, host: str = "127.0.0.1", port: int = 0) -> asyncio.Server: + """Serve the JSON-RPC side channel: ``url`` / ``contract`` / ``reset`` / ``result``.""" + return await asyncio.start_server(self._handle_control, host, port) + + async def _handle_control( + self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + with contextlib.suppress(ConnectionResetError, asyncio.IncompleteReadError): + while (msg := await read_frame(reader)) is not None: + try: + result = await self._dispatch_control(msg["method"], msg.get("params") or {}) + await send_frame(writer, reply(msg["id"], result)) + except Exception as exc: # surface to the caller, keep serving the link + await send_frame(writer, error(msg["id"], -32000, str(exc))) + writer.close() + with contextlib.suppress(Exception): + await writer.wait_closed() + + async def _dispatch_control(self, method: str, params: dict[str, Any]) -> dict[str, Any]: + if method == "url": + return {"url": self.url} + if method == "contract": + return {"contract": self.contract} + if method == "reset": + return {"prompt": await self._reset(**params)} + if method == "result": + return self.result() + raise ValueError(f"unknown method {method!r}") + class GymBridge(RobotBridge): - """Serve any gym-style env factory over the ``robot`` protocol, generically. + """Serve any gym-style env over the ``robot`` protocol, generically. - Task args are partitioned by the factory's signature: args the factory - accepts define the env (a change rebuilds it — so ``num_envs`` in the - factory signature is the vectorization declaration); everything else is - episodic and flows to ``env.reset(seed=..., options=...)``. + ``target`` is how this process builds the env: a factory callable, or a + string — a gymnasium registry id, or a declared env's spec as JSON + (:func:`gym_command` reduces an env instance to the latter). + + Task args are partitioned into env-defining **build args** (a change + rebuilds the env) and episodic args flowing to ``env.reset(seed=..., + options=...)``. For a factory, its signature is the partition — so + ``num_envs`` in the signature is the vectorization declaration; for a + registry target, ``num_envs`` is the one build arg (``gym.make_vec``). + + ``fps`` overrides the detected control rate; ``contract`` is the + ``contract.json`` round-trip path (load beats derive, so user edits stick; + ``None`` disables the file). """ - def __init__(self, factory: Any, **kwargs: Any) -> None: + def __init__( + self, + target: Any, + *, + fps: int | None = None, + contract: str | Path | None = "contract.json", + **kwargs: Any, + ) -> None: super().__init__(**kwargs) - self._factory = factory - self._factory_params = { - n - for n, p in inspect.signature(factory).parameters.items() - if p.kind in (p.POSITIONAL_OR_KEYWORD, p.KEYWORD_ONLY) - } + self._target = target + self._fps = fps + self._contract_path = contract + # Task args that define the env build (everything else is episodic). + if callable(target): + self._build_params = { + n + for n, p in inspect.signature(target).parameters.items() + if p.kind in (p.POSITIONAL_OR_KEYWORD, p.KEYWORD_ONLY) + } + else: + self._build_params = {"num_envs"} # registry targets: vectorization only self.env: Any = None self.batched = False # env carries a leading [N] dim (any env exposing num_envs) self._obs: Any = None # latest observation (reset or step) @@ -224,6 +310,22 @@ async def ensure_env(self, **task_args: Any) -> None: if self.env is None: await self._sim.call(self._sync_reset, task_args) + async def start(self) -> None: + """Build the env and derive/load the contract, then bring up the wire.""" + await self.ensure_env() + state, frames = self.sample_observation() + existed = self._contract_path is not None and Path(self._contract_path).exists() + self.contract = load_or_write_contract( + self._contract_path, + state, + frames, + action_dim_of(self.env, batched=self.batched), + self._fps or detect_fps(self.env), + ) + if not existed and self._contract_path is not None: + print(f"[env] wrote {self._contract_path} (edit names to relabel plots)", flush=True) + await super().start() + async def reset(self, **task_args: Any) -> str: return await self._sim.call(self._sync_reset, task_args) @@ -234,18 +336,17 @@ async def stop(self) -> None: self.env = None def _sync_reset(self, task_args: dict[str, Any]) -> str: - build = {k: v for k, v in task_args.items() if k in self._factory_params} - episodic = {k: v for k, v in task_args.items() if k not in self._factory_params} + build = {k: v for k, v in task_args.items() if k in self._build_params} + episodic = {k: v for k, v in task_args.items() if k not in self._build_params} key = tuple(sorted(build.items())) if self.env is None or key != self._instance: if self.env is not None: self.env.close() - self.env = self._factory(**build) + self.env = self._build_env(build) self._instance = key - base = getattr(self.env, "unwrapped", self.env) - n = getattr(self.env, "num_envs", getattr(base, "num_envs", None)) + n = num_envs_of(self.env) self.batched = n is not None - self.num_envs = int(n or 1) + self.num_envs = n or 1 seed = episodic.pop("seed", None) obs, _ = self.env.reset(seed=seed, options=episodic or None) self._obs = obs # the first frame an agent sees on connect/reset @@ -256,6 +357,21 @@ def _sync_reset(self, task_args: dict[str, Any]) -> str: self._seen_success = False return self._prompt(task_args) + def _build_env(self, build: dict[str, Any]) -> Any: + """Build the env from the target: factory call, or registry make/make_vec.""" + if callable(self._target): + return self._target(**build) + import gymnasium + from gymnasium.envs.registration import EnvSpec + + env_id, kwargs = self._target, dict(build) + if env_id.startswith("{"): # a declared env's spec, re-made in this process + spec = EnvSpec.from_json(env_id) + env_id, kwargs = spec.id, {**spec.kwargs, **kwargs} + if "num_envs" in kwargs: + return gymnasium.make_vec(env_id, **kwargs) + return gymnasium.make(env_id, **kwargs) + def _prompt(self, task_args: dict[str, Any]) -> str: base = getattr(self.env, "unwrapped", self.env) for attr in ("task_description", "instruction"): @@ -347,4 +463,89 @@ def _first_leaf(obs: Any) -> Any: return obs -__all__ = ["GymBridge", "RobotBridge"] +# ── the sim program (what env.gym() spawns; custom sims call serve_bridge) ──── + + +def serve_bridge(bridge: RobotBridge, *, host: str = "127.0.0.1", port: int = 0) -> None: + """Serve *bridge*, blocking for the process's lifetime — the sim program's last call. + + Brings up the robot WebSocket and the control side channel (announcing its + port on stdout as ``HUD_SIM_PORT=``), with the sim owning the main + thread. SIGTERM / Ctrl-C tear the bridge down on the way out. + """ + + async def _serve() -> None: + await bridge.start() # WS bound first, so `url` is concrete before any control call + server = await bridge.serve_control(host, port) + bound = server.sockets[0].getsockname()[1] + print(f"{PORT_ANNOUNCEMENT}{bound}", flush=True) + try: + await asyncio.Event().wait() # until SIGTERM / Ctrl-C cancels + finally: + server.close() + await bridge.stop() + + run_with_sim(_serve) + + +def gym_command( + target: Any, *, fps: int | None = None, contract: str | None = "contract.json" +) -> list[str]: + """Build the command line that spawns *target* as a sim process (what ``env.gym`` runs). + + ``env.py`` only declares the sim, it never runs it — this builds the argv + for the child process that will. A live env object can't cross that + boundary, only text can, so *target* is reduced to something the child can + rebuild itself: a factory's source path, a registry id as-is, or a + constructed env's spec (id + kwargs) — closed here since only the spec + crosses over; the child calls ``gym.make`` again for its own instance. + ``fps`` / ``contract`` flow through to the :class:`GymBridge` the CLI builds. + """ + if callable(target): + name = getattr(target, "__qualname__", "") + if not name or "." in name or "<" in name: + raise ValueError(f"env.gym factory must be a module-level callable, got {target!r}") + target = f"{inspect.getfile(target)}:{name}" + elif not isinstance(target, str): + spec = getattr(target, "spec", None) # registry-made envs carry their EnvSpec + if spec is None: + raise ValueError( + f"env.gym: {target!r} has no registry spec; pass a gymnasium id " + "or a module-level factory instead" + ) + target_env, target = target, spec.to_json() + target_env.close() # only the spec crosses; free the declaration-time instance + cmd = [sys.executable, "-m", "hud.environment.robot.bridge", target] + if fps is not None: + cmd += ["--fps", str(fps)] + if contract is not None: + cmd += ["--contract", str(contract)] + return cmd + + +def main() -> None: + import argparse + + from hud.utils.modules import load_module + + parser = argparse.ArgumentParser(description="Serve a gym-style env as a robot sim.") + parser.add_argument("target", help="'path/to/module.py:factory', a registry id, or spec JSON.") + parser.add_argument("--fps", type=int, default=None, help="Control-rate override.") + parser.add_argument("--contract", default=None, help="contract.json round-trip path.") + parser.add_argument("--host", default="127.0.0.1", help="Control-channel interface.") + parser.add_argument("--port", type=int, default=0, help="Control-channel port (0 = ephemeral).") + args = parser.parse_args() + + target: Any = args.target + if ".py:" in target: # factory by source path; ids / spec JSON pass through + path, _, attr = target.rpartition(":") + target = getattr(load_module(path), attr) + bridge = GymBridge(target, fps=args.fps, contract=args.contract) + serve_bridge(bridge, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() + + +__all__ = ["PORT_ANNOUNCEMENT", "GymBridge", "RobotBridge", "gym_command", "serve_bridge"] diff --git a/hud/environment/robot/endpoint.py b/hud/environment/robot/endpoint.py index e20138425..df0580e7a 100644 --- a/hud/environment/robot/endpoint.py +++ b/hud/environment/robot/endpoint.py @@ -1,21 +1,27 @@ -"""``RobotEndpoint`` — the env-side control handle for a :class:`RobotBridge`. +"""``RobotEndpoint`` — the env server's handle on a sim process. -The single surface an env uses to drive a bridge through an episode (``start`` / -``stop`` / ``reset`` / ``result`` / ``url``). Its whole point is to make *where the -bridge runs* irrelevant — the env code is identical either way: +A bridge always lives in the sim's own process (see :mod:`~.bridge`); the +endpoint is the JSON-RPC client that drives it through episodes (``reset`` / +``result``) and mints its capability — and, when it spawned the process, owns +its lifecycle. Two ways to build one, identical methods either way: -- **Same process** — ``RobotEndpoint(bridge)``: calls go straight through. -- **Different process** — ``RobotEndpoint.remote(host, port)`` on the env side, - ``RobotEndpoint(bridge).serve(host, port)`` in the process that owns the sim (e.g. - Isaac/Omniverse, which pins the main thread); calls are forwarded over JSON-RPC. +- **Spawned** — :meth:`RobotEndpoint.spawn`: fork the sim program, read its + announced control port, tear the process down on :meth:`stop` + (``env.gym(...)`` builds this). +- **Attached** — :meth:`RobotEndpoint.remote`: dial a sim process something + else runs (another container, a warm Isaac kept alive across env-server + restarts); :meth:`stop` only drops the link, never the process. -Control plane only: the agent's step/observation loop tunnels straight to the bridge's -``robot`` WebSocket, and the wire contract stays env-side. +Control plane only: the agent's step/observation loop tunnels straight to the +bridge's ``robot`` WebSocket, and templates drive episodes through the handle:: - async def my_task(task_id: int, seed: int = 0): - prompt = await endpoint.reset(task_id=task_id, seed=seed) - yield {"prompt": prompt} - yield await endpoint.result() + sim = env.gym(make_env) + + + @env.template(id="pawn_lift") + async def pawn_lift(task: str = "solo_pawn_lift", seed: int = 0, num_envs: int = 1): + yield {"prompt": await sim.reset(task=task, seed=seed, num_envs=num_envs)} + yield await sim.result() """ from __future__ import annotations @@ -24,171 +30,107 @@ async def my_task(task_id: int, seed: int = 0): import contextlib from typing import TYPE_CHECKING, Any -from hud.environment.utils import error, read_frame, reply, send_frame +from hud.environment.utils import read_frame, send_frame +from hud.utils.process import create_process_group_exec + +from .bridge import PORT_ANNOUNCEMENT if TYPE_CHECKING: - from hud.capabilities import Capability + from collections.abc import Sequence - from .bridge import RobotBridge + from hud.capabilities import Capability + from hud.utils.process import ProcessGroup class RobotEndpoint: - """Drive a simulation bridge - even if it's in another process. + """Drive a simulation bridge living in another process. - Build it one of two ways and use the *identical* methods either way: - ``RobotEndpoint(bridge)`` (local) or ``RobotEndpoint.remote(host, port)`` (a handle - on a bridge that another process exposes via :meth:`serve` defined here). + Build with :meth:`spawn` (own the sim process) or :meth:`remote` (attach to + one served elsewhere); :meth:`start` brings the link up either way. """ def __init__( self, - bridge: RobotBridge | None = None, *, + cmd: Sequence[str] | None = None, host: str | None = None, port: int | None = None, + connect_timeout_s: float = 240.0, ) -> None: - self._bridge = bridge # set => local; None => remote (dial host:port) + self._cmd = list(cmd) if cmd is not None else None # set => spawned mode self._host = host self._port = port + self._connect_timeout_s = connect_timeout_s + self._proc: ProcessGroup | None = None + self._forward: asyncio.Task[None] | None = None self._reader: asyncio.StreamReader | None = None self._writer: asyncio.StreamWriter | None = None @classmethod - def remote(cls, host: str, port: int) -> RobotEndpoint: - """A handle on a bridge served by another process; :meth:`connect` once it's up.""" - return cls(host=host, port=port) - - @property - def _is_remote(self) -> bool: - return self._bridge is None - - def _local_bridge(self) -> RobotBridge: - bridge = self._bridge - if bridge is None: - raise RuntimeError("local bridge required") - return bridge + def spawn(cls, cmd: Sequence[str], *, connect_timeout_s: float = 240.0) -> RobotEndpoint: + """An endpoint that forks *cmd* (a sim program; see :mod:`~.bridge`) and owns it.""" + return cls(cmd=cmd, connect_timeout_s=connect_timeout_s) - # ── control surface (same whether local or remote) ─────────────────── - async def url(self) -> str: - """The bridge's ``ws://`` address — publish it as the robot capability.""" - if self._is_remote: - return (await self._call("url"))["url"] - return self._local_bridge().url - - async def capability(self, *, name: str = "robot", contract: dict[str, Any]) -> Capability: - """The ``robot`` capability for this bridge — mirrors ``Workspace.capability()``. - - Publish it from an ``@env.initialize`` hook after :meth:`start` (the URL only - exists once the bridge has bound its socket):: - - @env.initialize - async def _up(): - await endpoint.start() - env.add_capability(await endpoint.capability(contract=CONTRACT)) - """ - from hud.capabilities import Capability + @classmethod + def remote(cls, host: str, port: int, *, connect_timeout_s: float = 240.0) -> RobotEndpoint: + """An endpoint attached to a sim process something else runs.""" + return cls(host=host, port=port, connect_timeout_s=connect_timeout_s) - return Capability.robot(name=name, url=await self.url(), contract=contract) + # ── lifecycle ───────────────────────────────────────────────────────── async def start(self) -> None: - if self._is_remote: - await self._call("start") - else: - await self._local_bridge().start() + """Bring the link up: fork the sim program (spawned mode) and connect.""" + if self._cmd is not None and self._proc is None: + self._proc = await create_process_group_exec( + *self._cmd, + term_timeout=10.0, + stdout=asyncio.subprocess.PIPE, # for the port announcement; stderr inherits + ) + self._host = "127.0.0.1" + self._port = await asyncio.wait_for( + self._read_announced_port(), self._connect_timeout_s + ) + # Keep passing the sim's stdout through so its logs stay visible + # (and the pipe never fills and blocks the child). + assert self._proc.stdout is not None + self._forward = asyncio.create_task(_forward_lines(self._proc.stdout)) + await self._connect() async def stop(self) -> None: - if self._is_remote: - await self._call("stop") - else: - await self._local_bridge().stop() - - async def reset(self, **task_args: Any) -> str: - """Start a new episode; return the task prompt.""" - if self._is_remote: - return (await self._call("reset", task_args))["prompt"] - return await self._local_bridge()._reset(**task_args) - - async def result(self, **extra: Any) -> dict[str, Any]: - """The episode score dict, merged with any caller ``extra`` metadata.""" - res = await self._call("result") if self._is_remote else self._local_bridge().result() - res = {**res, **extra} - print( - f"[env] result: success={res.get('success')} " - f"total_reward={res.get('total_reward', 0.0):.3f}", - flush=True, - ) - return res - - # ── serving: expose a local bridge so a remote endpoint can drive it ── - async def serve(self, host: str = "127.0.0.1", port: int = 9100) -> asyncio.AbstractServer: - """Serve this (local) bridge's control surface over JSON-RPC. - - The process that owns the sim calls this; a ``remote()`` endpoint elsewhere then - drives the bridge through it. Await the returned server's ``wait_closed()`` to run - for the process's lifetime. Calls dispatch on *this* loop — the sim's — so e.g. - ``reset`` runs inline on the sim thread. - """ - if self._bridge is None: - raise RuntimeError("serve() needs a local bridge: RobotEndpoint(bridge)") - server = await asyncio.start_server(self._handle, host, port) - print(f"[env] control endpoint listening on {host}:{port}", flush=True) - return server - - def serve_blocking(self, host: str = "0.0.0.0", port: int = 9100) -> None: # noqa: S104 — split-process sims serve cross-host - """Serve this (local) bridge for the process's lifetime, with the sim owning - the main thread — the entry a split-process sim program calls last. - - Same shape as ``hud.environment.server``: the control endpoint runs on a - background loop thread; every sim touch drains here on main. - """ - from .sim_thread import run_with_sim - - async def _serve() -> None: - server = await self.serve(host, port) - try: - await asyncio.Event().wait() # until SIGTERM / Ctrl-C cancels - finally: - server.close() - await self._local_bridge().stop() - - run_with_sim(_serve) - - async def _handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: - with contextlib.suppress(ConnectionResetError, asyncio.IncompleteReadError): - while (msg := await read_frame(reader)) is not None: - try: - result = await self._dispatch(msg["method"], msg.get("params") or {}) - await send_frame(writer, reply(msg["id"], result)) - except Exception as exc: # surface to the caller, keep serving the link - await send_frame(writer, error(msg["id"], -32000, str(exc))) - writer.close() - with contextlib.suppress(Exception): - await writer.wait_closed() - - async def _dispatch(self, method: str, params: dict[str, Any]) -> dict[str, Any]: - b = self._local_bridge() - if method == "url": - return {"url": b.url} - if method == "reset": - return {"prompt": await b._reset(**params)} - if method == "result": - return b.result() - if method == "start": - await b.start() - return {} - if method == "stop": - await b.stop() - return {} - raise ValueError(f"unknown method {method!r}") - - # ── remote link (no-ops when local) ────────────────────────────────── - async def connect(self, *, connect_timeout_s: float = 240.0, retry_every: float = 2.0) -> None: - """Dial the serving process, retrying until it's up. No-op for a local endpoint.""" - if not self._is_remote: - return + """Drop the link; tear the sim process down when this endpoint spawned it.""" + if self._forward is not None: + self._forward.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._forward + self._forward = None + if self._writer is not None: + self._writer.close() + with contextlib.suppress(Exception): + await self._writer.wait_closed() + self._reader = self._writer = None + if self._proc is not None: # owned: SIGTERM the sim's whole process group + await self._proc.terminate() + self._proc = None + + async def _read_announced_port(self) -> int: + """The sim program's ``HUD_SIM_PORT=`` line, passing boot logs through.""" + assert self._proc is not None and self._proc.stdout is not None + while line := await self._proc.stdout.readline(): + text = line.decode("utf-8", "replace").rstrip() + if text.startswith(PORT_ANNOUNCEMENT): + return int(text.removeprefix(PORT_ANNOUNCEMENT)) + print(text, flush=True) + code = await self._proc.wait() + raise RuntimeError(f"sim process exited with code {code} before announcing its port") + + async def _connect(self, retry_every: float = 2.0) -> None: + """Dial the control channel, retrying until the sim serves (it may boot slowly).""" + assert self._host is not None and self._port is not None try: - async with asyncio.timeout(connect_timeout_s): + async with asyncio.timeout(self._connect_timeout_s): while True: + if self._proc is not None and self._proc.returncode is not None: + raise RuntimeError(f"sim process exited with code {self._proc.returncode}") try: self._reader, self._writer = await asyncio.open_connection( self._host, self._port @@ -198,21 +140,44 @@ async def connect(self, *, connect_timeout_s: float = 240.0, retry_every: float await asyncio.sleep(retry_every) except TimeoutError as exc: raise TimeoutError( - f"timed out connecting to {self._host}:{self._port} after {connect_timeout_s}s" + f"timed out connecting to sim control at {self._host}:{self._port} " + f"after {self._connect_timeout_s}s" ) from exc - async def close(self) -> None: - """Drop the link (no-op when local; does not stop the bridge).""" - if self._writer is not None: - self._writer.close() - with contextlib.suppress(Exception): - await self._writer.wait_closed() - self._reader = self._writer = None + # ── the control surface ─────────────────────────────────────────────── + + async def url(self) -> str: + """The bridge's ``ws://`` address — the robot capability's url.""" + return (await self._call("url"))["url"] + + async def contract(self) -> dict[str, Any]: + """The env's self-describing wire contract, read from the bridge.""" + return (await self._call("contract"))["contract"] + + async def capability(self, name: str = "robot") -> Capability: + """The concrete ``robot`` capability — publish it from an ``@env.initialize`` hook.""" + from hud.capabilities import Capability + + return Capability.robot(name=name, url=await self.url(), contract=await self.contract()) + + async def reset(self, **task_args: Any) -> str: + """Start a new episode; return the task prompt.""" + return (await self._call("reset", task_args))["prompt"] + + async def result(self, **extra: Any) -> dict[str, Any]: + """The episode score dict, merged with any caller ``extra`` metadata.""" + res = {**(await self._call("result")), **extra} + print( + f"[env] result: success={res.get('success')} " + f"total_reward={res.get('total_reward', 0.0):.3f}", + flush=True, + ) + return res async def _call(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]: # Strictly request/reply, one call at a time, so a constant id is enough. if self._writer is None or self._reader is None: - raise RuntimeError("not connected; call connect() first") + raise RuntimeError("not connected; call start() first") await send_frame( self._writer, {"jsonrpc": "2.0", "id": 1, "method": method, "params": params or {}} ) @@ -224,4 +189,9 @@ async def _call(self, method: str, params: dict[str, Any] | None = None) -> dict return msg["result"] +async def _forward_lines(stream: asyncio.StreamReader) -> None: + while line := await stream.readline(): + print(line.decode("utf-8", "replace").rstrip(), flush=True) + + __all__ = ["RobotEndpoint"] diff --git a/hud/environment/robot/gym.py b/hud/environment/robot/gym.py index 5f8e53dcc..47f67d3df 100644 --- a/hud/environment/robot/gym.py +++ b/hud/environment/robot/gym.py @@ -1,124 +1,318 @@ -"""``Gym`` — the declarative sim handle over a gym-style env factory. +"""Gym-style env integration: introspection + ``hud.wrap`` trace streaming. -The robotics analog of :class:`~hud.environment.workspace.Workspace`: -construction is pure data, ``start()`` materializes everything (env build, -contract derivation, the ``robot`` WebSocket, a JSON-RPC control endpoint for -split-process setups), and ``capability()`` mints the wire capability. -``env.gym(make_env)`` wires the lifecycle. +Two consumers share the introspection here (split observations, probe success, +derive a minimal contract) and stay in lockstep because of it: -Templates drive episodes through the handle:: +- :class:`~.bridge.GymBridge` — the served path (``env.gym(...)``). +- :func:`hud.wrap` / :class:`TracedEnv` — the loop-owning path: wrap any + ``gym.Env``, ``gym.vector.VectorEnv``, or batched-tensor Isaac env you drive + yourself and every episode streams to the platform as a trace (numeric + state, per-camera H.264 video, actions, reward/success) under one job:: - sim = env.gym(make_env) + env = hud.wrap(make_env(...), job="chess-eval") - - @env.template(id="pawn_lift") - async def pawn_lift(task: str = "solo_pawn_lift", seed: int = 0, num_envs: int = 1): - yield {"prompt": await sim.reset(task=task, seed=seed, num_envs=num_envs)} - yield await sim.result() +On first reset a minimal ``contract.json`` is written next to your script +describing how the observation/action spaces were interpreted; edit its +``names`` to relabel the platform's plots. An existing file is loaded instead, +so edits stick. """ from __future__ import annotations +import atexit +import json import logging from pathlib import Path -from typing import TYPE_CHECKING, Any - -from .bridge import GymBridge -from .endpoint import RobotEndpoint -from .introspect import action_dim_of, detect_fps, load_or_write_contract +from typing import Any, Self -if TYPE_CHECKING: - import asyncio +import numpy as np - from hud.capabilities import Capability +from hud.telemetry.robot import JobRecorder, to_numpy logger = logging.getLogger(__name__) +#: Conventional info keys carrying env-reported success, probed in order. +SUCCESS_KEYS = ("success", "is_success", "task_success") + + +# ── env introspection (pure; no telemetry, no wrapper state) ────────────────── + + +def num_envs_of(env: Any) -> int | None: + """``num_envs`` from the env or its unwrapped core (gymnasium 1.x wrappers + no longer forward attributes, so an Isaac env inside ``gym.make`` hides it).""" + n = getattr(env, "num_envs", None) + if n is None: + n = getattr(getattr(env, "unwrapped", env), "num_envs", None) + return None if n is None else int(n) + -class Gym: - """One gym-style sim as a HUD building block: contract, capability, episode control. +def is_batched(env: Any) -> bool: + """Vectorized (gym VectorEnv) or batched-tensor (Isaac) envs carry a leading [N] dim. - Nothing is built at import — the factory runs at :meth:`start` (serve time). - All control goes through the bridge's public surface via a local - :class:`~.endpoint.RobotEndpoint`. + Any env exposing ``num_envs`` is batched — Isaac keeps batch semantics even at + ``num_envs == 1``; plain ``gym.Env``s don't have the attribute at all. + """ + if num_envs_of(env) is not None: + return True + try: + import gymnasium + + return isinstance(getattr(env, "unwrapped", env), gymnasium.vector.VectorEnv) + except ImportError: + return False + + +def flatten_observation(obs: Any, prefix: str = "") -> dict[str, Any]: + """Flatten nested dict observations to slash-keyed leaves (non-dicts -> ``{"obs": x}``).""" + if not isinstance(obs, dict): + return {prefix or "obs": obs} + flat: dict[str, Any] = {} + for k, v in obs.items(): + key = f"{prefix}/{k}" if prefix else str(k) + flat.update(flatten_observation(v, key) if isinstance(v, dict) else {key: v}) + return flat + + +def split_observation(obs: Any, *, batched: bool = False) -> tuple[dict[str, Any], dict[str, Any]]: + """Split an observation into ``(state, frames)``, both name -> array. + + A camera frame is a channel-last image (rank 3, or rank 4 when batched); state is + any flat numeric vector. Anything else is dropped. Batched arrays keep their + leading ``[N]`` dim — the recorder slices per slot. + """ + state: dict[str, Any] = {} + frames: dict[str, Any] = {} + img_rank, vec_rank = (4, 2) if batched else (3, 1) + for name, val in flatten_observation(obs).items(): + arr = to_numpy(val) + if arr.ndim == img_rank and arr.shape[-1] in (1, 3, 4): + frames[name] = arr + elif arr.ndim <= vec_rank and np.issubdtype(arr.dtype, np.number): + state[name] = arr if batched else np.atleast_1d(arr) + return state, frames + + +def probe_success(info: Any, *, num_envs: int = 1) -> np.ndarray | None: + """Per-env success bools from conventional info keys; ``None`` when the env reports none.""" + if not isinstance(info, dict): + return None + for key in SUCCESS_KEYS: + if info.get(key) is not None: + return np.broadcast_to(to_numpy(info[key]).astype(bool).ravel(), (num_envs,)) + return None + + +def detect_fps(env: Any) -> int: + """Control rate from env metadata (``render_fps``) or Isaac's ``step_dt``; default 10.""" + fps = (getattr(env, "metadata", None) or {}).get("render_fps") + if fps: + return round(fps) + dt = getattr(getattr(env, "unwrapped", env), "step_dt", None) + return round(1 / dt) if dt else 10 + + +def capture_task_params(kwargs: dict[str, Any]) -> dict[str, Any]: + """Reset parametrization as json-safe trace metadata — the full kwargs, never a label.""" + + def safe(v: Any) -> Any: + if v is None or isinstance(v, (str, int, float, bool)): + return v + if isinstance(v, dict): + return {str(k): safe(x) for k, x in v.items()} + if isinstance(v, (list, tuple)): + return [safe(x) for x in v] + return str(v) + + return {k: safe(v) for k, v in kwargs.items() if v is not None} + + +def action_dim_of(env: Any, *, batched: bool) -> int: + """Per-env action size from the env's (possibly batched) action space.""" + space = getattr(env, "single_action_space", None) or getattr(env, "action_space", None) + shape = tuple(getattr(space, "shape", None) or ()) + if batched and len(shape) > 1: + shape = shape[1:] # batched Isaac spaces carry the [N] dim + return int(np.prod(shape)) if shape else 1 + + +def load_or_write_contract( + path: str | Path | None, + state: dict[str, Any], + frames: dict[str, Any], + action_dim: int, + fps: int, +) -> dict[str, Any]: + """The contract round-trip: load the user-edited file if present, else derive a + minimal contract from one (per-env) sample observation and write it once. + + Derived: just enough to structure the spaces and label plots — cameras as rgb + observations, state vectors with positional names, one action feature. Users + edit the written file to rename dimensions; nothing else is derived on purpose. + """ + file = Path(path) if path else None + if file is not None and file.exists(): + return json.loads(file.read_text()) + features: dict[str, Any] = {name: {"role": "observation", "type": "rgb"} for name in frames} + for name, vec in state.items(): + leaf = name.split("/")[-1] + features[name] = { + "role": "observation", + "names": [f"{leaf}_{i}" for i in range(int(np.asarray(vec).size))], + } + features["action"] = {"role": "action", "names": [f"act_{i}" for i in range(action_dim)]} + contract = {"control_rate": fps, "features": features} + if file is not None: + file.write_text(json.dumps(contract, indent=2) + "\n") + return contract + + +# ── hud.wrap: trace streaming for a loop you own ────────────────────────────── + + +class TracedEnv: + """The observing wrapper: same ``reset``/``step``/``close`` surface, plus telemetry. + + Plain envs are recorded as a batch of one; vectorized/batched envs fan out into one + trace per episode per slot (``done[i]`` closes slot ``i``'s trace and opens the next). + + - ``job`` — job name on the platform (default: the env's spec id / class name). + - ``job_id`` — share one job across several wrapped envs (multi-task suites). + - ``task`` — optional instruction/label shown on each trace's timeline. + - ``fps`` — control rate override (default: detected from the env). + - ``contract`` — path for the derived contract round-trip; ``None`` disables it. + - ``record_indices`` — which env slots get rich traces (default: first 4). """ def __init__( self, - factory: Any, + env: Any, *, + job: str | None = None, + job_id: str | None = None, + task: str | None = None, fps: int | None = None, contract: str | Path | None = "contract.json", - host: str = "127.0.0.1", - port: int = 0, - control_port: int = 9100, + record_indices: list[int] | None = None, ) -> None: - self._bridge = GymBridge(factory, host=host, port=port) - self._endpoint = RobotEndpoint(self._bridge) - self._fps = fps - self._contract_path = contract - self._contract: dict[str, Any] = {} - self._control_host = host - self._control_port = control_port - self._control_server: asyncio.AbstractServer | None = None - - # ── lifecycle (driven by env.gym()'s hooks, or directly) ──────────────────── - - async def start(self) -> None: - """Build the env (factory defaults), derive/load the contract, bring up the wire. - - Also serves the bridge's JSON-RPC control endpoint so a split-process - env can dial this process directly. - """ - await self._bridge.ensure_env() - state, frames = self._bridge.sample_observation() - existed = self._contract_path is not None and Path(self._contract_path).exists() - self._contract = load_or_write_contract( - self._contract_path, - state, - frames, - action_dim_of(self._bridge.env, batched=self._bridge.batched), - self._fps or detect_fps(self._bridge.env), - ) - if not existed and self._contract_path is not None: - logger.info("gym: wrote %s (edit names to relabel plots)", self._contract_path) - await self._bridge.start() - if self._control_server is None: - self._control_server = await self._endpoint.serve( - self._control_host, self._control_port - ) + self.env = env + self._batched = is_batched(env) + self._n = num_envs_of(env) or 1 + self._fps = fps or detect_fps(env) + self._job = job or getattr(getattr(env, "spec", None), "id", None) or type(env).__name__ + self._job_id = job_id + self._task = task + self._contract_path = Path(contract) if contract else None + self._record_indices = record_indices + self._rec: JobRecorder | None = None + self._closed = False + atexit.register(self.close) # flush traces even without an explicit close() - async def stop(self) -> None: - if self._control_server is not None: - self._control_server.close() - self._control_server = None - await self._bridge.stop() # also closes the built env + def __getattr__(self, name: str) -> Any: + return getattr(self.env, name) + + # ── the observed surface ────────────────────────────────────────────────── + + def reset(self, **kwargs: Any) -> Any: + result = self.env.reset(**kwargs) + obs = result[0] if isinstance(result, tuple) else result + if self._rec is None: + self._rec = self._start(obs) + else: + self._rec.close_slots() # an explicit mid-run reset ends open episodes + # The episode's full parametrization (reset kwargs + options), never a label. + params = capture_task_params( + {k: v for k, v in kwargs.items() if k != "options"} | (kwargs.get("options") or {}) + ) + self._rec.extra_metadata = {"task_params": params} if params else {} + return result - def capability(self, name: str = "robot") -> Capability: - """The concrete ``robot`` capability — mirrors ``Workspace.capability()``.""" - from hud.capabilities import Capability + def step(self, action: Any) -> Any: + result = self.env.step(action) + obs, reward, terminated, truncated = result[:4] + info = result[4] if len(result) > 4 else {} + if self._rec is not None: + state, frames = split_observation(obs, batched=self._batched) + done = np.atleast_1d(to_numpy(terminated)).astype(bool) | np.atleast_1d( + to_numpy(truncated) + ).astype(bool) + if not self._batched: # record a plain env as a batch of one + state = {k: v[None] for k, v in state.items()} + frames = {k: v[None] for k, v in frames.items()} + action = to_numpy(action)[None] + reward = np.atleast_1d(to_numpy(reward)) + self._rec.record( + obs=state or None, + frames=frames or None, + action=action, + reward=reward, + done=done, + success=probe_success(info, num_envs=self._n), + ) + return result - return Capability.robot(name=name, url=self._bridge.url, contract=self._contract) + def close(self) -> None: + if self._closed: + return + self._closed = True + if self._rec is not None: + self._rec.close() + self.env.close() - # ── the template surface ───────────────────────────────────────────────────── + def _start(self, obs: Any) -> JobRecorder: + """First reset: derive/load the contract, then open the job's recorder.""" + state, frames = split_observation(obs, batched=self._batched) + sample = {k: to_numpy(v)[0] if self._batched else v for k, v in state.items()} + existed = self._contract_path is not None and self._contract_path.exists() + contract = load_or_write_contract( + self._contract_path, + sample, + frames, + action_dim_of(self.env, batched=self._batched), + self._fps, + ) + if not existed and self._contract_path is not None: + logger.info("hud.wrap: wrote %s (edit names to relabel plots)", self._contract_path) + feats = contract.get("features", {}) + return JobRecorder( + self._job, + self._n, + record_indices=self._record_indices, + fps=int(contract.get("control_rate") or self._fps), + job_id=self._job_id, + prompt=self._task, + action_names=next( + (f.get("names") for f in feats.values() if f.get("role") == "action"), None + ), + state_names={ + k: f["names"] + for k, f in feats.items() + if f.get("role") == "observation" and f.get("names") + }, + ) - async def reset(self, **task_args: Any) -> str: - """Start an episode (rebuilding the env if an env-defining arg changed); - returns the task prompt.""" - return await self._endpoint.reset(**task_args) + def __enter__(self) -> Self: + return self - async def result(self) -> dict[str, Any]: - """The episode grade: per-slot dicts under ``"slots"``, means at the top.""" - return await self._endpoint.result() + def __exit__(self, *exc: object) -> None: + self.close() - @property - def env(self) -> Any: - """The live env — privileged sim access for custom grading in templates.""" - return self._bridge.env - @property - def contract(self) -> dict[str, Any]: - return self._contract +# The public verb: ``hud.wrap(env, job=...)`` — construction is the wrapping. +wrap = TracedEnv -__all__ = ["Gym"] +__all__ = [ + "SUCCESS_KEYS", + "TracedEnv", + "action_dim_of", + "capture_task_params", + "detect_fps", + "flatten_observation", + "is_batched", + "load_or_write_contract", + "num_envs_of", + "probe_success", + "split_observation", + "wrap", +] diff --git a/hud/environment/robot/introspect.py b/hud/environment/robot/introspect.py deleted file mode 100644 index 337cd60a9..000000000 --- a/hud/environment/robot/introspect.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Gym-env introspection: split observations, probe success, derive a minimal contract. - -Pure functions shared by :func:`hud.wrap` (in-process trace streaming), the -:class:`~.bridge.GymBridge`, and the :class:`~.gym.Gym` handle — no telemetry, -no wrapper state. -""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any - -import numpy as np - -from hud.telemetry.robot import to_numpy - -#: Conventional info keys carrying env-reported success, probed in order. -SUCCESS_KEYS = ("success", "is_success", "task_success") - - -def flatten_observation(obs: Any, prefix: str = "") -> dict[str, Any]: - """Flatten nested dict observations to slash-keyed leaves (non-dicts -> ``{"obs": x}``).""" - if not isinstance(obs, dict): - return {prefix or "obs": obs} - flat: dict[str, Any] = {} - for k, v in obs.items(): - key = f"{prefix}/{k}" if prefix else str(k) - flat.update(flatten_observation(v, key) if isinstance(v, dict) else {key: v}) - return flat - - -def split_observation(obs: Any, *, batched: bool = False) -> tuple[dict[str, Any], dict[str, Any]]: - """Split an observation into ``(state, frames)``, both name -> array. - - A camera frame is a channel-last image (rank 3, or rank 4 when batched); state is - any flat numeric vector. Anything else is dropped. Batched arrays keep their - leading ``[N]`` dim — the recorder slices per slot. - """ - state: dict[str, Any] = {} - frames: dict[str, Any] = {} - img_rank, vec_rank = (4, 2) if batched else (3, 1) - for name, val in flatten_observation(obs).items(): - arr = to_numpy(val) - if arr.ndim == img_rank and arr.shape[-1] in (1, 3, 4): - frames[name] = arr - elif arr.ndim <= vec_rank and np.issubdtype(arr.dtype, np.number): - state[name] = arr if batched else np.atleast_1d(arr) - return state, frames - - -def probe_success(info: Any, *, num_envs: int = 1) -> np.ndarray | None: - """Per-env success bools from conventional info keys; ``None`` when the env reports none.""" - if not isinstance(info, dict): - return None - for key in SUCCESS_KEYS: - if info.get(key) is not None: - return np.broadcast_to(to_numpy(info[key]).astype(bool).ravel(), (num_envs,)) - return None - - -def detect_fps(env: Any) -> int: - """Control rate from env metadata (``render_fps``) or Isaac's ``step_dt``; default 10.""" - fps = (getattr(env, "metadata", None) or {}).get("render_fps") - if fps: - return round(fps) - dt = getattr(getattr(env, "unwrapped", env), "step_dt", None) - return round(1 / dt) if dt else 10 - - -def capture_task_params(kwargs: dict[str, Any]) -> dict[str, Any]: - """Reset parametrization as json-safe trace metadata — the full kwargs, never a label.""" - - def safe(v: Any) -> Any: - if v is None or isinstance(v, (str, int, float, bool)): - return v - if isinstance(v, dict): - return {str(k): safe(x) for k, x in v.items()} - if isinstance(v, (list, tuple)): - return [safe(x) for x in v] - return str(v) - - return {k: safe(v) for k, v in kwargs.items() if v is not None} - - -def derive_contract( - state: dict[str, Any], frames: dict[str, Any], action_dim: int, fps: int -) -> dict[str, Any]: - """Minimal contract from one (per-env) sample observation + the action size. - - Just enough to structure the spaces and label plots — cameras as rgb observations, - state vectors with positional names, one action feature. Users edit the written - ``contract.json`` to rename dimensions; nothing else is derived on purpose. - """ - features: dict[str, Any] = {} - for name in frames: - features[name] = {"role": "observation", "type": "rgb"} - for name, vec in state.items(): - leaf = name.split("/")[-1] - features[name] = { - "role": "observation", - "names": [f"{leaf}_{i}" for i in range(int(np.asarray(vec).size))], - } - features["action"] = { - "role": "action", - "names": [f"act_{i}" for i in range(action_dim)], - } - return {"control_rate": fps, "features": features} - - -def load_or_write_contract( - path: str | Path | None, - state: dict[str, Any], - frames: dict[str, Any], - action_dim: int, - fps: int, -) -> dict[str, Any]: - """The contract round-trip: load the user-edited file if present, else derive a - minimal contract from the sample observation and write it once for inspection.""" - file = Path(path) if path else None - if file is not None and file.exists(): - return json.loads(file.read_text()) - contract = derive_contract(state, frames, action_dim, fps) - if file is not None: - file.write_text(json.dumps(contract, indent=2) + "\n") - return contract - - -def action_dim_of(env: Any, *, batched: bool) -> int: - """Per-env action size from the env's (possibly batched) action space.""" - space = getattr(env, "single_action_space", None) or getattr(env, "action_space", None) - shape = tuple(getattr(space, "shape", None) or ()) - if batched and len(shape) > 1: - shape = shape[1:] # batched Isaac spaces carry the [N] dim - return int(np.prod(shape)) if shape else 1 - - -__all__ = [ - "SUCCESS_KEYS", - "action_dim_of", - "capture_task_params", - "derive_contract", - "detect_fps", - "flatten_observation", - "load_or_write_contract", - "probe_success", - "split_observation", -] diff --git a/hud/environment/robot/sim_thread.py b/hud/environment/robot/sim.py similarity index 83% rename from hud/environment/robot/sim_thread.py rename to hud/environment/robot/sim.py index 67b2bcc46..03c17db82 100644 --- a/hud/environment/robot/sim_thread.py +++ b/hud/environment/robot/sim.py @@ -1,15 +1,15 @@ -"""One way to run any simulation: the sim owns the process main thread. - -A simulator is usually *thread-affine* (every touch must run on the thread that -created its GL/device context) and some — Isaac/Omniverse — must own the process -main thread outright: Kit drives its own main-thread loop and ``env.reset()`` -nests ``run_until_complete``, which cannot run inside an asyncio task. - -So HUD serves every sim with one process shape: serving (control channel + -robot WebSocket) runs on a background loop thread, and every sim touch is -queued to the main thread via a :class:`SimThread`. :func:`run_with_sim` is -that shape — cheap CPU envs just block on the queue; when Kit is loaded, the -idle hook pumps it between touches. +"""The sim-process shape every sim program runs: the sim owns the main thread. + +There is one way to serve a simulator (see :mod:`~.bridge`), whatever the sim: +serving — the robot WebSocket and the control side channel — runs on a +background loop thread, and every sim touch is queued to the process main +thread through the shared :class:`SimThread`. One shape because the hardest +sims demand it: a simulator is usually *thread-affine* (every touch must run +on the thread that created its GL/device context), and Isaac/Omniverse must +own the process main thread outright — Kit drives its own main-thread loop and +``env.reset()`` nests ``run_until_complete``, which cannot run inside an +asyncio task. Cheap CPU envs pay ~nothing: the main thread just blocks on the +queue. Before :meth:`SimThread.run` starts (tests, in-loop use) calls execute inline on the caller — the degenerate single-thread case. @@ -97,11 +97,7 @@ def _execute(fn: Callable[[], Any], fut: Future) -> None: fut.set_exception(exc) -def run_with_sim( - serve: Callable[[], Coroutine[Any, Any, Any]], - *, - sim: SimThread | None = None, -) -> None: +def run_with_sim(serve: Callable[[], Coroutine[Any, Any, Any]]) -> None: """THE process shape for serving a sim, blocking for the serve's lifetime. ``await serve()`` runs on a background loop thread while the shared @@ -109,7 +105,7 @@ def run_with_sim( cancel the serve coroutine; the sim keeps draining through its teardown (``env.stop()`` touches the sim too), then this returns. """ - sim = sim or SimThread.shared() + sim = SimThread.shared() sim.bind() # claim main before serving starts, so no touch runs inline elsewhere loop = asyncio.new_event_loop() done = threading.Event() diff --git a/hud/environment/robot/wrap.py b/hud/environment/robot/wrap.py deleted file mode 100644 index 20042f280..000000000 --- a/hud/environment/robot/wrap.py +++ /dev/null @@ -1,222 +0,0 @@ -"""``hud.wrap`` — one-line trace streaming for gym-style envs. - -Wrap any ``gym.Env``, ``gym.vector.VectorEnv``, or batched-tensor Isaac env and keep -using it exactly as before; every episode streams to the platform as a trace (numeric -state, per-camera H.264 video, actions, reward/success) under one job. The user — or -``lerobot-eval`` — keeps owning the loop; the wrapper only observes ``reset``/``step``. - - env = hud.wrap(make_env(...), job="chess-eval") - -On first reset a minimal ``contract.json`` is written next to your script describing how -the observation/action spaces were interpreted; edit its ``names`` to relabel the -platform's plots. An existing file is loaded instead, so edits stick. -""" - -from __future__ import annotations - -import atexit -import logging -from pathlib import Path -from typing import Any, Self - -import numpy as np - -from hud.telemetry.robot import JobRecorder, to_numpy - -from .introspect import ( - action_dim_of, - capture_task_params, - detect_fps, - load_or_write_contract, - probe_success, - split_observation, -) - -logger = logging.getLogger(__name__) - - -def wrap( - env: Any, - *, - job: str | None = None, - job_id: str | None = None, - task: str | None = None, - fps: int | None = None, - contract: str | Path | None = "contract.json", - record_indices: list[int] | None = None, -) -> TracedEnv: - """Stream a gym-style env's episodes to the platform; returns the env, traced. - - - ``job`` — job name on the platform (default: the env's spec id / class name). - - ``job_id`` — share one job across several wrapped envs (multi-task suites). - - ``task`` — optional instruction/label shown on each trace's timeline. - - ``fps`` — control rate override (default: detected from the env). - - ``contract`` — path for the derived contract round-trip; ``None`` disables it. - - ``record_indices`` — which env slots get rich traces (default: first 4). - """ - return TracedEnv( - env, - job=job, - job_id=job_id, - task=task, - fps=fps, - contract=contract, - record_indices=record_indices, - ) - - -class TracedEnv: - """The observing wrapper: same ``reset``/``step``/``close`` surface, plus telemetry. - - Plain envs are recorded as a batch of one; vectorized/batched envs fan out into one - trace per episode per slot (``done[i]`` closes slot ``i``'s trace and opens the next). - """ - - def __init__( - self, - env: Any, - *, - job: str | None, - job_id: str | None, - task: str | None, - fps: int | None, - contract: str | Path | None, - record_indices: list[int] | None, - ) -> None: - self.env = env - self._batched = _is_batched(env) - self._n = int(_num_envs(env) or 1) - self._fps = fps or detect_fps(env) - self._job = job or getattr(getattr(env, "spec", None), "id", None) or type(env).__name__ - self._job_id = job_id - self._task = task - self._contract_path = Path(contract) if contract else None - self._record_indices = record_indices - self._rec: JobRecorder | None = None - self._closed = False - atexit.register(self.close) # flush traces even without an explicit close() - - def __getattr__(self, name: str) -> Any: - return getattr(self.env, name) - - # ── the observed surface ────────────────────────────────────────────────── - - def reset(self, **kwargs: Any) -> Any: - result = self.env.reset(**kwargs) - obs = result[0] if isinstance(result, tuple) else result - if self._rec is None: - self._rec = self._start(obs) - else: - self._rec.close_slots() # an explicit mid-run reset ends open episodes - # The episode's full parametrization (reset kwargs + options), never a label. - params = capture_task_params( - {k: v for k, v in kwargs.items() if k != "options"} | (kwargs.get("options") or {}) - ) - self._rec.extra_metadata = {"task_params": params} if params else {} - return result - - def step(self, action: Any) -> Any: - result = self.env.step(action) - obs, reward, terminated, truncated = result[:4] - info = result[4] if len(result) > 4 else {} - if self._rec is not None: - state, frames = split_observation(obs, batched=self._batched) - done = np.atleast_1d(to_numpy(terminated)).astype(bool) | np.atleast_1d( - to_numpy(truncated) - ).astype(bool) - if not self._batched: # record a plain env as a batch of one - state = {k: v[None] for k, v in state.items()} - frames = {k: v[None] for k, v in frames.items()} - action = to_numpy(action)[None] - reward = np.atleast_1d(to_numpy(reward)) - self._rec.record( - obs=state or None, - frames=frames or None, - action=action, - reward=reward, - done=done, - success=probe_success(info, num_envs=self._n), - ) - return result - - def close(self) -> None: - if self._closed: - return - self._closed = True - if self._rec is not None: - self._rec.close() - self.env.close() - - # ── setup ──────────────────────────────────────────────────────────────── - - def _start(self, obs: Any) -> JobRecorder: - """First reset: derive/load the contract, then open the job's recorder.""" - state, frames = split_observation(obs, batched=self._batched) - sample = {k: to_numpy(v)[0] if self._batched else v for k, v in state.items()} - contract = self._contract(sample, frames) - feats = contract.get("features", {}) - rec = JobRecorder( - self._job, - self._n, - record_indices=self._record_indices, - fps=int(contract.get("control_rate") or self._fps), - job_id=self._job_id, - prompt=self._task, - action_names=next( - (f.get("names") for f in feats.values() if f.get("role") == "action"), None - ), - state_names={ - k: f["names"] - for k, f in feats.items() - if f.get("role") == "observation" and f.get("names") - }, - ) - return rec - - def _contract(self, state: dict[str, Any], frames: dict[str, Any]) -> dict[str, Any]: - """Load the user-edited contract if present; else derive one and write it once.""" - existed = self._contract_path is not None and self._contract_path.exists() - contract = load_or_write_contract( - self._contract_path, - state, - frames, - action_dim_of(self.env, batched=self._batched), - self._fps, - ) - if not existed and self._contract_path is not None: - logger.info("hud.wrap: wrote %s (edit names to relabel plots)", self._contract_path) - return contract - - def __enter__(self) -> Self: - return self - - def __exit__(self, *exc: object) -> None: - self.close() - - -def _num_envs(env: Any) -> int | None: - """``num_envs`` from the env or its unwrapped core (gymnasium 1.x wrappers - no longer forward attributes, so an Isaac env inside ``gym.make`` hides it).""" - n = getattr(env, "num_envs", None) - if n is None: - n = getattr(getattr(env, "unwrapped", env), "num_envs", None) - return None if n is None else int(n) - - -def _is_batched(env: Any) -> bool: - """Vectorized (gym VectorEnv) or batched-tensor (Isaac) envs carry a leading [N] dim. - - Any env exposing ``num_envs`` is batched — Isaac keeps batch semantics even at - ``num_envs == 1``; plain ``gym.Env``s don't have the attribute at all. - """ - if _num_envs(env) is not None: - return True - try: - import gymnasium as gym - - return isinstance(getattr(env, "unwrapped", env), gym.vector.VectorEnv) - except ImportError: - return False - - -__all__ = ["TracedEnv", "wrap"] diff --git a/hud/environment/server.py b/hud/environment/server.py index 50c3a5830..e4ec9f838 100644 --- a/hud/environment/server.py +++ b/hud/environment/server.py @@ -462,19 +462,11 @@ async def _serve_until_terminated(env: Environment, host: str, port: int) -> Non def serve_blocking(env: Environment, host: str, port: int) -> None: """Serve *env*, blocking until terminated — the one entry every CLI path uses. - An env with attached sims (``env.gym(...)``) runs the sim-main shape: the - sim owns this (main) thread and serving runs on a background loop thread, - with every sim touch queued back here (see - :mod:`hud.environment.robot.sim_thread`). Anything else serves on a plain - ``asyncio.run``. + Sims never run here: an env with attached sims (``env.gym(...)``) spawns + each one as its own process from an ``@env.initialize`` hook (see + :mod:`hud.environment.robot.bridge`), so every env serves the same way. """ - if not getattr(env, "_sims", None): - asyncio.run(_serve_until_terminated(env, host, port)) - return - # Submodule import: only sim-serving needs it (keeps non-robot envs free of the extra). - from hud.environment.robot.sim_thread import run_with_sim - - run_with_sim(lambda: serve(env, host, port)) + asyncio.run(_serve_until_terminated(env, host, port)) def main() -> None: From 748ed4c42af316c66cb8d4653675f14ec239e4b4 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 01:39:53 +0000 Subject: [PATCH 18/43] feat(eval): add Run.started carrying the full tasks.start reply Generic hook for per-episode data an env hands back on start (e.g. a robot bridge slot token), without hud/eval learning robot concepts. Co-authored-by: Cursor --- hud/eval/run.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/hud/eval/run.py b/hud/eval/run.py index 41592cd9c..66140c8d5 100644 --- a/hud/eval/run.py +++ b/hud/eval/run.py @@ -117,6 +117,9 @@ def __init__(self, client: HudClient | None, task_id: str, args: dict[str, Any]) #: chat-style / multi-turn prompts. Agents consume the normalized #: views: :attr:`prompt_messages` / :attr:`prompt_text`. self.prompt: str | list[Any] | None = None + #: Full ``tasks.start`` reply frame — prompt plus any per-episode data + #: the env handed the agent (e.g. a robot slot token under ``robot``). + self.started: dict[str, Any] = {} #: The structured grading result (all-default until graded on exit). self.grade = Grade() self.trace = Trace() @@ -195,6 +198,7 @@ def record(self, step: Step) -> None: async def __aenter__(self) -> Self: started_at = now_iso() started = await self.client.start_task(self._task_id, self._args) + self.started = started self.prompt = started.get("prompt") self.record( Step( From b5258cd9cc59a56bbb5c67bb4b9aff49373cbe9e Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 01:39:58 +0000 Subject: [PATCH 19/43] feat(environment/robot): lock RobotEndpoint; reset/result carry a token N sessions now share one control-channel TCP link, so _call needs a lock around send+read to keep replies from crossing. reset() returns the full {prompt, token} reply and result(token=...) passes it through, matching the bridge's per-slot claim/release surface. Co-authored-by: Cursor --- hud/environment/robot/endpoint.py | 42 +++++++++++++++++-------------- 1 file changed, 23 insertions(+), 19 deletions(-) diff --git a/hud/environment/robot/endpoint.py b/hud/environment/robot/endpoint.py index df0580e7a..9290ed82b 100644 --- a/hud/environment/robot/endpoint.py +++ b/hud/environment/robot/endpoint.py @@ -19,9 +19,10 @@ @env.template(id="pawn_lift") - async def pawn_lift(task: str = "solo_pawn_lift", seed: int = 0, num_envs: int = 1): - yield {"prompt": await sim.reset(task=task, seed=seed, num_envs=num_envs)} - yield await sim.result() + async def pawn_lift(task: str = "solo_pawn_lift", seed: int = 0): + ep = await sim.reset(task=task, seed=seed) # {prompt, token} + yield {"prompt": ep["prompt"], "robot": {"token": ep["token"]}} + yield await sim.result(token=ep["token"]) """ from __future__ import annotations @@ -65,6 +66,8 @@ def __init__( self._forward: asyncio.Task[None] | None = None self._reader: asyncio.StreamReader | None = None self._writer: asyncio.StreamWriter | None = None + # N sessions share this one TCP link; serialize send+read so replies don't cross. + self._lock = asyncio.Lock() @classmethod def spawn(cls, cmd: Sequence[str], *, connect_timeout_s: float = 240.0) -> RobotEndpoint: @@ -160,13 +163,13 @@ async def capability(self, name: str = "robot") -> Capability: return Capability.robot(name=name, url=await self.url(), contract=await self.contract()) - async def reset(self, **task_args: Any) -> str: - """Start a new episode; return the task prompt.""" - return (await self._call("reset", task_args))["prompt"] + async def reset(self, **task_args: Any) -> dict[str, Any]: + """Claim a slot for a new episode; return ``{"prompt", "token"}``.""" + return await self._call("reset", task_args) - async def result(self, **extra: Any) -> dict[str, Any]: - """The episode score dict, merged with any caller ``extra`` metadata.""" - res = {**(await self._call("result")), **extra} + async def result(self, *, token: str, **extra: Any) -> dict[str, Any]: + """This slot's score dict (frees the slot), merged with any caller ``extra``.""" + res = {**(await self._call("result", {"token": token})), **extra} print( f"[env] result: success={res.get('success')} " f"total_reward={res.get('total_reward', 0.0):.3f}", @@ -175,18 +178,19 @@ async def result(self, **extra: Any) -> dict[str, Any]: return res async def _call(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]: - # Strictly request/reply, one call at a time, so a constant id is enough. + # One in-flight RPC: N sessions share this link; constant id is enough under the lock. if self._writer is None or self._reader is None: raise RuntimeError("not connected; call start() first") - await send_frame( - self._writer, {"jsonrpc": "2.0", "id": 1, "method": method, "params": params or {}} - ) - msg = await read_frame(self._reader) - if msg is None: - raise ConnectionError(f"connection closed awaiting {method!r} reply") - if "error" in msg: - raise RuntimeError(f"{method} failed: {msg['error']['message']}") - return msg["result"] + async with self._lock: + await send_frame( + self._writer, {"jsonrpc": "2.0", "id": 1, "method": method, "params": params or {}} + ) + msg = await read_frame(self._reader) + if msg is None: + raise ConnectionError(f"connection closed awaiting {method!r} reply") + if "error" in msg: + raise RuntimeError(f"{method} failed: {msg['error']['message']}") + return msg["result"] async def _forward_lines(stream: asyncio.StreamReader) -> None: From fded695047f247a76ed08d4ff224e59e20b35794 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 01:40:09 +0000 Subject: [PATCH 20/43] feat(environment/robot): slot registry + barrier stepping in the bridge Replaces the single-agent, batched-wire bridge with N claimable slots behind one scalar openpi WebSocket per connection: - _SlotRegistry: reset-on-first-claim (global reset + claim slot 0 when all free), claim a free slot otherwise, error when none are free. A released slot stays unreclaimable until the next global reset (v1 keeps whole-batch episodes). - Each connection's first frame is {"claim": token}; the bridge binds it to that slot and fans out that slot's scalar observation only. - _tick_loop barriers claimed, connected slots: gather pending actions (or step_timeout -> hold for stragglers), step once as [N, A], fan scalar frames back out. One slow agent can no longer stall the batch. - Control surface: reset() claims + returns {prompt, token}; result() reads one slot's grade by token and frees it. The old aggregate "slots" grade list is gone. GymBridge's build-arg partition now also takes env.gym(..., num_envs=8) style defaults, so num_envs is sim build config, not an eval parameter. Co-authored-by: Cursor --- hud/environment/robot/__init__.py | 4 +- hud/environment/robot/bridge.py | 342 +++++++++++++++++++++++------- 2 files changed, 266 insertions(+), 80 deletions(-) diff --git a/hud/environment/robot/__init__.py b/hud/environment/robot/__init__.py index 91ac461b1..c3ecf5ff0 100644 --- a/hud/environment/robot/__init__.py +++ b/hud/environment/robot/__init__.py @@ -5,8 +5,8 @@ thread and a bridge serves the wire. The env server holds a :class:`~.endpoint.RobotEndpoint` on it: -- :class:`~.bridge.RobotBridge` — the batched-first bridge base (``num_envs`` - slots in lockstep; a plain single env is a batch of one): the agent's +- :class:`~.bridge.RobotBridge` — bridge base (``num_envs`` slots in lockstep + internally; scalar openpi wire per claimed connection): the agent's ``robot`` WebSocket plus the JSON-RPC control side channel. - :class:`~.bridge.GymBridge` / :func:`~.bridge.gym_command` — the generic gym path (``env.gym(...)`` — a factory, registry id, or declared env): diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py index a4f7d8eca..e3141c929 100644 --- a/hud/environment/robot/bridge.py +++ b/hud/environment/robot/bridge.py @@ -5,10 +5,11 @@ there. A bridge lives in the sim's own process and serves two channels: the agent's obs/action **WebSocket** and the env server's JSON-RPC **control side channel** (a :class:`~.endpoint.RobotEndpoint` drives episodes through it). -Bridges are **batched-first**: one bridge serves ``num_envs`` slots in lockstep -(``[N, ...]`` obs frames, ``[N, A]`` actions, an ``[N]`` ``terminated`` mask). -``num_envs == 1`` — a plain single-env sim — speaks the scalar framing -(per-env arrays, scalar ``terminated``) on the same code path. + +Internally the sim may be vectorized (``num_envs`` slots in lockstep). On the +wire every connection is scalar openpi: one claim token binds a connection to +one slot; the bridge barriers actions across claimed slots, steps once, and +fans scalar obs back. Single-env is the one-slot degenerate case. - :class:`RobotBridge` — subclass with your sim: ``reset`` / ``step`` / ``get_observation``, and set ``self.contract``. The base owns both channels. @@ -34,8 +35,10 @@ import asyncio import contextlib import inspect +import secrets import sys from abc import ABC, abstractmethod +from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -64,6 +67,59 @@ PORT_ANNOUNCEMENT = "HUD_SIM_PORT=" +@dataclass +class _Slot: + """One sim slot: claimed by a control-plane token, driven by one WS connection.""" + + index: int + token: str | None = None + ws: Any = None + action: np.ndarray | None = None + idle: bool = False # timed out waiting for an action; hold next step + # Touched this episode; after release stays unreclaimable until the next global reset. + used: bool = False + + +@dataclass +class _SlotRegistry: + """Free/claimed slots for one bridge. All-free → next reset is global.""" + + slots: list[_Slot] = field(default_factory=list) + + def configure(self, n: int) -> None: + self.slots = [_Slot(index=i) for i in range(n)] + + @property + def all_free(self) -> bool: + return all(s.token is None for s in self.slots) + + def free_slot(self) -> _Slot | None: + # v1: a freed slot is only reclaimable after all slots free (whole-batch episodes). + return next((s for s in self.slots if s.token is None and not s.used), None) + + def by_token(self, token: str) -> _Slot | None: + return next((s for s in self.slots if s.token == token), None) + + def claimed(self) -> list[_Slot]: + return [s for s in self.slots if s.token is not None] + + def claim(self, slot: _Slot) -> str: + token = f"slot-{slot.index}-{secrets.token_hex(4)}" + slot.token = token + slot.used = True + slot.action = None + slot.idle = False + slot.ws = None + return token + + def release(self, slot: _Slot) -> None: + slot.token = None + slot.ws = None + slot.action = None + slot.idle = False + # keep used=True until configure() on the next global reset + + class RobotBridge(ABC): """Serves a sim over ``robot`` WebSocket + a JSON-RPC control side channel. @@ -74,22 +130,21 @@ class RobotBridge(ABC): (default 1). - :meth:`reset` initialises the sim for a new episode and returns the task - prompt. The base resets scoring state and pushes the first frame. - - :meth:`step` advances the sim by one action (``[A]``, or ``[N, A]`` batched). - - :meth:`get_observation` returns ``(data, terminated)`` — per-env arrays + - scalar bool for ``num_envs == 1``, ``[N, ...]`` arrays + ``[N]`` mask - otherwise — or ``None`` if not ready. - - :meth:`result_slots` returns one score dict per slot. The default covers - the common single-env binary-success case from ``self.success`` / - ``self.total_reward``; batched bridges override it. + prompt. The base resets scoring state. + - :meth:`step` advances the sim by one batched action ``[N, A]`` (N=1 ok). + - :meth:`get_observation` returns ``(data, terminated)`` with ``[N, ...]`` + arrays and an ``[N]`` terminated mask (N=1 ok), or ``None`` if not ready. + - :meth:`result_slots` returns one score dict per slot. """ + #: Seconds to wait for every claimed slot's action before stepping with holds. + step_timeout: float = 30.0 + def __init__(self, *, host: str = "127.0.0.1", port: int = 0) -> None: # Loopback + ephemeral by default; the concrete address is published in the # manifest post-``start()`` and tunneled, so no env manages bridge ports. self._host = host self._port = port - self._client: Any = None # robot serves a single agent at a time self._server: Any = None # Connect-time metadata frame (sent first on each connection); subclasses may set it. self.metadata: dict[str, Any] = {} @@ -99,22 +154,41 @@ def __init__(self, *, host: str = "127.0.0.1", port: int = 0) -> None: # Every sim touch runs on the process sim thread (see sim.py). self._sim = SimThread.shared() self.num_envs: int = 1 - # Episode scoring read by ``result()``; single-env subclasses update these + self._registry = _SlotRegistry() + self._registry.configure(1) + self._tick_task: asyncio.Task[None] | None = None + self._action_event = asyncio.Event() + # Episode scoring read by ``result_slots``; single-env subclasses update these # in ``reset``/``step`` (batched bridges override result_slots instead). self.task_description: str = "" self.total_reward: float = 0.0 self.success: bool = False self.terminated: bool = False - async def _reset(self, **kwargs: Any) -> str: - """Internal reset entry (called by the endpoint): reset scoring, run the - author's :meth:`reset`, push the first frame.""" - self.total_reward = 0.0 - self.success = False - self.terminated = False - self.task_description = await self.reset(**kwargs) - await self._send_observation() # first frame for an already-connected agent - return self.task_description + async def _claim_episode(self, **kwargs: Any) -> dict[str, Any]: + """Control-plane reset: global sim reset when all free, else claim a free slot.""" + if self._registry.all_free: + self.total_reward = 0.0 + self.success = False + self.terminated = False + self.task_description = await self.reset(**kwargs) + self._registry.configure(self.num_envs) + slot = self._registry.slots[0] + else: + slot = self._registry.free_slot() + if slot is None: + raise RuntimeError(f"all {self.num_envs} slots are claimed") + token = self._registry.claim(slot) + return {"prompt": self.task_description, "token": token} + + def _release_episode(self, token: str) -> dict[str, Any]: + """Control-plane result: this slot's score, then free it.""" + slot = self._registry.by_token(token) + if slot is None: + raise ValueError(f"unknown episode token: {token!r}") + grade = self.result_slots()[slot.index] + self._registry.release(slot) + return grade @abstractmethod async def reset(self, **kwargs: Any) -> str: @@ -122,11 +196,11 @@ async def reset(self, **kwargs: Any) -> str: @abstractmethod def step(self, action: np.ndarray) -> None: - """Advance the sim by one action.""" + """Advance the sim by one batched action ``[N, A]``.""" @abstractmethod - def get_observation(self) -> tuple[dict[str, np.ndarray], Any] | None: - """Return ``(data, terminated)`` for the current state, or ``None`` if not ready.""" + def get_observation(self) -> tuple[dict[str, np.ndarray], np.ndarray] | None: + """Return ``(data[N, ...], terminated[N])``, or ``None`` if not ready.""" def result_slots(self) -> list[dict[str, Any]]: """One score dict per slot. Default: single-env binary success.""" @@ -138,19 +212,14 @@ def result_slots(self) -> list[dict[str, Any]]: } ] - def result(self) -> dict[str, Any]: - """The episode grade: per-slot dicts under ``"slots"``, means at the top. - - The vectorized eval path grades each trace from ``slots[i]``; single-result - consumers read the aggregate ``score``/``success`` unchanged. - """ - slots = self.result_slots() - return { - "score": float(np.mean([s["score"] for s in slots])), - "success": float(np.mean([float(s["success"]) for s in slots])), - "total_reward": float(np.mean([s.get("total_reward", 0.0) for s in slots])), - "slots": slots, - } + def hold_action(self) -> np.ndarray: + """Action used for idle/stalled slots at a barrier step (zeros).""" + action = next( + (f for f in self.contract.get("features", {}).values() if f.get("role") == "action"), + {}, + ) + shape = action.get("shape") or (1,) + return np.zeros(shape, dtype=np.float32) @property def url(self) -> str: @@ -168,58 +237,134 @@ async def start(self) -> None: # second run against the same sim) is a no-op rather than an EADDRINUSE rebind. if self._server is not None: return + self._registry.configure(self.num_envs) + # No keepalive: a lockstep step can block the sim (and the GIL) for + # minutes during heavy resets; ping timeouts would sever a healthy run. self._server = await websockets.serve( - self._handle_client, self._host, self._port, max_size=None, reuse_address=True + self._handle_client, + self._host, + self._port, + max_size=None, + reuse_address=True, + ping_interval=None, ) if self._port == 0: self._port = self._server.sockets[0].getsockname()[1] + self._tick_task = asyncio.create_task(self._tick_loop()) print(f"[env] robot listening on ws://{self._host}:{self._port}", flush=True) async def stop(self) -> None: + if self._tick_task is not None: + self._tick_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._tick_task + self._tick_task = None if self._server is not None: self._server.close() await self._server.wait_closed() self._server = None async def _handle_client(self, ws: Any) -> None: - # A later connection replaces the previous one (only one agent at a time). - self._client = ws + """One agent connection: claim a slot, then feed actions into the barrier.""" + slot: _Slot | None = None try: await ws.send(_packb(self.metadata)) # connect-time metadata frame - await self._send_observation() # current obs on connect (if ready) - async for raw in ws: - action = _unpackb(raw)["actions"] # codec already returns an ndarray - await self._sim.call(self.step, action) # on the sim thread - await self._send_observation() + raw = await ws.recv() + if isinstance(raw, str): + raise RuntimeError(raw) + claim = _unpackb(raw) + token = claim.get("claim") + if not isinstance(token, str): + raise ValueError("first frame must be {\"claim\": }") + slot = self._registry.by_token(token) + if slot is None: + raise ValueError(f"unknown claim token: {token!r}") + if slot.ws is not None: + raise RuntimeError(f"slot {slot.index} already has a live connection") + slot.ws = ws + slot.action = None + slot.idle = False + await self._send_slot_observation(slot) + async for frame in ws: + action = _unpackb(frame)["actions"] # codec already returns an ndarray + slot.action = np.asarray(action, dtype=np.float32) + slot.idle = False + self._action_event.set() except websockets.exceptions.ConnectionClosed: pass except Exception: - # Surface failures as a string frame (a traceback) instead of a silent close. import traceback with contextlib.suppress(Exception): await ws.send(traceback.format_exc()) raise finally: - if self._client is ws: - self._client = None - - async def _send_observation(self) -> None: - """Send the current observation to the connected agent (if any). - - Framing follows ``num_envs``: scalar ``terminated`` for a single env, an - ``[N]`` mask for a batch — the one place the two wire shapes meet. - """ - if self._client is None: + if slot is not None and slot.ws is ws: + slot.ws = None + slot.action = None + self._action_event.set() # wake the barrier so it doesn't wait on us + + async def _tick_loop(self) -> None: + """Gather claimed slots' actions (or timeout → hold), step once, fan out.""" + while True: + try: + await asyncio.wait_for(self._action_event.wait(), timeout=self.step_timeout) + except TimeoutError: + pass + self._action_event.clear() + claimed = [s for s in self._registry.claimed() if s.ws is not None] + if not claimed: + continue + # Wait until every live claimed slot has an action, or step_timeout elapses. + deadline = asyncio.get_running_loop().time() + self.step_timeout + while True: + pending = [s for s in claimed if s.action is None and not s.idle] + if not pending: + break + remaining = deadline - asyncio.get_running_loop().time() + if remaining <= 0: + for s in pending: + s.idle = True # hold this step + break + self._action_event.clear() + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(self._action_event.wait(), timeout=remaining) + claimed = [s for s in self._registry.claimed() if s.ws is not None] + if not claimed: + break + if not claimed: + continue + hold = self.hold_action() + # Stack one row per sim slot (idle/unconnected claimed → hold; free → hold). + rows = [] + for s in self._registry.slots: + if s in claimed and s.action is not None: + rows.append(np.asarray(s.action, dtype=np.float32).reshape(-1)) + else: + rows.append(np.asarray(hold, dtype=np.float32).reshape(-1)) + s.action = None + actions = np.stack(rows) + await self._sim.call(self.step, actions) + for s in claimed: + await self._send_slot_observation(s) + + async def _send_slot_observation(self, slot: _Slot) -> None: + """Fan one scalar obs frame to a claimed connection.""" + if slot.ws is None: return out = await self._sim.call(self.get_observation) if out is None: return data, terminated = out - done = np.asarray(terminated, dtype=bool) - msg = {**data, "terminated": bool(done.ravel()[0]) if self.num_envs == 1 else done} + i = slot.index + # Slice the [N, ...] batch down to this slot's scalar row. + msg = { + **{k: (v[i] if getattr(v, "ndim", 0) >= 1 and len(v) == self.num_envs else v) + for k, v in data.items()}, + "terminated": bool(np.asarray(terminated).reshape(-1)[i]), + } with contextlib.suppress(websockets.exceptions.ConnectionClosed): - await self._client.send(_packb(msg)) + await slot.ws.send(_packb(msg)) # ── the control side channel (driven by a RobotEndpoint) ──────────────────── @@ -247,9 +392,12 @@ async def _dispatch_control(self, method: str, params: dict[str, Any]) -> dict[s if method == "contract": return {"contract": self.contract} if method == "reset": - return {"prompt": await self._reset(**params)} + return await self._claim_episode(**params) if method == "result": - return self.result() + token = params.get("token") + if not isinstance(token, str): + raise ValueError("result: 'token' must be a string") + return self._release_episode(token) raise ValueError(f"unknown method {method!r}") @@ -265,6 +413,8 @@ class GymBridge(RobotBridge): options=...)``. For a factory, its signature is the partition — so ``num_envs`` in the signature is the vectorization declaration; for a registry target, ``num_envs`` is the one build arg (``gym.make_vec``). + Default build kwargs (e.g. ``num_envs=8`` from ``env.gym(..., num_envs=8)``) + fill in when the task does not pass them. ``fps`` overrides the detected control rate; ``contract`` is the ``contract.json`` round-trip path (load beats derive, so user edits stick; @@ -277,12 +427,15 @@ def __init__( *, fps: int | None = None, contract: str | Path | None = "contract.json", - **kwargs: Any, + **defaults: Any, ) -> None: - super().__init__(**kwargs) + host = defaults.pop("host", "127.0.0.1") + port = int(defaults.pop("port", 0)) + super().__init__(host=host, port=port) self._target = target self._fps = fps self._contract_path = contract + self._defaults = defaults # e.g. num_envs from env.gym(..., num_envs=8) # Task args that define the env build (everything else is episodic). if callable(target): self._build_params = { @@ -301,6 +454,7 @@ def __init__( self._done: np.ndarray = np.zeros(1, dtype=bool) self._success: np.ndarray = np.zeros(1, dtype=bool) self._acc_reward: np.ndarray = np.zeros(1) + self._step_reward: np.ndarray = np.zeros(1) # last env.step reward (RL wire sibling) self._seen_success = False # any env-reported success signal this episode # ── env lifecycle (all sim touches on the sim thread) ─────────────────────── @@ -336,8 +490,10 @@ async def stop(self) -> None: self.env = None def _sync_reset(self, task_args: dict[str, Any]) -> str: - build = {k: v for k, v in task_args.items() if k in self._build_params} - episodic = {k: v for k, v in task_args.items() if k not in self._build_params} + # Defaults from env.gym(..., num_envs=N) fill in when the task omits them. + merged = {**self._defaults, **task_args} + build = {k: v for k, v in merged.items() if k in self._build_params} + episodic = {k: v for k, v in merged.items() if k not in self._build_params} key = tuple(sorted(build.items())) if self.env is None or key != self._instance: if self.env is not None: @@ -354,8 +510,9 @@ def _sync_reset(self, task_args: dict[str, Any]) -> str: self._done = np.zeros(self.num_envs, dtype=bool) self._success = np.zeros(self.num_envs, dtype=bool) self._acc_reward = np.zeros(self.num_envs) + self._step_reward = np.zeros(self.num_envs, dtype=np.float32) self._seen_success = False - return self._prompt(task_args) + return self._prompt(merged) def _build_env(self, build: dict[str, Any]) -> Any: """Build the env from the target: factory call, or registry make/make_vec.""" @@ -395,7 +552,7 @@ def step(self, action: np.ndarray) -> None: if not self.batched: act = act[0] if act.ndim > 1 else act # single plain env: drop the batch dim elif act.ndim == 1: - act = act[None] # batched-of-one served scalar-framed: restore the [N] dim + act = act[None] # The wire carries floats; discrete/int action spaces need their dtype + shape back. space = getattr(self.env, "action_space", None) dtype = getattr(space, "dtype", None) @@ -411,7 +568,9 @@ def step(self, action: np.ndarray) -> None: done = np.atleast_1d(to_numpy(terminated)).astype(bool) | np.atleast_1d( to_numpy(truncated) ).astype(bool) - self._acc_reward += np.atleast_1d(to_numpy(reward)) * ~self._done + # Per-step reward for the RL wire (get_observation attaches it as a sibling). + self._step_reward = np.atleast_1d(to_numpy(reward)).astype(np.float32) + self._acc_reward += self._step_reward * ~self._done newly = done & ~self._done if newly.any(): success = self._resolve_success(info) @@ -434,14 +593,20 @@ def _resolve_success(self, info: Any) -> np.ndarray | None: return None return None - def get_observation(self) -> tuple[dict[str, np.ndarray], Any] | None: + def get_observation(self) -> tuple[dict[str, np.ndarray], np.ndarray] | None: + """Always ``[N, ...]`` arrays + ``[N]`` terminated for the barrier fan-out.""" if self.env is None or self._obs is None: return None state, frames = split_observation(self._obs, batched=self.batched) data = {k: to_numpy(v) for k, v in {**state, **frames}.items()} - if self.batched and self.num_envs == 1: - data = {k: v[0] for k, v in data.items()} # batched-of-one: squeeze to scalar framing - return data, self._done if self.num_envs > 1 else bool(self._done[0]) + # Last env.step reward rides along for RL collection (client lifts it + # out of "data" to a top-level sibling, like "terminated"). + data["reward"] = self._step_reward + if not self.batched: + # Plain single env: lift scalars to a batch of one for uniform slicing. + data = {k: (v if getattr(v, "ndim", 0) >= 1 and v.shape[:1] == (1,) else np.asarray(v)[None]) + for k, v in data.items()} + return data, np.asarray(self._done, dtype=bool).reshape(self.num_envs) def result_slots(self) -> list[dict[str, Any]]: """Per-slot grades: env-reported success when available, else accumulated reward.""" @@ -489,7 +654,11 @@ async def _serve() -> None: def gym_command( - target: Any, *, fps: int | None = None, contract: str | None = "contract.json" + target: Any, + *, + fps: int | None = None, + contract: str | None = "contract.json", + **defaults: Any, ) -> list[str]: """Build the command line that spawns *target* as a sim process (what ``env.gym`` runs). @@ -499,7 +668,8 @@ def gym_command( rebuild itself: a factory's source path, a registry id as-is, or a constructed env's spec (id + kwargs) — closed here since only the spec crosses over; the child calls ``gym.make`` again for its own instance. - ``fps`` / ``contract`` flow through to the :class:`GymBridge` the CLI builds. + ``fps`` / ``contract`` / build defaults (e.g. ``num_envs=``) flow through + to the :class:`GymBridge` the CLI builds. """ if callable(target): name = getattr(target, "__qualname__", "") @@ -520,6 +690,9 @@ def gym_command( cmd += ["--fps", str(fps)] if contract is not None: cmd += ["--contract", str(contract)] + # Build defaults (num_envs, etc.) as --key value pairs the CLI re-applies. + for key, value in defaults.items(): + cmd += [f"--{key.replace('_', '-')}", str(value)] return cmd @@ -534,13 +707,26 @@ def main() -> None: parser.add_argument("--contract", default=None, help="contract.json round-trip path.") parser.add_argument("--host", default="127.0.0.1", help="Control-channel interface.") parser.add_argument("--port", type=int, default=0, help="Control-channel port (0 = ephemeral).") - args = parser.parse_args() + parser.add_argument("--num-envs", type=int, default=None, help="Vectorized slot count.") + args, unknown = parser.parse_known_args() target: Any = args.target if ".py:" in target: # factory by source path; ids / spec JSON pass through path, _, attr = target.rpartition(":") target = getattr(load_module(path), attr) - bridge = GymBridge(target, fps=args.fps, contract=args.contract) + defaults: dict[str, Any] = {} + if args.num_envs is not None: + defaults["num_envs"] = args.num_envs + # Accept further --key value pairs as build defaults. + i = 0 + while i < len(unknown): + flag = unknown[i] + if flag.startswith("--") and i + 1 < len(unknown): + defaults[flag[2:].replace("-", "_")] = unknown[i + 1] + i += 2 + else: + i += 1 + bridge = GymBridge(target, fps=args.fps, contract=args.contract, **defaults) serve_bridge(bridge, host=args.host, port=args.port) From 2c698d442d453b46a225b65a7be48824edd6a213 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 01:40:15 +0000 Subject: [PATCH 21/43] feat(capabilities/robot): RobotClient.connect claims a slot by token After the connect-time metadata frame, send {"claim": token} when the caller supplies one, binding this WebSocket to the bridge's matching slot. terminated is now always scalar (each connection is one slot). Co-authored-by: Cursor --- hud/capabilities/robot.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/hud/capabilities/robot.py b/hud/capabilities/robot.py index 5a0c88960..714021a3a 100644 --- a/hud/capabilities/robot.py +++ b/hud/capabilities/robot.py @@ -68,13 +68,19 @@ def spaces(self) -> tuple[dict[str, Any], dict[str, Any]]: return action, observations @classmethod - async def connect(cls, cap: Capability) -> Self: - ws = await websockets.connect(cap.url, max_size=None) + async def connect(cls, cap: Capability, *, token: str | None = None) -> Self: + """Dial the robot WebSocket; ``token`` claims a sim slot after the metadata frame.""" + # No keepalive: heavy sims (Isaac resets) legitimately stall for + # minutes between frames; a ping timeout would kill a healthy rollout. + ws = await websockets.connect(cap.url, max_size=None, ping_interval=None) # Consume the connect-time metadata frame (always first); a string frame # is the env's error convention. raw = await ws.recv() if isinstance(raw, str): raise RuntimeError(f"robot env error on connect:\n{raw}") + # Bind this connection to a claimed episode slot (scalar openpi from here). + if token is not None: + await ws.send(_packb({"claim": token})) return cls(cap, ws) async def get_observation(self) -> dict[str, Any]: @@ -95,11 +101,13 @@ async def get_observation(self) -> dict[str, Any]: msg = await self._queue.get() if "error" in msg: raise RuntimeError(f"robot env error:\n{msg['error']}") - # Passthrough (no bool() coercion): a scalar bool for single-env bridges, an [N] - # mask for vectorized ones (RobotClient serves both; the wire decides the shape). + # Scalar openpi: each connection is one slot; terminated is always a bool. terminated = msg.pop("terminated", False) meta = msg.pop("meta", None) + reward = msg.pop("reward", None) out: dict[str, Any] = {"data": msg, "terminated": terminated} + if reward is not None: + out["reward"] = reward # per-step reward sibling (RL collection) if meta is not None: out["meta"] = meta return out From 278376f7e9e14ae1d657892844018419f68ea8f9 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 01:40:22 +0000 Subject: [PATCH 22/43] refactor(agents/robot): drop drive() for a single scalar rollout loop RobotAgent goes back to the main-branch contract: __call__(run) only. Deletes drive(runs, client), the single-vs-batched wire-shape sniffing, the [None] lifts, per-slot recorder fan-out, and the n==1 vs asyncio.to_thread(model.infer) fork - ainfer is the only inference path now, so BatchedModel is the sole batching seam on the agent side. The loop claims its slot token from run.started["robot"]["token"] and connects one scalar RobotClient, same as any single-env run. Co-authored-by: Cursor --- hud/agents/robot/__init__.py | 5 +- hud/agents/robot/agent.py | 155 +++++++++++------------------------ 2 files changed, 52 insertions(+), 108 deletions(-) diff --git a/hud/agents/robot/__init__.py b/hud/agents/robot/__init__.py index 57ccfa7b0..f89d88748 100644 --- a/hud/agents/robot/__init__.py +++ b/hud/agents/robot/__init__.py @@ -1,8 +1,9 @@ """Agent-side robot harness: drive a ``robot`` env with a VLA policy. - :class:`~.agent.RobotAgent` — the harness: connects to the ``robot`` - capability, reads the contract, drives N >= 1 env slots with one open-loop - chunk queue. Subclass and set ``self.model`` + ``self.adapter``. + 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` / diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py index 5d17b3e37..ace2d15c0 100644 --- a/hud/agents/robot/agent.py +++ b/hud/agents/robot/agent.py @@ -1,17 +1,11 @@ -"""``RobotAgent`` — the one robot harness, driving N >= 1 env slots. +"""``RobotAgent`` — drive one ``robot`` env with an open-loop chunk queue. Subclass, set ``self.model`` and ``self.adapter`` in ``__init__``, and the base -owns the rest: connect to the ``robot`` capability, read the contract, run the -open-loop chunk queue until the env terminates. A plain single env is a batch -of one — the same loop drives a vectorized env's whole ``[N, ...]`` batch over -one connection (one batched forward per refill). +owns the rest: connect to the ``robot`` capability (claiming a slot token from +``run.started`` when present), read the contract, run until the env terminates. -Two entries, one loop: - -- ``__call__(run)`` — the generic rollout contract (one run, one trace). -- ``drive(runs, client)`` — the vectorized-eval entry - (:func:`hud.eval.run.vec_rollout`): N runs sharing one env instance, spans - recorded per slot onto each run's trace. +Vectorized sims still look like one scalar connection per rollout; concurrent +rollouts coalesce GPU forwards through :class:`~.batching.BatchedModel`. Most policies use :class:`~.adapter.LeRobotAdapter`; a policy whose spaces match the env natively can set ``adapter = None`` (raw pass-through). @@ -19,7 +13,6 @@ from __future__ import annotations -import asyncio from collections import deque from typing import TYPE_CHECKING, Any, ClassVar @@ -30,7 +23,6 @@ from hud.telemetry.robot import TraceRecorder if TYPE_CHECKING: - from hud.clients.client import HudClient from hud.eval.run import Run from .adapter import ActionArray, Adapter @@ -40,19 +32,18 @@ class RobotAgent(Agent): - """Drive a ``robot`` env — single or vectorized — with one open-loop chunk queue. + """Drive a ``robot`` env with one open-loop chunk queue. **Subclass contract:** in ``__init__`` set ``self.model`` (a :class:`~.model.Model`) and ``self.adapter`` (an :class:`~.adapter.Adapter`, - or ``None`` for raw pass-through). ``model.infer`` is batch-shaped - (``[N, ...] -> [N, T, A]``), so the same subclass drives both shapes. + or ``None`` for raw pass-through). """ 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 - #: (single-env runs only). Telemetry streams regardless; see :mod:`.dataset`. + #: Opt-in: also save a LeRobot v3 dataset of every (obs, action) pair. + #: Telemetry streams regardless; see :mod:`.dataset`. save: bool = False #: Runs the policy (preprocess -> forward -> postprocess). Subclasses set this. @@ -61,26 +52,20 @@ class RobotAgent(Agent): adapter: Adapter | None = None async def __call__(self, run: Run, *, max_steps: int = 520) -> None: - """The generic rollout contract: one run, one trace.""" - await self.drive([run], run.client, max_steps=max_steps) - run.trace.status = "completed" - run.trace.content = "done" - - async def drive( - self, runs: list[Run], client: HudClient, *, max_steps: int = 520 - ) -> None: - """Drive every env slot to termination, recording onto ``runs[i]``'s trace. - - ``len(runs)`` must match the env's batch size (the wire framing tells us: - scalar ``terminated`` = 1, an ``[N]`` mask = N). - """ + """The generic rollout contract: one run, one scalar robot connection.""" if self.model is None: raise RuntimeError(f"{type(self).__name__} must set self.model in __init__") - prompt = runs[0].prompt + prompt = run.prompt if not isinstance(prompt, str): raise TypeError(f"run.prompt must be a str, got {type(prompt).__name__}: {prompt!r}") - robot = await RobotClient.connect(client.binding(self.robot_protocol)) + # Per-episode slot token from tasks.start (opaque; env put it under "robot"). + robot_info = run.started.get("robot") if isinstance(run.started.get("robot"), dict) else {} + token = robot_info.get("token") + if not isinstance(token, str): + token = None + + robot = await RobotClient.connect(run.client.binding(self.robot_protocol), token=token) try: _, obs_space = robot.spaces() if self.adapter is not None: @@ -88,109 +73,67 @@ async def drive( self.adapter.reset() obs = await robot.get_observation() - single = np.ndim(obs["terminated"]) == 0 # wire framing: scalar vs [N] mask - n = 1 if single else int(np.asarray(obs["terminated"]).shape[0]) - if len(runs) != n: - raise ValueError(f"got {len(runs)} runs for an env batch of {n}") - fps = robot.get_control_rate() - recorders = [ - # A live run (single path) records through it so steps land on - # run.trace for training; vectorized receipts emit by trace id. - TraceRecorder(run=r, fps=fps, obs_space=obs_space) - if r._client is not None # pyright: ignore[reportPrivateUsage] - else TraceRecorder(trace_id=r.trace_id, fps=fps, obs_space=obs_space) - for r in runs - ] + recorder = TraceRecorder(run=run, fps=fps, obs_space=obs_space) writer = None if self.save: - if n == 1: - from .dataset import DatasetWriter - - writer = DatasetWriter(robot.contract, fps=fps) - else: - print("[agent] save=True is single-env only; streaming telemetry", flush=True) - - print(f"[agent] episode started: {prompt!r} (n={n})", flush=True) - await self._loop( - robot, obs, prompt, recorders, writer, single=single, max_steps=max_steps - ) - for rec in recorders: - rec.close() + from .dataset import DatasetWriter + + writer = DatasetWriter(robot.contract, fps=fps) + + print(f"[agent] episode started: {prompt!r}", flush=True) + await self._loop(robot, obs, prompt, recorder, writer, max_steps=max_steps) + recorder.close() if writer is not None: writer.end_episode() finally: await robot.close() + run.trace.status = "completed" + run.trace.content = "done" async def _loop( self, robot: RobotClient, obs: dict[str, Any], prompt: str, - recorders: list[TraceRecorder], + recorder: TraceRecorder, writer: Any, *, - single: bool, max_steps: int = 520, ) -> None: - """One batched forward per refill; execute chunks open-loop per slot.""" + """Open-loop chunk queue: ainfer refills, then execute one action per tick.""" model = self.model - assert model is not None # checked in drive() + assert model is not None adapter = self.adapter - n = len(recorders) - chunks: list[deque[ActionArray]] = [deque() for _ in range(n)] - ever_done = np.zeros(n, dtype=bool) + chunk: deque[ActionArray] = deque() for step in range(max_steps): - done = np.atleast_1d(np.asarray(obs["terminated"], dtype=bool)).reshape(-1) - for c, d in zip(chunks, done, strict=True): - if d: # a reset slot re-infers for its new episode - c.clear() - ever_done |= done - if step and ever_done.all(): - print(f"[agent] all slots terminated at step {step}", flush=True) + terminated = bool(np.asarray(obs["terminated"]).reshape(-1)[0]) + if step and terminated: + print(f"[agent] terminated at step {step}", flush=True) break + if terminated: + chunk.clear() - # Batched view of the observation: single framing lifts to a batch of one. - data = obs["data"] if not single else {k: v[None] for k, v in obs["data"].items()} - for i, rec in enumerate(recorders): - if not ever_done[i]: - rec.record_observation({k: v[i] for k, v in data.items()}, tick=step) + recorder.record_observation(obs["data"], tick=step) - if any(not c for c in chunks): # refill spent slots with a fresh forward - # The adapter sees the wire framing (unbatched when single), so a - # BatchedModel can still stack samples across concurrent rollouts. + if not chunk: # refill with a fresh forward (BatchedModel coalesces ainfer) batch = adapter.adapt_observation(obs, prompt) if adapter else obs - if n == 1: - # ainfer is the coalescing point for cross-rollout batching - # (BatchedModel), so the single slot goes through it. - chunk = np.atleast_2d(await model.ainfer(batch))[None] # [1, T, A] - else: - chunk = np.asarray(await asyncio.to_thread(model.infer, batch)) - for i, c in enumerate(chunks): - if not c: - rows = chunk[i] - if adapter is not None: # e.g. deltas -> absolute vs the query obs - slot = {"data": {k: v[i] for k, v in data.items()}} - rows = adapter.adapt_chunk(rows, slot) - c.extend(rows) - if not ever_done[i]: # finished slots still act, but stop recording - recorders[i].record_inference(rows, tick=step) - - raw: list[ActionArray] = [chunks[i].popleft() for i in range(n)] + rows = np.atleast_2d(await model.ainfer(batch)) + if adapter is not None: # e.g. deltas -> absolute vs the query obs + rows = adapter.adapt_chunk(rows, obs) + chunk.extend(rows) + recorder.record_inference(rows, tick=step) + + action = chunk.popleft() if adapter is not None: # per-step execution-time hook (default identity) - raw = [ - adapter.adapt_action(a, {"data": {k: v[i] for k, v in data.items()}}) - for i, a in enumerate(raw) - ] - action = raw[0] if single else np.stack(raw) + action = adapter.adapt_action(action, obs) if writer is not None: - writer.add(obs["data"], np.asarray(raw[0]), task=prompt) + writer.add(obs["data"], np.asarray(action), task=prompt) await robot.send_action(action) if self.log_every and step % self.log_every == 0: - live = int((~ever_done).sum()) - print(f"[agent] step {step}/{max_steps} live={live}/{n}", flush=True) + print(f"[agent] step {step}/{max_steps}", flush=True) obs = await robot.get_observation() else: print(f"[agent] reached max_steps={max_steps}", flush=True) From 880b902ff6f2f03dfd1c75a48f45dee3ab8ddb21 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 01:40:33 +0000 Subject: [PATCH 23/43] refactor(eval): drop vec_rollout/num_envs; add Shared provider hud/eval loses the vectorized-specific surface: vec_rollout, its "slots"-grade contract, and Taskset.run(num_envs=) all go, leaving group as the only rollout multiplier and rollout() as the only execution atom - both closer to main than before. Shared(provider, width=N) replaces it: a refcounting Provider wrapper that provisions the substrate on the first acquire and forwards that same address to later acquires, tearing down on the last release. Pairs with group=width, max_concurrent=width so N ordinary rollouts land as N sessions on one shared substrate (e.g. a vectorized robot sim) instead of each getting a fresh one. Taskset.run validates that pairing eagerly. Co-authored-by: Cursor --- hud/eval/__init__.py | 5 ++- hud/eval/run.py | 94 +------------------------------------------- hud/eval/runtime.py | 41 +++++++++++++++++++ hud/eval/taskset.py | 55 +++++++++++--------------- 4 files changed, 69 insertions(+), 126 deletions(-) diff --git a/hud/eval/__init__.py b/hud/eval/__init__.py index e27896f80..ee53010fa 100644 --- a/hud/eval/__init__.py +++ b/hud/eval/__init__.py @@ -32,7 +32,7 @@ from .chat import Chat from .job import Job -from .run import Grade, Run, rollout, vec_rollout +from .run import Grade, Run, rollout from .runtime import ( DaytonaRuntime, DockerRuntime, @@ -46,6 +46,7 @@ RuntimeGPU, RuntimeLimits, RuntimeResources, + Shared, SubprocessRuntime, ) from .sync import SyncPlan @@ -69,11 +70,11 @@ "RuntimeGPU", "RuntimeLimits", "RuntimeResources", + "Shared", "SubprocessRuntime", "SyncPlan", "Task", "Taskset", "Trace", "rollout", - "vec_rollout", ] diff --git a/hud/eval/run.py b/hud/eval/run.py index 66140c8d5..62728b012 100644 --- a/hud/eval/run.py +++ b/hud/eval/run.py @@ -419,96 +419,4 @@ async def _drive() -> None: return run -async def vec_rollout( - task: Task, - agent: Any, - *, - runtime: Provider, - num_envs: int, - job_id: str | None = None, - group_id: str | None = None, - rollout_timeout: float | None = None, -) -> list[Run]: - """Drive one vectorized env instance to ``num_envs`` graded runs. - - The vectorized counterpart of :func:`rollout`, still domain-agnostic: one - runtime, one task start (``num_envs`` injected into the task args; an env - whose template doesn't accept it fails loudly), the agent's - ``drive(runs, client)`` entry drives every slot over one connection, and - the grade's ``"slots"`` list grades run ``i``. The runs are *receipts* — - this atom owns trace lifecycle and grading. Its task -> runs shape also - permits a future scheduler packing different compatible tasks into one - instance's slots. - """ - if not hasattr(agent, "drive"): - raise TypeError( - f"{type(agent).__name__} cannot drive a vectorized rollout " - "(needs a drive(runs, client) entry, e.g. hud.agents.robot.RobotAgent)" - ) - if job_id is None: # a lone vectorized rollout is a job of one instance - job_id = uuid.uuid4().hex - await job_enter(job_id, name=task.id, group=1) - group_id = group_id or uuid.uuid4().hex - - runs = [Run(None, task.id, task.args) for _ in range(num_envs)] - for run in runs: - run.trace.trace_id = uuid.uuid4().hex - run.job_id = job_id - run.group_id = group_id - run.slug = task.slug or task.default_slug() - - loop = asyncio.get_running_loop() - deadline = None if rollout_timeout is None else loop.time() + rollout_timeout - - async def _bounded(awaitable: Any) -> Any: - # One shared wall-clock deadline across provision, start, and the agent - # loop (see rollout._bounded for why a read-timeout is not enough). - if deadline is None: - return await awaitable - return await asyncio.wait_for(awaitable, max(deadline - loop.time(), 0.0)) - - _phase = "provisioning" - try: - async with contextlib.AsyncExitStack() as stack: - addr = cast("Runtime", await _bounded(stack.enter_async_context(runtime(task)))) - _phase = "starting task" - client = cast("HudClient", await _bounded(stack.enter_async_context(connect(addr)))) - args = {**task.args, "num_envs": num_envs} - prompt = (await _bounded(client.start_task(task.id, args))).get("prompt") - for run in runs: - run.prompt = prompt - run._runtime = addr.url - tid = run.trace.trace_id - assert tid is not None # minted above - await trace_enter(tid, job_id=job_id, group_id=group_id, model=None) - with set_trace_context(tid): # opening user step per trace - run.record(Step(source="user", messages=run.prompt_messages)) - _phase = "agent loop" - await _bounded(agent.drive(runs, client)) - _phase = "grading" - evaluation = await client.grade({"answer": None}) - slots = evaluation.get("slots") or [] - if len(slots) != num_envs: - raise ValueError( - f"grade returned {len(slots)} slots for {num_envs} envs " - "(vectorized envs must return one 'slots' entry per slot)" - ) - for run, slot in zip(runs, slots, strict=True): - run.grade = Grade.from_dict(slot) - run.trace.status = "completed" - except Exception as exc: # isolate: one bad instance never kills the batch - detail = ( - f"timed out after {rollout_timeout:.0f}s" - if isinstance(exc, TimeoutError) and rollout_timeout - else str(exc) - ) - logger.warning("vectorized rollout failed (%s): %s", _phase, detail) - for run in runs: - run.trace.status = "error" - run.record(Step(source="system", error=f"[{_phase}] {detail}")) - for run in runs: - await trace_exit(run) - return runs - - -__all__ = ["Grade", "Run", "rollout", "vec_rollout"] +__all__ = ["Grade", "Run", "rollout"] diff --git a/hud/eval/runtime.py b/hud/eval/runtime.py index a05e16ab1..2e2d0eadd 100644 --- a/hud/eval/runtime.py +++ b/hud/eval/runtime.py @@ -150,6 +150,46 @@ def __call__(self, task: Task) -> AbstractAsyncContextManager[Runtime]: return nullcontext(self) +class Shared: + """Refcounting provider: N concurrent rollouts share one substrate. + + The first ``__call__`` provisions via *inner*; later callers reuse that + address until the last exit tears it down. ``width`` is the intended + concurrent occupancy (e.g. a robot sim's ``num_envs``); when callers pass + a smaller ``max_concurrent`` at the scheduler, slots are underfilled — + that is allowed. Use with ``Taskset.run(..., group=width, max_concurrent=width)``. + """ + + def __init__(self, inner: Provider, *, width: int) -> None: + if width < 1: + raise ValueError("Shared width must be >= 1") + self.inner = inner + self.width = width + self._addr: Runtime | None = None + self._refs = 0 + self._stack: contextlib.AsyncExitStack | None = None + self._lock = asyncio.Lock() + + @asynccontextmanager + async def __call__(self, task: Task) -> Any: + async with self._lock: + if self._addr is None: + self._stack = contextlib.AsyncExitStack() + await self._stack.__aenter__() + self._addr = await self._stack.enter_async_context(self.inner(task)) + self._refs += 1 + addr = self._addr + try: + yield addr + finally: + async with self._lock: + self._refs -= 1 + if self._refs == 0 and self._stack is not None: + await self._stack.aclose() + self._stack = None + self._addr = None + + def _modal_image_from_uri(modal: Any, image_uri: str) -> Any: modal_uri_prefix = "modal://" if image_uri.startswith(modal_uri_prefix): @@ -1216,5 +1256,6 @@ async def ws_to_tcp() -> None: "RuntimeGPU", "RuntimeLimits", "RuntimeResources", + "Shared", "SubprocessRuntime", ] diff --git a/hud/eval/taskset.py b/hud/eval/taskset.py index 5251ef783..2ee7ae021 100644 --- a/hud/eval/taskset.py +++ b/hud/eval/taskset.py @@ -22,8 +22,15 @@ from hud.utils.platform import PlatformClient from .job import Job, job_enter -from .run import rollout, vec_rollout -from .runtime import HostedRuntime, HUDRuntime, LocalRuntime, _declared_env, _declared_names +from .run import rollout +from .runtime import ( + HostedRuntime, + HUDRuntime, + LocalRuntime, + Shared, + _declared_env, + _declared_names, +) from .sync import fetch_taskset_tasks, resolve_taskset_id if TYPE_CHECKING: @@ -245,7 +252,6 @@ async def run( *, runtime: Provider | HostedRuntime | None = None, group: int | None = None, - num_envs: int | None = None, max_concurrent: int | None = None, job: Job | None = None, rollout_timeout: float | None = None, @@ -269,13 +275,10 @@ async def run( one id. Returned ``job.runs`` preserves expansion order (task-major, then group). - ``num_envs`` selects vectorized execution (:func:`~hud.eval.run.vec_rollout`): - each task instance runs one *vectorized* env whose N slots (each its own - seeded perturbation) become N graded traces sharing a group_id. Distinct - from ``group`` — statistical repeats as separate instances — and they - compose: ``group=3, num_envs=5`` is 3 instances x 5 traces per task. - ``max_concurrent`` always counts instances. Requires a self-managed - runtime and a ``drive``-capable agent (e.g. ``RobotAgent``). + ``group`` is the statistical-repeat multiplier (one GRPO group_id per + task's repeats). To land those repeats on one shared substrate (e.g. a + vectorized robot sim), pass a :class:`~hud.eval.runtime.Shared` provider + with ``width`` matching ``group`` / ``max_concurrent``. ``rollout_timeout`` is a hard per-rollout wall-clock cap (seconds) for the local (Provider) path: a rollout that exceeds it is cancelled and recorded @@ -288,19 +291,14 @@ async def run( raise ValueError("group must be >= 1") if max_concurrent is not None and max_concurrent < 1: raise ValueError("max_concurrent must be >= 1") - if num_envs is not None and num_envs < 1: - raise ValueError("num_envs must be >= 1") # Tasks are pure rows, shared across rollouts. The ``group`` repeats of one - # task share a group_id (the GRPO group) — except under ``num_envs``, where - # each instance's N slot-traces are their own group. + # task share a group_id (the GRPO group). expanded: list[tuple[Task, str]] = [] task_list = list(self) for task in task_list: group_id = uuid.uuid4().hex - expanded.extend( - (task, uuid.uuid4().hex if num_envs else group_id) for _ in range(group) - ) + expanded.extend((task, group_id) for _ in range(group)) if job is None: job = Job( @@ -321,24 +319,20 @@ async def run( # an error naming the forms to pass. # An empty taskset schedules nothing, so it needs no placement. placement = runtime if runtime is not None or not task_list else self._resolve_placement() - if num_envs is not None and isinstance(placement, HostedRuntime): - raise ValueError("num_envs (vectorized rollouts) requires a self-managed runtime") + # Shared substrates need every slot occupied concurrently — a smaller cap starves the barrier. + if isinstance(placement, Shared) and ( + max_concurrent is None or max_concurrent < placement.width + ): + raise ValueError( + f"Shared(width={placement.width}) requires max_concurrent>={placement.width} " + f"(got {max_concurrent!r})" + ) sem = asyncio.Semaphore(max_concurrent) if max_concurrent else None async def _run(task: Task, group_id: str) -> list[Run]: assert placement is not None # only reached when tasks were expanded if isinstance(placement, HostedRuntime): return [await placement.run(task, agent, job_id=job_id, group_id=group_id)] - if num_envs is not None: # vectorized: one instance -> num_envs graded runs - return await vec_rollout( - task, - agent, - runtime=placement, - num_envs=num_envs, - job_id=job_id, - group_id=group_id, - rollout_timeout=rollout_timeout, - ) return [ await rollout( task, @@ -357,11 +351,10 @@ async def _one(task: Task, group_id: str) -> list[Run]: return await _run(task, group_id) logger.info( - "running %d rollouts (%d tasks x %d group)%s%s", + "running %d rollouts (%d tasks x %d group)%s", len(expanded), len(task_list), group, - f" x {num_envs} envs" if num_envs else "", f", max_concurrent={max_concurrent}" if max_concurrent else "", ) waves = await asyncio.gather(*(_one(t, gid) for t, gid in expanded)) From aef1196534ac8e99e36d091bfc062dca761c36a8 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 01:40:44 +0000 Subject: [PATCH 24/43] docs(robots): document sessions-based vectorized evals and Shared Rewrites the "Vectorized envs and evals" section around the new mechanism: num_envs as sim build config, the agent claiming its slot token from run.started, and group + Shared(provider, width=N) as the run-side knobs - instead of the deleted num_envs= / "slots" contract. Adds a Shared reference entry to runtime.mdx (table row + constructor section) so the robot docs can link instead of repeat, and updates the endpoint reset/result examples in robots.mdx and the robot-benchmark cookbook to the {prompt, token} shape. Co-authored-by: Cursor --- docs/v6/advanced/robots.mdx | 99 +++++++++++++-------------- docs/v6/cookbooks/robot-benchmark.mdx | 8 +-- docs/v6/reference/runtime.mdx | 18 +++++ 3 files changed, 71 insertions(+), 54 deletions(-) diff --git a/docs/v6/advanced/robots.mdx b/docs/v6/advanced/robots.mdx index 1f6ec5766..35cc47cf4 100644 --- a/docs/v6/advanced/robots.mdx +++ b/docs/v6/advanced/robots.mdx @@ -103,8 +103,9 @@ sim = env.gym(make_env) # make_env: any callable returning a gym-style env @env.template() async def pick_and_place(task: str = "default", seed: int = 0): - yield {"prompt": await sim.reset(task=task, seed=seed)} - yield await sim.result() + ep = await sim.reset(task=task, seed=seed) # {prompt, token} + yield {"prompt": ep["prompt"], "robot": {"token": ep["token"]}} + yield await sim.result(token=ep["token"]) ``` Task args are partitioned by the factory's signature: args the factory accepts define the env (a @@ -117,14 +118,14 @@ from hud.environment.robot import RobotBridge class MySimBridge(RobotBridge): async def reset(self, task_id: str, seed: int = 0) -> str: ... # build the episode - await self._send_observation() # push the first frame return self.task_description # becomes the task prompt def step(self, action) -> None: ... # advance one tick; set success / terminated def get_observation(self): - return {"agentview_image": frame, "state": vec}, self.terminated + # Always [N, ...] arrays + [N] terminated (N=1 is fine). + return {"agentview_image": frames, "state": vecs}, done_mask ``` @@ -199,13 +200,13 @@ async def _down(): @env.template() async def pick_and_place(task_id: str, seed: int = 0): - prompt = yield {"prompt": await endpoint.reset(task_id=task_id, seed=seed)} - yield await endpoint.result() # {"score", "success", "total_reward", "slots": [...]} + ep = await endpoint.reset(task_id=task_id, seed=seed) # {prompt, token} + yield {"prompt": ep["prompt"], "robot": {"token": ep["token"]}} + yield await endpoint.result(token=ep["token"]) # this slot's {score, success, total_reward} ``` -The result dict carries a `"slots"` list with one score dict per env slot (a single env is a batch -of one). Aggregates sit at the top; the [vectorized eval](#vectorized-envs-and-evals) path -grades each trace from its slot. +Each session grades one slot. A vectorized sim fans N concurrent sessions across its slots; see +[vectorized evals](#vectorized-envs-and-evals). ## Agent side @@ -316,47 +317,44 @@ instance rather than one you also use for unbatched runs. ## Vectorized envs and evals -A GPU sim like Isaac runs many env copies in one instance, so batching at the model isn't enough - -the *environment itself* is vectorized. Three roles stay decoupled: - -- **The environment declares what it is.** A sim factory accepting `num_envs` is the declaration; - the template passes it through (`await sim.reset(task=task, seed=seed, num_envs=num_envs)`). Any - vectorized env still runs as a plain single env when `num_envs` is 1. -- **The agent declares what it can drive.** A stock `RobotAgent` drives N >= 1 slots over one - connection with one batched forward per refill - nothing extra to implement. -- **The run selects how.** `num_envs` on `Taskset.run` turns each task into one vectorized - instance whose N slots (each its own seeded perturbation) become N graded traces sharing a - group: +A GPU sim like Isaac amortizes its cost by stepping many env copies together in one process +(`num_envs`). Getting N graded rollouts by booting N separate containers throws that away - each +one pays for its own physics from scratch. Instead, N rollouts share one vectorized instance, and +every layer keeps the plain single-env contract described above: + +- **Vectorization is sim config, not an eval parameter.** A factory that accepts `num_envs` makes + `env.gym(make_env, num_envs=8)` build an eight-way batched sim. Internally the bridge steps + every claimed slot together each tick, but each rollout still gets its own ordinary WebSocket - + one observation in, one action out, no batch dimension on the wire. +- **The agent is unchanged.** A stock `RobotAgent` drives one connection per rollout, same as a + single-env run. It claims its slot with the token the environment handed back on `tasks.start` + (`run.started["robot"]["token"]`); concurrent rollouts still coalesce GPU forwards through + [`BatchedModel`](#batching). +- **`group` picks how many rollouts; `Shared` puts them on one substrate.** Left alone, `group` + rollouts would each provision a fresh container. Wrapping the provider in + [`Shared(provider, width=N)`](/v6/reference/runtime#shared) instead provisions one and lets + all `N` reuse it: ```python +from hud.eval import Shared, DockerRuntime + job = await taskset.run( MyAgent(), - runtime=Runtime("tcp://127.0.0.1:9100"), - num_envs=5, # slots per instance -> 5 graded traces per task - max_concurrent=2, # instances in flight (each has its own runtime) + runtime=Shared(DockerRuntime("hud-isaac-env"), width=8), + group=8, + max_concurrent=8, ) ``` -Distinct from `group` - statistical repeats as separate instances - and they compose: -`group=3, num_envs=5` is 3 instances x 5 traces per task. Grading is per slot: the engine reads -the template result's `"slots"` list and grades trace `i` from `slots[i]`. - -### Run parameters - -Everything about *how* an eval executes lives on `Taskset.run` (the future `hud eval` CLI maps -onto the same parameters): - -| Parameter | What it determines | -|-----------|--------------------| -| `runtime` | *Where* each instance runs: a provider (`Runtime(url)`, `LocalRuntime`, `DockerRuntime`, ...) acquired once per instance, or `HostedRuntime` to delegate rollouts to the platform. Unset: local source if the tasks share one, else the HUD tunnel. | -| `group` | Statistical repeats: each task runs as `group` *independent instances* (own runtime, reset, connection). Without `num_envs` the repeats share a `group_id`. | -| `num_envs` | Slots per instance: each task runs as *one vectorized instance* whose N seeded perturbations become N graded traces sharing a `group_id`. Composes with `group` (instances x slots). Needs a `drive`-capable agent and a self-managed runtime. | -| `max_concurrent` | How many *instances* are in flight at once - it never counts slots. | -| `rollout_timeout` | Wall-clock cap (seconds) per instance on the local path; a breach errors that instance's run(s) without stalling the batch. | -| `job` | An open `Job` to accumulate into, so one platform job spans several `run(...)` calls (e.g. a training arc). Unset: a fresh job per call. | +Each rollout opens its own control-channel connection and claims a slot on it: the first +`tasks.start` resets the sim and claims slot 0, later starts claim whatever is free, and a start +once every slot is taken errors instead of queuing silently. Grading is the same call every +rollout already makes - the environment resolves it to that rollout's slot and returns one +`Grade`, exactly as a single-env run would. -Runs produced per task = `group` x `num_envs` (each defaulting to 1). Everything about *what* runs -stays on the task rows (task args) and the agent - the run parameters never change task content. +`group` and `max_concurrent` are the ordinary [`Taskset.run` parameters](/v6/reference/tasks#running); +set `max_concurrent` to `width` so every slot fills and none sit idle waiting for a rollout that +never starts. ## Contract @@ -511,8 +509,9 @@ async def _down(): @env.template() async def pick_and_place(task_id: str, seed: int = 0): - prompt = yield {"prompt": await endpoint.reset(task_id=task_id, seed=seed)} - yield await endpoint.result() + ep = await endpoint.reset(task_id=task_id, seed=seed) + yield {"prompt": ep["prompt"], "robot": {"token": ep["token"]}} + yield await endpoint.result(token=ep["token"]) ``` **Sim process** - your program builds the bridge and serves its control surface, blocking for the @@ -535,15 +534,15 @@ downstream (`hud eval`, tasksets, the agent) is unchanged; only *where the bridg |--------|-------|------| | `env.gym(make_env)` / `Gym` | `hud.environment.robot` | Declarative sim: contract, capability, and serving derived from a gym factory | | `hud.wrap(env)` | `hud` | One-line trace streaming for any gym-style env you drive yourself | -| `RobotBridge` / `GymBridge` | `hud.environment.robot` | Env-side serve loop (batched-first); subclass `RobotBridge` for a custom sim | -| `RobotEndpoint` | `hud.environment.robot` | Episode bookkeeping + results (local or `.remote()`); `serve_blocking()` for a split process | -| `RobotEndpoint.capability(contract=...)` | `hud.environment.robot` | Build the `openpi/0` capability after `start()` | +| `RobotBridge` / `GymBridge` | `hud.environment.robot` | Env-side serve loop (slot barrier + scalar wire); subclass `RobotBridge` for a custom sim | +| `RobotEndpoint` | `hud.environment.robot` | Episode bookkeeping + results (`reset` → `{prompt, token}`, `result(token=...)`) | +| `RobotEndpoint.capability(...)` | `hud.environment.robot` | Build the `openpi/0` capability after `start()` | | `Capability.robot(name, url, contract)` | `hud.capabilities` | Lower-level constructor (usually via `endpoint.capability`) | -| `RobotClient` | `hud.capabilities.robot` | Agent-side wire client (`spaces`, `get_observation`, `send_action`, `send_chunk`) | -| `RobotAgent` | `hud.agents.robot` | The harness: one open-loop chunk queue driving N >= 1 env slots | +| `RobotClient` | `hud.capabilities.robot` | Agent-side wire client (`connect(..., token=)`, `spaces`, `get_observation`, `send_action`) | +| `RobotAgent` | `hud.agents.robot` | The harness: one open-loop chunk queue per scalar connection | | `Model` / `LeRobotModel`, `Adapter` / `LeRobotAdapter` | `hud.agents.robot` | Policy + space-translation seams | | `BatchedAgent` / `BatchedModel` | `hud.agents.robot.batching` | Many concurrent rollouts against one shared, batched model | -| `Taskset.run(..., num_envs=N)` | `hud.eval` | Grouped eval: one vectorized instance per task, N graded traces | +| `Shared(provider, width=N)` | `hud.eval` | Refcounting provider: N concurrent rollouts share one substrate | ## See also diff --git a/docs/v6/cookbooks/robot-benchmark.mdx b/docs/v6/cookbooks/robot-benchmark.mdx index a598d2141..151e72e12 100644 --- a/docs/v6/cookbooks/robot-benchmark.mdx +++ b/docs/v6/cookbooks/robot-benchmark.mdx @@ -37,10 +37,10 @@ async def _down(): @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} - yield await endpoint.result() + ep = await endpoint.reset(task_suite="libero_spatial", + task_id=libero_task_id, init_state_id=init_state_id) + yield {"prompt": ep["prompt"], "robot": {"token": ep["token"]}} + yield await endpoint.result(token=ep["token"]) ``` The image's CMD serves it with the standard entry point (`hud serve env.py --host 0.0.0.0 --port 8765`). This env lives in HUD's `demos/` examples tree, a sibling of the `hud-python` SDK; build it from the parent directory that holds **both** `demos/` and `hud-python/` so the image can install the SDK from local source: diff --git a/docs/v6/reference/runtime.mdx b/docs/v6/reference/runtime.mdx index 102760d83..829b7f1e7 100644 --- a/docs/v6/reference/runtime.mdx +++ b/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 From 121d945cc28f761f96405ace1896866877cd7d93 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 02:37:44 +0000 Subject: [PATCH 25/43] docs(agents): tighten BatchedModel and BatchedAgent docstrings Co-authored-by: Cursor --- hud/agents/robot/batching.py | 39 ++++++++++++++---------------------- 1 file changed, 15 insertions(+), 24 deletions(-) diff --git a/hud/agents/robot/batching.py b/hud/agents/robot/batching.py index b300274bd..1e9ec40a4 100644 --- a/hud/agents/robot/batching.py +++ b/hud/agents/robot/batching.py @@ -25,19 +25,15 @@ class BatchedModel(Model): """Coalesce concurrent ``ainfer`` calls into one stacked ``inner.infer``. - A lazily-started worker waits up to ``max_wait_s`` for callers to queue, stacks - them into one ``[N, ...]`` batch, runs a single forward, and scatters the - ``[N, T, A]`` rows back to each caller. With no ``batch_size`` the batch sizes - itself to whatever is in flight — the scheduler's ``max_concurrent`` already - bounds that, so the two never need manual pairing. Pass ``batch_size`` only to - cap the forward below the live concurrency (e.g. VRAM headroom); a set cap also - flushes the window early once reached, saving the tail of ``max_wait_s``. - - ``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. + Waits up to ``max_wait_s`` for callers, stacks to ``[N, ...]``, one forward, + scatters ``[N, T, A]`` rows back. Omit ``batch_size`` to size to in-flight + callers (``max_concurrent`` already bounds that); set it only to cap below + concurrency (e.g. VRAM) - that also flushes early when full. + + ``inner`` must batch the leading ``N`` in one in-process forward + (e.g. :class:`~hud.agents.robot.model.LeRobotModel`). Not + :class:`~hud.agents.robot.model.RemoteModel` (OpenPI has no batched request; + use one agent per rollout). """ def __init__( @@ -102,18 +98,13 @@ async def _batch_loop(self) -> None: class BatchedAgent(Agent): """Drive many rollouts concurrently against one shared, batched model. - Per run: a shallow clone of ``agent`` sharing a per-run adapter copy and the - single :class:`BatchedModel`, so concurrent ``ainfer`` calls coalesce into one - forward. The adapter copy keeps per-env bindings isolated; the model is - stateless by contract, so sharing it across clones is safe. + Per run: shallow-clone ``agent`` with a per-run adapter copy and the shared + :class:`BatchedModel` (stateless by contract; adapter copy keeps env bindings + isolated). Not for :class:`~hud.agents.robot.model.RemoteModel`. - 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. + Takes ownership: wraps ``agent.model`` in place with :class:`BatchedModel`, + shared by every clone. Pass a dedicated agent - do not also use that instance + for unbatched :class:`RobotAgent` rollouts. """ def __init__( From 8b42d5e88b543493ed4606bfc7d053f0ec239b21 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 02:37:44 +0000 Subject: [PATCH 26/43] feat(agents): share one LeRobot dataset across concurrent rollouts Buffer per-rollout frames and commit whole episodes into a process-shared dataset so BatchedAgent clones write one dataset instead of per-trace shards. Co-authored-by: Cursor --- docs/v6/advanced/robots.mdx | 4 +- hud/agents/robot/dataset.py | 101 ++++++++++++++++++++---------------- 2 files changed, 58 insertions(+), 47 deletions(-) diff --git a/docs/v6/advanced/robots.mdx b/docs/v6/advanced/robots.mdx index 35cc47cf4..5ad840bca 100644 --- a/docs/v6/advanced/robots.mdx +++ b/docs/v6/advanced/robots.mdx @@ -450,8 +450,8 @@ it already produces, so it runs in *your* process - not the environment containe sims (e.g. Isaac/RoboLab) whose dependency stack conflicts with `lerobot`; only your machine needs `pip install 'lerobot[dataset]'`. -Each rollout writes its own dataset (tagged with its trace id, so a shard maps back to its -trace), finalized at process exit. Saving is single-env; a vectorized run streams telemetry only. +All rollouts in a process record into one shared dataset: each buffers its episode and commits +it whole, so concurrent runs (e.g. `BatchedAgent`) interleave safely. Finalized at process exit. Destination and Hub push come from the environment: | Env var | Effect | diff --git a/hud/agents/robot/dataset.py b/hud/agents/robot/dataset.py index 1076e2ad1..ace77220a 100644 --- a/hud/agents/robot/dataset.py +++ b/hud/agents/robot/dataset.py @@ -1,10 +1,12 @@ """Opt-in LeRobot v3 dataset writing for robot rollouts. -One :class:`DatasetWriter` spans an agent's lifetime: every episode it drives -appends ``(observation, executed action)`` frames, and the dataset is finalized -at process exit (or an explicit :meth:`finalize`), optionally pushed to the HF -Hub. The contract drives the schema with no extra wiring. Destination + push -come from the environment: +Each rollout holds a :class:`DatasetWriter` that buffers its ``(observation, +executed action)`` frames and commits whole episodes into one process-shared +dataset — so concurrent rollouts (e.g. :class:`~hud.agents.robot.batching.BatchedAgent` +clones) record into a single dataset instead of one shard each. Finalized at +process exit (or an explicit :meth:`finalize`), optionally pushed to the HF Hub. +The contract drives the schema with no extra wiring. 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``) @@ -20,12 +22,10 @@ import time import uuid from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, ClassVar import numpy as np -from hud.telemetry.context import get_current_trace_id - if TYPE_CHECKING: from numpy.typing import NDArray @@ -88,36 +88,38 @@ def _names(feature: dict[str, Any], base: str) -> list[str]: class DatasetWriter: - """Appends robot frames to a LeRobot v3 dataset; a no-op shell when lerobot - is missing (warned once) so telemetry-only runs never break.""" + """Buffers one rollout's frames; commits whole episodes to a process-shared + LeRobot v3 dataset. A no-op shell when lerobot is missing (warned once) so + telemetry-only runs never break.""" + + # One dataset per process: concurrent rollouts (e.g. BatchedAgent clones) each + # buffer their own episode but commit into the same root. + _ds: ClassVar[Any | None] = None + _root: ClassVar[Path | None] = None + _repo_id: ClassVar[str] = "" def __init__(self, contract: dict[str, Any], *, fps: int) -> None: self._contract = contract self._fps = fps self._features, self._key_map = _lerobot_features(contract) - self._ds: Any | None = None # created lazily on the first frame - self._root: Path | None = None - self._repo_id = "" + self._frames: list[dict[str, Any]] = [] # this rollout's pending episode self._enabled = importlib.util.find_spec("lerobot") is not None if not self._enabled: logger.warning( "save=True but lerobot is not installed; streaming telemetry only " "(pip install 'lerobot[dataset]')" ) - else: - atexit.register(self.finalize) # keep the dataset loadable on abrupt exits def add(self, data: dict[str, Any], action: NDArray[Any], *, task: str) -> None: """One frame: the wire observation dict + the executed env-space action.""" if not self._enabled: return - if self._ds is None: # derived contracts carry no shapes; fill from the first frame - for wire, key in self._key_map.items(): - if not self._features[key]["shape"] and wire in data: - self._features[key]["shape"] = tuple(np.shape(data[wire])) - if not self._features["action"]["shape"]: - self._features["action"]["shape"] = tuple(np.shape(action)) - ds = self._ensure_dataset() + # Derived contracts carry no shapes; fill from the first real frame (no-op after). + for wire, key in self._key_map.items(): + if not self._features[key]["shape"] and wire in data: + self._features[key]["shape"] = tuple(np.shape(data[wire])) + if not self._features["action"]["shape"]: + self._features["action"]["shape"] = tuple(np.shape(action)) row: dict[str, Any] = {} for wire, key in self._key_map.items(): value = data.get(wire) @@ -133,55 +135,64 @@ def add(self, data: dict[str, Any], action: NDArray[Any], *, task: str) -> None: act_ft = self._features["action"] row["action"] = np.asarray(action, dtype=act_ft["dtype"]).reshape(act_ft["shape"]) row["task"] = task - ds.add_frame(row) + self._frames.append(row) def end_episode(self) -> None: - """Commit the episode's pending frames.""" - if self._ds is not None and self._ds.has_pending_frames(): - self._ds.save_episode() + """Commit the buffered episode to the shared dataset in one synchronous + pass (no awaits), so concurrent rollouts' episodes never interleave.""" + if not self._frames: + return + ds = self._ensure_dataset() + for row in self._frames: + ds.add_frame(row) + ds.save_episode() + self._frames.clear() def finalize(self) -> None: - """Write the parquet footer + optionally push to the Hub. Idempotent.""" - if self._ds is None: + """Flush, write the parquet footer + optionally push to the Hub. Idempotent.""" + self.end_episode() + ds, DatasetWriter._ds = DatasetWriter._ds, None + if ds is None: return - ds, self._ds = self._ds, None ds.finalize() - print(f"[agent] saved LeRobot dataset -> {self._root}", flush=True) + print(f"[agent] saved LeRobot dataset -> {DatasetWriter._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 ds.push_to_hub(private=private) - print(f"[agent] pushed -> https://huggingface.co/datasets/{self._repo_id}", flush=True) + print( + f"[agent] pushed -> https://huggingface.co/datasets/{DatasetWriter._repo_id}", + flush=True, + ) except Exception as exc: - logger.exception("HF push failed for %s", self._repo_id) + logger.exception("HF push failed for %s", DatasetWriter._repo_id) print(f"[agent] WARNING: HF push failed: {exc!r} (dataset still on disk)", flush=True) def _ensure_dataset(self) -> Any: - if self._ds is not None: - return self._ds + if DatasetWriter._ds is not None: + return DatasetWriter._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 writer so concurrent rollouts never share a root; tied 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] + # Stamp + random tag: unique root per process even across simultaneous launches. + tag = f"{time.strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}" 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}" + DatasetWriter._root = record_dir / f"{name}_{tag}" + DatasetWriter._repo_id = f"{os.environ.get('HF_REPO') or 'hud'}/{name}_{tag}" # LeRobotDataset.create requires a fresh root; images encode to per-episode video. - self._ds = lerobot_dataset.LeRobotDataset.create( - repo_id=self._repo_id, + DatasetWriter._ds = lerobot_dataset.LeRobotDataset.create( + repo_id=DatasetWriter._repo_id, fps=self._fps, features=self._features, - root=self._root, + root=DatasetWriter._root, robot_type=self._contract.get("robot_type"), use_videos=True, ) - print(f"[agent] recording LeRobot dataset -> {self._root}", flush=True) - return self._ds + atexit.register(self.finalize) # keep the dataset loadable on abrupt exits + print(f"[agent] recording LeRobot dataset -> {DatasetWriter._root}", flush=True) + return DatasetWriter._ds __all__ = ["DatasetWriter"] From b521d96c35bea0121585fa60129542cb9ef2058e Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 02:37:44 +0000 Subject: [PATCH 27/43] docs(robot): tighten RobotClient and Environment.gym docstrings Co-authored-by: Cursor --- hud/capabilities/robot.py | 21 ++++++--------------- hud/environment/env.py | 15 ++++----------- 2 files changed, 10 insertions(+), 26 deletions(-) diff --git a/hud/capabilities/robot.py b/hud/capabilities/robot.py index 714021a3a..885ac83be 100644 --- a/hud/capabilities/robot.py +++ b/hud/capabilities/robot.py @@ -70,11 +70,9 @@ def spaces(self) -> tuple[dict[str, Any], dict[str, Any]]: @classmethod async def connect(cls, cap: Capability, *, token: str | None = None) -> Self: """Dial the robot WebSocket; ``token`` claims a sim slot after the metadata frame.""" - # No keepalive: heavy sims (Isaac resets) legitimately stall for - # minutes between frames; a ping timeout would kill a healthy rollout. ws = await websockets.connect(cap.url, max_size=None, ping_interval=None) - # Consume the connect-time metadata frame (always first); a string frame - # is the env's error convention. + # Consume initial metadata; string means env error. + raw = await ws.recv() if isinstance(raw, str): raise RuntimeError(f"robot env error on connect:\n{raw}") @@ -86,17 +84,10 @@ async def connect(cls, cap: Capability, *, token: str | None = None) -> Self: async def get_observation(self) -> dict[str, Any]: """Await the latest observation: ``{"data": {name: ndarray}, "terminated": bool}``. - On the wire the env sends an openpi-style *flat* dict (``{name: ndarray, ...}``) - with ``terminated`` (and, for realtime bridges, ``meta``) as sibling keys; we - regroup the array fields under ``"data"`` for the agent harness. Arrays — nested - anywhere, including inside ``meta`` (e.g. ``unexecuted_chunk``) — are already - decoded by the codec. - - Realtime (free-running) bridges attach a ``"meta"`` block carrying the realtime - control state used for async/RTC inference (``obs_index``, ``queue_remaining``, - ``delay``, ``unexecuted_chunk``); sync bridges omit it. - - Raises if the env reported an error (a string traceback frame). + Wire format is a flat openpi dict (``{name: ndarray, ...}``) with ``terminated`` + as a sibling; array fields are regrouped under ``"data"``. Realtime bridges also + attach ``"meta"`` (``obs_index``, ``queue_remaining``, ``delay``, + ``unexecuted_chunk``); sync bridges omit it. Raises on an env error frame. """ msg = await self._queue.get() if "error" in msg: diff --git a/hud/environment/env.py b/hud/environment/env.py index 8c77a5bf6..9ba6362bc 100644 --- a/hud/environment/env.py +++ b/hud/environment/env.py @@ -308,17 +308,10 @@ async def _down() -> None: def gym(self, target: Any, *, name: str = "robot", **kwargs: Any) -> Any: """Attach a gym-style sim serving ``name`` over the ``robot`` protocol. - ``target`` declares the env: a module-level factory (the same - ``make_env`` an EnvHub repo exposes), a gymnasium registry id - (``"CartPole-v1"``), or a constructed registry env — reduced here to - its spec and closed, since only the declaration crosses the fork. The - sim runs in its own process (spawned at serve time; see - :mod:`hud.environment.robot.bridge`); this registers the spawn → - publish → teardown lifecycle on this env's hooks — nothing runs until - the env serves. Extra kwargs go to the sim program (``fps=``, - ``contract=``). Returns the - :class:`~hud.environment.robot.RobotEndpoint` templates drive episodes - through (``sim.reset`` / ``sim.result``). + ``target`` is a factory, gymnasium id (``"CartPole-v1"``), or constructed + registry env (reduced to its spec). Registers spawn → publish → teardown + on this env's hooks; nothing runs until serve. Returns a + :class:`~hud.environment.robot.RobotEndpoint` (``sim.reset`` / ``sim.result``). """ from hud.environment.robot import RobotEndpoint, gym_command From 42f6e28ee8e8d5e9a188de6efeee056f9bd7bd51 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 02:37:44 +0000 Subject: [PATCH 28/43] refactor(environment/robot): fold SimThread into bridge; move GymBridge to gym Dissolve the separate SimThread/run_with_sim surface into RobotBridge._run_on_sim and serve_bridge, and relocate GymBridge/gym_command/main next to the gym introspection helpers so each module has one job. Co-authored-by: Cursor --- hud/environment/robot/__init__.py | 13 +- hud/environment/robot/bridge.py | 446 +++++++----------------------- hud/environment/robot/gym.py | 392 +++++++++++++++++++++++--- hud/environment/robot/sim.py | 144 ---------- hud/environment/server.py | 2 +- 5 files changed, 457 insertions(+), 540 deletions(-) delete mode 100644 hud/environment/robot/sim.py diff --git a/hud/environment/robot/__init__.py b/hud/environment/robot/__init__.py index c3ecf5ff0..975a58088 100644 --- a/hud/environment/robot/__init__.py +++ b/hud/environment/robot/__init__.py @@ -1,14 +1,14 @@ """Env-side robot runtime: bridges, the sim program, and the control endpoint. A simulator always runs in its own process — the **sim program** -(``python -m hud.environment.robot.bridge``), where the sim owns the main +(``python -m hud.environment.robot.gym``), where the sim owns the main thread and a bridge serves the wire. The env server holds a :class:`~.endpoint.RobotEndpoint` on it: - :class:`~.bridge.RobotBridge` — bridge base (``num_envs`` slots in lockstep internally; scalar openpi wire per claimed connection): the agent's ``robot`` WebSocket plus the JSON-RPC control side channel. -- :class:`~.bridge.GymBridge` / :func:`~.bridge.gym_command` — the generic +- :class:`~.gym.GymBridge` / :func:`~.gym.gym_command` — the generic gym path (``env.gym(...)`` — a factory, registry id, or declared env): contract derivation, capability, episode control. - :class:`~.endpoint.RobotEndpoint` — the env server's handle: spawn (owned) @@ -17,8 +17,6 @@ bridges call it last). - :func:`hud.wrap` (:mod:`~.gym`) — one-line trace streaming for any gym env you drive yourself, plus the shared gym introspection. -- :class:`~.sim.SimThread` — the sim-process shape: the sim owns the main - thread, serving runs on a background loop thread. The agent-side counterpart, :class:`~hud.capabilities.robot.RobotClient`, lives under :mod:`hud.capabilities`; both ends share the wire codec defined there. @@ -26,19 +24,16 @@ from __future__ import annotations -from .bridge import GymBridge, RobotBridge, gym_command, serve_bridge +from .bridge import RobotBridge, serve_bridge from .endpoint import RobotEndpoint -from .gym import TracedEnv, wrap -from .sim import SimThread, run_with_sim +from .gym import GymBridge, TracedEnv, gym_command, wrap __all__ = [ "GymBridge", "RobotBridge", "RobotEndpoint", - "SimThread", "TracedEnv", "gym_command", - "run_with_sim", "serve_bridge", "wrap", ] diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py index e3141c929..8974cafc4 100644 --- a/hud/environment/robot/bridge.py +++ b/hud/environment/robot/bridge.py @@ -1,4 +1,4 @@ -"""Sim-process ``robot`` bridges — and the sim program that serves them. +"""The sim-process ``robot`` bridge base — and the sim program shape. The *server* side of the ``robot`` protocol (agent-side client: :class:`~hud.capabilities.robot.RobotClient`); both share the wire codec defined @@ -13,33 +13,32 @@ - :class:`RobotBridge` — subclass with your sim: ``reset`` / ``step`` / ``get_observation``, and set ``self.contract``. The base owns both channels. -- :class:`GymBridge` — the generic bridge over any gym-style target (factory, - registry id, or a declared env's spec); users get one via ``env.gym(...)`` - and never subclass it. - -This module is also **the sim program** — the one process shape every sim runs -(the env server never hosts one, it spawns or attaches through an endpoint): - -- ``python -m hud.environment.robot.bridge `` — what ``env.gym()`` - spawns; :func:`gym_command` builds the argv, so both ends of the format live - here. -- :func:`serve_bridge` — the blocking entry a custom sim program calls last. + (The generic gym-style bridge, :class:`~.gym.GymBridge`, lives in :mod:`~.gym`.) +- :func:`serve_bridge` — the blocking entry every sim program calls last. Inside the sim process the sim owns the main thread and serving runs on a -background loop; every sim touch routes through the shared -:class:`~.sim.SimThread`, so bridges stay thread-naive (see :mod:`~.sim`). +background loop thread; every sim touch is queued to main through +:meth:`RobotBridge._run_on_sim`. One shape because the hardest sims demand it: +sims are usually *thread-affine* (every touch must run on the thread that owns +their GL/device context), and Isaac/Omniverse must own the process main thread +outright — Kit drives its own main-thread loop and ``env.reset()`` nests +``run_until_complete``. Cheap CPU envs pay ~nothing: main just blocks on the +queue. """ from __future__ import annotations import asyncio import contextlib -import inspect +import queue import secrets +import signal import sys +import threading from abc import ABC, abstractmethod +from collections.abc import Callable +from concurrent.futures import Future from dataclasses import dataclass, field -from pathlib import Path from typing import Any import numpy as np @@ -50,17 +49,6 @@ # ends of the protocol stay in lockstep (env -> capabilities is the correct direction). from hud.capabilities.robot import _packb, _unpackb from hud.environment.utils import error, read_frame, reply, send_frame -from hud.telemetry.robot import to_numpy - -from .gym import ( - action_dim_of, - detect_fps, - load_or_write_contract, - num_envs_of, - probe_success, - split_observation, -) -from .sim import SimThread, run_with_sim #: Line the sim program prints once its control channel is bound; the spawning #: RobotEndpoint reads it from this process's stdout. @@ -151,8 +139,9 @@ def __init__(self, *, host: str = "127.0.0.1", port: int = 0) -> None: #: The env's self-describing wire contract (features, control_rate); the #: endpoint fetches it over the control channel to mint the capability. self.contract: dict[str, Any] = {} - # Every sim touch runs on the process sim thread (see sim.py). - self._sim = SimThread.shared() + # Sim touches queue to the thread that owns the simulator (see _run_on_sim). + self._sim_q: queue.Queue[tuple[Callable[[], Any], Future[Any]]] = queue.Queue() + self._sim_ident: int | None = None self.num_envs: int = 1 self._registry = _SlotRegistry() self._registry.configure(1) @@ -165,6 +154,18 @@ def __init__(self, *, host: str = "127.0.0.1", port: int = 0) -> None: self.success: bool = False self.terminated: bool = False + async def _run_on_sim(self, fn: Callable[..., Any], *args: Any) -> Any: + """Run ``fn(*args)`` on the sim thread, awaited on the caller's loop. + + Before :func:`serve_bridge` binds a sim thread (tests, in-loop use) calls + execute inline on the caller — the degenerate single-thread case. + """ + if self._sim_ident is None or threading.get_ident() == self._sim_ident: + return fn(*args) # not serving yet, or already on the sim thread + fut: Future[Any] = Future() + self._sim_q.put((lambda: fn(*args), fut)) + return await asyncio.wrap_future(fut) + async def _claim_episode(self, **kwargs: Any) -> dict[str, Any]: """Control-plane reset: global sim reset when all free, else claim a free slot.""" if self._registry.all_free: @@ -304,6 +305,10 @@ async def _handle_client(self, ws: Any) -> None: slot.action = None self._action_event.set() # wake the barrier so it doesn't wait on us + def _live_slots(self) -> list[_Slot]: + """Claimed slots with a live connection — the barrier's participants.""" + return [s for s in self._registry.claimed() if s.ws is not None] + async def _tick_loop(self) -> None: """Gather claimed slots' actions (or timeout → hold), step once, fan out.""" while True: @@ -312,7 +317,7 @@ async def _tick_loop(self) -> None: except TimeoutError: pass self._action_event.clear() - claimed = [s for s in self._registry.claimed() if s.ws is not None] + claimed = self._live_slots() if not claimed: continue # Wait until every live claimed slot has an action, or step_timeout elapses. @@ -329,7 +334,7 @@ async def _tick_loop(self) -> None: self._action_event.clear() with contextlib.suppress(TimeoutError): await asyncio.wait_for(self._action_event.wait(), timeout=remaining) - claimed = [s for s in self._registry.claimed() if s.ws is not None] + claimed = self._live_slots() if not claimed: break if not claimed: @@ -338,13 +343,11 @@ async def _tick_loop(self) -> None: # Stack one row per sim slot (idle/unconnected claimed → hold; free → hold). rows = [] for s in self._registry.slots: - if s in claimed and s.action is not None: - rows.append(np.asarray(s.action, dtype=np.float32).reshape(-1)) - else: - rows.append(np.asarray(hold, dtype=np.float32).reshape(-1)) + row = s.action if s in claimed and s.action is not None else hold + rows.append(np.asarray(row, dtype=np.float32).reshape(-1)) s.action = None actions = np.stack(rows) - await self._sim.call(self.step, actions) + await self._run_on_sim(self.step, actions) for s in claimed: await self._send_slot_observation(s) @@ -352,7 +355,7 @@ async def _send_slot_observation(self, slot: _Slot) -> None: """Fan one scalar obs frame to a claimed connection.""" if slot.ws is None: return - out = await self._sim.call(self.get_observation) + out = await self._run_on_sim(self.get_observation) if out is None: return data, terminated = out @@ -401,337 +404,76 @@ async def _dispatch_control(self, method: str, params: dict[str, Any]) -> dict[s raise ValueError(f"unknown method {method!r}") -class GymBridge(RobotBridge): - """Serve any gym-style env over the ``robot`` protocol, generically. - - ``target`` is how this process builds the env: a factory callable, or a - string — a gymnasium registry id, or a declared env's spec as JSON - (:func:`gym_command` reduces an env instance to the latter). - - Task args are partitioned into env-defining **build args** (a change - rebuilds the env) and episodic args flowing to ``env.reset(seed=..., - options=...)``. For a factory, its signature is the partition — so - ``num_envs`` in the signature is the vectorization declaration; for a - registry target, ``num_envs`` is the one build arg (``gym.make_vec``). - Default build kwargs (e.g. ``num_envs=8`` from ``env.gym(..., num_envs=8)``) - fill in when the task does not pass them. - - ``fps`` overrides the detected control rate; ``contract`` is the - ``contract.json`` round-trip path (load beats derive, so user edits stick; - ``None`` disables the file). - """ - - def __init__( - self, - target: Any, - *, - fps: int | None = None, - contract: str | Path | None = "contract.json", - **defaults: Any, - ) -> None: - host = defaults.pop("host", "127.0.0.1") - port = int(defaults.pop("port", 0)) - super().__init__(host=host, port=port) - self._target = target - self._fps = fps - self._contract_path = contract - self._defaults = defaults # e.g. num_envs from env.gym(..., num_envs=8) - # Task args that define the env build (everything else is episodic). - if callable(target): - self._build_params = { - n - for n, p in inspect.signature(target).parameters.items() - if p.kind in (p.POSITIONAL_OR_KEYWORD, p.KEYWORD_ONLY) - } - else: - self._build_params = {"num_envs"} # registry targets: vectorization only - self.env: Any = None - self.batched = False # env carries a leading [N] dim (any env exposing num_envs) - self._obs: Any = None # latest observation (reset or step) - self._instance: Any = None # current env-defining args; a mismatch rebuilds - self._is_torch = False - # Per-slot episode scoring, sticky until the next reset. - self._done: np.ndarray = np.zeros(1, dtype=bool) - self._success: np.ndarray = np.zeros(1, dtype=bool) - self._acc_reward: np.ndarray = np.zeros(1) - self._step_reward: np.ndarray = np.zeros(1) # last env.step reward (RL wire sibling) - self._seen_success = False # any env-reported success signal this episode - - # ── env lifecycle (all sim touches on the sim thread) ─────────────────────── - - async def ensure_env(self, **task_args: Any) -> None: - """Build the env (factory defaults unless task args say otherwise). Idempotent.""" - if self.env is None: - await self._sim.call(self._sync_reset, task_args) - - async def start(self) -> None: - """Build the env and derive/load the contract, then bring up the wire.""" - await self.ensure_env() - state, frames = self.sample_observation() - existed = self._contract_path is not None and Path(self._contract_path).exists() - self.contract = load_or_write_contract( - self._contract_path, - state, - frames, - action_dim_of(self.env, batched=self.batched), - self._fps or detect_fps(self.env), - ) - if not existed and self._contract_path is not None: - print(f"[env] wrote {self._contract_path} (edit names to relabel plots)", flush=True) - await super().start() - - async def reset(self, **task_args: Any) -> str: - return await self._sim.call(self._sync_reset, task_args) - - async def stop(self) -> None: - await super().stop() - if self.env is not None: - await self._sim.call(self.env.close) - self.env = None - - def _sync_reset(self, task_args: dict[str, Any]) -> str: - # Defaults from env.gym(..., num_envs=N) fill in when the task omits them. - merged = {**self._defaults, **task_args} - build = {k: v for k, v in merged.items() if k in self._build_params} - episodic = {k: v for k, v in merged.items() if k not in self._build_params} - key = tuple(sorted(build.items())) - if self.env is None or key != self._instance: - if self.env is not None: - self.env.close() - self.env = self._build_env(build) - self._instance = key - n = num_envs_of(self.env) - self.batched = n is not None - self.num_envs = n or 1 - seed = episodic.pop("seed", None) - obs, _ = self.env.reset(seed=seed, options=episodic or None) - self._obs = obs # the first frame an agent sees on connect/reset - self._is_torch = "torch" in type(_first_leaf(obs)).__module__ - self._done = np.zeros(self.num_envs, dtype=bool) - self._success = np.zeros(self.num_envs, dtype=bool) - self._acc_reward = np.zeros(self.num_envs) - self._step_reward = np.zeros(self.num_envs, dtype=np.float32) - self._seen_success = False - return self._prompt(merged) - - def _build_env(self, build: dict[str, Any]) -> Any: - """Build the env from the target: factory call, or registry make/make_vec.""" - if callable(self._target): - return self._target(**build) - import gymnasium - from gymnasium.envs.registration import EnvSpec - - env_id, kwargs = self._target, dict(build) - if env_id.startswith("{"): # a declared env's spec, re-made in this process - spec = EnvSpec.from_json(env_id) - env_id, kwargs = spec.id, {**spec.kwargs, **kwargs} - if "num_envs" in kwargs: - return gymnasium.make_vec(env_id, **kwargs) - return gymnasium.make(env_id, **kwargs) - - def _prompt(self, task_args: dict[str, Any]) -> str: - base = getattr(self.env, "unwrapped", self.env) - for attr in ("task_description", "instruction"): - text = getattr(getattr(base, "cfg", None), attr, None) or getattr(base, attr, None) - if isinstance(text, str) and text: - return text - return ", ".join(f"{k}={v}" for k, v in sorted(task_args.items())) or "run the task" - - def sample_observation(self) -> tuple[dict[str, Any], dict[str, Any]]: - """One per-env ``(state, frames)`` sample for contract derivation (post-build).""" - state, frames = split_observation(self._obs, batched=self.batched) - if self.batched: - state = {k: to_numpy(v)[0] for k, v in state.items()} - frames = {k: to_numpy(v)[0] for k, v in frames.items()} - return state, frames - - # ── bridge protocol ────────────────────────────────────────────────────────── - - def step(self, action: np.ndarray) -> None: - act: Any = np.array(action, dtype=np.float32) # wire buffer is read-only - if not self.batched: - act = act[0] if act.ndim > 1 else act # single plain env: drop the batch dim - elif act.ndim == 1: - act = act[None] - # The wire carries floats; discrete/int action spaces need their dtype + shape back. - space = getattr(self.env, "action_space", None) - dtype = getattr(space, "dtype", None) - if dtype is not None and np.issubdtype(dtype, np.integer): - act = act.astype(dtype).reshape(getattr(space, "shape", act.shape) or ()) - if self._is_torch: - import torch - - base = getattr(self.env, "unwrapped", self.env) - act = torch.as_tensor(act, device=getattr(base, "device", None)) - obs, reward, terminated, truncated, info = self.env.step(act) - self._obs = obs - done = np.atleast_1d(to_numpy(terminated)).astype(bool) | np.atleast_1d( - to_numpy(truncated) - ).astype(bool) - # Per-step reward for the RL wire (get_observation attaches it as a sibling). - self._step_reward = np.atleast_1d(to_numpy(reward)).astype(np.float32) - self._acc_reward += self._step_reward * ~self._done - newly = done & ~self._done - if newly.any(): - success = self._resolve_success(info) - if success is not None: - self._seen_success = True - self._success |= success & newly - self._done |= done - - def _resolve_success(self, info: Any) -> np.ndarray | None: - """Env-reported success at a done step: info keys, else Isaac's termination term.""" - found = probe_success(info, num_envs=self.num_envs) - if found is not None: - return found - base = getattr(self.env, "unwrapped", self.env) - manager = getattr(base, "termination_manager", None) - if manager is not None: - try: - return np.atleast_1d(to_numpy(manager.get_term("success"))).astype(bool) - except Exception: - return None - return None - - def get_observation(self) -> tuple[dict[str, np.ndarray], np.ndarray] | None: - """Always ``[N, ...]`` arrays + ``[N]`` terminated for the barrier fan-out.""" - if self.env is None or self._obs is None: - return None - state, frames = split_observation(self._obs, batched=self.batched) - data = {k: to_numpy(v) for k, v in {**state, **frames}.items()} - # Last env.step reward rides along for RL collection (client lifts it - # out of "data" to a top-level sibling, like "terminated"). - data["reward"] = self._step_reward - if not self.batched: - # Plain single env: lift scalars to a batch of one for uniform slicing. - data = {k: (v if getattr(v, "ndim", 0) >= 1 and v.shape[:1] == (1,) else np.asarray(v)[None]) - for k, v in data.items()} - return data, np.asarray(self._done, dtype=bool).reshape(self.num_envs) - - def result_slots(self) -> list[dict[str, Any]]: - """Per-slot grades: env-reported success when available, else accumulated reward.""" - # The env's own success check outranks accumulated shaped reward. - scores = self._success if self._seen_success else self._acc_reward - return [ - { - "score": float(scores[i]), - "success": bool(self._success[i]), - "total_reward": float(self._acc_reward[i]), - } - for i in range(self.num_envs) - ] - - -def _first_leaf(obs: Any) -> Any: - while isinstance(obs, dict): - obs = next(iter(obs.values())) - return obs - - -# ── the sim program (what env.gym() spawns; custom sims call serve_bridge) ──── +# ── the sim program shape (custom sims call serve_bridge last) ──────────────── def serve_bridge(bridge: RobotBridge, *, host: str = "127.0.0.1", port: int = 0) -> None: """Serve *bridge*, blocking for the process's lifetime — the sim program's last call. - Brings up the robot WebSocket and the control side channel (announcing its - port on stdout as ``HUD_SIM_PORT=``), with the sim owning the main - thread. SIGTERM / Ctrl-C tear the bridge down on the way out. + Serving — the robot WebSocket and the control side channel (port announced + on stdout as ``HUD_SIM_PORT=``) — runs on a background loop thread; + the calling (main) thread becomes the sim thread, executing every touch the + bridge queues (see :meth:`RobotBridge._run_on_sim`). SIGTERM / Ctrl-C cancel + serving; the sim keeps draining through teardown (``stop()`` touches it too). """ + # Claim main as the sim thread before serving starts, so no touch slips + # through inline on the wrong thread. + bridge._sim_ident = threading.get_ident() + loop = asyncio.new_event_loop() + done = threading.Event() + serve_task: asyncio.Task[Any] | None = None async def _serve() -> None: + nonlocal serve_task + serve_task = asyncio.current_task() await bridge.start() # WS bound first, so `url` is concrete before any control call server = await bridge.serve_control(host, port) - bound = server.sockets[0].getsockname()[1] - print(f"{PORT_ANNOUNCEMENT}{bound}", flush=True) + print(f"{PORT_ANNOUNCEMENT}{server.sockets[0].getsockname()[1]}", flush=True) try: await asyncio.Event().wait() # until SIGTERM / Ctrl-C cancels finally: server.close() await bridge.stop() - run_with_sim(_serve) - - -def gym_command( - target: Any, - *, - fps: int | None = None, - contract: str | None = "contract.json", - **defaults: Any, -) -> list[str]: - """Build the command line that spawns *target* as a sim process (what ``env.gym`` runs). - - ``env.py`` only declares the sim, it never runs it — this builds the argv - for the child process that will. A live env object can't cross that - boundary, only text can, so *target* is reduced to something the child can - rebuild itself: a factory's source path, a registry id as-is, or a - constructed env's spec (id + kwargs) — closed here since only the spec - crosses over; the child calls ``gym.make`` again for its own instance. - ``fps`` / ``contract`` / build defaults (e.g. ``num_envs=``) flow through - to the :class:`GymBridge` the CLI builds. - """ - if callable(target): - name = getattr(target, "__qualname__", "") - if not name or "." in name or "<" in name: - raise ValueError(f"env.gym factory must be a module-level callable, got {target!r}") - target = f"{inspect.getfile(target)}:{name}" - elif not isinstance(target, str): - spec = getattr(target, "spec", None) # registry-made envs carry their EnvSpec - if spec is None: - raise ValueError( - f"env.gym: {target!r} has no registry spec; pass a gymnasium id " - "or a module-level factory instead" - ) - target_env, target = target, spec.to_json() - target_env.close() # only the spec crosses; free the declaration-time instance - cmd = [sys.executable, "-m", "hud.environment.robot.bridge", target] - if fps is not None: - cmd += ["--fps", str(fps)] - if contract is not None: - cmd += ["--contract", str(contract)] - # Build defaults (num_envs, etc.) as --key value pairs the CLI re-applies. - for key, value in defaults.items(): - cmd += [f"--{key.replace('_', '-')}", str(value)] - return cmd - - -def main() -> None: - import argparse - - from hud.utils.modules import load_module - - parser = argparse.ArgumentParser(description="Serve a gym-style env as a robot sim.") - parser.add_argument("target", help="'path/to/module.py:factory', a registry id, or spec JSON.") - parser.add_argument("--fps", type=int, default=None, help="Control-rate override.") - parser.add_argument("--contract", default=None, help="contract.json round-trip path.") - parser.add_argument("--host", default="127.0.0.1", help="Control-channel interface.") - parser.add_argument("--port", type=int, default=0, help="Control-channel port (0 = ephemeral).") - parser.add_argument("--num-envs", type=int, default=None, help="Vectorized slot count.") - args, unknown = parser.parse_known_args() - - target: Any = args.target - if ".py:" in target: # factory by source path; ids / spec JSON pass through - path, _, attr = target.rpartition(":") - target = getattr(load_module(path), attr) - defaults: dict[str, Any] = {} - if args.num_envs is not None: - defaults["num_envs"] = args.num_envs - # Accept further --key value pairs as build defaults. - i = 0 - while i < len(unknown): - flag = unknown[i] - if flag.startswith("--") and i + 1 < len(unknown): - defaults[flag[2:].replace("-", "_")] = unknown[i + 1] - i += 2 - else: - i += 1 - bridge = GymBridge(target, fps=args.fps, contract=args.contract, **defaults) - serve_bridge(bridge, host=args.host, port=args.port) - + def _thread() -> None: + asyncio.set_event_loop(loop) + try: + with contextlib.suppress(asyncio.CancelledError): + loop.run_until_complete(_serve()) + finally: + done.set() -if __name__ == "__main__": - main() + def _cancel(*_: object) -> None: + if serve_task is not None: + loop.call_soon_threadsafe(serve_task.cancel) + with contextlib.suppress(ValueError): # signals are main-thread-only; tests may not be + signal.signal(signal.SIGTERM, _cancel) + thread = threading.Thread(target=_thread, name="hud-serve", daemon=True) + thread.start() -__all__ = ["PORT_ANNOUNCEMENT", "GymBridge", "RobotBridge", "gym_command", "serve_bridge"] + # The sim loop: execute queued touches until serving exits, pumping Kit when + # Omniverse is loaded (it needs continuous main-thread updates). + while not done.is_set(): + try: + # Re-checked every pass: Kit may finish loading after serving starts. + kit = sys.modules.get("omni.kit.app") + touches: list[tuple[Callable[[], Any], Future[Any]]] = [] + with contextlib.suppress(queue.Empty): + touches.append(bridge._sim_q.get(timeout=0.002 if kit else 0.05)) + while True: # drain the backlog behind the first touch + touches.append(bridge._sim_q.get_nowait()) + for fn, fut in touches: + if fut.set_running_or_notify_cancel(): + try: + fut.set_result(fn()) + except BaseException as exc: # propagate to the awaiting caller + fut.set_exception(exc) + if kit: + kit.get_app().update() + except KeyboardInterrupt: # Ctrl-C: cancel serving, keep draining through teardown + _cancel() + thread.join(timeout=30) + + +__all__ = ["PORT_ANNOUNCEMENT", "RobotBridge", "serve_bridge"] diff --git a/hud/environment/robot/gym.py b/hud/environment/robot/gym.py index 47f67d3df..f5b701d66 100644 --- a/hud/environment/robot/gym.py +++ b/hud/environment/robot/gym.py @@ -1,10 +1,14 @@ -"""Gym-style env integration: introspection + ``hud.wrap`` trace streaming. +"""Gym-style env integration: introspection, the served path, and ``hud.wrap``. Two consumers share the introspection here (split observations, probe success, derive a minimal contract) and stay in lockstep because of it: -- :class:`~.bridge.GymBridge` — the served path (``env.gym(...)``). -- :func:`hud.wrap` / :class:`TracedEnv` — the loop-owning path: wrap any +- :class:`GymBridge` - the served path (``env.gym(...)``): a generic + :class:`~.bridge.RobotBridge` over any gym-style target, spawned as + ``python -m hud.environment.robot.gym ``. This module is that sim + program; :func:`gym_command` builds the argv, so both ends of the format + live here. +- :func:`hud.wrap` / :class:`TracedEnv` - the loop-owning path: wrap any ``gym.Env``, ``gym.vector.VectorEnv``, or batched-tensor Isaac env you drive yourself and every episode streams to the platform as a trace (numeric state, per-camera H.264 video, actions, reward/success) under one job:: @@ -20,8 +24,10 @@ from __future__ import annotations import atexit +import inspect import json import logging +import sys from pathlib import Path from typing import Any, Self @@ -29,6 +35,8 @@ from hud.telemetry.robot import JobRecorder, to_numpy +from .bridge import RobotBridge, serve_bridge + logger = logging.getLogger(__name__) #: Conventional info keys carrying env-reported success, probed in order. @@ -41,16 +49,17 @@ def num_envs_of(env: Any) -> int | None: """``num_envs`` from the env or its unwrapped core (gymnasium 1.x wrappers no longer forward attributes, so an Isaac env inside ``gym.make`` hides it).""" - n = getattr(env, "num_envs", None) - if n is None: - n = getattr(getattr(env, "unwrapped", env), "num_envs", None) - return None if n is None else int(n) + for obj in (env, getattr(env, "unwrapped", env)): + n = getattr(obj, "num_envs", None) + if n is not None: + return int(n) + return None def is_batched(env: Any) -> bool: """Vectorized (gym VectorEnv) or batched-tensor (Isaac) envs carry a leading [N] dim. - Any env exposing ``num_envs`` is batched — Isaac keeps batch semantics even at + Any env exposing ``num_envs`` is batched - Isaac keeps batch semantics even at ``num_envs == 1``; plain ``gym.Env``s don't have the attribute at all. """ if num_envs_of(env) is not None: @@ -79,7 +88,7 @@ def split_observation(obs: Any, *, batched: bool = False) -> tuple[dict[str, Any A camera frame is a channel-last image (rank 3, or rank 4 when batched); state is any flat numeric vector. Anything else is dropped. Batched arrays keep their - leading ``[N]`` dim — the recorder slices per slot. + leading ``[N]`` dim - the recorder slices per slot. """ state: dict[str, Any] = {} frames: dict[str, Any] = {} @@ -113,7 +122,7 @@ def detect_fps(env: Any) -> int: def capture_task_params(kwargs: dict[str, Any]) -> dict[str, Any]: - """Reset parametrization as json-safe trace metadata — the full kwargs, never a label.""" + """Reset parametrization as json-safe trace metadata - the full kwargs, never a label.""" def safe(v: Any) -> Any: if v is None or isinstance(v, (str, int, float, bool)): @@ -146,7 +155,7 @@ def load_or_write_contract( """The contract round-trip: load the user-edited file if present, else derive a minimal contract from one (per-env) sample observation and write it once. - Derived: just enough to structure the spaces and label plots — cameras as rgb + Derived: just enough to structure the spaces and label plots - cameras as rgb observations, state vectors with positional names, one action feature. Users edit the written file to rename dimensions; nothing else is derived on purpose. """ @@ -167,6 +176,239 @@ def load_or_write_contract( return contract +# ── GymBridge: the served path (what env.gym() spawns) ──────────────────────── + + +class GymBridge(RobotBridge): + """Serve any gym-style env over the ``robot`` protocol, generically. + + ``target`` is how this process builds the env: a factory callable, or a + string — a gymnasium registry id, or a declared env's spec as JSON + (:func:`gym_command` reduces an env instance to the latter). + + Task args are partitioned into env-defining **build args** (a change + rebuilds the env) and episodic args flowing to ``env.reset(seed=..., + options=...)``. For a factory, its signature is the partition — so + ``num_envs`` in the signature is the vectorization declaration; for a + registry target, ``num_envs`` is the one build arg (``gym.make_vec``). + Default build kwargs (e.g. ``num_envs=8`` from ``env.gym(..., num_envs=8)``) + fill in when the task does not pass them. + + ``fps`` overrides the detected control rate; ``contract`` is the + ``contract.json`` round-trip path (load beats derive, so user edits stick; + ``None`` disables the file). + """ + + def __init__( + self, + target: Any, + *, + fps: int | None = None, + contract: str | Path | None = "contract.json", + **defaults: Any, + ) -> None: + host = defaults.pop("host", "127.0.0.1") + port = int(defaults.pop("port", 0)) + super().__init__(host=host, port=port) + self._target = target + self._fps = fps + self._contract_path = contract + self._defaults = defaults # e.g. num_envs from env.gym(..., num_envs=8) + # Task args that define the env build (everything else is episodic). + if callable(target): + self._build_params = { + n + for n, p in inspect.signature(target).parameters.items() + if p.kind in (p.POSITIONAL_OR_KEYWORD, p.KEYWORD_ONLY) + } + else: + self._build_params = {"num_envs"} # registry targets: vectorization only + self.env: Any = None + self.batched = False # env carries a leading [N] dim (any env exposing num_envs) + self._obs: Any = None # latest observation (reset or step) + self._instance: Any = None # current env-defining args; a mismatch rebuilds + self._is_torch = False + # Per-slot episode scoring, sticky until the next reset. + self._done: np.ndarray = np.zeros(1, dtype=bool) + self._success: np.ndarray = np.zeros(1, dtype=bool) + self._acc_reward: np.ndarray = np.zeros(1) + self._step_reward: np.ndarray = np.zeros(1) # last env.step reward (RL wire sibling) + self._seen_success = False # any env-reported success signal this episode + + @property + def _unwrapped(self) -> Any: + """The env's unwrapped core (gymnasium wrappers hide Isaac attributes).""" + return getattr(self.env, "unwrapped", self.env) + + # ── env lifecycle (all sim touches on the sim thread) ─────────────────────── + + async def ensure_env(self, **task_args: Any) -> None: + """Build the env (factory defaults unless task args say otherwise). Idempotent.""" + if self.env is None: + await self._run_on_sim(self._sync_reset, task_args) + + async def start(self) -> None: + """Build the env and derive/load the contract, then bring up the wire.""" + await self.ensure_env() + state, frames = self.sample_observation() + existed = self._contract_path is not None and Path(self._contract_path).exists() + self.contract = load_or_write_contract( + self._contract_path, + state, + frames, + action_dim_of(self.env, batched=self.batched), + self._fps or detect_fps(self.env), + ) + if not existed and self._contract_path is not None: + print(f"[env] wrote {self._contract_path} (edit names to relabel plots)", flush=True) + await super().start() + + async def reset(self, **task_args: Any) -> str: + return await self._run_on_sim(self._sync_reset, task_args) + + async def stop(self) -> None: + await super().stop() + if self.env is not None: + await self._run_on_sim(self.env.close) + self.env = None + + def _sync_reset(self, task_args: dict[str, Any]) -> str: + # Defaults from env.gym(..., num_envs=N) fill in when the task omits them. + merged = {**self._defaults, **task_args} + build = {k: v for k, v in merged.items() if k in self._build_params} + episodic = {k: v for k, v in merged.items() if k not in self._build_params} + key = tuple(sorted(build.items())) + if self.env is None or key != self._instance: + if self.env is not None: + self.env.close() + self.env = self._build_env(build) + self._instance = key + n = num_envs_of(self.env) + self.batched = n is not None + self.num_envs = n or 1 + seed = episodic.pop("seed", None) + obs, _ = self.env.reset(seed=seed, options=episodic or None) + self._obs = obs # the first frame an agent sees on connect/reset + self._is_torch = "torch" in type(_first_leaf(obs)).__module__ + self._done = np.zeros(self.num_envs, dtype=bool) + self._success = np.zeros(self.num_envs, dtype=bool) + self._acc_reward = np.zeros(self.num_envs) + self._step_reward = np.zeros(self.num_envs, dtype=np.float32) + self._seen_success = False + return self._prompt(merged) + + def _build_env(self, build: dict[str, Any]) -> Any: + """Build the env from the target: factory call, or registry make/make_vec.""" + if callable(self._target): + return self._target(**build) + import gymnasium + from gymnasium.envs.registration import EnvSpec + + env_id, kwargs = self._target, dict(build) + if env_id.startswith("{"): # a declared env's spec, re-made in this process + spec = EnvSpec.from_json(env_id) + env_id, kwargs = spec.id, {**spec.kwargs, **kwargs} + if "num_envs" in kwargs: + return gymnasium.make_vec(env_id, **kwargs) + return gymnasium.make(env_id, **kwargs) + + def _prompt(self, task_args: dict[str, Any]) -> str: + base = self._unwrapped + for attr in ("task_description", "instruction"): + text = getattr(getattr(base, "cfg", None), attr, None) or getattr(base, attr, None) + if isinstance(text, str) and text: + return text + return ", ".join(f"{k}={v}" for k, v in sorted(task_args.items())) or "run the task" + + def sample_observation(self) -> tuple[dict[str, Any], dict[str, Any]]: + """One per-env ``(state, frames)`` sample for contract derivation (post-build).""" + state, frames = split_observation(self._obs, batched=self.batched) + if self.batched: + state = {k: to_numpy(v)[0] for k, v in state.items()} + frames = {k: to_numpy(v)[0] for k, v in frames.items()} + return state, frames + + # ── bridge protocol ────────────────────────────────────────────────────────── + + def step(self, action: np.ndarray) -> None: + act: Any = np.array(action, dtype=np.float32) # wire buffer is read-only + if not self.batched: + act = act[0] if act.ndim > 1 else act # single plain env: drop the batch dim + elif act.ndim == 1: + act = act[None] + # The wire carries floats; discrete/int action spaces need their dtype + shape back. + space = getattr(self.env, "action_space", None) + dtype = getattr(space, "dtype", None) + if dtype is not None and np.issubdtype(dtype, np.integer): + act = act.astype(dtype).reshape(getattr(space, "shape", act.shape) or ()) + if self._is_torch: + import torch + + act = torch.as_tensor(act, device=getattr(self._unwrapped, "device", None)) + obs, reward, terminated, truncated, info = self.env.step(act) + self._obs = obs + done = np.atleast_1d(to_numpy(terminated)).astype(bool) | np.atleast_1d( + to_numpy(truncated) + ).astype(bool) + # Per-step reward for the RL wire (get_observation attaches it as a sibling). + self._step_reward = np.atleast_1d(to_numpy(reward)).astype(np.float32) + self._acc_reward += self._step_reward * ~self._done + newly = done & ~self._done + if newly.any(): + success = self._resolve_success(info) + if success is not None: + self._seen_success = True + self._success |= success & newly + self._done |= done + + def _resolve_success(self, info: Any) -> np.ndarray | None: + """Env-reported success at a done step: info keys, else Isaac's termination term.""" + found = probe_success(info, num_envs=self.num_envs) + if found is not None: + return found + manager = getattr(self._unwrapped, "termination_manager", None) + if manager is not None: + try: + return np.atleast_1d(to_numpy(manager.get_term("success"))).astype(bool) + except Exception: + return None + return None + + def get_observation(self) -> tuple[dict[str, np.ndarray], np.ndarray] | None: + """Always ``[N, ...]`` arrays + ``[N]`` terminated for the barrier fan-out.""" + if self.env is None or self._obs is None: + return None + state, frames = split_observation(self._obs, batched=self.batched) + data = {k: to_numpy(v) for k, v in {**state, **frames}.items()} + # Last env.step reward rides along for RL collection (client lifts it + # out of "data" to a top-level sibling, like "terminated"). + data["reward"] = self._step_reward + if not self.batched: + # Plain single env: lift scalars to a batch of one for uniform slicing. + data = {k: (v if getattr(v, "ndim", 0) >= 1 and v.shape[:1] == (1,) else np.asarray(v)[None]) + for k, v in data.items()} + return data, np.asarray(self._done, dtype=bool).reshape(self.num_envs) + + def result_slots(self) -> list[dict[str, Any]]: + """Per-slot grades: env-reported success when available, else accumulated reward.""" + # The env's own success check outranks accumulated shaped reward. + scores = self._success if self._seen_success else self._acc_reward + return [ + { + "score": float(scores[i]), + "success": bool(self._success[i]), + "total_reward": float(self._acc_reward[i]), + } + for i in range(self.num_envs) + ] + + +def _first_leaf(obs: Any) -> Any: + while isinstance(obs, dict): + obs = next(iter(obs.values())) + return obs + + # ── hud.wrap: trace streaming for a loop you own ────────────────────────────── @@ -176,12 +418,12 @@ class TracedEnv: Plain envs are recorded as a batch of one; vectorized/batched envs fan out into one trace per episode per slot (``done[i]`` closes slot ``i``'s trace and opens the next). - - ``job`` — job name on the platform (default: the env's spec id / class name). - - ``job_id`` — share one job across several wrapped envs (multi-task suites). - - ``task`` — optional instruction/label shown on each trace's timeline. - - ``fps`` — control rate override (default: detected from the env). - - ``contract`` — path for the derived contract round-trip; ``None`` disables it. - - ``record_indices`` — which env slots get rich traces (default: first 4). + - ``job`` - job name on the platform (default: the env's spec id / class name). + - ``job_id`` - share one job across several wrapped envs (multi-task suites). + - ``task`` - optional instruction/label shown on each trace's timeline. + - ``fps`` - control rate override (default: detected from the env). + - ``contract`` - path for the derived contract round-trip; ``None`` disables it. + - ``record_indices`` - which env slots get rich traces (default: first 4). """ def __init__( @@ -197,7 +439,7 @@ def __init__( ) -> None: self.env = env self._batched = is_batched(env) - self._n = num_envs_of(env) or 1 + self._num_envs = num_envs_of(env) or 1 self._fps = fps or detect_fps(env) self._job = job or getattr(getattr(env, "spec", None), "id", None) or type(env).__name__ self._job_id = job_id @@ -233,21 +475,19 @@ def step(self, action: Any) -> Any: info = result[4] if len(result) > 4 else {} if self._rec is not None: state, frames = split_observation(obs, batched=self._batched) - done = np.atleast_1d(to_numpy(terminated)).astype(bool) | np.atleast_1d( - to_numpy(truncated) - ).astype(bool) - if not self._batched: # record a plain env as a batch of one + # Episode ends on either gym flag; coerce to a 1-D bool mask. + done = np.atleast_1d(to_numpy(terminated) | to_numpy(truncated)).astype(bool) + if not self._batched: # plain env → batch of one for the recorder state = {k: v[None] for k, v in state.items()} frames = {k: v[None] for k, v in frames.items()} action = to_numpy(action)[None] - reward = np.atleast_1d(to_numpy(reward)) self._rec.record( obs=state or None, frames=frames or None, action=action, - reward=reward, + reward=np.atleast_1d(to_numpy(reward)), done=done, - success=probe_success(info, num_envs=self._n), + success=probe_success(info, num_envs=self._num_envs), ) return result @@ -262,31 +502,34 @@ def close(self) -> None: def _start(self, obs: Any) -> JobRecorder: """First reset: derive/load the contract, then open the job's recorder.""" state, frames = split_observation(obs, batched=self._batched) + # Contract naming uses one per-env sample (slot 0 when batched). sample = {k: to_numpy(v)[0] if self._batched else v for k, v in state.items()} - existed = self._contract_path is not None and self._contract_path.exists() + path = self._contract_path + wrote = path is not None and not path.exists() contract = load_or_write_contract( - self._contract_path, + path, sample, frames, action_dim_of(self.env, batched=self._batched), self._fps, ) - if not existed and self._contract_path is not None: - logger.info("hud.wrap: wrote %s (edit names to relabel plots)", self._contract_path) - feats = contract.get("features", {}) + if wrote: + logger.info("hud.wrap: wrote %s (edit names to relabel plots)", path) + features = contract.get("features", {}) return JobRecorder( self._job, - self._n, + self._num_envs, record_indices=self._record_indices, fps=int(contract.get("control_rate") or self._fps), job_id=self._job_id, prompt=self._task, action_names=next( - (f.get("names") for f in feats.values() if f.get("role") == "action"), None + (f.get("names") for f in features.values() if f.get("role") == "action"), + None, ), state_names={ k: f["names"] - for k, f in feats.items() + for k, f in features.items() if f.get("role") == "observation" and f.get("names") }, ) @@ -298,17 +541,98 @@ def __exit__(self, *exc: object) -> None: self.close() -# The public verb: ``hud.wrap(env, job=...)`` — construction is the wrapping. +# The public verb: ``hud.wrap(env, job=...)`` - construction is the wrapping. wrap = TracedEnv +# ── the sim program (python -m hud.environment.robot.gym) ───────────────────── + + +def gym_command( + target: Any, + *, + fps: int | None = None, + contract: str | None = "contract.json", + **defaults: Any, +) -> list[str]: + """Build the command line that spawns *target* as a sim process (what ``env.gym`` runs). + + ``env.py`` only declares the sim, it never runs it — this builds the argv + for the child process that will. A live env object can't cross that + boundary, only text can, so *target* is reduced to something the child can + rebuild itself: a factory's source path, a registry id as-is, or a + constructed env's spec (id + kwargs) — closed here since only the spec + crosses over; the child calls ``gym.make`` again for its own instance. + ``fps`` / ``contract`` / build defaults (e.g. ``num_envs=``) flow through + to the :class:`GymBridge` the CLI builds. + """ + if callable(target): + name = getattr(target, "__qualname__", "") + if not name or "." in name or "<" in name: + raise ValueError(f"env.gym factory must be a module-level callable, got {target!r}") + target = f"{inspect.getfile(target)}:{name}" + elif not isinstance(target, str): + spec = getattr(target, "spec", None) # registry-made envs carry their EnvSpec + if spec is None: + raise ValueError( + f"env.gym: {target!r} has no registry spec; pass a gymnasium id " + "or a module-level factory instead" + ) + target_env, target = target, spec.to_json() + target_env.close() # only the spec crosses; free the declaration-time instance + cmd = [sys.executable, "-m", "hud.environment.robot.gym", target] + if fps is not None: + cmd += ["--fps", str(fps)] + if contract is not None: + cmd += ["--contract", str(contract)] + # Build defaults (num_envs, etc.) as --key value pairs the CLI re-applies. + for key, value in defaults.items(): + cmd += [f"--{key.replace('_', '-')}", str(value)] + return cmd + + +def main() -> None: + import argparse + + from hud.utils.modules import load_module + + parser = argparse.ArgumentParser(description="Serve a gym-style env as a robot sim.") + parser.add_argument("target", help="'path/to/module.py:factory', a registry id, or spec JSON.") + parser.add_argument("--fps", type=int, default=None, help="Control-rate override.") + parser.add_argument("--contract", default=None, help="contract.json round-trip path.") + parser.add_argument("--host", default="127.0.0.1", help="Control-channel interface.") + parser.add_argument("--port", type=int, default=0, help="Control-channel port (0 = ephemeral).") + parser.add_argument("--num-envs", type=int, default=None, help="Vectorized slot count.") + args, unknown = parser.parse_known_args() + + target: Any = args.target + if ".py:" in target: # factory by source path; ids / spec JSON pass through + path, _, attr = target.rpartition(":") + target = getattr(load_module(path), attr) + defaults: dict[str, Any] = {} + if args.num_envs is not None: + defaults["num_envs"] = args.num_envs + # Further --key value pairs from gym_command become build defaults. + for flag, value in zip(unknown, unknown[1:]): + if flag.startswith("--"): + defaults[flag[2:].replace("-", "_")] = value + bridge = GymBridge(target, fps=args.fps, contract=args.contract, **defaults) + serve_bridge(bridge, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() + + __all__ = [ "SUCCESS_KEYS", + "GymBridge", "TracedEnv", "action_dim_of", "capture_task_params", "detect_fps", "flatten_observation", + "gym_command", "is_batched", "load_or_write_contract", "num_envs_of", diff --git a/hud/environment/robot/sim.py b/hud/environment/robot/sim.py deleted file mode 100644 index 03c17db82..000000000 --- a/hud/environment/robot/sim.py +++ /dev/null @@ -1,144 +0,0 @@ -"""The sim-process shape every sim program runs: the sim owns the main thread. - -There is one way to serve a simulator (see :mod:`~.bridge`), whatever the sim: -serving — the robot WebSocket and the control side channel — runs on a -background loop thread, and every sim touch is queued to the process main -thread through the shared :class:`SimThread`. One shape because the hardest -sims demand it: a simulator is usually *thread-affine* (every touch must run -on the thread that created its GL/device context), and Isaac/Omniverse must -own the process main thread outright — Kit drives its own main-thread loop and -``env.reset()`` nests ``run_until_complete``, which cannot run inside an -asyncio task. Cheap CPU envs pay ~nothing: the main thread just blocks on the -queue. - -Before :meth:`SimThread.run` starts (tests, in-loop use) calls execute inline -on the caller — the degenerate single-thread case. -""" - -from __future__ import annotations - -import asyncio -import contextlib -import queue -import signal -import sys -import threading -from concurrent.futures import Future -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from collections.abc import Callable, Coroutine - - -class SimThread: - """Executes sim touches on the thread that owns the simulator. - - ``call`` is awaitable from any event loop and routes the touch to the - owning thread; :meth:`run` is the blocking loop that thread drives. - Re-entrant calls (a sim touch calling back in) and calls made before - :meth:`run` starts execute inline. - """ - - _shared: SimThread | None = None - - def __init__(self) -> None: - self._q: queue.Queue[tuple[Callable[[], Any], Future]] = queue.Queue() - self._ident: int | None = None - - @classmethod - def shared(cls) -> SimThread: - """The process-wide instance (one main thread, one sim thread).""" - if cls._shared is None: - cls._shared = cls() - return cls._shared - - async def call(self, fn: Callable[..., Any], *args: Any) -> Any: - """Run ``fn(*args)`` on the sim thread, awaited on the caller's loop.""" - if self._ident is None or threading.get_ident() == self._ident: - return fn(*args) # not serving yet, or already on the sim thread - fut: Future = Future() - self._q.put((lambda: fn(*args), fut)) - return await asyncio.wrap_future(fut) - - def bind(self) -> None: - """Claim the calling thread as the sim thread (before serving starts, - so no touch can slip through inline on the wrong thread).""" - self._ident = threading.get_ident() - - def run(self, until: Callable[[], bool]) -> None: - """Own the calling thread: execute queued touches until ``until()``. - - Each pass drains the queue, then runs the idle hook — a Kit update when - Omniverse is loaded (it needs continuous pumping), else nothing (the - queue ``get`` timeout is the wait). - """ - self.bind() - while not until(): - kit = sys.modules.get("omni.kit.app") - try: - item = self._q.get(timeout=0.002 if kit else 0.05) - except queue.Empty: - pass - else: - self._execute(*item) - with contextlib.suppress(queue.Empty): # drain the backlog - while True: - self._execute(*self._q.get_nowait()) - if kit: - kit.get_app().update() - - @staticmethod - def _execute(fn: Callable[[], Any], fut: Future) -> None: - if not fut.set_running_or_notify_cancel(): - return - try: - fut.set_result(fn()) - except BaseException as exc: # propagate to the awaiting caller - fut.set_exception(exc) - - -def run_with_sim(serve: Callable[[], Coroutine[Any, Any, Any]]) -> None: - """THE process shape for serving a sim, blocking for the serve's lifetime. - - ``await serve()`` runs on a background loop thread while the shared - :class:`SimThread` owns the calling (main) thread. SIGTERM and Ctrl-C - cancel the serve coroutine; the sim keeps draining through its teardown - (``env.stop()`` touches the sim too), then this returns. - """ - sim = SimThread.shared() - sim.bind() # claim main before serving starts, so no touch runs inline elsewhere - loop = asyncio.new_event_loop() - done = threading.Event() - task_box: dict[str, asyncio.Task[Any]] = {} - - async def _main() -> None: - task_box["task"] = asyncio.current_task() # type: ignore[assignment] - with contextlib.suppress(asyncio.CancelledError): - await serve() - - def _thread() -> None: - asyncio.set_event_loop(loop) - try: - loop.run_until_complete(_main()) - finally: - done.set() - - def _cancel(*_: object) -> None: - task = task_box.get("task") - if task is not None: - loop.call_soon_threadsafe(task.cancel) - - with contextlib.suppress(ValueError): # signals are main-thread-only; tests may not be - signal.signal(signal.SIGTERM, _cancel) - - thread = threading.Thread(target=_thread, name="hud-serve", daemon=True) - thread.start() - try: - sim.run(until=done.is_set) - except KeyboardInterrupt: - _cancel() - sim.run(until=done.is_set) # keep draining through teardown - thread.join(timeout=30) - - -__all__ = ["SimThread", "run_with_sim"] diff --git a/hud/environment/server.py b/hud/environment/server.py index e4ec9f838..a20732739 100644 --- a/hud/environment/server.py +++ b/hud/environment/server.py @@ -464,7 +464,7 @@ def serve_blocking(env: Environment, host: str, port: int) -> None: Sims never run here: an env with attached sims (``env.gym(...)``) spawns each one as its own process from an ``@env.initialize`` hook (see - :mod:`hud.environment.robot.bridge`), so every env serves the same way. + :mod:`hud.environment.robot.gym`), so every env serves the same way. """ asyncio.run(_serve_until_terminated(env, host, port)) From d1d7e23a0fdc3c2d6b1c668db63131347c6b4890 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 02:38:01 +0000 Subject: [PATCH 29/43] refactor(hud): drop lazy hud.wrap export from package root Co-authored-by: Cursor --- hud/__init__.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/hud/__init__.py b/hud/__init__.py index 6ebf8a276..fddc41fa8 100644 --- a/hud/__init__.py +++ b/hud/__init__.py @@ -61,15 +61,6 @@ ] -def __getattr__(name: str) -> object: - # Lazy: hud.wrap pulls in the robot extra (numpy, websockets) only when used. - if name == "wrap": - from .environment.robot import wrap - - return wrap - raise AttributeError(f"module 'hud' has no attribute {name!r}") - - try: from .version import __version__ except ImportError: From 1df6832df5aedfe6adfaa8c096daa124a049f78d Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 02:41:53 +0000 Subject: [PATCH 30/43] docs(robot): point wrap at hud.environment.robot, not hud.wrap Co-authored-by: Cursor --- docs/v6/advanced/robots.mdx | 2 +- hud/environment/robot/__init__.py | 4 ++-- hud/environment/robot/gym.py | 13 +++++++------ hud/telemetry/robot/__init__.py | 2 +- hud/telemetry/robot/recorder.py | 2 +- 5 files changed, 12 insertions(+), 11 deletions(-) diff --git a/docs/v6/advanced/robots.mdx b/docs/v6/advanced/robots.mdx index 5ad840bca..fc422c6ac 100644 --- a/docs/v6/advanced/robots.mdx +++ b/docs/v6/advanced/robots.mdx @@ -533,7 +533,7 @@ downstream (`hud eval`, tasksets, the agent) is unchanged; only *where the bridg | Symbol | Where | Role | |--------|-------|------| | `env.gym(make_env)` / `Gym` | `hud.environment.robot` | Declarative sim: contract, capability, and serving derived from a gym factory | -| `hud.wrap(env)` | `hud` | One-line trace streaming for any gym-style env you drive yourself | +| `wrap(env)` | `hud.environment.robot` | One-line trace streaming for any gym-style env you drive yourself | | `RobotBridge` / `GymBridge` | `hud.environment.robot` | Env-side serve loop (slot barrier + scalar wire); subclass `RobotBridge` for a custom sim | | `RobotEndpoint` | `hud.environment.robot` | Episode bookkeeping + results (`reset` → `{prompt, token}`, `result(token=...)`) | | `RobotEndpoint.capability(...)` | `hud.environment.robot` | Build the `openpi/0` capability after `start()` | diff --git a/hud/environment/robot/__init__.py b/hud/environment/robot/__init__.py index 975a58088..f77c0b32b 100644 --- a/hud/environment/robot/__init__.py +++ b/hud/environment/robot/__init__.py @@ -15,8 +15,8 @@ or remote (attached), same control surface either way. - :func:`~.bridge.serve_bridge` — the sim program's blocking entry (custom bridges call it last). -- :func:`hud.wrap` (:mod:`~.gym`) — one-line trace streaming for any gym env - you drive yourself, plus the shared gym introspection. +- :func:`~.gym.wrap` — one-line trace streaming for any gym env you drive + yourself, plus the shared gym introspection. The agent-side counterpart, :class:`~hud.capabilities.robot.RobotClient`, lives under :mod:`hud.capabilities`; both ends share the wire codec defined there. diff --git a/hud/environment/robot/gym.py b/hud/environment/robot/gym.py index f5b701d66..259d64eb9 100644 --- a/hud/environment/robot/gym.py +++ b/hud/environment/robot/gym.py @@ -1,4 +1,4 @@ -"""Gym-style env integration: introspection, the served path, and ``hud.wrap``. +"""Gym-style env integration: introspection, the served path, and ``wrap``. Two consumers share the introspection here (split observations, probe success, derive a minimal contract) and stay in lockstep because of it: @@ -8,12 +8,13 @@ ``python -m hud.environment.robot.gym ``. This module is that sim program; :func:`gym_command` builds the argv, so both ends of the format live here. -- :func:`hud.wrap` / :class:`TracedEnv` - the loop-owning path: wrap any +- :func:`wrap` / :class:`TracedEnv` - the loop-owning path: wrap any ``gym.Env``, ``gym.vector.VectorEnv``, or batched-tensor Isaac env you drive yourself and every episode streams to the platform as a trace (numeric state, per-camera H.264 video, actions, reward/success) under one job:: - env = hud.wrap(make_env(...), job="chess-eval") + from hud.environment.robot import wrap + env = wrap(make_env(...), job="chess-eval") On first reset a minimal ``contract.json`` is written next to your script describing how the observation/action spaces were interpreted; edit its @@ -409,7 +410,7 @@ def _first_leaf(obs: Any) -> Any: return obs -# ── hud.wrap: trace streaming for a loop you own ────────────────────────────── +# ── wrap: trace streaming for a loop you own ────────────────────────────────── class TracedEnv: @@ -514,7 +515,7 @@ def _start(self, obs: Any) -> JobRecorder: self._fps, ) if wrote: - logger.info("hud.wrap: wrote %s (edit names to relabel plots)", path) + logger.info("wrap: wrote %s (edit names to relabel plots)", path) features = contract.get("features", {}) return JobRecorder( self._job, @@ -541,7 +542,7 @@ def __exit__(self, *exc: object) -> None: self.close() -# The public verb: ``hud.wrap(env, job=...)`` - construction is the wrapping. +# The public verb: ``from hud.environment.robot import wrap`` - construction is the wrapping. wrap = TracedEnv diff --git a/hud/telemetry/robot/__init__.py b/hud/telemetry/robot/__init__.py index 5f4288902..67c0482e2 100644 --- a/hud/telemetry/robot/__init__.py +++ b/hud/telemetry/robot/__init__.py @@ -1,6 +1,6 @@ """Shared robot telemetry: trace/job recorders + per-camera H.264 video streaming. -Imported by both sides of the robot stack (agent harness, ``hud.wrap``, gym +Imported by both sides of the robot stack (agent harness, ``wrap``, gym bridges); needs the ``robot`` extra (numpy, PyAV). """ diff --git a/hud/telemetry/robot/recorder.py b/hud/telemetry/robot/recorder.py index 24ca66fb1..d94a3b4f2 100644 --- a/hud/telemetry/robot/recorder.py +++ b/hud/telemetry/robot/recorder.py @@ -1,6 +1,6 @@ """Robot trace recording: numeric state, per-camera H.264 video, action chunks. -Shared by both sides of the robot stack (the agent harness, ``hud.wrap``, the +Shared by both sides of the robot stack (the agent harness, ``wrap``, the gym bridges), so it lives under telemetry rather than either side: - :class:`TraceRecorder` — one trace. Emits spans with an explicit ``trace_id`` From 08d424e54837c20082755d227277c36a2f98cd82 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 02:45:53 +0000 Subject: [PATCH 31/43] docs(environment/robot): shorten the bridge module docstring Co-authored-by: Cursor --- hud/environment/robot/bridge.py | 34 ++++++++------------------------- 1 file changed, 8 insertions(+), 26 deletions(-) diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py index 8974cafc4..b2b431d88 100644 --- a/hud/environment/robot/bridge.py +++ b/hud/environment/robot/bridge.py @@ -1,29 +1,11 @@ -"""The sim-process ``robot`` bridge base — and the sim program shape. - -The *server* side of the ``robot`` protocol (agent-side client: -:class:`~hud.capabilities.robot.RobotClient`); both share the wire codec defined -there. A bridge lives in the sim's own process and serves two channels: the -agent's obs/action **WebSocket** and the env server's JSON-RPC **control side -channel** (a :class:`~.endpoint.RobotEndpoint` drives episodes through it). - -Internally the sim may be vectorized (``num_envs`` slots in lockstep). On the -wire every connection is scalar openpi: one claim token binds a connection to -one slot; the bridge barriers actions across claimed slots, steps once, and -fans scalar obs back. Single-env is the one-slot degenerate case. - -- :class:`RobotBridge` — subclass with your sim: ``reset`` / ``step`` / - ``get_observation``, and set ``self.contract``. The base owns both channels. - (The generic gym-style bridge, :class:`~.gym.GymBridge`, lives in :mod:`~.gym`.) -- :func:`serve_bridge` — the blocking entry every sim program calls last. - -Inside the sim process the sim owns the main thread and serving runs on a -background loop thread; every sim touch is queued to main through -:meth:`RobotBridge._run_on_sim`. One shape because the hardest sims demand it: -sims are usually *thread-affine* (every touch must run on the thread that owns -their GL/device context), and Isaac/Omniverse must own the process main thread -outright — Kit drives its own main-thread loop and ``env.reset()`` nests -``run_until_complete``. Cheap CPU envs pay ~nothing: main just blocks on the -queue. +"""Sim-process ``robot`` bridge: WebSocket obs/action + JSON-RPC control. + +Subclass :class:`RobotBridge` (``reset`` / ``step`` / ``get_observation``); +:func:`serve_bridge` is the blocking entry. Wire is scalar openpi per claimed +slot; the sim may be vectorized internally (barrier step, fan-out). Sim owns +main; serving runs on a background loop and queues touches via +:meth:`RobotBridge._run_on_sim` (thread-affine / Isaac). Gym path: +:class:`~.gym.GymBridge`. """ from __future__ import annotations From 0a2d2cf0b8155783e7e55fd2fb2d6f719ac284fd Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 02:45:53 +0000 Subject: [PATCH 32/43] fix(environment/robot): barrier over all claimed slots, not just connected ones A slot claimed via the control plane but whose agent is still dialing no longer gets stepped with hold actions; the tick barrier waits for it (up to step_timeout) so episodes cannot advance before the first observation. Co-authored-by: Cursor --- hud/environment/robot/bridge.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py index b2b431d88..a614d0b8a 100644 --- a/hud/environment/robot/bridge.py +++ b/hud/environment/robot/bridge.py @@ -287,10 +287,6 @@ async def _handle_client(self, ws: Any) -> None: slot.action = None self._action_event.set() # wake the barrier so it doesn't wait on us - def _live_slots(self) -> list[_Slot]: - """Claimed slots with a live connection — the barrier's participants.""" - return [s for s in self._registry.claimed() if s.ws is not None] - async def _tick_loop(self) -> None: """Gather claimed slots' actions (or timeout → hold), step once, fan out.""" while True: @@ -299,10 +295,13 @@ async def _tick_loop(self) -> None: except TimeoutError: pass self._action_event.clear() - claimed = self._live_slots() - if not claimed: + claimed = self._registry.claimed() + # No live connection at all → don't step; episodes must not advance unseen. + if not any(s.ws is not None for s in claimed): continue - # Wait until every live claimed slot has an action, or step_timeout elapses. + # Barrier over every *claimed* slot: one whose agent is still dialing + # (claimed, no WS yet) blocks too, so its episode can't advance before + # the first observation. step_timeout marks stragglers idle (hold). deadline = asyncio.get_running_loop().time() + self.step_timeout while True: pending = [s for s in claimed if s.action is None and not s.idle] @@ -316,10 +315,10 @@ async def _tick_loop(self) -> None: self._action_event.clear() with contextlib.suppress(TimeoutError): await asyncio.wait_for(self._action_event.wait(), timeout=remaining) - claimed = self._live_slots() - if not claimed: + claimed = self._registry.claimed() + if not any(s.ws is not None for s in claimed): break - if not claimed: + if not any(s.ws is not None for s in claimed): continue hold = self.hold_action() # Stack one row per sim slot (idle/unconnected claimed → hold; free → hold). From 9884ccc35b493bd34ff702dc77dcb5f12c8f0f0d Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 02:47:02 +0000 Subject: [PATCH 33/43] docs(agents/robot): clarify max_steps is a call-site kwarg with default 520 Co-authored-by: Cursor --- hud/agents/robot/agent.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py index ace2d15c0..e4a07f331 100644 --- a/hud/agents/robot/agent.py +++ b/hud/agents/robot/agent.py @@ -52,7 +52,10 @@ class RobotAgent(Agent): adapter: Adapter | None = None async def __call__(self, run: Run, *, max_steps: int = 520) -> None: - """The generic rollout contract: one run, one scalar robot connection.""" + """The generic rollout contract: one run, one scalar robot connection. + + ``max_steps`` caps control ticks (default 520); omit it for the default. + """ if self.model is None: raise RuntimeError(f"{type(self).__name__} must set self.model in __init__") prompt = run.prompt @@ -99,7 +102,7 @@ async def _loop( recorder: TraceRecorder, writer: Any, *, - max_steps: int = 520, + max_steps: int, ) -> None: """Open-loop chunk queue: ainfer refills, then execute one action per tick.""" model = self.model From 9f0cc5a45e6bc4a9d660179e7bfa021b2c27a75a Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 02:50:16 +0000 Subject: [PATCH 34/43] fix(robot): fail fast when the slot claim is missing instead of deadlocking The agent raises on a run without a robot token, and the bridge times out a connection that never sends its claim frame (both ends previously waited on each other forever). Co-authored-by: Cursor --- hud/agents/robot/agent.py | 11 ++++++++--- hud/capabilities/robot.py | 7 +++++-- hud/environment/robot/bridge.py | 4 +++- 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py index e4a07f331..cd5b9a0e4 100644 --- a/hud/agents/robot/agent.py +++ b/hud/agents/robot/agent.py @@ -63,10 +63,15 @@ async def __call__(self, run: Run, *, max_steps: int = 520) -> None: raise TypeError(f"run.prompt must be a str, got {type(prompt).__name__}: {prompt!r}") # Per-episode slot token from tasks.start (opaque; env put it under "robot"). - robot_info = run.started.get("robot") if isinstance(run.started.get("robot"), dict) else {} - token = robot_info.get("token") + # The bridge won't send observations until the connection claims a slot, so a + # missing token would deadlock connect — fail fast instead. + robot_info = run.started.get("robot") + token = robot_info.get("token") if isinstance(robot_info, dict) else None if not isinstance(token, str): - token = None + raise RuntimeError( + "run.started carries no robot slot token; the env template must yield " + '{"robot": {"token": ...}} from endpoint.reset() before the agent runs' + ) robot = await RobotClient.connect(run.client.binding(self.robot_protocol), token=token) try: diff --git a/hud/capabilities/robot.py b/hud/capabilities/robot.py index 885ac83be..681864ffe 100644 --- a/hud/capabilities/robot.py +++ b/hud/capabilities/robot.py @@ -69,10 +69,13 @@ def spaces(self) -> tuple[dict[str, Any], dict[str, Any]]: @classmethod async def connect(cls, cap: Capability, *, token: str | None = None) -> Self: - """Dial the robot WebSocket; ``token`` claims a sim slot after the metadata frame.""" + """Dial the robot WebSocket; ``token`` claims a sim slot after the metadata frame. + + HUD bridges require the claim and send nothing until it arrives — pass the + token from ``endpoint.reset()``. Omit it only for servers without slots. + """ ws = await websockets.connect(cap.url, max_size=None, ping_interval=None) # Consume initial metadata; string means env error. - raw = await ws.recv() if isinstance(raw, str): raise RuntimeError(f"robot env error on connect:\n{raw}") diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py index a614d0b8a..cddc5bd46 100644 --- a/hud/environment/robot/bridge.py +++ b/hud/environment/robot/bridge.py @@ -252,7 +252,9 @@ async def _handle_client(self, ws: Any) -> None: slot: _Slot | None = None try: await ws.send(_packb(self.metadata)) # connect-time metadata frame - raw = await ws.recv() + # Fail fast on a client that never claims (it would otherwise deadlock: + # we wait for the claim frame, it waits for an observation). + raw = await asyncio.wait_for(ws.recv(), timeout=self.step_timeout) if isinstance(raw, str): raise RuntimeError(raw) claim = _unpackb(raw) From 5ae0641ea5d167014e5129540c442fb745075e40 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 02:51:57 +0000 Subject: [PATCH 35/43] fix(environment/robot): reject mismatched claim kwargs; fix hold shape and per-slot default grades Three vectorized-bridge fixes: a later claim with different task kwargs now raises instead of silently running the first claim's task (lockstep sims share one batch reset); hold_action derives the action dim from contract names when a derived contract omits shape; the default result_slots grades every slot so releasing an index > 0 no longer raises. Co-authored-by: Cursor --- hud/environment/robot/bridge.py | 31 ++++++++++++++++++++++--------- 1 file changed, 22 insertions(+), 9 deletions(-) diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py index cddc5bd46..c1ccce7d1 100644 --- a/hud/environment/robot/bridge.py +++ b/hud/environment/robot/bridge.py @@ -135,6 +135,8 @@ def __init__(self, *, host: str = "127.0.0.1", port: int = 0) -> None: self.total_reward: float = 0.0 self.success: bool = False self.terminated: bool = False + # kwargs of the current batch's global reset; later claims must match. + self._episode_kwargs: dict[str, Any] = {} async def _run_on_sim(self, fn: Callable[..., Any], *args: Any) -> Any: """Run ``fn(*args)`` on the sim thread, awaited on the caller's loop. @@ -155,9 +157,19 @@ async def _claim_episode(self, **kwargs: Any) -> dict[str, Any]: self.success = False self.terminated = False self.task_description = await self.reset(**kwargs) + self._episode_kwargs = kwargs self._registry.configure(self.num_envs) slot = self._registry.slots[0] else: + # Lockstep sim: one global reset serves the whole batch, so a later claim + # cannot get its own task/seed — reject differing kwargs instead of + # silently running the first claim's task. + if kwargs and kwargs != self._episode_kwargs: + raise ValueError( + f"slots share one batch reset ({self._episode_kwargs}); a concurrent " + f"claim cannot use different task kwargs ({kwargs}) — group tasks with " + "identical args, or use one sim per distinct task/seed" + ) slot = self._registry.free_slot() if slot is None: raise RuntimeError(f"all {self.num_envs} slots are claimed") @@ -186,14 +198,14 @@ def get_observation(self) -> tuple[dict[str, np.ndarray], np.ndarray] | None: """Return ``(data[N, ...], terminated[N])``, or ``None`` if not ready.""" def result_slots(self) -> list[dict[str, Any]]: - """One score dict per slot. Default: single-env binary success.""" - return [ - { - "score": 1.0 if self.success else 0.0, - "success": bool(self.success), - "total_reward": float(self.total_reward), - } - ] + """One score dict per slot. Default: the scalar episode grade for every slot + (vectorized bridges with per-slot scoring override this).""" + grade = { + "score": 1.0 if self.success else 0.0, + "success": bool(self.success), + "total_reward": float(self.total_reward), + } + return [dict(grade) for _ in range(self.num_envs)] def hold_action(self) -> np.ndarray: """Action used for idle/stalled slots at a barrier step (zeros).""" @@ -201,7 +213,8 @@ def hold_action(self) -> np.ndarray: (f for f in self.contract.get("features", {}).values() if f.get("role") == "action"), {}, ) - shape = action.get("shape") or (1,) + # Derived contracts carry only per-dim names, no shape; the label count is the dim. + shape = action.get("shape") or (len(action.get("names") or []) or 1,) return np.zeros(shape, dtype=np.float32) @property From 9e1df7038188347f127edcf81e920828e973aea6 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 03:10:44 +0000 Subject: [PATCH 36/43] fix(agents/robot): honor subclass max_steps and skip acting when already terminated Co-authored-by: Cursor --- hud/agents/robot/agent.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py index cd5b9a0e4..66cc2ecc1 100644 --- a/hud/agents/robot/agent.py +++ b/hud/agents/robot/agent.py @@ -40,6 +40,8 @@ class RobotAgent(Agent): """ robot_protocol: ClassVar[str] = ROBOT_PROTOCOL + #: Max control ticks before the episode is cut off. Subclasses may override. + max_steps: ClassVar[int] = 520 #: 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. @@ -51,10 +53,10 @@ class RobotAgent(Agent): #: Translates env<->policy spaces. Subclasses set this; ``None`` = raw pass-through. adapter: Adapter | None = None - async def __call__(self, run: Run, *, max_steps: int = 520) -> None: + async def __call__(self, run: Run, *, max_steps: int | None = None) -> None: """The generic rollout contract: one run, one scalar robot connection. - ``max_steps`` caps control ticks (default 520); omit it for the default. + ``max_steps`` caps control ticks; omit it to use the class ``max_steps`` (520). """ if self.model is None: raise RuntimeError(f"{type(self).__name__} must set self.model in __init__") @@ -90,7 +92,14 @@ async def __call__(self, run: Run, *, max_steps: int = 520) -> None: writer = DatasetWriter(robot.contract, fps=fps) print(f"[agent] episode started: {prompt!r}", flush=True) - await self._loop(robot, obs, prompt, recorder, writer, max_steps=max_steps) + await self._loop( + robot, + obs, + prompt, + recorder, + writer, + max_steps=self.max_steps if max_steps is None else max_steps, + ) recorder.close() if writer is not None: writer.end_episode() @@ -117,11 +126,11 @@ async def _loop( for step in range(max_steps): terminated = bool(np.asarray(obs["terminated"]).reshape(-1)[0]) - if step and terminated: - print(f"[agent] terminated at step {step}", flush=True) - break + # Already done (including a pre-terminated first obs) → don't act. if terminated: - chunk.clear() + if step: + print(f"[agent] terminated at step {step}", flush=True) + break recorder.record_observation(obs["data"], tick=step) From 083122fcbf5d31a7f5df245c36fd7709e881b1e6 Mon Sep 17 00:00:00 2001 From: lukass16 Date: Mon, 20 Jul 2026 03:18:07 +0000 Subject: [PATCH 37/43] fix(robot): barrier, claim, telemetry, and gym-argv robustness from review - bridge: a dropped connection marks its slot idle so the lockstep barrier doesn't stall the other slots for step_timeout - bridge/client: metadata advertises claim_required; connecting without a token fails fast instead of blocking forever in get_observation - agent: recorder.close()/writer.end_episode() run in finally, so video tails flush and buffered episodes commit even when the rollout raises - agent: the terminal observation is recorded before breaking (traces and datasets keep the final frame) - gym: build defaults cross the argv as JSON, so bools/ints/floats reach the child process with their real types Co-authored-by: Cursor --- hud/agents/robot/agent.py | 33 ++++++++++++++++++--------------- hud/capabilities/robot.py | 9 +++++++++ hud/environment/robot/bridge.py | 7 ++++++- hud/environment/robot/gym.py | 11 ++++++++--- 4 files changed, 41 insertions(+), 19 deletions(-) diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py index 66cc2ecc1..65dfc80b8 100644 --- a/hud/agents/robot/agent.py +++ b/hud/agents/robot/agent.py @@ -92,17 +92,21 @@ async def __call__(self, run: Run, *, max_steps: int | None = None) -> None: writer = DatasetWriter(robot.contract, fps=fps) print(f"[agent] episode started: {prompt!r}", flush=True) - await self._loop( - robot, - obs, - prompt, - recorder, - writer, - max_steps=self.max_steps if max_steps is None else max_steps, - ) - recorder.close() - if writer is not None: - writer.end_episode() + try: + await self._loop( + robot, + obs, + prompt, + recorder, + writer, + max_steps=self.max_steps if max_steps is None else max_steps, + ) + finally: + # Flush video tails / commit the buffered episode even when the + # rollout raises mid-loop. + recorder.close() + if writer is not None: + writer.end_episode() finally: await robot.close() run.trace.status = "completed" @@ -125,15 +129,14 @@ async def _loop( chunk: deque[ActionArray] = deque() for step in range(max_steps): - terminated = bool(np.asarray(obs["terminated"]).reshape(-1)[0]) + # Record every frame, including the terminal one. + recorder.record_observation(obs["data"], tick=step) # Already done (including a pre-terminated first obs) → don't act. - if terminated: + if bool(np.asarray(obs["terminated"]).reshape(-1)[0]): if step: print(f"[agent] terminated at step {step}", flush=True) break - recorder.record_observation(obs["data"], tick=step) - if not chunk: # refill with a fresh forward (BatchedModel coalesces ainfer) batch = adapter.adapt_observation(obs, prompt) if adapter else obs rows = np.atleast_2d(await model.ainfer(batch)) diff --git a/hud/capabilities/robot.py b/hud/capabilities/robot.py index 681864ffe..5fefc0abd 100644 --- a/hud/capabilities/robot.py +++ b/hud/capabilities/robot.py @@ -82,6 +82,15 @@ async def connect(cls, cap: Capability, *, token: str | None = None) -> Self: # Bind this connection to a claimed episode slot (scalar openpi from here). if token is not None: await ws.send(_packb({"claim": token})) + else: + # A slotted bridge sends nothing until a claim arrives - fail fast + # instead of blocking forever in get_observation(). + meta = _unpackb(raw) + if isinstance(meta, dict) and meta.get("claim_required"): + await ws.close() + raise RuntimeError( + "this robot env requires a slot claim; pass token= from endpoint.reset()" + ) return cls(cap, ws) async def get_observation(self) -> dict[str, Any]: diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py index c1ccce7d1..bd364b713 100644 --- a/hud/environment/robot/bridge.py +++ b/hud/environment/robot/bridge.py @@ -264,7 +264,9 @@ async def _handle_client(self, ws: Any) -> None: """One agent connection: claim a slot, then feed actions into the barrier.""" slot: _Slot | None = None try: - await ws.send(_packb(self.metadata)) # connect-time metadata frame + # Connect-time metadata frame; claim_required lets clients without a + # token fail fast instead of waiting forever for an observation. + await ws.send(_packb({**self.metadata, "claim_required": True})) # Fail fast on a client that never claims (it would otherwise deadlock: # we wait for the claim frame, it waits for an observation). raw = await asyncio.wait_for(ws.recv(), timeout=self.step_timeout) @@ -300,6 +302,9 @@ async def _handle_client(self, ws: Any) -> None: if slot is not None and slot.ws is ws: slot.ws = None slot.action = None + # Idle, not pending: a dropped connection must not stall the barrier + # for the other slots (a reconnect clears idle again). + slot.idle = True self._action_event.set() # wake the barrier so it doesn't wait on us async def _tick_loop(self) -> None: diff --git a/hud/environment/robot/gym.py b/hud/environment/robot/gym.py index 259d64eb9..78afbaeb1 100644 --- a/hud/environment/robot/gym.py +++ b/hud/environment/robot/gym.py @@ -587,8 +587,9 @@ def gym_command( if contract is not None: cmd += ["--contract", str(contract)] # Build defaults (num_envs, etc.) as --key value pairs the CLI re-applies. + # JSON-encoded so the child gets real bools/ints/strings back (see main()). for key, value in defaults.items(): - cmd += [f"--{key.replace('_', '-')}", str(value)] + cmd += [f"--{key.replace('_', '-')}", json.dumps(value)] return cmd @@ -613,10 +614,14 @@ def main() -> None: defaults: dict[str, Any] = {} if args.num_envs is not None: defaults["num_envs"] = args.num_envs - # Further --key value pairs from gym_command become build defaults. + # Further --key value pairs from gym_command become build defaults. JSON + # round-trip restores real types; bare strings (hand-written argv) pass as-is. for flag, value in zip(unknown, unknown[1:]): if flag.startswith("--"): - defaults[flag[2:].replace("-", "_")] = value + try: + defaults[flag[2:].replace("-", "_")] = json.loads(value) + except json.JSONDecodeError: + defaults[flag[2:].replace("-", "_")] = value bridge = GymBridge(target, fps=args.fps, contract=args.contract, **defaults) serve_bridge(bridge, host=args.host, port=args.port) From e878d307e322ae06ac4b193ec110966271b7e378 Mon Sep 17 00:00:00 2001 From: Lukass Kellijs Date: Mon, 20 Jul 2026 04:44:02 +0000 Subject: [PATCH 38/43] feat(robot): make the slot token optional for single-env bridges A None claim (connect or result) resolves to the sole claimed slot, so single-env templates drop the token plumbing; ambiguous vectorized claims still error. Templates: yield {"prompt": ...} and sim.result() suffice. --- docs/v6/advanced/robots.mdx | 9 ++- hud/agents/robot/agent.py | 12 ++-- hud/capabilities/robot.py | 20 +++---- hud/environment/robot/bridge.py | 43 ++++++++----- hud/environment/robot/endpoint.py | 17 ++++-- hud/environment/tests/test_robot_bridge.py | 70 ++++++++++++++++++++++ 6 files changed, 128 insertions(+), 43 deletions(-) create mode 100644 hud/environment/tests/test_robot_bridge.py diff --git a/docs/v6/advanced/robots.mdx b/docs/v6/advanced/robots.mdx index fc422c6ac..0ca4386f8 100644 --- a/docs/v6/advanced/robots.mdx +++ b/docs/v6/advanced/robots.mdx @@ -104,12 +104,17 @@ sim = env.gym(make_env) # make_env: any callable returning a gym-style env @env.template() async def pick_and_place(task: str = "default", seed: int = 0): ep = await sim.reset(task=task, seed=seed) # {prompt, token} - yield {"prompt": ep["prompt"], "robot": {"token": ep["token"]}} - yield await sim.result(token=ep["token"]) + yield {"prompt": ep["prompt"]} + yield await sim.result() ``` Task args are partitioned by the factory's signature: args the factory accepts define the env (a change rebuilds it), everything else is episodic and flows to `env.reset(seed=..., options=...)`. + +A single-env sim has one slot, so the template may omit the slot token as above. A vectorized sim +(`num_envs > 1`) must thread it through so each session grades its own slot — yield +`{"prompt": ep["prompt"], "robot": {"token": ep["token"]}}` and call `sim.result(token=ep["token"])`; +see [vectorized evals](#vectorized-envs-and-evals). For a sim that isn't gym-shaped, subclass the **bridge** instead: ```python diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py index 65dfc80b8..0e433d447 100644 --- a/hud/agents/robot/agent.py +++ b/hud/agents/robot/agent.py @@ -65,15 +65,11 @@ async def __call__(self, run: Run, *, max_steps: int | None = None) -> None: raise TypeError(f"run.prompt must be a str, got {type(prompt).__name__}: {prompt!r}") # Per-episode slot token from tasks.start (opaque; env put it under "robot"). - # The bridge won't send observations until the connection claims a slot, so a - # missing token would deadlock connect — fail fast instead. + # Single-env templates may omit it — a None claim binds the sole claimed + # slot; vectorized bridges reject the ambiguity at connect. robot_info = run.started.get("robot") - token = robot_info.get("token") if isinstance(robot_info, dict) else None - if not isinstance(token, str): - raise RuntimeError( - "run.started carries no robot slot token; the env template must yield " - '{"robot": {"token": ...}} from endpoint.reset() before the agent runs' - ) + raw_token = robot_info.get("token") if isinstance(robot_info, dict) else None + token = raw_token if isinstance(raw_token, str) else None robot = await RobotClient.connect(run.client.binding(self.robot_protocol), token=token) try: diff --git a/hud/capabilities/robot.py b/hud/capabilities/robot.py index 5fefc0abd..54244f656 100644 --- a/hud/capabilities/robot.py +++ b/hud/capabilities/robot.py @@ -71,26 +71,20 @@ def spaces(self) -> tuple[dict[str, Any], dict[str, Any]]: async def connect(cls, cap: Capability, *, token: str | None = None) -> Self: """Dial the robot WebSocket; ``token`` claims a sim slot after the metadata frame. - HUD bridges require the claim and send nothing until it arrives — pass the - token from ``endpoint.reset()``. Omit it only for servers without slots. + A ``None`` token binds the sole claimed slot on a single-env bridge; + vectorized bridges (ambiguous slots) reject it — pass the token from + ``endpoint.reset()`` there. """ ws = await websockets.connect(cap.url, max_size=None, ping_interval=None) # Consume initial metadata; string means env error. raw = await ws.recv() if isinstance(raw, str): raise RuntimeError(f"robot env error on connect:\n{raw}") - # Bind this connection to a claimed episode slot (scalar openpi from here). - if token is not None: + # Bind this connection to an episode slot (scalar openpi from here). HUD + # bridges require the claim frame; plain openpi servers have no slots. + meta = _unpackb(raw) + if token is not None or (isinstance(meta, dict) and meta.get("claim_required")): await ws.send(_packb({"claim": token})) - else: - # A slotted bridge sends nothing until a claim arrives - fail fast - # instead of blocking forever in get_observation(). - meta = _unpackb(raw) - if isinstance(meta, dict) and meta.get("claim_required"): - await ws.close() - raise RuntimeError( - "this robot env requires a slot claim; pass token= from endpoint.reset()" - ) return cls(cap, ws) async def get_observation(self) -> dict[str, Any]: diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py index bd364b713..be00c4ddb 100644 --- a/hud/environment/robot/bridge.py +++ b/hud/environment/robot/bridge.py @@ -67,8 +67,20 @@ def free_slot(self) -> _Slot | None: # v1: a freed slot is only reclaimable after all slots free (whole-batch episodes). return next((s for s in self.slots if s.token is None and not s.used), None) - def by_token(self, token: str) -> _Slot | None: - return next((s for s in self.slots if s.token == token), None) + def resolve(self, token: str | None) -> _Slot: + """Token → its claimed slot; ``None`` binds the sole claimed slot (single-env).""" + if token is None: + claimed = self.claimed() + if len(claimed) != 1: + raise ValueError( + f"tokenless claim needs exactly one claimed slot, found {len(claimed)}; " + "pass the token from reset()" + ) + return claimed[0] + slot = next((s for s in self.slots if s.token == token), None) + if slot is None: + raise ValueError(f"unknown episode token: {token!r}") + return slot def claimed(self) -> list[_Slot]: return [s for s in self.slots if s.token is not None] @@ -176,11 +188,9 @@ async def _claim_episode(self, **kwargs: Any) -> dict[str, Any]: token = self._registry.claim(slot) return {"prompt": self.task_description, "token": token} - def _release_episode(self, token: str) -> dict[str, Any]: + def _release_episode(self, token: str | None) -> dict[str, Any]: """Control-plane result: this slot's score, then free it.""" - slot = self._registry.by_token(token) - if slot is None: - raise ValueError(f"unknown episode token: {token!r}") + slot = self._registry.resolve(token) grade = self.result_slots()[slot.index] self._registry.release(slot) return grade @@ -273,12 +283,11 @@ async def _handle_client(self, ws: Any) -> None: if isinstance(raw, str): raise RuntimeError(raw) claim = _unpackb(raw) - token = claim.get("claim") - if not isinstance(token, str): - raise ValueError("first frame must be {\"claim\": }") - slot = self._registry.by_token(token) - if slot is None: - raise ValueError(f"unknown claim token: {token!r}") + if not isinstance(claim, dict) or "claim" not in claim: + raise ValueError('first frame must be {"claim": }') + # A None claim binds the sole claimed slot — single-env agents skip + # the token plumbing; ambiguous (vectorized) claims error instead. + slot = self._registry.resolve(claim["claim"]) if slot.ws is not None: raise RuntimeError(f"slot {slot.index} already has a live connection") slot.ws = ws @@ -363,8 +372,10 @@ async def _send_slot_observation(self, slot: _Slot) -> None: i = slot.index # Slice the [N, ...] batch down to this slot's scalar row. msg = { - **{k: (v[i] if getattr(v, "ndim", 0) >= 1 and len(v) == self.num_envs else v) - for k, v in data.items()}, + **{ + k: (v[i] if getattr(v, "ndim", 0) >= 1 and len(v) == self.num_envs else v) + for k, v in data.items() + }, "terminated": bool(np.asarray(terminated).reshape(-1)[i]), } with contextlib.suppress(websockets.exceptions.ConnectionClosed): @@ -399,8 +410,8 @@ async def _dispatch_control(self, method: str, params: dict[str, Any]) -> dict[s return await self._claim_episode(**params) if method == "result": token = params.get("token") - if not isinstance(token, str): - raise ValueError("result: 'token' must be a string") + if token is not None and not isinstance(token, str): + raise ValueError("result: 'token' must be a string when given") return self._release_episode(token) raise ValueError(f"unknown method {method!r}") diff --git a/hud/environment/robot/endpoint.py b/hud/environment/robot/endpoint.py index 9290ed82b..21a0c14d7 100644 --- a/hud/environment/robot/endpoint.py +++ b/hud/environment/robot/endpoint.py @@ -21,8 +21,13 @@ @env.template(id="pawn_lift") async def pawn_lift(task: str = "solo_pawn_lift", seed: int = 0): ep = await sim.reset(task=task, seed=seed) # {prompt, token} - yield {"prompt": ep["prompt"], "robot": {"token": ep["token"]}} - yield await sim.result(token=ep["token"]) + yield {"prompt": ep["prompt"]} + yield await sim.result() + +A single-env sim has one slot, so the template may omit the token as above. A +vectorized sim (``num_envs > 1``) must thread it through so each session grades +its own slot: ``yield {"prompt": ..., "robot": {"token": ep["token"]}}`` and +``sim.result(token=ep["token"])``. """ from __future__ import annotations @@ -167,8 +172,12 @@ async def reset(self, **task_args: Any) -> dict[str, Any]: """Claim a slot for a new episode; return ``{"prompt", "token"}``.""" return await self._call("reset", task_args) - async def result(self, *, token: str, **extra: Any) -> dict[str, Any]: - """This slot's score dict (frees the slot), merged with any caller ``extra``.""" + async def result(self, *, token: str | None = None, **extra: Any) -> dict[str, Any]: + """This slot's score dict (frees the slot), merged with any caller ``extra``. + + ``token`` may be omitted on a single-env bridge (one claimed slot); + vectorized envs must pass the token from :meth:`reset`. + """ res = {**(await self._call("result", {"token": token})), **extra} print( f"[env] result: success={res.get('success')} " diff --git a/hud/environment/tests/test_robot_bridge.py b/hud/environment/tests/test_robot_bridge.py new file mode 100644 index 000000000..4208e98c3 --- /dev/null +++ b/hud/environment/tests/test_robot_bridge.py @@ -0,0 +1,70 @@ +"""Slot-token contract on the bridge control plane (reset/result).""" + +from __future__ import annotations + +from typing import Any + +import pytest + +pytest.importorskip("numpy") +pytest.importorskip("openpi_client") + +from hud.environment.robot.bridge import RobotBridge + + +class StubBridge(RobotBridge): + """Minimal concrete bridge: no sim, fixed prompt.""" + + def __init__(self, num_envs: int = 1) -> None: + super().__init__() + self.num_envs = num_envs + + async def reset(self, **kwargs: Any) -> str: + return "do the task" + + def step(self, action: Any) -> None: + pass + + def get_observation(self) -> None: + return None + + +async def test_result_without_token_grades_the_single_slot() -> None: + bridge = StubBridge() + ep = await bridge._dispatch_control("reset", {}) + assert ep["prompt"] == "do the task" + assert isinstance(ep["token"], str) + + grade = await bridge._dispatch_control("result", {}) + assert {"score", "success", "total_reward"} <= set(grade) + + # The slot was freed: a new episode claims cleanly. + await bridge._dispatch_control("reset", {}) + + +async def test_result_with_token_still_resolves_its_slot() -> None: + bridge = StubBridge() + ep = await bridge._dispatch_control("reset", {}) + grade = await bridge._dispatch_control("result", {"token": ep["token"]}) + assert {"score", "success", "total_reward"} <= set(grade) + + +async def test_tokenless_result_rejects_ambiguous_slots() -> None: + bridge = StubBridge(num_envs=2) + await bridge._dispatch_control("reset", {}) + await bridge._dispatch_control("reset", {}) + with pytest.raises(ValueError, match="exactly one claimed slot"): + await bridge._dispatch_control("result", {}) + + +async def test_tokenless_result_rejects_no_claimed_slot() -> None: + bridge = StubBridge() + with pytest.raises(ValueError, match="exactly one claimed slot"): + await bridge._dispatch_control("result", {}) + + +async def test_unknown_token_still_errors() -> None: + bridge = StubBridge() + await bridge._dispatch_control("reset", {}) + with pytest.raises(ValueError, match="unknown episode token"): + await bridge._dispatch_control("result", {"token": "slot-0-deadbeef"}) From cc47fff08af098c808fbbc022a50330505af989c Mon Sep 17 00:00:00 2001 From: Lukass Kellijs Date: Mon, 20 Jul 2026 04:44:22 +0000 Subject: [PATCH 39/43] =?UTF-8?q?feat(environment/robot):=20lazy=20Isaac?= =?UTF-8?q?=20spawn=20=E2=80=94=20serve=20the=20wire=20before=20the=20sim?= =?UTF-8?q?=20boots?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GymBridge.start announces the port and loads a pre-written contract immediately; the env builds on first reset. Endpoint connect timeout raised to 900s for cold Isaac boots. --- hud/environment/robot/endpoint.py | 6 +++--- hud/environment/robot/gym.py | 32 +++++++++++++++++++++++++++---- 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/hud/environment/robot/endpoint.py b/hud/environment/robot/endpoint.py index 21a0c14d7..7f8cb579d 100644 --- a/hud/environment/robot/endpoint.py +++ b/hud/environment/robot/endpoint.py @@ -61,7 +61,7 @@ def __init__( cmd: Sequence[str] | None = None, host: str | None = None, port: int | None = None, - connect_timeout_s: float = 240.0, + connect_timeout_s: float = 900.0, ) -> None: self._cmd = list(cmd) if cmd is not None else None # set => spawned mode self._host = host @@ -75,12 +75,12 @@ def __init__( self._lock = asyncio.Lock() @classmethod - def spawn(cls, cmd: Sequence[str], *, connect_timeout_s: float = 240.0) -> RobotEndpoint: + def spawn(cls, cmd: Sequence[str], *, connect_timeout_s: float = 900.0) -> RobotEndpoint: """An endpoint that forks *cmd* (a sim program; see :mod:`~.bridge`) and owns it.""" return cls(cmd=cmd, connect_timeout_s=connect_timeout_s) @classmethod - def remote(cls, host: str, port: int, *, connect_timeout_s: float = 240.0) -> RobotEndpoint: + def remote(cls, host: str, port: int, *, connect_timeout_s: float = 900.0) -> RobotEndpoint: """An endpoint attached to a sim process something else runs.""" return cls(host=host, port=port, connect_timeout_s=connect_timeout_s) diff --git a/hud/environment/robot/gym.py b/hud/environment/robot/gym.py index 78afbaeb1..1eca75bbf 100644 --- a/hud/environment/robot/gym.py +++ b/hud/environment/robot/gym.py @@ -248,9 +248,19 @@ async def ensure_env(self, **task_args: Any) -> None: if self.env is None: await self._run_on_sim(self._sync_reset, task_args) - async def start(self) -> None: - """Build the env and derive/load the contract, then bring up the wire.""" - await self.ensure_env() + def _load_contract_if_present(self) -> None: + """Load a pre-written contract.json without touching Isaac (lazy spawn).""" + if self.contract or self._contract_path is None: + return + path = Path(self._contract_path) + if path.exists(): + self.contract = json.loads(path.read_text()) + self._fps = self._fps or int(self.contract.get("control_rate") or 15) + + def _ensure_contract_from_env(self) -> None: + """Derive/load contract once the env exists (first reset under lazy spawn).""" + if self.contract or self.env is None: + return state, frames = self.sample_observation() existed = self._contract_path is not None and Path(self._contract_path).exists() self.contract = load_or_write_contract( @@ -262,10 +272,24 @@ async def start(self) -> None: ) if not existed and self._contract_path is not None: print(f"[env] wrote {self._contract_path} (edit names to relabel plots)", flush=True) + + async def start(self) -> None: + """Bring up the wire immediately; Isaac builds on first reset (lazy spawn). + + Cold Isaac boots can take several minutes. Announcing ``HUD_SIM_PORT`` + before ``make_env`` keeps the parent server alive while the child cooks. + """ + self._load_contract_if_present() + if "num_envs" in self._defaults: + self.num_envs = int(self._defaults["num_envs"]) + self.batched = self.num_envs > 0 + print("[env] lazy spawn: announcing port before Isaac boot", flush=True) await super().start() async def reset(self, **task_args: Any) -> str: - return await self._run_on_sim(self._sync_reset, task_args) + prompt = await self._run_on_sim(self._sync_reset, task_args) + self._ensure_contract_from_env() + return prompt async def stop(self) -> None: await super().stop() From 16104f32d1adaa1b12915d7ee57291482bf6827e Mon Sep 17 00:00:00 2001 From: Lukass Kellijs Date: Tue, 21 Jul 2026 02:45:45 +0000 Subject: [PATCH 40/43] fix(environment/robot): probe env at start when contract.json is missing GymBridge.start builds with factory defaults to mint the contract before the capability is published, so env.initialize never binds an empty manifest. Existing contract files still skip the probe (lazy spawn). Co-authored-by: Cursor --- hud/environment/robot/gym.py | 30 ++++--- .../tests/test_gym_bridge_probe.py | 90 +++++++++++++++++++ 2 files changed, 109 insertions(+), 11 deletions(-) create mode 100644 hud/environment/tests/test_gym_bridge_probe.py diff --git a/hud/environment/robot/gym.py b/hud/environment/robot/gym.py index 1eca75bbf..f9adb0c33 100644 --- a/hud/environment/robot/gym.py +++ b/hud/environment/robot/gym.py @@ -16,10 +16,10 @@ from hud.environment.robot import wrap env = wrap(make_env(...), job="chess-eval") -On first reset a minimal ``contract.json`` is written next to your script -describing how the observation/action spaces were interpreted; edit its -``names`` to relabel the platform's plots. An existing file is loaded instead, -so edits stick. +If no contract file exists, ``GymBridge.start`` builds the env once (factory +defaults) and writes a minimal ``contract.json`` before the capability is +published — the agent binds from that manifest. Edit ``names`` to relabel +plots; an existing file is loaded instead, so edits stick. """ from __future__ import annotations @@ -249,7 +249,7 @@ async def ensure_env(self, **task_args: Any) -> None: await self._run_on_sim(self._sync_reset, task_args) def _load_contract_if_present(self) -> None: - """Load a pre-written contract.json without touching Isaac (lazy spawn).""" + """Load a pre-written contract.json when present (skips the start-time probe).""" if self.contract or self._contract_path is None: return path = Path(self._contract_path) @@ -258,7 +258,7 @@ def _load_contract_if_present(self) -> None: self._fps = self._fps or int(self.contract.get("control_rate") or 15) def _ensure_contract_from_env(self) -> None: - """Derive/load contract once the env exists (first reset under lazy spawn).""" + """Derive/load contract from a sample observation once the env exists.""" if self.contract or self.env is None: return state, frames = self.sample_observation() @@ -274,21 +274,29 @@ def _ensure_contract_from_env(self) -> None: print(f"[env] wrote {self._contract_path} (edit names to relabel plots)", flush=True) async def start(self) -> None: - """Bring up the wire immediately; Isaac builds on first reset (lazy spawn). + """Bind the wire with a complete contract so ``env.initialize`` can publish it. - Cold Isaac boots can take several minutes. Announcing ``HUD_SIM_PORT`` - before ``make_env`` keeps the parent server alive while the child cooks. + A contract file is loaded when present. Otherwise the env is built with + factory defaults and probed here — ``@env.initialize`` awaits this — so + the capability manifest is never empty when the agent binds. """ self._load_contract_if_present() if "num_envs" in self._defaults: self.num_envs = int(self._defaults["num_envs"]) self.batched = self.num_envs > 0 - print("[env] lazy spawn: announcing port before Isaac boot", flush=True) + if not self.contract: + # Manifest is minted from bridge.contract right after start returns. + print( + "[env] no contract found; building env to probe observation/action structure", + flush=True, + ) + await self._run_on_sim(self._sync_reset, {}) + self._ensure_contract_from_env() await super().start() async def reset(self, **task_args: Any) -> str: prompt = await self._run_on_sim(self._sync_reset, task_args) - self._ensure_contract_from_env() + self._ensure_contract_from_env() # no-op when start() already probed return prompt async def stop(self) -> None: diff --git a/hud/environment/tests/test_gym_bridge_probe.py b/hud/environment/tests/test_gym_bridge_probe.py new file mode 100644 index 000000000..d784204ac --- /dev/null +++ b/hud/environment/tests/test_gym_bridge_probe.py @@ -0,0 +1,90 @@ +"""GymBridge probes the env at start when no contract file exists.""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import numpy as np +import pytest + +pytest.importorskip("gymnasium") + +import gymnasium as gym +from gymnasium import spaces + +from hud.environment.robot.gym import GymBridge + + +class _ToyEnv(gym.Env): + """Minimal HWC camera + state env for contract probing.""" + + metadata = {"render_modes": ["rgb_array"], "render_fps": 10} + + def __init__(self) -> None: + super().__init__() + self.observation_space = spaces.Dict( + { + "pixels": spaces.Box(0, 255, shape=(4, 4, 3), dtype=np.uint8), + "state": spaces.Box(-1.0, 1.0, shape=(2,), dtype=np.float32), + } + ) + self.action_space = spaces.Box(-1.0, 1.0, shape=(3,), dtype=np.float32) + + def reset(self, *, seed: int | None = None, options: dict | None = None): + super().reset(seed=seed) + obs = { + "pixels": np.zeros((4, 4, 3), dtype=np.uint8), + "state": np.zeros(2, dtype=np.float32), + } + return obs, {} + + def step(self, action): + obs, _ = self.reset() + return obs, 0.0, False, False, {} + + +def make_toy_env() -> gym.Env: + return _ToyEnv() + + +def test_start_probes_env_when_contract_missing(tmp_path: Path) -> None: + contract_path = tmp_path / "contract.json" + bridge = GymBridge(make_toy_env, contract=contract_path) + + assert not contract_path.exists() + + async def _run() -> None: + await bridge.start() + await bridge.stop() + + asyncio.run(_run()) + + assert contract_path.exists() + assert "pixels" in bridge.contract["features"] + assert "state" in bridge.contract["features"] + assert bridge.contract["features"]["pixels"]["type"] == "rgb" + assert bridge.contract["control_rate"] == 10 + + +def test_start_loads_existing_contract_without_rebuild(tmp_path: Path) -> None: + contract_path = tmp_path / "contract.json" + contract_path.write_text( + '{"control_rate": 7, "features": {"state": {"role": "observation", "names": ["s0"]}}}\n' + ) + builds: list[int] = [] + + def make_counting_env() -> gym.Env: + builds.append(1) + return _ToyEnv() + + bridge = GymBridge(make_counting_env, contract=contract_path) + + async def _run() -> None: + await bridge.start() + await bridge.stop() + + asyncio.run(_run()) + + assert builds == [] # file present → no probe build + assert bridge.contract["control_rate"] == 7 From 21efad73f66a03fa595ad31c499545c000302949 Mon Sep 17 00:00:00 2001 From: Lukass Kellijs Date: Tue, 21 Jul 2026 04:21:21 +0000 Subject: [PATCH 41/43] docs(robots): split env-side docs into gym-style envs and custom bridges Separates env.gym (with its three target forms and EnvHub/Isaac Lab Arena examples) from the manual RobotBridge path, and links the contract on first mention. Co-authored-by: Cursor --- docs/v6/advanced/robots.mdx | 159 ++++++++++++++++++++++++++++++++---- 1 file changed, 141 insertions(+), 18 deletions(-) diff --git a/docs/v6/advanced/robots.mdx b/docs/v6/advanced/robots.mdx index 0ca4386f8..4a801768d 100644 --- a/docs/v6/advanced/robots.mdx +++ b/docs/v6/advanced/robots.mdx @@ -85,15 +85,24 @@ drives model until env terminates. - **`Model`** - the actual stateless checkpoint of the model (includes pre-/post-processing) - **`Adapter`** - translates the env's observation space to the model's, and the model's action space to the env's -**The contract** (of the environment) - the one artifact both sides share: a self-describing JSON schema of the -embodiment's control rate, observation and action spaces, carried in the capability's manifest params. +[**The contract**](#contract) (of the environment) - the one artifact both sides share: a self-describing JSON +schema of the environment including the embodiment's control rate, observation and action spaces, +carried in the capability's manifest params. The agent wires observations to policy inputs purely from the manifest; there is no shared config. +In addition, the contract serves to define useful metadata that you can easily toggle to improve +trace visualization on the platform. ## Environment side -For a gym-style sim you implement nothing: `env.gym(make_env)` takes any callable returning a -gym-style env, derives the [contract](#contract) from a sample observation, serves the sim over the -``robot`` WebSocket, and returns the handle templates drive episodes through: +HUD has two ways to build the environment side: [`env.gym`](#gym-style-envs) for anything +gym-shaped, which derives the contract, capability, and serving for you, or a +[custom `RobotBridge`](#custom-bridges) for a sim that isn't. + +## Gym-style envs + +Calling `env.gym(make_env)` derives the [contract](#contract) from a sample observation, serves +the env over the `robot` WebSocket, and returns a [`RobotEndpoint`](#custom-bridges) - the handle +templates drive episodes through: ```python env.py from hud import Environment @@ -108,14 +117,132 @@ async def pick_and_place(task: str = "default", seed: int = 0): yield await sim.result() ``` -Task args are partitioned by the factory's signature: args the factory accepts define the env (a -change rebuilds it), everything else is episodic and flows to `env.reset(seed=..., options=...)`. +`env.gym` takes one argument, in any of three forms: + +- **A factory callable**, like `make_env` above - a module-level function that returns a + gym-style env. Its signature splits task args into build args (a change rebuilds the env) and + episodic args, which flow to `env.reset(seed=..., options=...)`. +- **A gymnasium registry id**, a string such as `"CartPole-v1"` - built with `gym.make` (or + `gym.make_vec` when `num_envs` is passed). +- **An already-constructed env** made through the gymnasium registry - `env.gym` reduces it to + its spec and rebuilds it in the sim process. -A single-env sim has one slot, so the template may omit the slot token as above. A vectorized sim -(`num_envs > 1`) must thread it through so each session grades its own slot — yield +`env.gym` serves either a standard, single-instance env or a vectorized one - one gym env that +internally holds `num_envs` copies of the environment, stepped together in the same process for +GPU efficiency. Each copy is a **slot**, and a vectorized sim hands back a slot token on +`reset` so the rollout that claimed it can step and grade that copy alone, leaving the rest +untouched. + +A single-env sim has only one slot, so the template may omit the token, as above. A vectorized +sim (`num_envs > 1`) must thread it through so each session grades its own slot - yield `{"prompt": ep["prompt"], "robot": {"token": ep["token"]}}` and call `sim.result(token=ep["token"])`; -see [vectorized evals](#vectorized-envs-and-evals). -For a sim that isn't gym-shaped, subclass the **bridge** instead: +see [vectorized envs and evals](#vectorized-envs-and-evals) for the full pattern. + +You can find gym-shaped environments in various places - [LeRobot EnvHub](https://huggingface.co/docs/lerobot/en/envhub), +including Isaac Lab Arena environments, or environments built into LeRobot directly. + + +Vectorized (`num_envs=4`), with cameras enabled for RTX frames. `make_env` downloads the +[EnvHub](https://huggingface.co/docs/lerobot/en/envhub_isaaclab_arena) repo's module and calls its +`make_env` the same way LeRobot's loader does. + +```python env.py +from __future__ import annotations +import os +from hud import Environment + +HUB = "nvidia/isaaclab-arena-envs" + +def make_env(num_envs: int = 4): + from lerobot.envs.configs import IsaaclabArenaEnv + from lerobot.envs.utils import ( + _call_make_env, + _download_hub_file, + _import_hub_module, + _normalize_hub_result, + ) + + cfg = IsaaclabArenaEnv( + hub_path=HUB, + environment="gr1_microwave", + embodiment="gr1_pink", + object="mustard_bottle", + headless=True, + enable_cameras=True, # required with headless for RTX frames + camera_keys="robot_pov_cam_rgb", + num_envs=num_envs, + episode_length=50, + state_keys="robot_joint_pos", + state_dim=54, + action_dim=36, # GR1 + task="Reach out to the microwave and open it.", + ) + repo_id, _, local_file, _ = _download_hub_file(cfg.hub_path, True, None) + module = _import_hub_module(local_file, repo_id) + suites = _normalize_hub_result(_call_make_env(module, num_envs, False, cfg)) + return next(iter(next(iter(suites.values())).values())) + + +env = Environment(name="arena") +sim = env.gym(make_env, num_envs=4, fps=30) + + +@env.template() +async def episode(seed: int = 0): + ep = await sim.reset(seed=seed) + # Vectorized: token pins this rollout's slot. + yield {"prompt": ep["prompt"], "robot": {"token": ep["token"]}} + yield await sim.result(token=ep["token"]) +``` + + + +Meta-World ships inside LeRobot itself rather than as an EnvHub repo, so `make_env` calls +LeRobot's own factory directly. GymBridge keeps the env's `pixels` and `agent_pos` keys as-is. + +```python metaworld_sim.py +from __future__ import annotations + +import os + +os.environ.setdefault("MUJOCO_GL", "egl") + + +def make_env(task: str = "assembly-v3"): + from lerobot.envs.configs import MetaworldEnv + from lerobot.envs.factory import make_env as lerobot_make_env + + # Unwrap vec-of-one; pixels_agent_pos = camera + 4-D eef/gripper state. + suites = lerobot_make_env(MetaworldEnv(task=task, obs_type="pixels_agent_pos"), n_envs=1) + suite = next(iter(suites)) + task_id = next(iter(suites[suite])) + return suites[suite][task_id].envs[0] +``` + +```python env.py +from metaworld_sim import make_env + +from hud import Environment + +env = Environment(name="metaworld") +# LeRobot hardcodes render_fps=80; Meta-World control is slower - pin playback. +sim = env.gym(make_env, fps=20) + + +@env.template() +async def episode(task: str = "assembly-v3", seed: int = 0): + """One Meta-World episode; task name is a build arg (e.g. assembly-v3, dial-turn-v3).""" + ep = await sim.reset(task=task, seed=seed) + yield {"prompt": ep["prompt"]} + yield await sim.result() +``` + +Needs `metaworld` / `lerobot[metaworld]` installe. + + +## Custom bridges + +For a sim that isn't gym-shaped, subclass `RobotBridge` instead: ```python from hud.environment.robot import RobotBridge @@ -133,9 +260,8 @@ class MySimBridge(RobotBridge): return {"agentview_image": frames, "state": vecs}, done_mask ``` - -Those three methods are all you write. Under the hood the framework takes care of communication -with the agent and starting/stopping as well as stepping of the simulator at the *control rate*. +Those three methods are all you write. Under the hood the framework takes care of communication +with the agent and starting/stopping as well as stepping of the simulator at the *control rate*. - **`reset`** starts a fresh episode for a task and returns its prompt (the text the agent is given). - **`step`** applies one action and advances the sim a tick, setting `success` / `terminated` as the @@ -185,7 +311,7 @@ through. `start` / `stop` bring the bridge's socket up and down; `capability` pu binding once that URL exists (call it after `start`); `reset` begins an episode and returns its prompt; `result` returns the episode's score. It's control-plane only - the agent's observe/act loop tunnels straight to the bridge's WebSocket - and the same calls work whether the bridge is local -(shown here) or [in another process](#running-a-sim-in-another-process). +(shown here) or in another process. ```python from hud import Environment @@ -210,9 +336,6 @@ async def pick_and_place(task_id: str, seed: int = 0): yield await endpoint.result(token=ep["token"]) # this slot's {score, success, total_reward} ``` -Each session grades one slot. A vectorized sim fans N concurrent sessions across its slots; see -[vectorized evals](#vectorized-envs-and-evals). - ## Agent side The harness lives in `hud.agents.robot`. From 37dc2ab6bb8f7c70c2412150e2cdfadced7543fe Mon Sep 17 00:00:00 2001 From: Lukass Kellijs Date: Tue, 21 Jul 2026 04:21:30 +0000 Subject: [PATCH 42/43] docs(robots): rewrite vectorized envs and evals section Leads with the motivation (why one vectorized sim beats N containers), moves the working Shared() example up front, and reduces the surrounding explanation to three short notes instead of dense inline bullets. Co-authored-by: Cursor --- docs/v6/advanced/robots.mdx | 43 +++++++++++++++---------------------- 1 file changed, 17 insertions(+), 26 deletions(-) diff --git a/docs/v6/advanced/robots.mdx b/docs/v6/advanced/robots.mdx index 4a801768d..046772e72 100644 --- a/docs/v6/advanced/robots.mdx +++ b/docs/v6/advanced/robots.mdx @@ -445,23 +445,12 @@ instance rather than one you also use for unbatched runs. ## Vectorized envs and evals -A GPU sim like Isaac amortizes its cost by stepping many env copies together in one process -(`num_envs`). Getting N graded rollouts by booting N separate containers throws that away - each -one pays for its own physics from scratch. Instead, N rollouts share one vectorized instance, and -every layer keeps the plain single-env contract described above: - -- **Vectorization is sim config, not an eval parameter.** A factory that accepts `num_envs` makes - `env.gym(make_env, num_envs=8)` build an eight-way batched sim. Internally the bridge steps - every claimed slot together each tick, but each rollout still gets its own ordinary WebSocket - - one observation in, one action out, no batch dimension on the wire. -- **The agent is unchanged.** A stock `RobotAgent` drives one connection per rollout, same as a - single-env run. It claims its slot with the token the environment handed back on `tasks.start` - (`run.started["robot"]["token"]`); concurrent rollouts still coalesce GPU forwards through - [`BatchedModel`](#batching). -- **`group` picks how many rollouts; `Shared` puts them on one substrate.** Left alone, `group` - rollouts would each provision a fresh container. Wrapping the provider in - [`Shared(provider, width=N)`](/v6/reference/runtime#shared) instead provisions one and lets - all `N` reuse it: +Grading N rollouts by booting N separate containers wastes a GPU sim's whole advantage - each +container pays for its own physics from scratch, when a sim like Isaac Sim could instead step all +N copies together in one [vectorized](#gym-style-envs) instance. Running an eval against a vectorized sim needs no changes to +the agent or the env's contract, only to how the eval provisions rollouts: wrap the runtime in +[`Shared(provider, width=N)`](/v6/reference/runtime#shared) so `group` rollouts share that one +substrate instead of provisioning `N` containers. ```python from hud.eval import Shared, DockerRuntime @@ -474,15 +463,17 @@ job = await taskset.run( ) ``` -Each rollout opens its own control-channel connection and claims a slot on it: the first -`tasks.start` resets the sim and claims slot 0, later starts claim whatever is free, and a start -once every slot is taken errors instead of queuing silently. Grading is the same call every -rollout already makes - the environment resolves it to that rollout's slot and returns one -`Grade`, exactly as a single-env run would. - -`group` and `max_concurrent` are the ordinary [`Taskset.run` parameters](/v6/reference/tasks#running); -set `max_concurrent` to `width` so every slot fills and none sit idle waiting for a rollout that -never starts. +`group` and `max_concurrent` are the ordinary [`Taskset.run` parameters](/v6/reference/tasks#running). +A few things to note: + +- **The agent doesn't change.** Each rollout still drives one ordinary WebSocket connection and + calls `reset` / `result` as usual - the environment resolves each call to that rollout's own + [slot](#gym-style-envs). +- **Slots fill in order, then stop.** The first `reset` claims slot 0; once every slot is taken, a + further `reset` errors instead of queuing, so set `max_concurrent` to `width` so every slot fills. +- **EnvHub sims work unchanged.** A [LeRobot EnvHub](https://huggingface.co/docs/lerobot/en/envhub) + package's `make_env` often already returns a vectorized sim, so `env.gym(make_env, num_envs=N)` + serves it directly. ## Contract From bb698af1379c6107624bd039e6856efb65e177a5 Mon Sep 17 00:00:00 2001 From: Lukass Kellijs Date: Tue, 21 Jul 2026 04:21:46 +0000 Subject: [PATCH 43/43] docs(robots): simplify overview and recording notes, drop sim internals Tightens the Overview's environment/agent split, tightens Recording datasets wording, and removes the How sims run and Running a sim in another process sections. Co-authored-by: Cursor --- docs/v6/advanced/robots.mdx | 91 ++++--------------------------------- 1 file changed, 8 insertions(+), 83 deletions(-) diff --git a/docs/v6/advanced/robots.mdx b/docs/v6/advanced/robots.mdx index 046772e72..2ebeb42ce 100644 --- a/docs/v6/advanced/robots.mdx +++ b/docs/v6/advanced/robots.mdx @@ -31,8 +31,7 @@ pip install 'hud[robot]' ## Overview -Like with other HUD workflows there's the environment side -(server - containerized, served on the runtime) and the agent side (client - swappable, model with harness). +Like with other HUD workflows there's the environment side and the agent side. For robotics the **environment side** translates incoming actions into changes in the digital or physical environment and serves observations. The **agent side** owns the policy: it reads those observations, runs @@ -545,18 +544,6 @@ spec - the closed symbol sets and known traps - lives outside the SDK alongside -## How sims run - -The loop is lockstep - the bridge steps the sim once per received action. A simulator is usually -**thread-affine** (every touch must run on the thread that created its GL/device context), and some -runtimes - Isaac/Omniverse - must own the process main thread outright. HUD therefore serves every -sim with one process shape: **the sim owns the main thread**, serving (the control channel and the -robot WebSocket) runs on a background loop thread, and the framework queues every sim touch back to -the main thread. A cheap CPU sim blocks on that queue; when Omniverse is loaded, the main thread -also pumps Kit between touches. - -There is nothing to configure. `hud serve` and `python -m hud.environment.server` detect an -attached sim and pick the shape; any fix to the loop applies to every sim the same way. ## Recording datasets @@ -564,14 +551,14 @@ Set `agent.save = True` (wire it to a `--save` flag on your runner) to also reco `(observation, executed action)` tick into a **LeRobot v3 dataset** - the rollouts you just ran, ready to finetune a policy on. Telemetry streams either way; saving is the opt-in extra. -Recording is **agent-side**: it consumes the observations the agent already receives and the actions -it already produces, so it runs in *your* process - not the environment container. That sidesteps -sims (e.g. Isaac/RoboLab) whose dependency stack conflicts with `lerobot`; only your machine needs -`pip install 'lerobot[dataset]'`. +Recording is **agent-side**: it consumes the observations and actions the agent already sees, so +it runs in *your* process, never the environment container. Only your machine needs +`pip install 'lerobot[dataset]'` - the sim's own dependency stack (e.g. Isaac/RoboLab) never +enters the picture. -All rollouts in a process record into one shared dataset: each buffers its episode and commits -it whole, so concurrent runs (e.g. `BatchedAgent`) interleave safely. Finalized at process exit. -Destination and Hub push come from the environment: +All rollouts in a process record into one shared dataset: each buffers its episode and commits it +whole, so concurrent runs (e.g. `BatchedAgent`) interleave safely, and the dataset finalizes at +process exit. Destination and Hub push come from the environment: | Env var | Effect | |---------|--------| @@ -585,68 +572,6 @@ The [contract](#contract) drives the schema with no extra wiring: image features `task`. -## Running a sim in another process - -Serving in one process is the default (see [how sims run](#how-sims-run)), but the sim can also -live in its own process - on another machine, or built against a dependency stack the env -container can't carry. The env code stays essentially unchanged. - -`RobotEndpoint` is built for exactly this: the same control surface (`start` / `reset` / `result` / -`stop`) works whether the bridge is local or remote. - -- **Env process** - publish a *remote* handle with `RobotEndpoint.remote(host, port)`. It dials the - sim process and forwards every control call over JSON-RPC. -- **Sim process** - wrap the real bridge and expose it with - `RobotEndpoint(bridge).serve_blocking(host, port)`, the same sim-owns-main shape `hud serve` uses. - -The two planes split cleanly, which is why the agent never knows the sim is remote: - -- **Control plane** (`start` / `reset` / `result`) - JSON-RPC between the remote endpoint and the - serving process. -- **Data plane** (the agent's `observe → act` loop) - tunnels straight to the bridge's `robot` - WebSocket; the contract stays env-side. - -**Env side** - identical to the local example, but the endpoint is remote and you `connect()` to it -first: - -```python env.py -from hud import Environment -from hud.environment.robot import RobotEndpoint - -env = Environment(name="isaac-sim") -endpoint = RobotEndpoint.remote("127.0.0.1", 9100) # a handle on the bridge in the sim process - -@env.initialize -async def _up(): - await endpoint.connect() # retries until the sim process is serving - await endpoint.start() - env.add_capability(await endpoint.capability(contract=CONTRACT)) - -@env.shutdown -async def _down(): - await endpoint.close() # drops the link; does not stop the sim - -@env.template() -async def pick_and_place(task_id: str, seed: int = 0): - ep = await endpoint.reset(task_id=task_id, seed=seed) - yield {"prompt": ep["prompt"], "robot": {"token": ep["token"]}} - yield await endpoint.result(token=ep["token"]) -``` - -**Sim process** - your program builds the bridge and serves its control surface, blocking for the -process's lifetime with the sim owning the main thread: - -```python sim_main.py -from hud.environment.robot import RobotEndpoint - -bridge = MySimBridge() -RobotEndpoint(bridge).serve_blocking("127.0.0.1", 9100) -``` - -Bring the two up together - the env's `connect()` retries until the sim is listening. Everything -downstream (`hud eval`, tasksets, the agent) is unchanged; only *where the bridge runs* moved. - - ## API summary | Symbol | Where | Role |