refactor(robot): unify the robot stack - one sim process shape, batched-first bridge/agent, grouped evals#481
Draft
lukass16 wants to merge 14 commits into
Draft
refactor(robot): unify the robot stack - one sim process shape, batched-first bridge/agent, grouped evals#481lukass16 wants to merge 14 commits into
lukass16 wants to merge 14 commits into
Conversation
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.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 3d341a1. Configure here.
Contributor
|
Preview deployment for your docs. Learn more about Mintlify Previews.
|
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.
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
1 participant
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 framework maintained two parallel implementations - single-env and vectorized - with duplicated agent loops, recorders, bridges, and eval paths. Sim execution had three separate threading mechanisms, and the vectorized eval path leaked robot-specific arguments (
contract=,vectorized=) into the genericTaskset.runAPI.Solution
SimThreadreplaces the threeSimRunnerstrategies: the sim always owns the main thread, serving runs on a background thread. The same mechanism covers CPU envs and Isaac.RobotBridgeis now batched-first;VecRobotBridgeandIsaacBridgeare removed.result()returns per-slot scores under"slots".hud.telemetry.robotmodule, replacing the three previous recorder classes. LeRobot dataset writing is split intoagents/robot/dataset.py.RobotAgentdrives any number of env slots with a single loop;VecRobotAgentis removed.Taskset.run(agent, runtime=..., num_envs=N)runs vectorized evals through a genericrollout_group; the robot-specific kwargs are removed andhud/evalno longer imports robot code.env.gym(make_env)serves a gym-style sim declaratively (contract and capability derived from the factory), andhud.wrap(env)streams traces from a loop the user owns.Outcome / Verification
Roughly 1,500 lines of duplicated code removed; the diff against
mainis net negative despite the new features. Test suite and lint pass. Verified end-to-end by serving a vectorized stub env in a child process and runningTaskset.run(num_envs=3)to three graded traces. Downstream consumers (assembly_bench,bench,rl) updated.Note
High Risk
Large refactor across sim threading, the robot wire protocol (scalar vs batched terminated), eval orchestration, and telemetry/recording; behavior changes for Isaac-style sims and vectorized grading.
Overview
This PR collapses the parallel single-env and vectorized robot paths into one stack:
SimThreadreplaces the threeSimRunnerstrategies so every sim uses the same shape (sim on main, serving on a background loop).hud serve/serve_blockingpick that automatically whenenv.gym(...)is attached.Environment:
env.gym(make_env)plusGymBridgederive the contract, capability, androbotWebSocket from a gym factory; bridgeresult()now returns per-slot scores under"slots".hud.wrap(env)is exposed lazily fromhudfor in-process trace streaming.Agent:
RobotAgentis the only harness—drive(runs, client)drives N slots with one open-loop chunk loop; telemetry goes through sharedTraceRecorderinhud.telemetry.robot, with LeRobot saving inDatasetWriter(per-trace shards; vectorized runs stream telemetry only). Adapters gainadapt_chunkand batchedLeRobotAdaptersupport.Eval:
Taskset.run(..., num_envs=N)routes to genericvec_rollout, grading each trace fromslots[i];BatchedAgentdefaultsbatch_sizeto live concurrency.Reviewed by Cursor Bugbot for commit 5736eb2. Bugbot is set up for automated code reviews on this repo. Configure here.