Skip to content

feat(training): train with Tinker on a bundle's tasks - #458

Open
polakamtejas wants to merge 4 commits into
feat/episode-run-indexfrom
feat/training-tinker
Open

polakamtejas wants to merge 4 commits into
feat/episode-run-indexfrom
feat/training-tinker

Conversation

@polakamtejas

@polakamtejas polakamtejas commented Nov 3, 2026 •

Copy link
Copy Markdown
Collaborator

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 main once it merges.

What

agent_env.training.tinker.AgentEnvDatasetBuilder is tinker-cookbook's rl.train dataset builder for a bundle folder's tasks:

  • Batches: each batch is groups of episodes, a group playing one task.
  • Splits: train_split and test_split are all or an eval of the bundle.
  • Reward: reward= names the scoring step when a task records more than one.
  • Concurrency: max_concurrent_episodes bounds the runs under way at once.
  • Tasks left out: a task that can't run as an episode, or that gives the agent the trainer plays skills (an add_skills step, or a deploy_agent step with skills), is logged and left out before training starts. No trainer shows its model skills yet, and a rollout refused at that turn would only be retried and refused again.

Each rollout is an EpisodeEnv, a cookbook MessageEnv played on the Episode API, in process:

  • The first observation starts the episode and renders the turn's system prompt, its tools (the envs' tools, and bash when the agent works in a sandbox) and its prompt with the cookbook's renderer. The system prompt and the tools are written once, at the start of the conversation.
  • Each step records the sampled message on the turn, then runs its tool calls. A reply without a call ends the turn. A later turn comes back as the next user message.
  • A later turn's new system prompt is written as text before that turn's prompt (_episodes.next_prompt). That is what the OpenEnv-based helpers already did; they now call the same function, so every trainer's model reads the same conversation. A later turn that offers other tools than the first fails its rollout, and a turn that gives the agent skills is refused whichever turn it is.
  • When the run ends, the reward is the task's scores under the rules build_app applies (agent_env.openenv.reward). Each score is a metric score/<step>, next to messages, tool_calls and tool_errors.
  • Tool-call arguments are typed by the tool's schema. A renderer that writes parameters as untyped text (Qwen3.5's <parameter=zip>19122</parameter>) parses each one with json.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's calls.parse does.
  • A rollout the cookbook cuts (the token limit, context overflow, a structural parse error, or a first prompt already over the token budget) is scored the moment it's cut, by _ScoredWhenCut over the cookbook's EnvFromMessageEnv. Its open turn ends with no reply, the verifiers score where it got to, and that score is added to the cookbook's failed_reward. Scoring at the cut also frees the rollout's slot under max_concurrent_episodes: a slot held until the group finished meant a group larger than the slots could never finish. The metric cut is 1 for such a rollout and 0 for the rest.
  • A failed episode raises, so the cookbook's RetryOnFailure retries it or drops its group. It never trains as 0. An episode that fails to start, or whose first prompt can't be written (a renderer without tool support raises there), gives back its slot at once.
  • Step logs: the model's message, each call's arguments and each tool's answer go into the logs the cookbook renders in its per-rollout report, each clipped at EpisodeEnv.LOGGED_CHARS (2000 characters). The full text stays in the episode's run record.
  • Run groups: each group's episodes are filed in the run index as <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_turn tool, probing each task's tools, and the server hop. It shares the same rules: the reward rules, score/<step> metrics, how a later turn's system prompt is written, and failures left out of training.

Packaging and CI

  • The tinker extra is openenv, for the reward rules. tinker-cookbook, Tinker and torch are the trainer's to pin, as with verl and SkyRL. The extra is in dev and adds nothing heavy, so uv.lock changes by a few lines.
  • agent_env.training.tinker imports only where tinker-cookbook is installed. It's listed with the other trainer helpers as not core.
  • agent_env.training._episodes holds what needs no trainer installed: argument typing, next_prompt, RunGroup and TrainingSteps. It's unit-tested in the regular tier.
  • agent_env.training._openenv: Session writes a later turn's changed system prompt through next_prompt instead of inline. What the TRL, NeMo Gym, verl, SkyRL, Miles and slime helpers' models read is unchanged; TRL's tier still asserts the same text.
  • The new tst/tinker tier plays EpisodeEnv on tinker-cookbook with a stand-in renderer, so nothing samples and no Tinker key is needed. It's covered by a new tinker.yml workflow (tinker-cookbook as released and its main, CPU torch; not a required check) and by make tinker-test.

