Add OpenEnv agent-training recipe: multi-turn Wordle GRPO (Hugging Face Jobs)

[behroozazarkhalili]

What does this PR do?

Adds a new Agents recipe: grpo_agent_wordle_hf_jobs — an end-to-end walkthrough that trains a small language model to play Wordle as a multi-turn agent inside a live OpenEnv/TextArena environment with GRPO (via TRL's environment_factory path), designed to run non-interactively on Hugging Face Jobs with papermill.

Each rollout spans up to six guesses with stateful per-turn feedback and an environment-driven done signal; the trained agent is evaluated the faithful way — by playing complete games and measuring win rate (the agentic analogue of held-out accuracy).

Why the "▶ Submit this notebook as an HF Job" panel? Unlike most recipes, it opens with a hf jobs run … panel — deliberate, since GRPO holds the policy and generates rollouts, so the notebook is meant to be submitted to HF's cloud rather than run locally cell-by-cell. The recipe clones the cookbook itself and papermills the notebook, so it works as-is once merged. An interactive (vLLM-off) path is included for reading/experimentation.

Checklist against the contribution guide:

• Lowercase filename; _Authored by: line right after the first header; "What you'll learn" + tiered References with links to every open-source resource.
• No pip install/non-informative outputs; no empty code cells.
• Open-source stack only (TRL, OpenEnv, TextArena, Qwen3).
• Executes without errors — tested end-to-end on Hugging Face Jobs (SMOKE=1, papermill exit 0, multi-turn GRPO → play-the-game eval).
• Added to _toctree.yml (Agents Recipes) and index.md.

A companion single-turn reasoning-gym recipe is submitted in a separate PR.

Who can review?

@merveenoyan @stevhliu — and @sergiopaniego, since this builds on the OpenEnv reasoning_gym Space and the SFT-warmup tutorial.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[behroozazarkhalili]

Ready for review 🙌 cc @merveenoyan @stevhliu — and @sergiopaniego, since this builds on the OpenEnv reasoning_gym Space and the SFT-warmup tutorial. Happy to adjust anything, including the "Submit as an HF Job" panel if you'd prefer a different framing.

[sergiopaniego]

Hi @behroozazarkhalili, thanks a lot for the careful work on this pair of PRs! The notebooks are clearly written, the SMOKE gating and portability plumbing are nicely done, and I really like the core idea. Since GRPO holds the policy and generates the rollouts, submitting the notebook to HF Jobs with papermill instead of running it interactively is a genuinely useful pattern that isn't documented anywhere else.

I also think that Jobs workflow is the real contribution here. The Wordle payload itself (system prompt, WordleEnv, reward_func, GRPO config, eval loop) closely follows TRL's official example (openenv_wordle_grpo.ipynb and the TRL OpenEnv guide). The same is true of #361 relative to the OpenEnv SFT-warmup and end-to-end walkthrough tutorials. The cookbook tries not to duplicate content that's maintained in the library docs, so rather than merging two recipes that differ mainly in payload, here's what I'd propose.

Consolidate #361 and #362 into this single recipe, reframed so the HF Jobs workflow is the star. Concretely:

• Lead with "how to run agent RL non-interactively on HF Jobs" as the thesis. Wordle stays as the worked example, since it uses the official openenv/wordle Space and needs no external API keys.
• Add explicit attribution up front, something like "the environment and training setup build on TRL's official Wordle example (link), and this recipe adds the non-interactive Jobs workflow around it".
• Make it self-contained. Drop the references and links to grpo_agent_reasoning_gym_hf_jobs.ipynb and chain_sum, since they'd 404 on hf.co/learn until that PR merges (relative links also shouldn't carry the .ipynb extension). For the single-turn plus SFT warm-start story, link to the OpenEnv tutorials instead.

A few concrete fixes on the current content, most of which apply regardless of the consolidation:

1. Committed outputs: all 31 cells have empty outputs, so readers never see a training log, reward delta, or win rate. Could you commit the outputs of a real execution? At least the smoke run, ideally the full one.
2. The vLLM job command is likely broken: the first hf jobs run variant installs ipykernel and registers the python3 kernel, but the vLLM variant doesn't, so papermill should fail there with "No such kernel".
3. SMOKE=0 sizing: 150 steps with 64 grad-accum is around 9,600 multi-turn rollouts (up to 6 sequential ~1k-token generations each) with plain generate() on an a10g-small inside a 3-hour timeout. That math doesn't seem to close. Since the PR notes only SMOKE=1 was tested, could you share a completed SMOKE=0 run, or adjust the recommended flavor, timeout, and steps?
4. Smoke-first UX: the flagship copy-paste command hardcodes SMOKE=0, so a first-time user gets the expensive run. I'd make the primary command SMOKE=1 and show the full run as the follow-up.
5. Cell 2 isn't idempotent: git clone into /tmp/wordle_env with check=True fails on re-run, and on a Job it duplicates the install already done in the bash command. Also worth double-checking the promisor-clone claim, since Add OpenEnv agent-training recipe: GRPO on Reasoning-Gym (Hugging Face Jobs) #361 installs its env Space directly with pip install git+https://huggingface.co/spaces/... without issue.
6. Train/eval mismatch: training injects the guess tool schema via the chat template, but the eval calls apply_chat_template without tools=, so the model is evaluated on a prompt format it wasn't trained with. The regex fallback masks this and likely understates the win rate.
7. Minor: torch_dtype="auto" is deprecated in transformers v5 (use dtype="auto"), import os appears in three cells, and the H1 title would be better aligned with the _toctree.yml title ("Train a multi-turn Wordle agent with GRPO on OpenEnv using Hugging Face Jobs").

