Repository navigation
feat(training): train with Tinker on a bundle's tasks - #458
polakamtejas wants to merge 1 commit into
Conversation
`agent_env.training.tinker.AgentEnvDatasetBuilder` is tinker-cookbook's `rl.train` dataset builder: each batch is groups of episodes of one task, and every rollout is an episode played in process on the Episode API, not through the OpenEnv server, since Tinker samples itself. The turn's tools are rendered by the cookbook's renderer, each sampled message is recorded on the turn before its calls run, a reply without a call ends the turn, and the reward is made by the rules `build_app` applies (`openenv/reward.py`). Tool-call arguments are typed by the tool's schema, since a renderer that writes parameters as untyped text reads digits as numbers. A rollout the cookbook cuts is scored once its group is done; a failed episode raises for the cookbook's retry. Each group is filed in the run index as `<run>/train|eval/step-<n>` (`_episodes.RunGroup`). The `tinker` extra is `openenv`; tinker-cookbook, Tinker and torch are the trainer's to pin. The `tst/tinker` tier plays episodes on tinker-cookbook with a stand-in renderer, run by `tinker.yml` and `make tinker-test`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| async def compute_group_rewards(self, trajectory_group: list[Trajectory], | ||
| env_group: Sequence[Env]) -> list[tuple[float, Metrics]]: | ||
| return list(await asyncio.gather(*(env.message_env.score_cut() for env in env_group))) |
There was a problem hiding this comment.
Cut rollouts keep their episode slots until compute_group_rewards, which runs after the group finishes. With group_size=16 and max_concurrent_episodes=8, if the first eight rollouts hit a token or parse limit, the remaining eight cannot start. The group never finishes, so scoring never releases the slots.
Finish and score each cut episode before waiting for its siblings, while still bounding active episodes.
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/agent_env/training/tinker.py
Line: 289-291
Comment:
**Cut rollouts can hang training**
Cut rollouts keep their episode slots until `compute_group_rewards`, which runs after the group finishes. With `group_size=16` and `max_concurrent_episodes=8`, if the first eight rollouts hit a token or parse limit, the remaining eight cannot start. The group never finishes, so scoring never releases the slots.
Finish and score each cut episode before waiting for its siblings, while still bounding active episodes.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.
| await self._open(turn) | ||
| self.history.append({"role": "user", "content": turn.prompt}) | ||
| return MessageStepResult(reward=0.0, episode_done=False, next_messages=self.history) |
There was a problem hiding this comment.
Later turns append only turn.prompt. A task that changes its system_prompt between prompts still trains the model under the first turn's instructions, even though EpisodeSpec supports these changes. This produces rollouts and rewards for the wrong model input.
Update the rendered instructions when opening a later turn, or reject such tasks explicitly.
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/agent_env/training/tinker.py
Line: 203-205
Comment:
**Later turns lose changed instructions**
Later turns append only `turn.prompt`. A task that changes its `system_prompt` between prompts still trains the model under the first turn's instructions, even though `EpisodeSpec` supports these changes. This produces rollouts and rewards for the wrong model input.
Update the rendered instructions when opening a later turn, or reject such tasks explicitly.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.| # A call whose arguments aren't an object can't run, so it isn't recorded: the turn runs every recorded call. | ||
| await self._turn.record_model_call( | ||
| _chat_message(text, [call for call in calls if isinstance(call[2], dict)]), model=self._model_name) | ||
| logs: dict[str, str | int | float] = {"task": self._name, "episode": self._episode.id, "assistant": text} |
There was a problem hiding this comment.
The rollout logs include the complete assistant response here and complete serialized tool arguments at line 133. This violates the repository directive: “Never log full response bodies, request bodies, or large serialized objects in production code.”
Keep summaries such as tool names, sizes, and error counts here, and leave full content in the episode record. This repository requirement must be satisfied before merging.
Rule Used: What: Never log full response bodies, request bodies, or large serialized objects in production code. Why: Large logs are expensive, noisy, hard to search, risk pipeline limits, and can leak sensitive data. Good: ``` logger.info('Search completed',... (source)
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/agent_env/training/tinker.py
Line: 125
Comment:
**Logs copy full message content**
The rollout logs include the complete assistant response here and complete serialized tool arguments at line 133. This violates the repository directive: “Never log full response bodies, request bodies, or large serialized objects in production code.”
Keep summaries such as tool names, sizes, and error counts here, and leave full content in the episode record. This repository requirement must be satisfied before merging.
**Rule Used:** What: Never log full response bodies, request bodies, or large serialized objects in production code. Why: Large logs are expensive, noisy, hard to search, risk pipeline limits, and can leak sensitive data. Good: ``` logger.info('Search completed',... ([source](https://app.greptile.com/-/switch-tenant?tenant=scale-ai&next=%2Fcustom-context%3Fmemory%3Dfa8d684f-4686-4f3e-b1ef-c27453f614ea))
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.| self._tool_errors += 1 | ||
| result = _tool_message(call, call_id, "error: the arguments aren't a JSON object") | ||
| self.history.append(result) | ||
| logs[f"tool_call_{n}"] = f"{call.function.name}({json.dumps(arguments)}) -> {str(result['content'])[:600]}" |
There was a problem hiding this comment.
The inline 600 sets the tool-result log limit. The repository requires magic numbers to be stored as descriptively named class or instance variables. Give this character limit a name so its purpose is clear and the logging policy is easy to change. This repository requirement must be satisfied before merging.
Rule Used: Store magic numbers as class or instance variables with descriptive names rather than using them inline in the code. (source)
Learned From
scaleapi/scaleapi#126388
Prompt To Fix With AI
This is a comment left during a code review.
Path: src/agent_env/training/tinker.py
Line: 133
Comment:
**The log limit is unnamed**
The inline `600` sets the tool-result log limit. The repository requires magic numbers to be stored as descriptively named class or instance variables. Give this character limit a name so its purpose is clear and the logging policy is easy to change. This repository requirement must be satisfied before merging.
**Rule Used:** Store magic numbers as class or instance variables with descriptive names rather than using them inline in the code. ([source](https://app.greptile.com/-/switch-tenant?tenant=scale-ai&next=%2Fcustom-context%3Fmemory%3D002e0051-41ad-46c1-9098-47433c580150))
**Learned From**
[scaleapi/scaleapi#126388](https://github.com/scaleapi/scaleapi/pull/126388)
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Stacked on #457, which records episodes in the run index. This PR's run groups need it. Review that one first; this one retargets to
mainonce it merges.What
agent_env.training.tinker.AgentEnvDatasetBuilderis tinker-cookbook'srl.traindataset builder for a bundle folder's tasks:train_splitandtest_splitareallor an eval of the bundle.reward=names the scoring step when a task records more than one.max_concurrent_episodesbounds the runs under way at once.Each rollout is an
EpisodeEnv, a cookbookMessageEnvplayed on the Episode API, in process:bashwhen the agent works in a sandbox) and its prompt with the cookbook's renderer.build_appapplies (agent_env.openenv.reward). Each score is a metricscore/<step>, next tomessages,tool_callsandtool_errors.<parameter=zip>19122</parameter>) parses each one withjson.loads, which reads digits as a number. So without this, a model could never pass an id or a zip code. Qwen's own parser types arguments by the schema, as SkyRL'scalls.parsedoes.compute_group_rewardsends its turn with no reply once its group is done, so the verifiers still score it. The metriccutmarks it.RetryOnFailureretries it or drops its group. It never trains as 0.<run_name>/train|eval/step-<n>(_episodes.RunGroup). The explorer then shows a task's rollouts by training step. The cookbook makes its eval groups once and replays them at every eval, so an eval group's step is read when it's played.Bundles are prepared as the OpenEnv server prepares them (
prepare_bundle(played=True)), without the agent the trainer plays.Why in process, not through the OpenEnv server
Tinker samples tokens itself, so the trainer owns the loop. The env needs only the Episode API, and skips the pieces the server-based helpers need: an
end_turntool, probing each task's tools, and the server hop. It shares the same rules: the reward rules,score/<step>metrics, and failures left out of training.Packaging and CI
tinkerextra isopenenv, for the reward rules. tinker-cookbook, Tinker and torch are the trainer's to pin, as with verl and SkyRL. The extra is indevand adds nothing heavy, souv.lockchanges by a few lines.agent_env.training.tinkerimports only where tinker-cookbook is installed. It's listed with the other trainer helpers as not core.agent_env.training._episodesholds what needs no trainer installed: argument typing,RunGroupandTrainingSteps. It's unit-tested in the regular tier.tst/tinkertier playsEpisodeEnvon tinker-cookbook with a stand-in renderer, so nothing samples and no Tinker key is needed. It's covered by a newtinker.ymlworkflow (tinker-cookbook as released and itsmain, CPU torch; not a required check) and bymake tinker-test.Tests
tst/tinker(8 tests), on tinker-cookbook 0.5.9:Breaking the cut metric fails the tier.
tst/unit/training/episodes_test.py: argument typing (plain, nullable andanyOfstrings; integers and lists left alone) and run-group naming, including an eval's group read at the step it's played.Unit tier, plain dev venv, macOS: 11,567 passed and 4 failed. The same 4 fail on feat(episode): record episodes in the run index #457 alone, without this change, so they come from this machine's environment, not from this PR:
test_mcp_cli_builder×2, andtest_docker_image_github_clone×2, whosegithangs on a local credential wrapper.End to end on local Docker, Qwen3.5-9B on Tinker: a bundle of tau2-bench's retail domain. It has the store database as the env's universe, its 15 tools, 89 tasks, and a verifier comparing the database the agent leaves with the gold actions'.
main, trained and filed its groups as…/eval/step-0000,…/train/step-0000and…/eval/step-0001.Not covered yet
[image content]placeholder.🤖 Generated with Claude Code
Fix the cut-rollout deadlock, preserve later-turn instructions, and satisfy the logging requirements before merging.
Fix with agent prompt
Summary
Adds a Tinker dataset builder that plays bundle tasks through the Episode API.
The assessment is based on source inspection; no runtime verification was performed.
Diagram
%%{init: {'theme': 'neutral'}}%% flowchart TD A[Prepare bundle splits] --> B[Create task groups] B --> C[Acquire episode slot] C --> D[Start Episode API run] D --> E[Render turn and sample message] E --> F{Rollout outcome} F -->|Tool call| G[Record and run tools] G --> E F -->|Reply| H[End turn or finish episode] H -->|Finished| I[Return reward and release slot] F -->|Cookbook cuts rollout| J[Keep open episode and slot] J --> K[Wait for group to finish] K --> L[Score cut episode and release slot] K -. Waiting siblings may need held slots .-> CReviews (1) · Last reviewed commit: "feat(training): train with Tinker on a b..." · Reviewed by Greptile