Tests

  • tst/tinker (12 tests), on tinker-cookbook 0.5.9:
    • a rollout's tool call runs on its turn, and its reply ends it with the task's reward; the run is in the index under its group;
    • a later turn's new system prompt comes before its prompt;
    • a turn that gives the agent skills is refused, whichever turn it is;
    • the turn bound ends a turn;
    • a step's logs are clipped;
    • a cut rollout has its open turn ended and is scored once;
    • a group larger than the slots finishes when every rollout in it is cut, and a rollout whose first prompt is over the token budget is scored and frees its slot. Both hang against an adapter that scores cut rollouts only once their group is done;
    • a rollout whose first prompt can't be written is abandoned and frees its slot. This fails with the slot still held when the prompt is built outside the cleanup block;
    • a bundle's splits are prepared without the played agent;
    • tasks that can't be played, or that give the agent skills, are left out, and a dataset whose tasks make no reward is refused.
  • tst/unit/training/episodes_test.py: argument typing (plain, nullable and anyOf strings; integers and lists left alone), next_prompt (a changed system prompt before the prompt, and kept in effect after), 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, and test_docker_image_github_clone ×2, whose git hangs on a local credential wrapper.

Training runs

Both on local Docker, with Tinker sampling and training.

GridSim, as the other training pages run it

The bundle the other trainers' docs pages train on, with their recipe: an SFT warm start on the train days' best schedules, then 40 RL steps of 32 rollouts (4 days, 8 rollouts each) at temperature 1.1, learning rate 2e-5, LoRA rank 32, and at most 4 model messages a rollout. The 200 eval days are played once each at steps 0, 20 and 40. Tinker has no Qwen3-4B, the model those runs train, so this one trains Qwen3-8B with qwen3_disable_thinking.

Qwen3-8B, 200 eval days Eval reward Beat idling (0.2) Best schedule (1.0)
Untrained 0.195 2.5% 0%
After SFT (step 0) 0.645 82.5% 12.5%
After 20 steps 0.697 93% 9.5%
After 40 steps 0.715 91.5% 13.5%
  • RL added 0.07, with standard errors about 0.02. The other trainers' Qwen3-4B runs score 0.51–0.57 after SFT and 0.57–0.69 after 40 steps, RL gains of 0 to 0.18. The higher start here is mostly the larger model, and each is one run.
  • 1,880 rollouts, no group skipped. RetryOnFailure retried 22: 18 after a dropped connection and 4 after a timeout, while another job shared the 4-CPU Docker. A step took about 70 seconds.
  • Cut rollouts are scored: the cookbook cut 0%, 2% and 3.5% of the eval rollouts at steps 0, 20 and 40, and each was scored where it stood.
  • Repeat calls: after its schedule is accepted, the model keeps calling submit_schedule until the 4-message limit. The env refuses the repeats, and the reward is the first schedule's. RL didn't change it: 2.7 refused calls a rollout at the first step, 2.9 at the last.

tau2-bench retail, Qwen3.5-9B

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'.

  • Replaying the gold actions through the Episode API scores 1.0; doing nothing scores 0.0.
  • Three training runs (8 tasks × 4 rollouts a step, 14 steps) took held-out reward on 32 unseen tasks from 0.16–0.21 to 0.83–0.92 by step 8.
  • Since cut rollouts are scored, every rollout has a verifier score.
  • The sampler and trainer stayed on-policy (KL ≤ 6e-4), across about 900 episodes with 0 episode errors.
  • A smoke run on this branch trained and filed its groups as …/eval/step-0000, …/train/step-0000 and …/eval/step-0001.
  • The later changes (scoring at the cut, the system prompt, skills, slots) are covered by the tst/tinker tier above. The retail tasks have one turn each, so a later turn's system prompt doesn't come up there.

Not covered yet

  • Skills: tasks that give the agent skills are left out, since no trainer shows them to the model yet. The OpenEnv-based helpers don't refuse them yet either; that's a separate change.
  • Images in tool results: they reach the model as a [image content] placeholder.
  • Remote sandboxes: only local Docker has been exercised.

🤖 Generated with Claude Code

RetriggerView in GreptileConfidence Score: 4/5

Do not merge until failed runs are excluded from training rewards.

Fix All in CursorFindings

  1. P1 Failed runs still earn rewards ▶
Fix with agent prompt
### Issue 1
src/agent_env/training/tinker.py:240-243
`EpisodeEnv._finish` treats returned scores as proof that the run succeeded. But `Episode.result()` also returns scores when a step failed with `fail_task_on_error=False`; that run is recorded as `FAILED`. A failed setup step alongside a successful verifier can therefore produce a positive training reward and bypass retries.