Happy to go back and forth on any of this. And again, the Jobs angle is a great addition, I'd just like it to land as one strong, self-contained recipe 🙌

[behroozazarkhalili]

Thanks for the thorough review, and for the consolidation call — I agree the Jobs/papermill workflow is the part worth keeping, and the recipe is much stronger framed that way. I've pushed a revision to #362 that does that. Everything below was verified against the notebook source rather than taken as read, and one item came back with a result worth flagging.

Consolidation

• Jobs-first thesis. H1 is now Train a multi-turn Wordle agent with GRPO on OpenEnv using Hugging Face Jobs (matching _toctree.yml), and the opening leads with why agent RL wants a batch job rather than an interactive kernel.
• Attribution up front. An explicit note now credits TRL's openenv_wordle_grpo.ipynb and the TRL OpenEnv guide for the environment and training setup, and states that what this recipe adds is the Jobs workflow around it.
• Self-contained. Dropped the relative grpo_agent_reasoning_gym_hf_jobs.ipynb link and the chain_sum references; the single-turn and SFT-warm-start story now links the OpenEnv SFT-warmup and end-to-end walkthrough tutorials. Every URL in the notebook returns 200.

The seven items on #362

2. vLLM command — confirmed and fixed. The a10g-small variant had ipykernel + registration; the vLLM variant had neither, so papermill would have failed with NoSuchKernel before running a cell. All three command variants now install and register it.

4. Smoke-first — fixed. The primary copy-paste command is now SMOKE=1 with a short timeout; the full run is the documented follow-up.

3. SMOKE=0 sizing — you were right, and I answered it with a measurement rather than another guess. Since rollout cost depends on GPU, policy size, and whether vLLM is on, any single recommended triple would be wrong across the flavors the notebook advertises. So the notebook now measures it. From the committed run (a10g-small, Qwen3-0.6B, no vLLM):

measured: 92.4s over 40 rollouts -> 2.31s per rollout
projected full run (150 steps x 64 grad-accum = 9,600 rollouts): 6.2h

Steps that actually fit a given --timeout on THIS hardware:
  --timeout   3600 ( 1.0h)  ->  GRPO_MAX_STEPS <= 18    [full run does NOT fit]
  --timeout  10800 ( 3.0h)  ->  GRPO_MAX_STEPS <= 54    [full run does NOT fit]
  --timeout  21600 ( 6.0h)  ->  GRPO_MAX_STEPS <= 109   [full run does NOT fit]
  --timeout  43200 (12.0h)  ->  GRPO_MAX_STEPS <= 219   [full run fits]

So the original recipe was off by roughly 2x on wall-clock: at the documented --timeout 10800 only 54 steps fit, not 150. Every knob is env-overridable, so this projection feeds straight back into the SMOKE=0 submission as -e GRPO_MAX_STEPS=….

5. Idempotency — confirmed and fixed. git clone with check=True and no guard did fail on re-run, and duplicated the job's own install. The helper now short-circuits when the env module already imports (which is the Job case) and guards the clone target.

5b. The promisor claim — I checked, and it turns out we're both right about our own Space. Reproduced directly:

$ git clone --filter=blob:none https://huggingface.co/spaces/openenv/wordle
fatal: expected 'packfile'
fatal: could not fetch a6344aac8c09253b3b630fb776ae94478aa0275b from promisor remote
warning: Clone succeeded, but checkout failed.          # exit 128, empty worktree

$ git clone --filter=blob:none https://huggingface.co/spaces/sergiopaniego/reasoning_gym
exit 0                                                   # 10 files, clean

So reasoning_gym genuinely installs fine via pip install git+… (which uses a partial clone), and openenv/wordle genuinely does not. wordle additionally vendors src/, envs/ and an egg-info dir, so it's a materially bigger object graph — though I haven't inspected HF's git backend and won't claim the mechanism. I've kept the full clone and replaced the over-broad comment with the exact command, the verbatim error, and an explicit note that this is per-Space behaviour, not a general property of HF Spaces.

6. Train/eval mismatch — confirmed, and it was bigger than it looks. Measured on Qwen3-0.6B with identical messages:

	rendered prompt
apply_chat_template(...) — what eval did	81 chars
apply_chat_template(..., tools=…) — what training did	845 chars

The 845-char version carries a full <tools> block with the schema auto-derived from guess's signature and docstring. Rather than hand-write that schema in the eval cell, the notebook now derives it with a small helper that mirrors GRPOTrainer's own rule (public bound methods, minus reset/get_reward) and passes the same callables TRL passes:

