Play a chunk in a layer and command the harness one round at a time - #679
Play a chunk in a layer and command the harness one round at a time#679vertix wants to merge 10 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5df7ecf333
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bcf78dd251
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9ac9ee6190
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1ddc8e056a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
`Session.__call__` answers `(commands, resume_at_ns)`: the commands to emit now, and the instant the session wants its next call. `ChunkedSchedule` becomes `ChunkPlayer`: it holds the anchored chunk, drops the waypoints already passed, answers the ones due, and asks the sessions below for a new chunk in the call that drains the last one. The harness sheds its waypoint deques, `_play`, `_reschedule` and `_assert_anchored`. It emits what the session answers and paces the next round by `resume_at_ns`, still capped at the poll period. Sessions under the player keep answering chunks, and `base.py` names that contract `ChunkSession` until PR 9 makes the chunk a served function's answer. `openpi` moves `ChangeEEFrame` under the player, where a codec belongs. `episode.steps` counts the control rounds the harness ran. The golden pipeline does not move.
…found A model session answers a chunk, so the six vendor sessions move to `ChunkSession`. They stayed on `Session`, which now means "answers commands", and `blocking()`'s guard refused them: no inference server could load a checkpoint or open a websocket. `basedpyright` excludes `positronic/vendors`, so nothing caught it. A recording tap picks its side when it is made, from the contract its inner answers, so a tap no longer claims to answer commands while it forwards a chunk. It also wraps a bare-dict chunk, which `ChunkSession` allows. `ChunkPlayer` converts the offset rather than the sum, so a waypoint at 0.0 lands on the call that loads it whatever the clock reads. `probe` asks the endpoint once and waits the round trip out, instead of calling the session every 10 ms and logging the observation again each time. It closes the runtime before it reports, so the wait is not a surprise. An episode that refuses to open closes the session it opened. The eval warmup closes its runtime before the session it served. Tests: waypoints due in one round keep every channel they name; the drained chunk's last command survives the call that loads the next; the harness sleeps to the moment the session asked for; a tap above the player logs a command per round. Each one fails against the code it pins.
A `match` sequence pattern matches a two-element list, so a chunk of two actions read as a `(commands, resume_at_ns)` pair: the tap plotted the first action as a command and dropped the second. Each tap now states what its own side answers, which is what it knew when it was made. `probe` logs nothing for an endpoint that commands nothing, instead of taking the joint-velocity branch on an empty list and raising `IndexError`. It waits on the round trip without a deadline of its own: the endpoint's `infer_timeout` is the one that bounds it, and a failed round trip raises out of the next call with its own cause.
The round tests the deadline before it calls the session, and the world can reach the deadline while that call runs. `_reschedule` used to drop a chunk answered past it; `_infer` now holds back the commands instead, and `_run` finishes the trial on the next round.
Finding out which contract a session answers takes opening it, and a model session holds a backend connection from that moment. The runtime still closes first, and the session it served follows.
`deadline_ns` replaced the float `_deadline` on main while this branch was open, and the guard the round runs after a session call was written against the old name.
215f531 to
a807baa
Compare
The call that drains a chunk also loads the next, and it merged what the drained chunk still held into what the new one commands. A channel only the drained chunk named then took a setpoint from the plan that had just been replaced, while every other channel followed the new one. A driver holds what it last took, so the new chunk alone says what to command. This is what `_reschedule` did, which cleared every channel before it refilled.
The harness sleeps to `resume_at_ns` instead of a fixed poll period, bounded by the episode deadline, a 1 ms floor and a 1 s ceiling. `_EpisodeInference` asserts the instant is in the future, so `ChunkPlayer` and `StopOnFault` each name a poll period of their own while they wait for something they cannot time.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e6ded06173
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if not self._waypoints or self._waypoints[-1].time_ns <= time_ns: | ||
| chunk = self._inner(obs, time_ns) | ||
| if chunk is not None: | ||
| self._load(chunk, time_ns) | ||
| commands: dict[str, Any] = {} |
There was a problem hiding this comment.
When a ChunkSession synchronously returns its next chunk and the current chunk ends with an actionable waypoint rather than a timestamp-only sentinel, this branch calls _load() before the draining loop, replacing _waypoints and silently discarding that final command. The same session emits the final waypoint when its next answer is asynchronous (None), so playback also differs between inline and remote execution; drain the due waypoint before replacing the deque, or make a closing sentinel an enforced part of the chunk contract.
Useful? React with 👍 / 👎.
`ChunkSession.__call__` answers an `Answer` instead of a chunk-or-`None`. The one-call-in-flight state machine leaves `RemoteSession` and `LerobotPolicy._Session`, which each become one line, and moves to `ChunkPlayer`: it holds the handle, decides when to ask and when to read, and could hold more than one. A source that runs the work inside the call answers a `Done`.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 34af1a8841
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| that receives it and asks for a call at each waypoint, so the walk follows the instants it names. An | ||
| endpoint that commands nothing gives an empty list. | ||
| """ | ||
| session(obs, time.time_ns()) |
There was a problem hiding this comment.
When the first endpoint call completes synchronously—such as a Done-backed inline policy, or an executor worker that finishes before ChunkPlayer checks done()—this call can already load the chunk and drain its time-zero command. _play discards the returned (commands, resume_at_ns), so the probe omits the first waypoint or reports no commands for a single-action chunk; capture the priming call's commands before waiting and continuing the walk.
Useful? React with 👍 / 👎.
| # The server's clock is not the rig's. | ||
| actions = await asyncio.to_thread(session, raw_obs, time.time_ns()) | ||
| answer = await asyncio.to_thread(session, raw_obs, time.time_ns()) | ||
| actions = answer.result() |
There was a problem hiding this comment.
When the remote half contains codecs or recording taps, answer carries lazy Answer.map callbacks, so result() performs action decoding and Rerun logging synchronously on the ASGI event-loop thread. A large chunk or slow recording write therefore stalls every connection, including handshakes and keepalives, even though the session itself was moved to a worker; resolve the answer in asyncio.to_thread as well.
Useful? React with 👍 / 👎.
`Session` and `ChunkSession` answered different types from one method name, so every caller narrowed a union, `DelegatingChunkSession` repeated its twin, and three `isinstance` asserts stood in for a type the code could not state. A policy now declares the work of one episode under `INFER`, and `Policy.episode` opens what that work holds and closes it after. `ChunkPlayer` holds `rt.fns[INFER]` and is the bottom session. `Codec` and the chunk tap become `ChunkLayer`s, which wrap that call rather than a session, so the encode and the JPEG leave the control thread. `ChunkSession`, `AnySession`, `DelegatingChunkSession`, `Done`, `Answer.map`, `RemoteSession`, `blocking`, `Session.meta` and `Policy.functions` go with it.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7390bb2ea5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| assert not (played and isinstance(component, ChunkLayer)), ( | ||
| f'compose {type(component).__name__} under the layer that plays the chunk, not above it' | ||
| ) | ||
| played = played or component.PLAYS_CHUNKS |
There was a problem hiding this comment.
Rule hidden-dependency violated:
_ComposedLayer.wrap() rejects a ChunkLayer above the player but accepts ordinary session layers below it, even though ChunkPlayer._Policy.new_session() bypasses their make_session(). This already breaks RemotePolicy(recording_dir=...), where rec.tap('server') sits below the declared player and is silently skipped, leaving the RRD without the server observation/action chunk. Reject normal layers below the player and make that recording layer a chunk_tap so the placement requirement is enforced.
AGENTS.md reference: AGENTS.md:L7-L8
Useful? React with 👍 / 👎.
1 participant
Stage 7 of #661. The session call answers
(commands, resume_at_ns), the chunk moves out of the harness into a layer, and the caller owns the calls that layer makes.What changes
Session.__call__(obs, time_ns)answers a pair: the commands to emit now, and the instant the session wants its next call.ChunkedSchedulebecomesChunkPlayer(wire namechunk_player). It holds the anchored chunk, drops the waypoints already passed, answers the ones due, and asks the sessions below for a new chunk in the call that drains the last one.The harness sheds
_schedules,_play,_reschedule,_assert_anchoredandMAX_ACTION_SKEW_SEC. It emits what the session answers and paces the next round byresume_at_ns. That instant is what drives the loop: the episode deadline, a 1 ms floor and a 1 s ceiling bound it, and the fixed poll period is gone._EpisodeInferenceasserts the instant is in the future, soChunkPlayerandStopOnFaulteach name a poll period of their own while they wait for something they cannot time.Sessions under the player keep answering chunks, and
base.pynames that contractChunkSession. A call on it answers anAnswer— the caller's handle on the work — rather than a chunk-or-None. The one-call-in-flight state machine leavesRemoteSessionandLerobotPolicy._Session, which each become one line, and moves toChunkPlayer: it holds the handle, decides when to ask and when to read, and could hold more than one. A source that runs the work inside the call answers aDone. A TODO says PR 9 deletes the contract when the chunk becomes a served function's answer. The wire protocol is untouched.Consequences
golden_pipeline.json.gzis byte-identical. The latency simulator in that test moves under the player, where a thing that holds chunks belongs.openpimovesChangeEEFrameunder the player. A codec decodes a chunk, so it belongs there; above a player it would silently pass the commands through unconverted._CodecSessionrefuses that composition rather than corrupting it.episode.stepscounts control rounds. The harness no longer sees chunks, so it counts the rounds it called the session.rawtap loses its 3D path; theservertap keeps it, and logs when the caller reads the chunk, against the observation and timelines of the call that started the work.probewalks the chunk. It asks the endpoint once, waits the round trip out, then steps the clock through what the player emits. It draws the predicted path itself.Note for reviewers
basedpyrightexcludespositronic/vendors, so nothing checked that the vendor model sessions matched the new contract. They did not, and every inference server would have failed to load a checkpoint. Caught by review, fixed inbcf78dd2. Turning that exclusion off is its own change.Test plan
uv run --locked pytest: 1691 passed, 8 skippeduv run --locked ruff check .andruff format .: cleanbasedpyright: 0 errors; the baseline shrank by 21 entriesgolden_pipeline.json.gzis byte-identical: md54ce319042d23a50c9c9726135cd5cc3atest_in_process_equals_remote_for_same_pipelinestays green, rewritten to compare what both sides playWhat the caller-owned calls buy
RestrictImageSizecost 3 ms per 1280x720 frame on every one of those rounds. It now runs on the call that starts the work.blockingis three lines. It calls, waits, and hands back a handle that is already done.