feat(robot): OpenPI policy harness, H.264 trace video, rollout batching against one agent#425
Open
lukass16 wants to merge 12 commits into
Open
feat(robot): OpenPI policy harness, H.264 trace video, rollout batching against one agent#425lukass16 wants to merge 12 commits into
lukass16 wants to merge 12 commits into
Conversation
Docker for slow envs like Isaac Sim publishes the port before @env.initialize finishes, so hello retries can exceed 120s on slow container boots.
Add a weightless Model that queries a remote policy server over the OpenPI msgpack/WebSocket protocol: the adapter builds the request dict, the server owns all pre/post-processing + the forward, and infer() ships it and returns the [T, A] chunk. connect() is lazy and idempotent (blocks until the server is up); response_key covers "actions" (stock OpenPI) vs "action" (Cosmos).
…erence BatchedModel wraps any Model and coalesces concurrent ainfer() calls into a single stacked forward: a lazily-started worker drains up to batch_size queued calls (or flushes after max_wait_s for the suite tail), runs one inner.infer, and scatters the [N, T, A] rows back to each caller. BatchedAgent wraps a RobotAgent and shallow-clones it per run so each rollout keeps isolated episode state while sharing the one batched model. Usage stays a one-liner: BatchedAgent(agent, batch_size=8) with max_concurrent set to match.
Migrate the robot harness to OpenPI-standard, slash-delimited observation keys end-to-end, and add a thin OpenPIAdapter so a generic OpenPI policy server drives the harness with no agent code changes.
Replace per-tick JPEG observation images with per-camera H.264/CMAF video streaming for robot traces: - Add hud/agents/robot/video.py (SegmentEncoder/VideoStreamer): encode each camera on a background thread, emitting CMAF fragments as VideoSegmentStep spans without blocking the act loop. - RobotAgent starts/finalizes the streamer at the env control rate; finalize in `finally` so a crashed run still leaves video. - ObservationStep.from_obs records only numeric state now; camera frames travel as video. - Step.emit accepts an explicit trace_id so the encoder thread (no contextvars trace context) attributes spans correctly. - Add RobotClient.get_control_rate(); add "video_segment" RobotStepSource; add PyAV (av>=12) to the robot extra.
Add ModalRuntime as a Provider alongside DockerRuntime: resolve image once (from_name or lazy build), create an isolated Sandbox per rollout, expose the env control channel over raw TCP, terminate on exit. Export from hud.eval and add optional [modal] extra.
…oxes Add DaytonaRuntime as a Provider alongside ModalRuntime: resolve snapshot once (build from image if missing), create an isolated sandbox per rollout, start the env server in a background session, reach it via an asyncssh local-forward (Daytona exposes only HTTPS previews, connect dials tcp://), delete on exit. workdir defaults to /app to match the scaffolded Dockerfile.hud. Export from hud.eval and add optional [daytona] extra.
Remove the per-episode model.reset() hook (Model/LeRobotModel/RemoteModel/ BatchedModel + agent.on_episode_start); per-episode state lives only on the agent, so a shared BatchedModel can no longer clear one rollout's policy state mid-episode. Document that RemoteModel is not batchable (OpenPI server has no batched-request shape) on RemoteModel, BatchedModel, and BatchedAgent.
…ship Spell out on Model.infer/ainfer that implementations must keep the leading batch dim N (ainfer indexes [0], BatchedModel scatters rows along it) and add a one-line assert in LeRobotModel.infer. Document that BatchedAgent mutates the passed-in agent in place, leaving it permanently batched. Co-authored-by: Cursor <cursoragent@cursor.com>
Clamp get_control_rate to max(1, round(...)) so sub-0.5 Hz contracts no longer emit 0 FPS on VideoSegmentStep. Init _hooks_done before add_capability in Environment.__init__. Load optional robot deps via importlib for pyright, add shim-test ignores, and ruff-format flagged files. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes 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 446a05b. Configure here.
| except Exception: # not found: build it under this name | ||
| await daytona.snapshot.create( | ||
| CreateSnapshotParams(name=self.snapshot_name, image=self._image) | ||
| ) |
There was a problem hiding this comment.
Daytona snapshot probe swallows errors
Medium Severity
DaytonaRuntime._ensure_snapshot treats any snapshot.get failure like a missing snapshot and always calls snapshot.create. Transient API or auth errors can trigger a redundant create attempt and mark the snapshot resolved, hiding the real failure until sandbox startup.
Reviewed by Cursor Bugbot for commit 446a05b. Configure here.
Wrap long lines, move NDArray to TYPE_CHECKING, noqa intentional 0.0.0.0 bind in LocalRuntime, and reformat legacy shim test imports. Co-authored-by: Cursor <cursoragent@cursor.com>
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
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 v6 robot harness needed to drive real OpenPI policy servers, run concurrent rollouts efficiently, and stream camera data to traces without bloating each step with JPEG frames. Slow sim boots (e.g. Isaac Sim) also exceeded the default env connect timeout.
Solution
RemoteModel— WebSocket/msgpack client for OpenPI policy servers (lazy connect, supportsactions/actionresponse keys).BatchedAgent/BatchedModel— coalesce concurrentainfer()calls into stacked forwards for parallel rollouts.OpenPIAdapterso a stock OpenPI server drives the harness with no agent changes.VideoStreamer(hud/agents/robot/video.py); numeric state stays onObservationStep, frames go asVideoSegmentStepspans.RobotClientconnectready_timeoutdefault to 240s for slow container boots.lukass/modal-daytona-runtimes.Outcome / Verification
RemoteModel+OpenPIAdapterBatchedAgent(batch_size=N)video_segmentspans with playable H.264 segmentsNote
Medium Risk
Substantial changes to the robot inference contract (batch shape, removed model reset/ensembler) and new threaded video encoding could affect existing LeRobot rollouts; cloud runtimes add operational dependencies but are isolated behind optional extras.
Overview
Robot policy harness gains OpenPI server support via
RemoteModel(WebSocket client, lazy connect, configurableactions/actionresponse key) andOpenPIAdapter(passesobs['data']+ prompt).Modelis now stateless with a fixed[N, T, A]batch contract; per-episode chunking stays onRobotAgent(nomodel.reset()).LeRobotModeldropsEnsembler/ standalonelerobot_inferand asserts the batch dimension is preserved.Concurrent eval:
BatchedModel/BatchedAgentcoalesce parallelainfercalls into one stacked forward (not compatible withRemoteModel).Traces: camera frames move off per-tick JPEGs in
ObservationStepto background H.264/CMAF encoding (video.py→VideoSegmentStepspans, optionaltrace_idonStep.emitfor encoder threads).RobotClient.get_control_rate()drives video FPS.Placement:
ModalRuntimeandDaytonaRuntimeproviders;connectdefaultready_timeoutraised to 240s.robotextra adds PyAV; optionalmodal/daytonaextras inpyproject.toml.Reviewed by Cursor Bugbot for commit 51786b6. Bugbot is set up for automated code reviews on this repo. Configure here.