Check the run's outcome before returning the reward, and raise for failed runs so `RetryOnFailure` can retry or drop the group.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Summary

Adds Tinker training on bundle tasks through in-process episodes.

  • Tinker rollouts play bundle tasks as Episode API runs.
  • Tinker batches bundle tasks and files runs by training step.
  • Training conversations carry changed prompts and schema-matched tool values.
  • Cut rollouts get scored and release their slots immediately.
  • Tasks without a playable episode or reward stay out of Tinker training.
  • The Tinker helper stays optional and has its own test workflow.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Bundle tasks] --> B[Dataset batches]
  B --> C[Episode groups]
  C --> D[EpisodeEnv]
  D --> E[Model messages and tool calls]
  E --> D
  D --> F[Task scores]
  G[Cookbook cuts rollout] --> F
  F --> H[Reward and metrics]
  F --> I[Release episode slot]
Loading

Reviews (15) · Last reviewed commit: "fix(training): give a rollout's slot bac..." · Reviewed by Greptile

@polakamtejas
polakamtejas requested a review from a team as a code owner November 3, 2026 00:21
Comment thread src/agent_env/training/tinker.py Resolved
Comment thread src/agent_env/training/tinker.py Resolved
Comment thread src/agent_env/training/tinker.py Resolved
Comment thread src/agent_env/training/tinker.py Resolved
Comment thread src/agent_env/training/tinker.py Resolved
Comment thread src/agent_env/training/tinker.py Resolved
polakamtejas and others added 4 commits November 10, 2026 22:44
`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>
… change the system prompt, and clip step logs

- A rollout tinker-cookbook's adapter ends itself (the token limit, a structural parse failure, context overflow,
  or a first prompt over the budget) is scored where it is cut (`_ScoredWhenCut`), which frees its slot under
  `max_concurrent_episodes`. Scored only once its group was done, such rollouts could hold every slot while their
  groups' other rollouts waited for one, and training stopped.
- `playable` leaves out a task whose turns don't share one system prompt: the renderer writes it once, at the start
  of the conversation, so later turns would be trained under the first turn's instructions.
- A step's logs clip the model's message, each call's arguments and each tool's answer at
  `EpisodeEnv.LOGGED_CHARS`.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…nd leave out tasks that give the agent skills

- `_episodes.next_prompt` is the one rule for a later turn's prompt: a system prompt the turn changes comes before it,
  as text, since a trainer's conversation writes its system prompt once. `_openenv.Session` (TRL, NeMo Gym, verl,
  SkyRL, Miles, slime) already wrote it this way and now calls it; Tinker calls it too instead of leaving such tasks
  out, so every trainer shows the model the same conversation.
- Tinker checks every turn for skills, not just the first, and `playable` leaves out a task that gives the agent it
  plays skills (an `add_skills` step for it, or its `deploy_agent` with `skills`), since no trainer shows its model
  skills yet and a rollout refused at that turn would be retried and refused again.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… be written

The renderer's conversation prefix was built after initial_observation's cleanup block, so a renderer that can't
write tools (the cookbook's role_colon raises NotImplementedError) left the episode running and its slot under
max_concurrent_episodes held until the group was cleaned up. The prefix and first user message are now built inside
the block, so such a rollout is abandoned and frees its slot at once.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Comment on lines +240 to +243
scores = await self._episode.result()
self._release()
self._scored = True
reward = self._reward(scores)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Failed runs still earn rewards

EpisodeEnv._finish treats returned scores as proof that the run succeeded. But Episode.result() also returns scores when a step failed with fail_task_on_error=False; that run is recorded as FAILED. A failed setup step alongside a successful verifier can therefore produce a positive training reward and bypass retries.

Check the run's outcome before returning the reward, and raise for failed runs so RetryOnFailure can retry or drop the group.

Prompt To Fix With AI
This is a comment left during a code review.
Path: src/agent_env/training/tinker.py
Line: 240-243

Comment:
**Failed runs still earn rewards**

`EpisodeEnv._finish` treats returned scores as proof that the run succeeded. But `Episode.result()` also returns scores when a step failed with `fail_task_on_error=False`; that run is recorded as `FAILED`. A failed setup step alongside a successful verifier can therefore produce a positive training reward and bypass retries.

Check the run's outcome before returning the reward, and raise for failed runs so `RetryOnFailure` can retry or drop the group.

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Fix in Cursor Fix in Claude Code Fix in Codex

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant