feat(robot,eval): unified robot stack + Shared runtime for concurrent rollouts#481
Open
lukass16 wants to merge 56 commits into
Open
feat(robot,eval): unified robot stack + Shared runtime for concurrent rollouts#481lukass16 wants to merge 56 commits into
lukass16 wants to merge 56 commits into
Conversation
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Contributor
|
Preview deployment for your docs. Learn more about Mintlify Previews.
|
lukass16
marked this pull request as draft
July 7, 2026 00:10
The sim always owns the process main thread; serving (control channel + robot WebSocket) runs on a background loop thread, with every sim touch queued back to main. Cheap CPU envs block on the queue; Isaac/Omniverse gets its Kit pump from the idle hook - no per-sim threading choices. serve_blocking is the single entry hud serve and the module runner use, replacing the serve_pumped/asyncio.run branching.
…y.robot Recording is shared infrastructure (agent harness, hud.wrap, gym bridges), not agent code. TraceRecorder (one trace, span emission only - lifecycle stays with the caller) and JobRecorder (a vectorized env as one Job of per-episode traces) replace Recorder/VecRecorder/EpisodeRecorder and their dual-mode flags; video streaming moves alongside. The LeRobot dataset half of record.py becomes agents/robot/dataset.py (DatasetWriter). InferenceStep gains contract-derived per-dim action names so the viewer can label plots.
…mBridge
One bridge serves num_envs slots in lockstep; a plain single env is a
batch of one speaking the scalar wire framing on the same code path, so
VecRobotBridge and IsaacBridge collapse into the base. GymBridge is the
generic bridge over any gym-style factory (arg partitioning by factory
signature, rebuild-on-instance-change, success probing, torch/int action
handling). Grading becomes result() -> {score, success, slots: [...]}
(one dict per slot), replacing the result_batch plumbing through the
endpoint's JSON-RPC. SimRunner strategies are gone - bridges route every
sim touch through the shared SimThread - and the endpoint gains
serve_blocking() for split-process sims (replacing serve_forever).
env.gym(make_env) turns any gym-style factory into a served robot sim: contract derived from a sample observation (round-tripped through an editable contract.json), capability minted, lifecycle wired to the env's hooks, episodes driven via the returned handle (sim.reset / sim.result). A factory accepting num_envs is the vectorization declaration. hud.wrap is the loop-owning counterpart: wrap an env you drive yourself and every episode streams to the platform as a trace under one job (lazy top-level attr so core hud stays free of the robot extra).
One open-loop chunk queue drives single and vectorized envs alike: the wire framing (scalar terminated vs [N] mask) sets the batch size, spent slots refill from one batched forward, adapt_chunk converts per slot at inference time. __call__(run) is the generic rollout contract (a group of one, still coalescing through ainfer so BatchedAgent works unchanged); drive(runs, client) is the grouped-eval entry recording spans per slot. VecRobotAgent and VecLeRobotAdapter are deleted - LeRobotAdapter maps batched and unbatched observations with the same wiring.
…vs=) One vectorized env instance becomes num_envs graded runs sharing a group_id, behind the same Taskset.run surface and with no robot concepts in hud/eval: the atom goes through the normal env client and template (num_envs injected into task args; an env that doesn't accept it fails loudly), the contract rides the capability manifest, the agent's drive(runs, client) entry drives the slots, and run i grades from the template result's slots[i]. Replaces vec_rollout and the vectorized=/ contract= kwargs; num_envs composes with group= (k instances x N traces, group ids per instance) and requires a self-managed runtime.
env.gym-first environment side, the one sim-serving process shape (SimRunner section replaced), slots-based grading, serve_blocking for split-process sims, a new vectorized-envs-and-grouped-evals section, and a refreshed API summary.
batch_size is now an optional cap: with no value the worker stacks every ainfer call queued within the max_wait_s window, so the forward is as wide as whatever is in flight and max_concurrent never needs manual pairing. A set cap still flushes the window early once reached.
"group" already means statistical repeats (group=k); the vectorized atom deserves its own word. Docstrings and error messages now say "vectorized" throughout and the atom's docstring is tightened. Also narrows the minted trace id before trace_enter (pyright).
hud/agents is strict-typed in CI: restore the stubless-torch Any alias in LeRobotAdapter, type the chunk queues as deque[ActionArray], narrow the optional model to a local before the loop, and suppress the deliberate private _client check when picking recorder mode.
Cameras in derived gym contracts are tagged type: rgb (no dtype: image), so frames mapped to non-video LeRobot columns. Image detection now accepts both conventions; missing dtypes default to float32 and missing shapes are filled from the first real frame before the dataset is created. Also imports importlib.util explicitly (pyright).
Finished slots keep acting so the [N, A] action frame stays complete, but their chunk refills no longer emit InferenceSteps - matching the existing observation guard, so a finished slot's trace gains no orphan inference.
Adds a Taskset.run parameter reference (runtime / group / num_envs / max_concurrent / rollout_timeout / job) under the vectorized-evals section, updates the batching example for the self-sizing BatchedModel, and aligns wording with the vec_rollout rename.
lukass16
force-pushed
the
lukass/phys-experimental
branch
from
July 18, 2026 01:15
7bd800c to
7710838
Compare
Key the control channel's suspended tasks by session id: each connection starts and grades its own task instead of the last start cancelling whatever any other connection had in flight. A connection drop parks the session's task; a later connection grades it by resuming the session (hello with its session_id) or, when exactly one session is parked, by a plain tasks.grade — so the split 'hud task start' / 'hud task grade' flow keeps working unchanged, and the ambiguous case errors loudly instead of guessing.
max_steps was configurable per-subclass but every override just repeated the same default; keep one constant (520) via the call-site argument instead of a class attribute agents had to remember to set. Co-authored-by: Cursor <cursoragent@cursor.com>
…ative gym targets The env server no longer forks itself around a sim; every sim runs as a separate process, always the same way. env.gym(...) spawns it (or a custom sim program calls serve_bridge directly), and the env server's serving path (server.serve_blocking) is a plain asyncio.run — no more sim-main special case. - bridge.py absorbs the sim program (formerly serve.py): serve_bridge, gym_command, and the CLI (`python -m hud.environment.robot.bridge`) now live next to RobotBridge/GymBridge, since both ends of the argv format belong together. - GymBridge / env.gym(target) take three declarative target forms: a module-level factory (unchanged), a gymnasium registry id, or a constructed registry env — reduced by gym_command to its EnvSpec JSON and closed, since only the declaration can cross the fork; the sim process rebuilds it via gym.make/gym.make_vec. - sim_thread.py -> sim.py: the sim-process shape (SimThread/run_with_sim) is unconditional infrastructure every sim uses, not a special case to fork around, so it's named and documented as the one shape, not an alternative. - introspect.py + wrap.py merge into gym.py: one module for the gym integration both GymBridge and hud.wrap share (observation splitting, contract derivation, TracedEnv). - endpoint.py (RobotEndpoint) and env.py (Environment.gym) updated for the merged module layout and broadened target type. Co-authored-by: Cursor <cursoragent@cursor.com>
Generic hook for per-episode data an env hands back on start (e.g. a robot bridge slot token), without hud/eval learning robot concepts. Co-authored-by: Cursor <cursoragent@cursor.com>
N sessions now share one control-channel TCP link, so _call needs a
lock around send+read to keep replies from crossing. reset() returns
the full {prompt, token} reply and result(token=...) passes it through,
matching the bridge's per-slot claim/release surface.
Co-authored-by: Cursor <cursoragent@cursor.com>
Replaces the single-agent, batched-wire bridge with N claimable slots
behind one scalar openpi WebSocket per connection:
- _SlotRegistry: reset-on-first-claim (global reset + claim slot 0 when
all free), claim a free slot otherwise, error when none are free.
A released slot stays unreclaimable until the next global reset
(v1 keeps whole-batch episodes).
- Each connection's first frame is {"claim": token}; the bridge binds
it to that slot and fans out that slot's scalar observation only.
- _tick_loop barriers claimed, connected slots: gather pending actions
(or step_timeout -> hold for stragglers), step once as [N, A], fan
scalar frames back out. One slow agent can no longer stall the batch.
- Control surface: reset() claims + returns {prompt, token}; result()
reads one slot's grade by token and frees it. The old aggregate
"slots" grade list is gone.
GymBridge's build-arg partition now also takes env.gym(..., num_envs=8)
style defaults, so num_envs is sim build config, not an eval parameter.
Co-authored-by: Cursor <cursoragent@cursor.com>
…eview - bridge: a dropped connection marks its slot idle so the lockstep barrier doesn't stall the other slots for step_timeout - bridge/client: metadata advertises claim_required; connecting without a token fails fast instead of blocking forever in get_observation - agent: recorder.close()/writer.end_episode() run in finally, so video tails flush and buffered episodes commit even when the rollout raises - agent: the terminal observation is recorded before breaking (traces and datasets keep the final frame) - gym: build defaults cross the argv as JSON, so bools/ints/floats reach the child process with their real types Co-authored-by: Cursor <cursoragent@cursor.com>
A None claim (connect or result) resolves to the sole claimed slot, so
single-env templates drop the token plumbing; ambiguous vectorized claims
still error. Templates: yield {"prompt": ...} and sim.result() suffice.
… sim boots GymBridge.start announces the port and loads a pre-written contract immediately; the env builds on first reset. Endpoint connect timeout raised to 900s for cold Isaac boots.
GymBridge.start builds with factory defaults to mint the contract before the capability is published, so env.initialize never binds an empty manifest. Existing contract files still skip the probe (lazy spawn). Co-authored-by: Cursor <cursoragent@cursor.com>
Separates env.gym (with its three target forms and EnvHub/Isaac Lab Arena examples) from the manual RobotBridge path, and links the contract on first mention. Co-authored-by: Cursor <cursoragent@cursor.com>
Leads with the motivation (why one vectorized sim beats N containers), moves the working Shared() example up front, and reduces the surrounding explanation to three short notes instead of dense inline bullets. Co-authored-by: Cursor <cursoragent@cursor.com>
Tightens the Overview's environment/agent split, tightens Recording datasets wording, and removes the How sims run and Running a sim in another process sections. Co-authored-by: Cursor <cursoragent@cursor.com>
- Only still-dialing slots time out to hold; connected agents mid-inference wait - Derive contract on the contract RPC so env.gym publishes a non-empty manifest - Make reset sync and hop it through _run_on_sim; fan out one batched obs read Co-authored-by: Cursor <cursoragent@cursor.com>
Lock shared create/add_frame/save_episode across BatchedAgent rollouts, and pass contract action names into TraceRecorder for InferenceStep plots. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Clean up a half-open AsyncExitStack when inner provisioning fails, refuse acquires beyond width, and require Taskset max_concurrent to equal width so shared substrates are neither underfilled nor over-claimed. Co-authored-by: Cursor <cursoragent@cursor.com>
The SimThread-era serve_blocking switch in the CLI is obsolete now that env.gym spawns its own sim process; drop the unrelated serve.py diff from the robot PR. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Format gym.py, tighten Shared/RobotAgent/RobotClient typing, and apply small ruff cleanups in the bridge so lint-ruff and lint-pyright pass. Co-authored-by: Cursor <cursoragent@cursor.com>
Pass a module-level bridge class/instance; start() forks python -m hud.environment.robot.bridge so authors need not call serve_bridge themselves. Docs drop the removed capability(contract=) kwarg. Co-authored-by: Cursor <cursoragent@cursor.com>
lukass16
changed the title
…rame Spawn encodes JSON-safe subclass __init__ attrs as --init so configured instances (e.g. use_delta=True) reach the child. DatasetWriter creates the shared LeRobot dataset and atexit flush on the first add(). Co-authored-by: Cursor <cursoragent@cursor.com>
result_slots is filled in step on the sim thread; hop the result RPC through _run_on_sim so grading cannot race a mid-step update. Co-authored-by: Cursor <cursoragent@cursor.com>
Let adapters declare how many predicted actions to execute before replanning, applied by RobotAgent after adapt_chunk so subclasses stay truncation-free. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 5928734. Configure here.
| 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 |
There was a problem hiding this comment.
Task cancel leaves sim slots claimed
High Severity
Robot episodes now claim a sim slot on control reset and only free it on result. tasks.cancel, bye, and rollout cancellation call TaskRunner.cancel, which closes the task generator without running the template’s result RPC, so claimed slots (and their tokens) stay occupied on the shared sim.
Additional Locations (2)
Reviewed by Cursor Bugbot for commit 5928734. Configure here.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Issue
The robot stack duplicated single-env vs vectorized paths (agents, bridges, recorders, eval), and vectorized eval leaked robot-specific knobs (
num_envs,vec_rollout,contract=) into generic APIs. Sims also had three threading strategies, custom bridges required hand-writtenserve_bridgeprograms, and concurrent rollouts needed a first-class way to share one substrate at a fixed width.Solution
Robot (env + agent)
serve_bridgeon main, wire on a side loop;_run_on_simfor thread-affine sims.RobotBridgewith slot tokens, barrier stepping, and optional tokenless single-env claims.env.gym(make_env)/GymBridgewith contract derive/load; lazy port announce and non-empty manifest whencontract.jsonis missing.RobotEndpoint(MyBridge):start()spawns and serves the bridge;capability()fetches the contract (nocontract=).RobotAgentloop; shared lockedDatasetWriter; telemetry inhud.telemetry.robot.step_timeoutonly holds still-dialing slots (slow inference does not get zeros).Eval / runtime
Shared(provider, width=N)— refcount-shares one substrate across concurrent rollouts; replacesnum_envs/vec_rollout.Taskset.runenforcesmax_concurrent == widthwhen using Shared.Run.startedcarries the fulltasks.startpayload (robot token metadata).hello(session_id=...).Outcome / Verification
One robot surface for gym and custom bridges; vectorized evals go through Shared instead of robot-specific eval kwargs. Docs updated accordingly. Stub-bridge smoke for auto-spawn endpoint. Lint/typecheck cleaned on the branch.
Note
High Risk
Large architectural change to robot wire protocol (slot claims), sim process lifecycle, and eval placement; concurrent shared substrates and session grading affect core rollout correctness.
Overview
This PR collapses the robot stack onto one env shape (sim in a child process,
serve_bridgeon main) and one agent shape (scalar WebSocket + open-loop chunks), with slot tokens threading throughendpoint.reset→ taskrobotmetadata →RobotClient.connect(token=...).Environment:
Environment.gym(make_env)spawnsGymBridge, derives or loadscontract.json, and publishes capabilities without a hand-passedcontract=. Custom bridges use batched[N,...]obs/actions, JSON-RPC control, barrier stepping, and claim/release per slot.RobotEndpointalways talks to a remote sim process (spawn or attach);SimRunnerand the old in-process / separate-process Isaac doc path are removed.Agent:
RobotAgentis simplified into a single_loopwithadapt_chunk/ optionalchunk_size, batchedLeRobotAdapter, and defaultmax_steps. Telemetry moves tohud.telemetry.robot.TraceRecorder; LeRobot recording is a new thread-safeDatasetWriter(replacingRecorder).BatchedAgent/BatchedModeldefault batch width to live concurrency.Eval: New
Shared(provider, width=N)refcount-shares one substrate;Taskset.runrequiresmax_concurrent == widthwhen using it.Run.startedholds the fulltasks.startpayload for robot tokens. Control-channel sessions are per-connection with park-on-disconnect andhello(session_id=...)resume.Docs and cookbooks are updated for
env.gym, vectorized evals, and the removed threading/remote-bridge patterns.Reviewed by Cursor Bugbot for commit 5928734. Bugbot is set up for automated code reviews on this repo. Configure here.