def env_tools(env) -> list:
    return [m for n, m in inspect.getmembers(env, predicate=inspect.ismethod)
            if n not in ("reset", "get_reward") and not n.startswith("_")]

It names no tool, so swapping the env class keeps eval correct for free. CHAT_TEMPLATE_KWARGS is now a single dict shared by GRPOConfig and eval so they can't drift either. And since you noted the regex fallback was masking this, the eval now reports how many turns came from real tool calls vs the fallback, so any residual mismatch shows up in the output instead of being absorbed.

7. Minor — all fixed. torch_dtype="auto" → dtype="auto" (confirmed deprecated: modeling_utils.py logs "`torch_dtype` is deprecated! Use `dtype` instead!"), import os consolidated from four cells to one, H1 aligned with _toctree.yml.

1. Committed outputs — done. The notebook now ships executed, PAPERMILL_EXIT=0, every cell carrying real output: the GRPO step tables with completions and advantages, the reward delta, the calibration projection above, and the eval.

The committed run turned up a bug your review nearly caught

You predicted the regex fallback was making the win rate understated. It was actually hiding a broken metric in the other direction. The first executed run reported:

win rate: 3/3 = 100%

for a 0.6B policy after five GRPO steps. It had not solved anything. won = env.done and env.reward > 0 was the culprit — TextArena ends an episode that keeps submitting invalid moves and still pays partial credit, so rule-breaking scored as success.

Verified by replaying the model's own guesses over a real session against the Space:

[apple]   done=False reward=0.0   valid
[banana]  done=False reward=0.0   INVALID — "must be exactly 5 letters"
[cherry]  done=True  reward=0.2   episode terminated

To find the real signal I scripted a solver and actually won games against the same Space: a genuine solve returns exactly reward=1.0, a plain loss 0.0. So the eval now tests reward >= ENV_SOLVED_REWARD, reports a solve rate, and prints an invalid-move rate beside it so a policy breaking the rules reads differently from one losing fairly. The honest committed numbers:

solve rate: 0/3 = 0%   mean reward=0.33
invalid moves: 3/9 turns (33%)
parse modes over 9 turns: tool_call=9 (100%)

Zero solves is the correct result for a smoke run — and the 33% invalid-move rate says why: the policy breaks the 5-letter rule, which is the first thing GRPO has to train away.

There was a second, related bug that only became reachable because of your item 6. The tool contract asks for [abcde], so once tools= was supplied and the model started emitting real tool calls, the JSON parse path returned an already-bracketed string and env.guess(f"[{word}]") re-wrapped it to [[apple]]. Previously unreachable, since without tools= the regex path (which strips brackets via group(1)) always ran. parse_guess now normalises brackets — and deliberately does not repair a wrong-length word, since that is a real invalid move the eval should report rather than launder.

One thing your review surfaced that I think extends further

You argued on #361 that defaulting ENV_BASE_URL to a low-concurrency Space is wrong, and that duplicating should be a requirement rather than an optional step. I think that's right, and it applies to the official openenv/wordle Space too — for a stronger reason than load. From the Space's own source:

• TextArenaEnvironment inherits SUPPORTS_CONCURRENT_SESSIONS = False from OpenEnv's Environment base
• the Space calls create_app() with no override, so the server defaults to max_concurrent_envs=1
• _validate_concurrency_safety() raises at startup if the cap is >1 for an env not marked concurrent

So the 1-session limit is by design and can't simply be raised. A previous smoke run of mine died at the GRPO cell with Server at capacity: 1/1 sessions active (code: CAPACITY_REACHED) — most likely sessions leaked by earlier crashed attempts, since reset() opens one and a crash never calls close(). Owning the Space is what lets you restart it and clear those.

The revision documents this as a required step for training runs (with the shared Space fine for the smoke run), because it also turned out to be the prerequisite for item 1 — the committed outputs above were produced against an owned duplicate, exactly as the recipe now instructs.

On #361

Happy to fold it in as you suggested. I've left it open for the moment so it doesn't go stale before this lands — glad to close it the moment you're happy with the consolidated recipe. For the record, its five items resolve as: (1) and (3) disappear with the SFT/teacher phase leaving the cookbook (and the open-model-first guideline is well taken — it would have applied), (2) is fixed in every surviving command variant, (4) is the point above, and (5) goes away with the cross-link.

Thanks again — the sizing and train/eval catches in particular made this a materially better recipe.

[sergiopaniego]

Thanks for the thorough revision! I verified everything against the notebook and it all checks out. The solve-rate fix and the calibration cell in particular make this a much stronger recipe.

Approving now. Note the failing doc build is unrelated to this PR: it is a doc-builder regression that breaks the build for every cookbook PR, fixed in huggingface/doc-builder#819. Once that merges we will re-run the check here and merge.

Tiny optional nit if you touch the notebook again: the %pip cell has the pip root-user warning committed as output, worth stripping.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.