Make Rollout the per-rollout unit and withhold its scene from the policy - #620
Make Rollout the per-rollout unit and withhold its scene from the policy#620vertix wants to merge 10 commits into
Rollout the per-rollout unit and withhold its scene from the policy#620Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bab8091ed8
ℹ️ 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".
| instruction = task.instruction | ||
| self._policy_inputs = {keys.TASK: instruction} if instruction is not None else {} | ||
| self.context = context | self._policy_inputs | ||
| self._policy_session = self.policy.new_session(self.context, clock.now) |
There was a problem hiding this comment.
Withhold scene metadata from session creation
When a policy or wrapper reads its new_session context, it receives eval.seed, task IDs, and every other scene field because self.context combines the scene with _policy_inputs. For example, SampledPolicy.new_session dispatches using this context, so policy selection can condition on the ground truth despite observations being filtered. Pass only _policy_inputs to new_session while retaining the full context for reset, recording, and completion bookkeeping.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Deferring this one, with the reasoning — the concern is right, but the fix as written regresses a live path.
new_session's context has exactly one consumer today: SampledPolicy.new_session, which reads the fields named in its group_fields to pick a sub-policy and to keep the counter's buckets balanced. No policy implementation forwards that context to a model, and RemotePolicy.new_session drops it entirely — so nothing in the session context crosses the wire. The observation dict is the only path to a model, and that is what this PR closed.
The catch is which fields the sampler needs. cfg/policy.py's production config groups by [keys.TASK, 'eval.object', 'eval.tote_placement', 'eval.external_camera'] — operator-staged scene fields that arrive on the RUN payload and now land in Task.scene. Narrowing new_session to _policy_inputs would leave the sampler with task alone and silently collapse phail's grouping into one bucket, which is a worse failure than the one being fixed: it is invisible in a run's output.
Deciding which scene fields are policy-facing is exactly #464's step 15 (routed policy — "the model is a per-trial input"), where the model field and the sampler's keys get a declared home on the rollout. That is where this belongs, and I would rather leave the leak visible and named than trade it for a silent sampling regression.
Leaving this thread open for @vertix to weigh in.
There was a problem hiding this comment.
I went and traced this properly rather than leaving it on the earlier deferral, and the conclusion is that the suggested fix would introduce a silent bug. Declining it, but the docstring it contradicts is corrected in 2663c2f.
Where the context actually goes. I read every new_session and wrap_session in the tree:
ChunkedSchedule,TemporalStack,RecordingTapandCodecall takecontextand ignore it, so nothing in the wrapper stack reads it.RemotePolicy.new_sessiondrops it entirely — it opens a websocket session carrying no payload. Nothing about the scene crosses the wire.SampledPolicy.new_sessionis the only reader, and it passes the context tosampler.sample(keys, ctx, counter.counts(keys, ctx)). BothUniformSampler.sampleandBalancedSampler.sampleignore theircontextargument; the sole consumer isEpisodeCounter._group(), which reduces it totuple(context.get(f) for f in group_fields)— the tally bucket the balancer reads.
So the scene selects no policy. It picks which per-condition tally the balancer balances within, which is stratification the experimenter configures — the deliberate use, not a leak.
Why the suggested split breaks. counts() is called with the new_session context; record() is called from _finalize_recording with self.context. Passing _policy_inputs to the first while "retaining the full context for completion bookkeeping" makes the two group by different keys, so counts() would read a bucket record() never writes — every tally reads 0, BalancedSampler weights go uniform, and the balancing silently stops working with nothing to notice it. The production config makes this concrete: cfg/policy.py groups by [task, eval.object, eval.tote_placement, eval.external_camera], and those three eval.* fields are operator-staged scene fields that now live in Rollout.scene.
What was fair in the finding. My own comment overstated the boundary — it said the scene was "withheld" full stop, when it does reach session construction. That is now stated accurately:
Only the rollout's ``instruction`` reaches an observation — its scene goes to ``reset``, to the
recording, and to ``new_session``, where a sampling policy groups by it, so no model is shown the
ground truth its rollout is scored against.
The invariant the PR actually establishes is that no model sees the scene, and that holds: the scene never enters an observation, and new_session's context reaches no model.
Leaving this open — if you want the sampler's grouping fed from a declared source rather than the recorded context, that is the routed-policy work in #464 step 15, and worth doing there with the counter's read and write paths moved together.
| reset: Callable[[dict[str, Any]], None] | None = None, | ||
| static_meta: dict[str, Any] | None = None, | ||
| on_episode_complete: Callable[[Session, dict[str, Any]], None] | None = None, | ||
| inference_latency: bool | float = False, |
There was a problem hiding this comment.
Use a domain type for inference latency
Rule primitive-type violated:
inference_latency: bool | float makes False mean disabled, True mean measured latency, and a float mean fixed latency, forcing consumers to rederive overloaded primitive semantics; introduce a domain type for these modes and convert the CLI value at the boundary.
AGENTS.md reference: AGENTS.md:L7-L8
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Declining as out of scope.
bool | float is not new here — it is the shape _inference_latency already had on main; this PR moved where it arrives (a Harness constructor argument instead of a key on every RUN context) without touching what the values mean.
The value comes off a configuronic CLI flag, which ast.literal_evals --inference_latency=True to a bool and --inference_latency=0.2 to a float. A domain type means building a conversion seam at that boundary for a sim-only debug knob, which is more than this change is for. Worth doing, but as its own change rather than folded into the eval-spine one.
Leaving open so it is not lost.
There was a problem hiding this comment.
Correcting my earlier reply first: I declined this as pre-existing, and that was only half right. The bool | float annotation and the is True check are both on main, but this PR promoted the value into a public constructor parameter, so it is code I touched and "pre-existing" is not the reason to leave it.
Still declining, on the merits this time.
The rule targets a primitive whose overloaded meaning every consumer re-derives. Here there is exactly one consumer — a single is True branch in _inference_delay — because False and 0.0 are already behaviourally identical (float(False) == 0.0, and the caller gates on delay > 0.0). So the domain is two modes, not three.
The reason a domain type does not pay for itself is that both boundaries are genuinely bool | float. The CLI flag is documented as --inference_latency=True and --inference_latency=0.2 (docs/evaluation.md), and the value is recorded into episode statics under keys.INFERENCE_LATENCY. A Fixed | Measured pair would be constructed at the first boundary and destructured at the second — for the recorded scalar there is no clean answer, since "measured" has no seconds to record — so the union moves rather than disappears, and a .delay() polymorphic call replaces one is True.
If this is worth doing, the version that actually pays is making the recorded fact honest — recording the latency each episode was charged rather than the knob it was configured with — which is a change to what the eval records, not a retyping. Happy to file that separately; leaving this open for your call.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0af5931190
ℹ️ 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".
Task the per-rollout unit and withhold its scene from the policyRollout the per-rollout unit and withhold its scene from the policy
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 23eb34a0c8
ℹ️ 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: 9bbf1da552
ℹ️ 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".
ccc26f5 to
63f270c
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 63f270c520
ℹ️ 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: 2663c2fa43
ℹ️ 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: e9e0564556
ℹ️ 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: 8166966882
ℹ️ 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: d3e5e4e3b6
ℹ️ 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".
| build_rollouts(lambda: proxy.meta[META_INSTRUCTION], timeout, seed, rollout_count, scenes), | ||
| reset=proxy.reset, |
There was a problem hiding this comment.
Couple instruction resolution to scene reset
Rule hidden-dependency violated:
The LIBERO rollout’s instruction lambda reads proxy.meta, which is unavailable until the separate reset=proxy.reset has run, so the Rollout works only when a caller knows and preserves the Harness’s ordering; direct inspection or another runner can hit meta read before the first reset. Have scene staging return the resolved instruction, or otherwise make the resolver explicitly consume the reset result, so the dependency is represented in the interface.
AGENTS.md reference: AGENTS.md:L14-L22
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Valid, and it is a regression this PR introduced rather than an inherited one — holding it for the author's call rather than landing an interface change unilaterally.
Before this change, Task carried instruction and reset together, so "the lambda resolves once the reset beside it has run" was one object's internal contract. Splitting the per-rollout unit from the eval's wiring put the resolver on Rollout and the reset on Eval, so a Rollout can no longer resolve its own instruction. That is the coupling you are pointing at, and I made it.
The mechanics work out for the fix you suggest: RemoteEnvControlSystem.reset assigns self._meta = self._frame['meta'] synchronously, so it can return the staged scene's meta rather than leaving it to be fetched later. That would make it
reset: Callable[[dict], dict] # stage the scene, return what staging reported
instruction: str | Callable[[dict], str] | None # resolved from that reportwith the Harness passing one to the other, and the ordering expressed in the types instead of in its call order.
Two notes on severity, neither of which makes the finding wrong. proxy.meta asserts meta read before the first reset, so a consumer that gets this wrong is broken loudly rather than silently — the safe side of the rule's own distinction. And it reaches Eval, Rollout, Harness, the proxy, and every eval config, including the ones whose reset reports nothing and would have to say so.
Leaving this open pending that decision.
An eval config now returns a list of `Eval`, one per embodiment, so a selection spanning several sims — `.sim.all` — is a single `--eval=` against one warm policy. The Worlds still rebuild one at a time, so only the env server of the eval in flight is ever alive. `Task` becomes the rollout: instruction, timeout and the `scene` payload the eval's `reset` stages from. The RUN-context dicts are gone, and the scene wiring that was task-invariant — `reset`, `privileged`, `done` — moves onto `Eval`. A RUN directive is read as one such rollout, so attended and self-driven episodes are the same thing downstream. The policy boundary is now structural: only the instruction crosses to the policy. `eval.seed`, `eval.task_id` and the rest of the scene are recorded in episode statics but never enter the observation dict — the old `inputs.update(self.context)` shipped them and relied on codecs to ignore them. Trial index and count are derived from the plan rather than stamped into each context, and `inference_latency` becomes a run-level Harness argument instead of riding every trial. The droid eval stops recording a random `eval.seed` for a scene that has none.
`eval.seed` was spelled in three modules and read with `.get`, so a typo on either side drew an unseeded scene in silence rather than failing. It becomes `keys.EVAL_SEED` and both readers now subscript. The backend-specific scene keys move onto the adapter that unpacks them — LIBERO's five, RoboLab's two — so the eval config writes the same names its reset token reads. Also: the policy-facing subset is now built once and merged into the recorded context, rather than each dict testing the instruction separately; `_run_world`'s wiring arguments are keyword-only, so a call site no longer reads as a row of bare positionals.
The tests and the two test env adapters still spelled the seed key, and both adapters read it with `.get`, which is the silent-miss the constant exists to prevent. `EVAL_SEED`'s comment now says what the key means rather than naming the actuators that read it.
`Task` named one execution attempt while `eval.task_id`, `eval.task` and `proxy.meta['task']` all name the scenario being attempted — three meanings in one constructor call. The type is `Rollout`, `Eval.rollouts` is the plan, and "task" goes back to meaning the benchmark task. The count and position follow the type: `--eval.trial_count` becomes `--eval.rollout_count`, and the recorded statics become `eval.rollout_index` / `eval.rollout_count`. Docs that spell the flag are updated with it.
…keys intrinsically `workflows/nebius/eval.sh`'s own help still printed the old `--eval.trial_count`, so copying the example it prints would have had configuronic reject the run. The scene-key and eval-statics comments said which code writes and reads them; they now say what the keys mean. The `eval.*` block also claimed every episode carries all of them, which is wrong twice over — an unbounded rollout has no timeout, and a directive-driven one has no plan position. `_directive_rollout` moves into `Harness` beside the directive dispatch that is its only caller, and the harness tests read the eval statics through their constants.
The integration tests read `eval.rollout_index` and its neighbours as literals two lines below an assertion that already went through `keys.EVAL_TERMINATED`, so the block was inconsistent rather than pinning anything.
A harness holding a rollout plan stamped every episode with the plan's current position, so an episode a RUN directive carried was recorded at index -1 before the plan started, or at the preceding rollout's index between planned ones. The position is now read where the plan advances and passed in, and a directive episode is given none. `_rollout_terminal` follows the unit it evaluates, and the boundary comments say what holds: the scene does reach `new_session`, where a sampling policy groups by it, and reaches no model.
The scene's docstring promised it is never fed to the policy, which the Harness contradicts by handing the scene-bearing context to `new_session` for sampler grouping. The scene is defined by what it stages and records; which of it a policy sees is stated where the observations are assembled.
The rename left locals and docstrings calling a `Rollout` a task, which is the collision the rename set out to remove.
The env server reports its language goal under a meta key the client read as a bare literal at three sites. The server runs in its own interpreter and cannot import positronic, so the name belongs in `protocol` — the module the launcher puts on its PYTHONPATH — and both sides now spell it from there. Resolving that import for the type checker needs the search-path entry `robolab` already had; adding libero's clears two baseline entries. Rollouts are called rollouts in the remaining plan docstrings.
d3e5e4e to
2c24d3d
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2c24d3d72d
ℹ️ 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".
| 'control_mode': context['eval.control_mode'], | ||
| 'seed': context.get('eval.seed'), | ||
| 'settle_steps': context['eval.settle_steps'], | ||
| 'suite': context[SUITE], |
There was a problem hiding this comment.
Centralize the LIBERO reset-token field names
Rule hardcoded-keys violated:
LiberoAdapter._reset_token still spells the env-server wire fields (suite, task_id, camera_resolution, control_mode, seed, and settle_steps) as literals even though LiberoEnv.reset must independently spell the same names; define these fields in the positronic-free protocol module and use those constants on both sides of the isolated-interpreter boundary.
AGENTS.md reference: AGENTS.md:L14-L22
Useful? React with 👍 / 👎.
| def _reset_token(self, context: dict[str, Any]) -> Any: | ||
| # No seed rides the token: RoboLab's eval path has no seed hook, so a recorded seed would only mislead. | ||
| return {'task': context['eval.task'], 'instruction_type': context['eval.instruction_type']} | ||
| return {'task': context[TASK], 'instruction_type': context[INSTRUCTION_TYPE]} |
There was a problem hiding this comment.
Centralize the RoboLab reset-token field names
Rule hardcoded-keys violated:
RobolabAdapter._reset_token writes task and instruction_type as bare wire-key literals while RobolabEnv.reset and the validation clients repeat them; place constants beside META_INSTRUCTION in the positronic-free protocol module and import them at every producer and consumer.
AGENTS.md reference: AGENTS.md:L14-L22
Useful? React with 👍 / 👎.
| # The rollout's model-facing half, assembled per episode: the instruction and nothing else. The | ||
| # scene, seed and plan position are recorded and reach ``new_session``, which a sampling policy | ||
| # groups episodes by, but they never enter an observation and so never reach a model. |
There was a problem hiding this comment.
Remove the deleted sampler from the context comment
Rule diff-comments violated:
The _policy_inputs comment explains why scene metadata reaches new_session by claiming that a sampling policy groups episodes with it, but fresh evidence relative to the earlier review is that the current parent removed SampledPolicy and a repo-wide search finds no sampling-policy consumer; describe the intrinsic distinction between episode context and observation inputs without narrating deleted code.
AGENTS.md reference: AGENTS.md:L14-L22
Useful? React with 👍 / 👎.
Summary
Steps 14's core from #464: the per-rollout unit becomes a typed
Rollout, and the policy boundary becomes structural.list[Eval], one per embodiment.--eval=.sim.allis a single selection covering LIBERO's 40 tasks, RoboLab's 120 and the native sim, against one warm policy. The Worlds still rebuild one at a time, so only the env server of the eval in flight is ever alive.Rolloutis the unit: instruction, timeout, and thescenepayload its eval'sresetstages from. The RUN-context dicts are gone. The wiring that was task-invariant —reset,privileged,done— moves ontoEval.Directive.RUN(**kwargs)is unchanged for its emitters;timeoutjoinstaskas a field of the rollout rather than the scene.eval.seed,eval.task_idand the rest of the scene are recorded in episode statics but never enter the observation dict. The oldinputs.update(self.context)shipped them and relied on every codec to ignore them — blindness by convention, now by structure.eval.seedwas spelled in three modules and read with.get, so a typo on either side drew an unseeded scene in silence; it iskeys.EVAL_SEEDand both readers subscript. The backend-specific keys live on the adapter that unpacks them.RolloutreplacesTask, whose name collided with the benchmark taskeval.task_idnames; the count and position follow it, so--eval.trial_countis now--eval.rollout_countand the recorded statics areeval.rollout_index/eval.rollout_count.Index and count are derived from the plan instead of stamped into each context, and
inference_latencybecomes a run-levelHarnessargument rather than riding every trial. The droid eval stops recording a randomeval.seedfor a scene that has none.Recorded episode statics are unchanged apart from that droid seed.
Test plan
uv run --locked pytest --no-cov— 1108 passed, 9 skippedruff check,ruff format,basedpyrightall cleantest_scene_is_recorded_but_withheld_from_the_policyis the new guard for the boundary; verified it fails when_policy_inputsis swapped back forself.contexttest_golden_pipelinepasses unchanged, so movinginference_latencyoff the RUN context did not shift any timing.sim.allinstantiates to 3 evals / 161 rollouts;.sim.libero.spatial --eval.seed=3 --eval.trial_count=2gives 20 tasks with the expected scenes