Add rollout traces to built-in envs#749
Conversation
derek-tml
changed the title
derek-tml
requested review from
dphuang2 and
joschu
and removed request for
dphuang2
|
@codex review |
![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: 6d34086f40
ℹ️ 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".
|
@claude review |
|
Claude encountered an error after 0s —— View job I'll analyze this and get back to you. |
dphuang2
left a comment
dphuang2
left a comment
There was a problem hiding this comment.
Nice, the structured trace is much easier to consume than detokenizing ob/ac. Few things below, mostly test coverage. Non-blocking except maybe the message_env error-path coverage.
Two broader questions on the trace being unconditional (the field lives in #748 but it's this PR that populates it everywhere):
- It's built on every step and always serialized into the summaries, plus the per-step
deepcopyinmessage_env, so long multi-turn rollouts pay CPU + a bigger summaries file. Is always-on intended or should it be gated? log_formatteralready stashesto_data()into*_logtree.json, so for thelog_formatterenvs the prompt/response is somewhat duplicated. Fine if it's deliberate (summaries as the machine-readable per-traj record vs the HTML logtree), just want to confirm it's a choice.
|
|
||
| if self._last_messages is None: | ||
| self._last_messages = await self.message_env.initial_observation() | ||
| prompt_messages = deepcopy(self._last_messages) |
There was a problem hiding this comment.
The pre-step deepcopy is the right call since MessageEnv.step() mutates history in place. test_trace_prompt_snapshots_messages_before_step_mutation only covers the happy path though. The parse-error early return and the context-overflow return both emit a trace too, can we get the snapshot assertion on those paths as well? Those are the ones most likely to regress quietly.
| response_messages, | ||
| is_valid_list, | ||
| ): | ||
| trajectory.transitions[0].trace = RolloutTrace( |
There was a problem hiding this comment.
This assumes every preference trajectory has exactly one transition, and that trajectory_group/response_messages/is_valid_list stay aligned. safezip catches a length mismatch but nothing asserts the right response lands on the right trajectory. Worth a small test, this is the one spot here with real back-fill logic.
| trace=RolloutTrace( | ||
| prompt=prompt_formatter.to_data(), | ||
| policy_response=response_formatter.to_data(), | ||
| reward_data=reward_terms, |
There was a problem hiding this comment.
Since reward_terms feeds both table_from_dict and reward_data, the trace ends up with display strings (f"{...:.3f}") rather than numbers. For a durable record I think we want raw floats so downstream can compute on them. Cheap fix: keep a numeric dict and format only for the table. Same thing in problem_env.py with reward_table.
2 participants
Summary
RolloutTracepayloads to built-in rollout envs that can decode prompt/response text.Stack from ghstack (oldest at bottom):
Co-authored-by: Cursor cursoragent@cursor.com