diff --git a/docs/v6/advanced/robots.mdx b/docs/v6/advanced/robots.mdx
index 40d78d03e..0208c6ed1 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
@@ -85,33 +84,185 @@ 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
-You implement one class - the **bridge**.
+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
+
+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):
+ ep = await sim.reset(task=task, seed=seed) # {prompt, token}
+ yield {"prompt": ep["prompt"]}
+ yield await sim.result()
+```
+
+`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.
+
+`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 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
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
+ def reset(self, task_id: str, seed: int = 0) -> str:
+ ... # build the episode (runs on the sim thread)
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
```
-
-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*.
+All three hooks run on the sim thread (main under `serve_bridge`), so thread-affine sims such as
+Isaac are safe without extra hopping in your subclass.
- **`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
@@ -157,23 +308,22 @@ Actions come back the same way: the agent sends them under openpi's `actions` ke
`RobotEndpoint` is the env's control handle on the bridge - the one surface it drives an episode
-through. `start` / `stop` bring the bridge's socket up and down; `capability` publishes the `robot`
-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).
+through. Pass your bridge class (or instance); `start` spawns the sim process and serves it,
+`capability` publishes the `robot` binding (contract comes from the bridge), `reset` /
+`result` run the episode. Control-plane only - the agent's observe/act loop tunnels straight
+to the bridge's WebSocket.
```python
from hud import Environment
from hud.environment.robot import RobotEndpoint
env = Environment(name="my-sim")
-endpoint = RobotEndpoint(MySimBridge()) # the env drives the bridge only through the endpoint
+endpoint = RobotEndpoint(MySimBridge()) # start() spawns + serves the bridge process
@env.initialize
async def _up():
await endpoint.start()
- env.add_capability(await endpoint.capability(contract=CONTRACT))
+ env.add_capability(await endpoint.capability())
@env.shutdown
async def _down():
@@ -181,8 +331,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() # {"score", "success", "total_reward"}
+ 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}
```
## Agent side
@@ -208,7 +359,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
@@ -275,15 +426,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
@@ -291,6 +443,38 @@ 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 evals
+
+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
+
+job = await taskset.run(
+ MyAgent(),
+ runtime=Shared(DockerRuntime("hud-isaac-env"), width=8),
+ group=8,
+ max_concurrent=8,
+)
+```
+
+`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
Embodiments and policies disagree on cameras, state layout, action semantics, and control rate, so
@@ -361,21 +545,6 @@ spec - the closed symbol sets and known traps - lives outside the SDK alongside
-## Sim threading
-
-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.
-
-Pass one to the bridge (`RobotBridge(sim_runner=ThreadSimRunner())`), or subclass `SimRunner` for an
-exotic topology.
## Recording datasets
@@ -383,13 +552,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.
-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:
+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 |
|---------|--------|
@@ -397,92 +567,27 @@ 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`.
-## 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.
-
-`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.
-
-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):
- prompt = yield {"prompt": await endpoint.reset(task_id=task_id, seed=seed)}
- yield await endpoint.result()
-```
-
-**Sim process** - your Isaac program builds the bridge and serves its control surface, then runs for
-the process's lifetime:
-
-```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()
-
-asyncio.run(main()) # launched on the main thread the sim owns
-```
-
-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 |
|--------|-------|------|
-| `RobotEndpoint.capability(contract=...)` | `hud.environment.robot` | Build the `openpi/0` capability after `start()` |
+| `env.gym(make_env)` / `Gym` | `hud.environment.robot` | Declarative sim: contract, capability, and serving derived from a gym factory |
+| `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()` |
| `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 |
+| `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 |
+| `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 f9e6521e1..b65ab42da 100644
--- a/docs/v6/cookbooks/robot-benchmark.mdx
+++ b/docs/v6/cookbooks/robot-benchmark.mdx
@@ -29,7 +29,7 @@ endpoint = RobotEndpoint(LiberoSimBridge(use_delta=True)) # drive the bridge th
@env.initialize
async def _up():
await endpoint.start()
- env.add_capability(await endpoint.capability(contract=CONTRACT))
+ env.add_capability(await endpoint.capability())
@env.shutdown
async def _down():
@@ -37,9 +37,9 @@ 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}
+ ep = await endpoint.reset(task_suite="libero_spatial",
+ task_id=libero_task_id, init_state_id=init_state_id)
+ yield {"prompt": ep["prompt"]}
yield await endpoint.result()
```
@@ -67,8 +67,6 @@ from hud.eval import DockerRuntime, Task, Taskset
CHECKPOINT = "lerobot/pi05_libero_finetuned"
class PI05Agent(RobotAgent):
- max_steps = 400
-
def __init__(self):
device = "cuda" if torch.cuda.is_available() else "cpu"
policy = PI05Policy.from_pretrained(CHECKPOINT).to(device).eval()
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
diff --git a/hud/__init__.py b/hud/__init__.py
index 79c520916..fddc41fa8 100644
--- a/hud/__init__.py
+++ b/hud/__init__.py
@@ -60,6 +60,7 @@
"instrument",
]
+
try:
from .version import __version__
except ImportError:
diff --git a/hud/agents/robot/__init__.py b/hud/agents/robot/__init__.py
index cba9db77b..f89d88748 100644
--- a/hud/agents/robot/__init__.py
+++ b/hud/agents/robot/__init__.py
@@ -1,24 +1,19 @@
-"""Robot agent harness: drive a ``robot`` capability with a policy.
-
-The harness splits a policy rollout into three seams, each replaceable on its own:
-
-- :class:`~hud.agents.robot.agent.RobotAgent` — the loop: connect to the env's
- ``robot`` capability, observe, act, stop.
-- :class:`~hud.agents.robot.model.Model` — *how to run* the policy (preprocess →
- forward → postprocess). :class:`~hud.agents.robot.model.LeRobotModel` ships the
- LeRobot checkpoint convention.
-- :class:`~hud.agents.robot.adapter.Adapter` — translate between the env's
- observation/action spaces (from the contract) and the policy's.
-
-Wrap an agent in :class:`~hud.agents.robot.batching.BatchedAgent` to run many rollouts
-concurrently off one batched GPU forward (``max_concurrent`` rollouts, shared model).
-
-Per-tick platform tracing is emitted by the loop itself: each step records an
-:class:`~hud.agents.types.ObservationStep`, and each re-inference an
-:class:`~hud.agents.types.InferenceStep`, so runs stream live into the HUD trace viewer.
-
-This subpackage needs the ``robot`` extra (``pip install 'hud[robot]'``) for
-``numpy`` + ``msgpack``; importing :mod:`hud.agents` alone never pulls them in.
+"""Agent-side robot harness: drive a ``robot`` env with a VLA policy.
+
+- :class:`~.agent.RobotAgent` — the harness: connects to the ``robot``
+ capability (claiming a slot token from ``run.started``), reads the contract,
+ drives one scalar connection with an open-loop chunk queue. Subclass and
+ set ``self.model`` + ``self.adapter``.
+- :class:`~.model.Model` / :class:`~.model.LeRobotModel` /
+ :class:`~.model.RemoteModel` — the policy and its inference mechanics.
+- :class:`~.adapter.Adapter` / :class:`~.adapter.LeRobotAdapter` /
+ :class:`~.adapter.OpenPIAdapter` — env <-> policy space translation.
+- :class:`~.batching.BatchedAgent` — many concurrent single-env rollouts
+ sharing one batched model.
+- :class:`~.dataset.DatasetWriter` — opt-in LeRobot v3 dataset recording
+ (``agent.save = True``).
+
+This subpackage needs the ``robot`` extra (``pip install 'hud[robot]'``).
"""
from __future__ import annotations
@@ -26,16 +21,19 @@
from .adapter import Adapter, LeRobotAdapter, OpenPIAdapter
from .agent import ROBOT_PROTOCOL, RobotAgent
from .batching import BatchedAgent, BatchedModel
-from .model import LeRobotModel, Model
+from .dataset import DatasetWriter
+from .model import LeRobotModel, Model, RemoteModel
__all__ = [
"ROBOT_PROTOCOL",
"Adapter",
"BatchedAgent",
"BatchedModel",
+ "DatasetWriter",
"LeRobotAdapter",
"LeRobotModel",
"Model",
"OpenPIAdapter",
+ "RemoteModel",
"RobotAgent",
]
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 08c5fca72..8020ffca6 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,38 +53,52 @@ def adapt_observation(self, obs: dict[str, Any], prompt: str) -> Any:
"""Translate an env observation + task prompt into the policy's input."""
raise NotImplementedError
+ def adapt_chunk(self, chunk: ActionArray, obs: dict[str, Any]) -> ActionArray:
+ """Translate a freshly-inferred ``[T, A]`` chunk to env space, given the
+ (per-slot) query-time observation it was inferred from (default identity).
+
+ Called once per slot at inference time, so a chunk expressed relative to
+ the query state — e.g. joint *deltas* to be added to the query-time
+ joints — converts in one shot.
+ """
+ return chunk
+
def adapt_action(self, action: ActionArray, obs: dict[str, Any]) -> ActionArray:
- """Translate a policy action into the env's action space (default identity)."""
+ """Per-step execution-time hook on the popped action (default identity)."""
return action
class LeRobotAdapter(Adapter):
- """Vanilla LeRobot adapter for a standard image/state env.
+ """Vanilla LeRobot adapter for a standard image/state env, single or batched.
- Maps env cameras onto the model's image slots in order, converts HWC ``uint8`` to
- CHW ``float`` in ``[0, 1]``, and passes state + prompt through. Actions are identity
- (postprocess already returns env-space actions); subclass for resize/pad/reshaping.
+ Maps env cameras onto the model's image slots in order and converts HWC
+ ``uint8`` to CHW ``float`` in ``[0, 1]``; state and prompt pass through.
+ A batched observation (state ``[N, S]``) maps in one go — cameras become
+ ``[N, C, H, W]`` and the shared task is repeated to ``N``. Actions are
+ identity (postprocess already returns env-space actions).
"""
def adapt_observation(self, obs: dict[str, Any], prompt: str) -> dict[str, Any]:
import torch # pyright: ignore[reportMissingImports]
- torch_mod: Any = torch
+ torch_mod: Any = torch # torch ships no stubs; keep strict mode quiet
data = obs["data"]
+ state = np.asarray(data[self.state_key], dtype=np.float32)
+ batched = state.ndim > 1 # [N, S] vs [S]
batch: dict[str, Any] = {
- "observation.state": torch_mod.from_numpy(data[self.state_key].astype(np.float32)),
- "task": prompt,
+ "observation.state": torch_mod.from_numpy(state),
+ "task": [prompt] * len(state) if batched else prompt,
}
for model_key, env_key in zip(self.model_image_keys, self.image_keys, strict=False):
- batch[model_key] = torch_mod.from_numpy(data[env_key]).permute(2, 0, 1).float() / 255.0
+ img = torch_mod.from_numpy(np.asarray(data[env_key]))
+ perm = (0, 3, 1, 2) if batched else (2, 0, 1)
+ batch[model_key] = img.permute(*perm).float() / 255.0
return batch
- def adapt_action(self, action: ActionArray, obs: dict[str, Any]) -> ActionArray:
- return action
-
class OpenPIAdapter(Adapter):
- """unwraps obs['data'] to OpenPI wire keys, attaches prompt; actions are passthrough"""
+ """Unwraps ``obs['data']`` to OpenPI wire keys and attaches the prompt;
+ actions are pass-through."""
def adapt_observation(self, obs: dict[str, Any], prompt: str) -> dict[str, Any]:
out = dict(obs["data"])
@@ -98,8 +106,4 @@ def adapt_observation(self, obs: dict[str, Any], prompt: str) -> dict[str, Any]:
return out
-__all__ = [
- "Adapter",
- "LeRobotAdapter",
- "OpenPIAdapter",
-]
+__all__ = ["ActionArray", "Adapter", "LeRobotAdapter", "OpenPIAdapter"]
diff --git a/hud/agents/robot/agent.py b/hud/agents/robot/agent.py
index 0b05310f4..ced8f7c0f 100644
--- a/hud/agents/robot/agent.py
+++ b/hud/agents/robot/agent.py
@@ -1,169 +1,167 @@
-"""Base v6 agent for any env that exposes a ``robot`` capability.
+"""``RobotAgent`` — drive one ``robot`` env with an open-loop chunk queue.
-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 (claiming a slot token from
+``run.started`` when present), read the contract, run until the env terminates.
-The base calls the adapter and model at the right moments::
+Vectorized sims still look like one scalar connection per rollout; concurrent
+rollouts coalesce GPU forwards through :class:`~.batching.BatchedModel`.
- setup_robot -> adapter.bind(spaces) # once after connect
- on_episode_start -> adapter.reset() # per episode; model is stateless
- select_action -> adapt_observation -> model.ainfer -> pop chunk -> adapt_action
-
-``model.ainfer`` always returns a ``[T, A]`` chunk; :meth:`RobotAgent.select_action`
-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
from collections import deque
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar, cast
import numpy as np
from hud.agents.base import Agent
from hud.capabilities.robot import RobotClient
-
-from .record import Recorder
+from hud.telemetry.robot import TraceRecorder
if TYPE_CHECKING:
from hud.eval.run import Run
- from ._types import ActionArray
- from .adapter import Adapter
+ from .adapter import ActionArray, Adapter
from .model import Model
ROBOT_PROTOCOL = "openpi/0"
class RobotAgent(Agent):
- """Drive a ``robot`` side-channel for one :class:`~hud.client.Run`.
+ """Drive a ``robot`` env 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).
"""
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 to disk
- #: (the ``--save`` flag). Telemetry streams regardless; see :mod:`.record`.
+ #: 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.
+ #: 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: Recorder | 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 = Recorder(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 scalar robot connection.
+
+ ``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__")
- 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)
-
- 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)
+ prompt = run.prompt
+ if not isinstance(prompt, str):
+ 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").
+ # 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: str | None = None
+ if isinstance(robot_info, dict):
+ raw_token = cast("dict[str, Any]", robot_info).get("token")
+ if isinstance(raw_token, str):
+ token = raw_token
+
+ robot = await RobotClient.connect(run.client.binding(self.robot_protocol), token=token)
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}"
+ action_space, obs_space = robot.spaces()
+ if self.adapter is not None:
+ self.adapter.bind(action_space, obs_space)
+ self.adapter.reset()
+
+ obs = await robot.get_observation()
+ fps = robot.get_control_rate()
+ # Contract labels: obs_space for state, action names for InferenceStep plots.
+ recorder = TraceRecorder(
+ run=run,
+ fps=fps,
+ obs_space=obs_space,
+ action_names=action_space.get("names"),
+ )
+ writer = None
+ if self.save:
+ from .dataset import DatasetWriter
+
+ writer = DatasetWriter(robot.contract, fps=fps)
+
+ print(f"[agent] episode started: {prompt!r}", flush=True)
+ try:
+ await self._loop(
+ robot,
+ obs,
+ prompt,
+ recorder,
+ writer,
+ max_steps=self.max_steps if max_steps is None else max_steps,
)
- 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"
+ 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:
- if self._recorder is not None:
- self._recorder.end() # flush video tails + commit the LeRobot episode
- await client.close()
+ await robot.close()
+ run.trace.status = "completed"
+ run.trace.content = "done"
+
+ async def _loop(
+ self,
+ robot: RobotClient,
+ obs: dict[str, Any],
+ prompt: str,
+ recorder: TraceRecorder,
+ writer: Any,
+ *,
+ max_steps: int,
+ ) -> None:
+ """Open-loop chunk queue: ainfer refills, then execute one action per tick."""
+ model = self.model
+ assert model is not None
+ adapter = self.adapter
+ chunk: deque[ActionArray] = deque()
+
+ for step in range(max_steps):
+ # 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 bool(np.asarray(obs["terminated"]).reshape(-1)[0]):
+ if step:
+ print(f"[agent] terminated at step {step}", flush=True)
+ break
+
+ 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))
+ 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)
+ action = adapter.adapt_action(action, obs)
+ if writer is not None:
+ writer.add(obs["data"], np.asarray(action), task=prompt)
+ await robot.send_action(action)
+
+ if self.log_every and step % self.log_every == 0:
+ 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)
__all__ = ["ROBOT_PROTOCOL", "RobotAgent"]
diff --git a/hud/agents/robot/batching.py b/hud/agents/robot/batching.py
index a24594488..1e9ec40a4 100644
--- a/hud/agents/robot/batching.py
+++ b/hud/agents/robot/batching.py
@@ -18,28 +18,29 @@
if TYPE_CHECKING:
from hud.eval.run import Run
- from ._types import ActionArray
+ from .adapter import ActionArray
from .agent import RobotAgent
class BatchedModel(Model):
"""Coalesce concurrent ``ainfer`` calls into one stacked ``inner.infer``.
- A lazily-started worker drains up to ``batch_size`` queued calls (or waits up to
- ``max_wait_s`` for stragglers — which avoids stalling when fewer rollouts are live,
- e.g. the tail of a suite), stacks them into one ``[N, ...]`` batch, runs a single
- forward, and scatters the ``[N, T, A]`` rows back to each caller.
+ 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 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.
+ ``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__(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 +65,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
@@ -97,22 +98,18 @@ 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: 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__(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
diff --git a/hud/agents/robot/dataset.py b/hud/agents/robot/dataset.py
new file mode 100644
index 000000000..a80b1f218
--- /dev/null
+++ b/hud/agents/robot/dataset.py
@@ -0,0 +1,218 @@
+"""Opt-in LeRobot v3 dataset writing for robot rollouts.
+
+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. The shared
+dataset (and an ``atexit`` finalizer that flushes open buffers) is created on
+the first frame. A class lock serializes commits so episodes stay contiguous.
+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.util
+import logging
+import os
+import threading
+import time
+import uuid
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar
+
+import numpy as np
+
+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 not _is_image(f) and f.get("dtype") != "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 _is_image(f):
+ key, dtype = f"observation.images.{leaf}", "video"
+ elif leaf == "state" or single_state:
+ key = "observation.state"
+ else:
+ key = f"observation.{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 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 _is_image(feature):
+ return ["height", "width", "channel"]
+ return [f"{base}_{i}" for i in range(int((feature.get("shape") or [1])[0]))]
+
+
+class DatasetWriter:
+ """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 under ``_lock``.
+ _ds: ClassVar[Any | None] = None
+ _root: ClassVar[Path | None] = None
+ _repo_id: ClassVar[str] = ""
+ # Serialize create / add_frame / save_episode / finalize across rollouts.
+ _lock: ClassVar[threading.RLock] = threading.RLock()
+
+ 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._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]')"
+ )
+
+ 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
+ # 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)
+ 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
+ # Open the shared dataset on the first frame so atexit can flush if we
+ # die before end_episode (still in-memory until then; kill -9 loses it).
+ with DatasetWriter._lock:
+ self._ensure_dataset()
+ self._frames.append(row)
+
+ def end_episode(self) -> None:
+ """Commit this rollout's buffered episode to the shared dataset.
+
+ Whole episodes stay contiguous: ``_lock`` serializes create / add_frame /
+ save_episode so concurrent BatchedAgent rollouts cannot interleave frames.
+ """
+ if not self._frames:
+ return
+ with DatasetWriter._lock:
+ ds = self._ensure_dataset()
+ for row in self._frames:
+ ds.add_frame(row)
+ ds.save_episode()
+ self._frames.clear()
+
+ def finalize(self) -> None:
+ """Flush, write the parquet footer + optionally push to the Hub. Idempotent."""
+ with DatasetWriter._lock:
+ # Re-entrant: end_episode takes the same lock when frames remain.
+ self.end_episode()
+ ds, DatasetWriter._ds = DatasetWriter._ds, None
+ if ds is None:
+ return
+ root, repo_id = DatasetWriter._root, DatasetWriter._repo_id
+ ds.finalize()
+ print(f"[agent] saved LeRobot dataset -> {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/{repo_id}", flush=True)
+ except Exception as exc:
+ logger.exception("HF push failed for %s", repo_id)
+ print(
+ f"[agent] WARNING: HF push failed: {exc!r} (dataset still on disk)", flush=True
+ )
+
+ def _ensure_dataset(self) -> Any:
+ """Return the process-shared dataset, creating it on first frame.
+
+ Caller must hold ``_lock`` (create is check-then-act on ``_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 + 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)
+ 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.
+ DatasetWriter._ds = lerobot_dataset.LeRobotDataset.create(
+ repo_id=DatasetWriter._repo_id,
+ fps=self._fps,
+ features=self._features,
+ root=DatasetWriter._root,
+ robot_type=self._contract.get("robot_type"),
+ use_videos=True,
+ )
+ 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"]
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/record.py b/hud/agents/robot/record.py
deleted file mode 100644
index 2a053710a..000000000
--- a/hud/agents/robot/record.py
+++ /dev/null
@@ -1,230 +0,0 @@
-"""Per-episode recording for robot rollouts — telemetry, plus an optional LeRobot dataset.
-
-The agent loop hands every tick to one :class:`Recorder`. It always streams the telemetry
-the HUD viewer needs (an :class:`~hud.agents.types.ObservationStep` of numeric state +
-per-camera H.264 video); when ``save`` is on it *also* appends each
-``(observation, executed action)`` pair to a LeRobot v3 dataset for offline
-training/finetuning.
-
-Saving is opt-in (the agent's ``save`` flag — the ``--save`` runner flag), so the heavy
-LeRobot/PyAV imports stay deferred until a dataset is actually built. One dataset spans the
-whole run (every episode the shared agent drives appends to it) and is finalized at process
-exit, optionally pushed to the HF Hub. Destination + push come from the environment:
-
-- ``RECORD_DIR`` — dataset root (default ``./data`` from where the rollout launched)
-- ``HF_REPO`` — HF namespace to also push to (needs ``HF_TOKEN``)
-- ``HF_PRIVATE`` — push the dataset private
-"""
-
-from __future__ import annotations
-
-import atexit
-import importlib.util
-import logging
-import os
-import time
-import uuid
-from pathlib import Path
-from typing import TYPE_CHECKING, Any
-
-import numpy as np
-
-from hud.agents.types import InferenceStep, ObservationStep
-from hud.telemetry.context import get_current_trace_id
-
-from .video import VideoStreamer
-
-if TYPE_CHECKING:
- from numpy.typing import NDArray
-
- from hud.capabilities.robot import RobotClient
- from hud.eval.run import Run
-
-logger = logging.getLogger(__name__)
-
-
-def _lerobot_features(contract: dict[str, Any]) -> tuple[dict[str, dict[str, Any]], dict[str, str]]:
- """Map a robot contract to LeRobot ``features`` + a wire-key -> LeRobot-key map.
-
- Image obs -> ``observation.images.`` (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 Recorder:
- """Records one agent's rollouts: always telemetry, optionally a LeRobot dataset.
-
- The agent owns a single instance for its lifetime and routes *all* recording through
- it: :meth:`begin`/:meth:`end` bracket each episode, :meth:`record_observation` /
- :meth:`record_inference` / :meth:`record_action` feed each tick (the first two write
- telemetry steps onto the run passed to :meth:`begin`; the last completes a LeRobot
- frame), and :meth:`save` (also an ``atexit`` hook) finalizes the cross-episode dataset.
- With ``save=False`` only the telemetry path runs and the LeRobot deps are never imported.
- """
-
- def __init__(self, client: RobotClient, *, save: bool = False) -> None:
- self._obs_space = client.spaces()[1]
- self._fps = client.get_control_rate()
- self._contract = client.contract
- # Telemetry is always on; saving also needs lerobot installed.
- if save and importlib.util.find_spec("lerobot") is None:
- logger.warning(
- "save=True but lerobot is not installed; streaming telemetry only "
- "(pip install 'lerobot[dataset]')"
- )
- save = False
- self._save = save
- self._features: dict[str, dict[str, Any]] = {}
- self._key_map: dict[str, str] = {}
- if save:
- self._features, self._key_map = _lerobot_features(self._contract)
-
- self._video: VideoStreamer | None = None # per-episode
- self._run: Run | None = None
- self._task = ""
- self._pending: dict[str, Any] | None = None # last obs awaiting its action
- # LeRobot dataset spans every episode; created lazily on the first frame.
- self._ds: Any | None = None
- self._root: Path | None = None
- self._repo_id = ""
- if save:
- atexit.register(self.save) # finalize even on an abrupt exit (parquet footer)
-
- # ── episode lifecycle (called from the agent harness) ─────────────────────
- def begin(self, run: Run, prompt: str) -> None:
- """Open an episode: fresh per-camera video stream + the task prompt."""
- self._run = run
- self._task = prompt
- self._pending = None
- self._video = VideoStreamer(fps=self._fps, trace_id=get_current_trace_id())
-
- def record_observation(self, obs: dict[str, Any], *, tick: int) -> None:
- """One observation: numeric-state span + per-camera video (always streamed)."""
- assert self._run is not None and self._video is not None # set in begin()
- self._run.record(ObservationStep.from_obs(obs, tick=tick, obs_space=self._obs_space))
- self._video.record(obs)
- self._pending = obs.get("data") # paired with the action in record_action()
-
- def record_inference(self, chunk: NDArray[Any], *, tick: int) -> None:
- """One re-inference: the freshly inferred ``[T, A]`` action chunk, onto the run."""
- assert self._run is not None # set in begin()
- self._run.record(InferenceStep(tick=tick, chunk=chunk.tolist(), chunk_length=len(chunk)))
-
- def record_action(self, action: NDArray[Any]) -> None:
- """The executed (env-space) action: completes the pending LeRobot frame."""
- if self._save and self._pending is not None:
- self._add_frame(self._pending, action)
- self._pending = None
-
- def end(self) -> None:
- """Close the episode: flush video tails; commit the LeRobot episode (if any frames)."""
- if self._video is not None:
- self._video.finalize()
- if self._ds is not None and self._ds.has_pending_frames():
- self._ds.save_episode()
-
- def save(self) -> None:
- """Finalize the dataset (writes the parquet footer) + optionally push to the Hub.
-
- Idempotent; registered with ``atexit`` so the dataset stays loadable even if the
- process exits without an explicit call.
- """
- if not self._save or self._ds is None:
- return
- self._save = False # idempotent across the explicit call + the atexit hook
- self._ds.finalize()
- print(f"[agent] saved LeRobot dataset -> {self._root}", flush=True)
- if not os.environ.get("HF_REPO"):
- return
- private = os.environ.get("HF_PRIVATE", "0") not in ("0", "", "false", "False")
- try: # best-effort: the on-disk dataset is the source of truth
- self._ds.push_to_hub(private=private)
- print(f"[agent] pushed -> https://huggingface.co/datasets/{self._repo_id}", flush=True)
- except Exception as exc:
- logger.exception("HF push failed for %s", self._repo_id)
- print(f"[agent] WARNING: HF push failed: {exc!r} (dataset still on disk)", flush=True)
-
- # ── LeRobot writing ───────────────────────────────────────────────────────
- def _add_frame(self, data: dict[str, Any], action: NDArray[Any]) -> None:
- ds = self._ensure_dataset()
- row: dict[str, Any] = {}
- for wire, key in self._key_map.items():
- value = data.get(wire)
- if value is None:
- logger.warning("obs missing contract feature %r; skipping frame", wire)
- return
- ft = self._features[key]
- row[key] = (
- np.ascontiguousarray(value, dtype=np.uint8) # bridge images are uint8 HWC
- if ft["dtype"] in ("video", "image")
- else np.asarray(value, dtype=ft["dtype"]).reshape(ft["shape"])
- )
- act_ft = self._features["action"]
- row["action"] = np.asarray(action, dtype=act_ft["dtype"]).reshape(act_ft["shape"])
- row["task"] = self._task
- ds.add_frame(row)
-
- def _ensure_dataset(self) -> Any:
- if self._ds is not None:
- return self._ds
- lerobot_dataset: Any = importlib.import_module("lerobot.datasets.lerobot_dataset")
-
- name = self._contract.get("robot_type") or "robot"
- stamp = time.strftime("%Y%m%d_%H%M%S")
- # Unique per recorder so concurrent (batched) rollouts never share a root;
- # tie it to the trace id when there is one so a shard maps back to its trace.
- tag = (get_current_trace_id() or uuid.uuid4().hex)[:8]
- # Default under ./data (relative to where the rollout was launched), created if absent.
- record_dir = Path(os.environ.get("RECORD_DIR", "data"))
- record_dir.mkdir(parents=True, exist_ok=True)
- self._root = record_dir / f"{name}_{stamp}_{tag}"
- self._repo_id = f"{os.environ.get('HF_REPO') or 'hud'}/{name}_{stamp}_{tag}"
- # LeRobotDataset.create requires a fresh root; images encode to per-episode video.
- self._ds = lerobot_dataset.LeRobotDataset.create(
- repo_id=self._repo_id,
- fps=self._fps,
- features=self._features,
- root=self._root,
- robot_type=self._contract.get("robot_type"),
- use_videos=True,
- )
- print(f"[agent] recording LeRobot dataset -> {self._root}", flush=True)
- return self._ds
-
-
-__all__ = ["Recorder"]
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/capabilities/robot.py b/hud/capabilities/robot.py
index 25af6a337..052ff5bc1 100644
--- a/hud/capabilities/robot.py
+++ b/hud/capabilities/robot.py
@@ -15,7 +15,7 @@
import asyncio
import contextlib
-from typing import Any, ClassVar, Self
+from typing import Any, ClassVar, Self, cast
import numpy as np
import websockets
@@ -68,36 +68,43 @@ 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)
- # Consume the connect-time metadata frame (always first); a string frame
- # is the env's error convention.
+ async def connect(cls, cap: Capability, *, token: str | None = None) -> Self:
+ """Dial the robot WebSocket; ``token`` claims a sim slot after the metadata frame.
+
+ 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 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(cast("Any", _packb({"claim": token})))
return cls(cap, ws)
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:
raise RuntimeError(f"robot env error:\n{msg['error']}")
- terminated = bool(msg.pop("terminated", False))
+ # 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
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/env.py b/hud/environment/env.py
index b3072e78c..4ba23bd2b 100644
--- a/hud/environment/env.py
+++ b/hud/environment/env.py
@@ -305,6 +305,30 @@ async def _down() -> None:
return ws
+ def gym(self, target: Any, *, name: str = "robot", **kwargs: Any) -> Any:
+ """Attach a gym-style sim serving ``name`` over the ``robot`` protocol.
+
+ ``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
+ from hud.environment.robot.gym import gym_command
+
+ sim = RobotEndpoint.spawn(gym_command(target, **kwargs))
+
+ @self.initialize
+ async def _up() -> None:
+ await sim.start()
+ self.add_capability(await 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 538e0c784..02c1bc19e 100644
--- a/hud/environment/robot/__init__.py
+++ b/hud/environment/robot/__init__.py
@@ -1,29 +1,38 @@
-"""Env-side robot runtime: the ``robot`` bridge + its building blocks.
+"""Env-side robot runtime: bridges, the sim program, and the control endpoint.
-This package holds 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.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:`~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` — 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:`~.gym.GymBridge` — 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
+ (``RobotEndpoint(MyBridge)`` spawns it; custom containers can call it too).
+- :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` (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 RobotBridge
+from .bridge import RobotBridge, serve_bridge
from .endpoint import RobotEndpoint
-from .sim_runner import InlineSimRunner, MainThreadSimRunner, SimRunner, ThreadSimRunner
+from .gym import GymBridge, TracedEnv, wrap
__all__ = [
- "InlineSimRunner",
- "MainThreadSimRunner",
+ "GymBridge",
"RobotBridge",
"RobotEndpoint",
- "SimRunner",
- "ThreadSimRunner",
+ "TracedEnv",
+ "serve_bridge",
+ "wrap",
]
diff --git a/hud/environment/robot/bridge.py b/hud/environment/robot/bridge.py
index 2fc5303c2..961b31fdf 100644
--- a/hud/environment/robot/bridge.py
+++ b/hud/environment/robot/bridge.py
@@ -1,176 +1,549 @@
-"""Env-side ``robot`` bridge: the base class users subclass to wrap their sim.
+"""Sim-process ``robot`` bridge: WebSocket obs/action + JSON-RPC control.
-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.
-
-An injected :class:`~.sim_runner.SimRunner` owns *which thread runs the
-(thread-affine) sim*, so subclasses stay thread-naive.
+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
+import asyncio
import contextlib
+import queue
+import secrets
+import signal
+import sys
+import threading
from abc import ABC, abstractmethod
+from concurrent.futures import Future
+from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
+import numpy as np
import websockets
import websockets.exceptions
# 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 .sim_runner import InlineSimRunner, SimRunner
+from hud.environment.utils import error, read_frame, reply, send_frame
if TYPE_CHECKING:
- import numpy as np
+ from collections.abc import Callable
+
+#: 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="
+
+
+@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 # hold next step (dialing timed out, or WS dropped)
+ # 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 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]
+
+ 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 ``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:`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`, 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). All three hooks run on the sim thread via :meth:`_run_on_sim`
+ (under :func:`serve_bridge`, that is main — safe for Isaac / thread-affine sims).
- :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.
+ - :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.
"""
- def __init__(
- self,
- *,
- host: str = "127.0.0.1",
- port: int = 0,
- sim_runner: SimRunner | None = None,
- ) -> None:
+ #: Seconds to wait for a *still-dialing* claimed slot before stepping with a
+ #: hold. Connected slots never hit this — slow inference must not advance the
+ #: sim with zeros.
+ 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] = {}
- # 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``.
+ #: 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] = {}
+ # 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)
+ 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
+ # 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, **kwargs: Any) -> Any:
+ """Run ``fn(*args, **kwargs)`` 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, **kwargs) # not serving yet, or already on the sim thread
+ fut: Future[Any] = Future()
+ self._sim_q.put((lambda: fn(*args, **kwargs), 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:
+ self.total_reward = 0.0
+ self.success = False
+ self.terminated = False
+ # Same sim-thread hop as step/obs — custom bridges must not touch Isaac here
+ # on the background serve loop.
+ self.task_description = await self._run_on_sim(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")
+ token = self._registry.claim(slot)
+ return {"prompt": self.task_description, "token": token}
+
+ async def _release_episode(self, token: str | None) -> dict[str, Any]:
+ """Control-plane result: this slot's score, then free it.
- 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
+ Scores are written in ``step`` on the sim thread — read them there too
+ so grading never races a mid-step update on the serve loop.
+ """
+ slot = self._registry.resolve(token)
+ grade = (await self._run_on_sim(self.result_slots))[slot.index]
+ self._registry.release(slot)
+ return grade
@abstractmethod
- async def reset(self, **kwargs: Any) -> str:
+ 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.
+ Always invoked on the sim thread (via :meth:`_run_on_sim`) — implement
+ this as ordinary sync code, same as :meth:`step` / :meth:`get_observation`.
"""
@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], bool] | 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(self) -> dict[str, Any]:
- """Return the episode score dict after the episode ends.
+ def result_slots(self) -> list[dict[str, Any]]:
+ """One score dict per slot. Default: the scalar episode grade for every slot
+ (vectorized bridges with per-slot scoring override this).
- 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.
+ Invoked on the sim thread via :meth:`_run_on_sim` (same as ``step``).
"""
- return {
+ 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)."""
+ action = next(
+ (f for f in self.contract.get("features", {}).values() if f.get("role") == "action"),
+ {},
+ )
+ # 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
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")
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._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_runner.call(self.step, action) # on the sim thread
- await self._send_observation()
+ # 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)
+ if isinstance(raw, str):
+ raise RuntimeError(raw)
+ claim = _unpackb(raw)
+ 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
+ slot.action = None
+ slot.idle = False
+ # First frame after claim: one sim read, then this slot's scalar row.
+ await self._send_slot_observation(slot, await self._run_on_sim(self.get_observation))
+ 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
+ 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 _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:
+ async def _tick_loop(self) -> None:
+ """Gather claimed slots' actions (or dialing-timeout → hold), step once, fan out."""
+ while True:
+ with contextlib.suppress(TimeoutError):
+ await asyncio.wait_for(self._action_event.wait(), timeout=self.step_timeout)
+ self._action_event.clear()
+ 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
+ # Barrier over every *claimed* slot: dialing (no WS yet) blocks too, so
+ # its episode can't advance before the first observation. step_timeout
+ # only holds still-dialing slots — a connected agent mid-inference waits.
+ 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
+ dialing = [s for s in pending if s.ws is None]
+ connected_pending = [s for s in pending if s.ws is not None]
+ if dialing:
+ remaining = deadline - asyncio.get_running_loop().time()
+ if remaining <= 0:
+ for s in dialing:
+ s.idle = True # never connected → hold so the batch can move
+ if not connected_pending:
+ break
+ # Connected agents still deciding — keep waiting without a deadline.
+ continue
+ self._action_event.clear()
+ with contextlib.suppress(TimeoutError):
+ await asyncio.wait_for(self._action_event.wait(), timeout=remaining)
+ else:
+ # Only live agents pending: wait for an action (or a disconnect).
+ self._action_event.clear()
+ await self._action_event.wait()
+ claimed = self._registry.claimed()
+ if not any(s.ws is not None for s in claimed):
+ break
+ 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).
+ rows = []
+ for s in self._registry.slots:
+ 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)
+ # One step + one batched obs on the sim thread, then fan-out (not N reads).
+ await self._run_on_sim(self.step, actions)
+ batch = await self._run_on_sim(self.get_observation)
+ for s in claimed:
+ await self._send_slot_observation(s, batch)
+
+ async def _send_slot_observation(
+ self,
+ slot: _Slot,
+ batch: tuple[dict[str, np.ndarray], np.ndarray] | None,
+ ) -> None:
+ """Fan one scalar obs frame to a claimed connection from a batched read."""
+ if slot.ws is None or batch is None:
return
- data, terminated = out
- # openpi-style flat obs dict: array fields at the top level, terminated alongside.
- msg = {**data, "terminated": bool(terminated)}
+ data, terminated = batch
+ 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) ────────────────────
+
+ 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 ensure_contract(self) -> dict[str, Any]:
+ """Return the wire contract, deriving it when a subclass can.
+
+ Default is the already-set ``self.contract``. Lazy-spawn bridges
+ (:class:`~.gym.GymBridge`) override to build the env and derive when
+ no pre-written ``contract.json`` was loaded at start — so
+ ``endpoint.capability()`` publishes a real manifest, not ``{}``.
+ """
+ return self.contract
+
+ async def _dispatch_control(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
+ if method == "url":
+ return {"url": self.url}
+ if method == "contract":
+ # May trigger env build under lazy spawn when the contract is still empty.
+ return {"contract": await self.ensure_contract()}
+ if method == "reset":
+ return await self._claim_episode(**params)
+ if method == "result":
+ token = params.get("token")
+ if token is not None and not isinstance(token, str):
+ raise ValueError("result: 'token' must be a string when given")
+ return await self._release_episode(token)
+ raise ValueError(f"unknown method {method!r}")
+
+
+# ── 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.
+
+ 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)
+ 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()
+
+ def _thread() -> None:
+ asyncio.set_event_loop(loop)
+ try:
+ with contextlib.suppress(asyncio.CancelledError):
+ loop.run_until_complete(_serve())
+ finally:
+ done.set()
+
+ 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()
+
+ # 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)
+
+
+def main() -> None:
+ """Child entry: ``python -m hud.environment.robot.bridge path.py:Class [--init JSON]``."""
+ import json
+
+ from hud.utils.modules import load_module
+
+ path, _, name = sys.argv[1].rpartition(":")
+ kwargs: dict[str, Any] = {}
+ if len(sys.argv) >= 4 and sys.argv[2] == "--init":
+ kwargs = json.loads(sys.argv[3])
+ serve_bridge(getattr(load_module(path), name)(**kwargs))
+
+
+if __name__ == "__main__":
+ main()
-__all__ = ["RobotBridge"]
+__all__ = ["PORT_ANNOUNCEMENT", "RobotBridge", "serve_bridge"]
diff --git a/hud/environment/robot/endpoint.py b/hud/environment/robot/endpoint.py
index 28517920d..e84cccaec 100644
--- a/hud/environment/robot/endpoint.py
+++ b/hud/environment/robot/endpoint.py
@@ -1,177 +1,184 @@
-"""``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):
+ ep = await sim.reset(task=task, seed=seed) # {prompt, 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
import asyncio
import contextlib
+import inspect
+import json
+import sys
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, RobotBridge
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
+
+
+def _bridge_init_kwargs(bridge: RobotBridge) -> dict[str, Any]:
+ """JSON-safe subclass ``__init__`` kwargs from a declaration instance (skip base host/port)."""
+ base = set(inspect.signature(RobotBridge.__init__).parameters) - {"self"}
+ out: dict[str, Any] = {}
+ for name, param in inspect.signature(type(bridge).__init__).parameters.items():
+ if name == "self" or name in base:
+ continue
+ if param.kind not in (param.POSITIONAL_OR_KEYWORD, param.KEYWORD_ONLY):
+ continue
+ attr = name if hasattr(bridge, name) else f"_{name}"
+ if not hasattr(bridge, attr):
+ continue
+ value = getattr(bridge, attr)
+ try:
+ json.dumps(value) # stream sinks / recorders are not child-portable
+ except TypeError:
+ continue
+ out[name] = value
+ return out
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 a bridge class/instance (``start()`` spawns its process),
+ :meth:`spawn` (explicit argv), or :meth:`remote` (attach).
"""
def __init__(
self,
- bridge: RobotBridge | None = None,
+ bridge: type[RobotBridge] | RobotBridge | None = None,
*,
+ cmd: Sequence[str] | None = None,
host: str | None = None,
port: int | None = None,
+ connect_timeout_s: float = 900.0,
) -> None:
- self._bridge = bridge # set => local; None => remote (dial host:port)
+ if bridge is not None:
+ # Child re-imports the class; instance ctor kwargs ride along as --init JSON.
+ cls = bridge if isinstance(bridge, type) else type(bridge)
+ name = getattr(cls, "__qualname__", "")
+ if not name or "." in name or "<" in name:
+ raise ValueError(f"bridge class must be module-level, got {bridge!r}")
+ cmd = [
+ sys.executable,
+ "-m",
+ "hud.environment.robot.bridge",
+ f"{inspect.getfile(cls)}:{name}",
+ ]
+ if isinstance(bridge, RobotBridge) and (kwargs := _bridge_init_kwargs(bridge)):
+ cmd += ["--init", json.dumps(kwargs)]
+ 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
+ # N sessions share this one TCP link; serialize send+read so replies don't cross.
+ self._lock = asyncio.Lock()
@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
-
- # ── 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)::
+ 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)
- @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 = 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)
- 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
-
- """ 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.
-
- 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
-
- 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
@@ -181,30 +188,67 @@ 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.
+
+ Under lazy spawn this may build the env and derive ``contract.json`` when
+ none was pre-written — so a capability published at initialize is complete.
+ """
+ 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) -> 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 | 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')} "
+ 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.
+ # 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 connect() 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"]
+ raise RuntimeError("not connected; call start() first")
+ 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:
+ 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
new file mode 100644
index 000000000..ad31b62bd
--- /dev/null
+++ b/hud/environment/robot/gym.py
@@ -0,0 +1,687 @@
+"""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:
+
+- :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:`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::
+
+ from hud.environment.robot import wrap
+
+ env = wrap(make_env(...), job="chess-eval")
+
+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
+
+import atexit
+import inspect
+import itertools
+import json
+import logging
+import sys
+from pathlib import Path
+from typing import Any, Self
+
+import numpy as np
+
+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.
+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)."""
+ 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
+ ``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
+
+
+# ── 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) ───────────────────────
+
+ def _load_contract_if_present(self) -> None:
+ """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)
+ 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 from a sample observation once the env exists."""
+ 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(
+ 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)
+
+ async def ensure_contract(self) -> dict[str, Any]:
+ """Return the wire contract, probing the env when start() did not already.
+
+ ``start()`` mints the contract for ``env.gym`` publish; this covers a
+ bare ``contract`` RPC (tests / custom callers) that skips that path.
+ """
+ if not self.contract:
+ await self._run_on_sim(self.reset)
+ return self.contract
+
+ async def start(self) -> None:
+ """Bind the wire with a complete contract so ``env.initialize`` can publish it.
+
+ 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
+ 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.reset)
+ await super().start()
+
+ def reset(self, **task_args: Any) -> str:
+ """Build/reset the gym env on the sim thread (base claim hops here)."""
+ # 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
+ self._ensure_contract_from_env() # no-op when start() already probed
+ return self._prompt(merged)
+
+ 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 _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
+
+
+# ── 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,
+ 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,
+ ) -> None:
+ self.env = env
+ self._batched = is_batched(env)
+ 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
+ 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)
+ # 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]
+ self._rec.record(
+ obs=state or None,
+ frames=frames or None,
+ action=action,
+ reward=np.atleast_1d(to_numpy(reward)),
+ done=done,
+ success=probe_success(info, num_envs=self._num_envs),
+ )
+ 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()
+
+ 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()}
+ path = self._contract_path
+ wrote = path is not None and not path.exists()
+ contract = load_or_write_contract(
+ path,
+ sample,
+ frames,
+ action_dim_of(self.env, batched=self._batched),
+ self._fps,
+ )
+ if wrote:
+ logger.info("wrap: wrote %s (edit names to relabel plots)", path)
+ features = contract.get("features", {})
+ return JobRecorder(
+ self._job,
+ 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 features.values() if f.get("role") == "action"),
+ None,
+ ),
+ state_names={
+ k: f["names"]
+ for k, f in features.items()
+ if f.get("role") == "observation" and f.get("names")
+ },
+ )
+
+ def __enter__(self) -> Self:
+ return self
+
+ def __exit__(self, *exc: object) -> None:
+ self.close()
+
+
+# The public verb: ``from hud.environment.robot import wrap`` - 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.
+ # JSON-encoded so the child gets real bools/ints/strings back (see main()).
+ for key, value in defaults.items():
+ cmd += [f"--{key.replace('_', '-')}", json.dumps(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. JSON
+ # round-trip restores real types; bare strings (hand-written argv) pass as-is.
+ for flag, value in itertools.pairwise(unknown):
+ if flag.startswith("--"):
+ 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)
+
+
+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",
+ "probe_success",
+ "split_observation",
+ "wrap",
+]
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"]
diff --git a/hud/environment/server.py b/hud/environment/server.py
index f4e264c22..a20732739 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(
@@ -418,6 +459,16 @@ 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.
+
+ 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.gym`), so every env serves the same way.
+ """
+ asyncio.run(_serve_until_terminated(env, host, port))
+
+
def main() -> None:
from hud.environment import load_environment
@@ -429,9 +480,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__":
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)
diff --git a/hud/eval/__init__.py b/hud/eval/__init__.py
index 2824645c7..ee53010fa 100644
--- a/hud/eval/__init__.py
+++ b/hud/eval/__init__.py
@@ -46,6 +46,7 @@
RuntimeGPU,
RuntimeLimits,
RuntimeResources,
+ Shared,
SubprocessRuntime,
)
from .sync import SyncPlan
@@ -69,6 +70,7 @@
"RuntimeGPU",
"RuntimeLimits",
"RuntimeResources",
+ "Shared",
"SubprocessRuntime",
"SyncPlan",
"Task",
diff --git a/hud/eval/run.py b/hud/eval/run.py
index 5949786bd..62728b012 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(
diff --git a/hud/eval/runtime.py b/hud/eval/runtime.py
index a05e16ab1..85671321c 100644
--- a/hud/eval/runtime.py
+++ b/hud/eval/runtime.py
@@ -150,6 +150,55 @@ 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 hard cap on
+ concurrent occupancy (e.g. a robot sim's ``num_envs``) — a further acquire
+ raises. Pair 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) -> AsyncIterator[Runtime]:
+ async with self._lock:
+ if self._refs >= self.width:
+ raise RuntimeError(
+ f"Shared(width={self.width}) already has {self._refs} concurrent "
+ "callers; raise max_concurrent no higher than width"
+ )
+ if self._addr is None:
+ stack = contextlib.AsyncExitStack()
+ await stack.__aenter__()
+ try:
+ self._addr = await stack.enter_async_context(self.inner(task))
+ except BaseException:
+ await stack.aclose() # failed boot — don't orphan a half-open stack
+ raise
+ self._stack = stack # publish only after inner succeeds
+ 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 +1265,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 2ef0d223a..0de511a7b 100644
--- a/hud/eval/taskset.py
+++ b/hud/eval/taskset.py
@@ -23,7 +23,14 @@
from .job import Job, job_enter
from .run import rollout
-from .runtime import HostedRuntime, HUDRuntime, LocalRuntime, _declared_env, _declared_names
+from .runtime import (
+ HostedRuntime,
+ HUDRuntime,
+ LocalRuntime,
+ Shared,
+ _declared_env,
+ _declared_names,
+)
from .sync import fetch_taskset_tasks, resolve_taskset_id
if TYPE_CHECKING:
@@ -268,6 +275,11 @@ async def run(
one id. Returned ``job.runs`` preserves expansion order (task-major,
then group).
+ ``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`` equal to both ``group`` and ``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
as a failed/errored run so one wedged rollout (e.g. a stuck sampling
@@ -280,8 +292,8 @@ async def run(
if max_concurrent is not None and max_concurrent < 1:
raise ValueError("max_concurrent 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).
expanded: list[tuple[Task, str]] = []
task_list = list(self)
for task in task_list:
@@ -307,22 +319,31 @@ 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()
+ # Shared substrates have exactly ``width`` slots: a smaller cap starves the
+ # barrier; a larger one over-claims and fails mid-batch on reset.
+ if isinstance(placement, Shared) and 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) -> 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)]
+ 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:
@@ -335,7 +356,8 @@ async def _one(task: Task, group_id: str) -> Run:
group,
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.
diff --git a/hud/telemetry/robot/__init__.py b/hud/telemetry/robot/__init__.py
new file mode 100644
index 000000000..67c0482e2
--- /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, ``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..d94a3b4f2
--- /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, ``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