Buckets:
| <meta charset="utf-8" /><meta name="hf:doc:metadata" content="{"title":"OpenEnv Wordle with GRPO using TRL","local":"openenv-wordle-with-grpo-using-trl","sections":[{"title":"Install dependencies","local":"install-dependencies","sections":[{"title":"Log in to Hugging Face","local":"log-in-to-hugging-face","sections":[],"depth":3}],"depth":2},{"title":"Define the system prompt","local":"define-the-system-prompt","sections":[],"depth":2},{"title":"Define the environment","local":"define-the-environment","sections":[],"depth":2},{"title":"Define the reward function","local":"define-the-reward-function","sections":[],"depth":2},{"title":"Create the dataset","local":"create-the-dataset","sections":[],"depth":2},{"title":"Set GRPO Config","local":"set-grpo-config","sections":[],"depth":2},{"title":"Create the GRPOTrainer and start training","local":"create-the-grpotrainer-and-start-training","sections":[],"depth":2},{"title":"Save and push to Hub","local":"save-and-push-to-hub","sections":[],"depth":2},{"title":"Load the fine-tuned model and run inference","local":"load-the-fine-tuned-model-and-run-inference","sections":[],"depth":2}],"depth":1}"> | |
| <link href="/docs/openenv/pr_749/en/_app/immutable/assets/0.e3b0c442.css" rel="modulepreload"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/entry/start.85477f45.js"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/chunks/scheduler.2b22cead.js"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/chunks/singletons.63566282.js"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/chunks/paths.dd876c7b.js"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/entry/app.51835dc5.js"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/chunks/preload-helper.0820fbc7.js"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/chunks/index.1a0e8013.js"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/nodes/0.167255c0.js"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/chunks/each.e59479a4.js"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/nodes/66.cae83d34.js"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/chunks/Heading.c0d3f116.js"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/chunks/MermaidChart.svelte_svelte_type_style_lang.21bcf336.js"> | |
| <link rel="modulepreload" href="/docs/openenv/pr_749/en/_app/immutable/chunks/CodeBlock.c8d73295.js"><!-- HEAD_svelte-u9bgzb_START --><meta name="hf:doc:metadata" content="{"title":"OpenEnv Wordle with GRPO using TRL","local":"openenv-wordle-with-grpo-using-trl","sections":[{"title":"Install dependencies","local":"install-dependencies","sections":[{"title":"Log in to Hugging Face","local":"log-in-to-hugging-face","sections":[],"depth":3}],"depth":2},{"title":"Define the system prompt","local":"define-the-system-prompt","sections":[],"depth":2},{"title":"Define the environment","local":"define-the-environment","sections":[],"depth":2},{"title":"Define the reward function","local":"define-the-reward-function","sections":[],"depth":2},{"title":"Create the dataset","local":"create-the-dataset","sections":[],"depth":2},{"title":"Set GRPO Config","local":"set-grpo-config","sections":[],"depth":2},{"title":"Create the GRPOTrainer and start training","local":"create-the-grpotrainer-and-start-training","sections":[],"depth":2},{"title":"Save and push to Hub","local":"save-and-push-to-hub","sections":[],"depth":2},{"title":"Load the fine-tuned model and run inference","local":"load-the-fine-tuned-model-and-run-inference","sections":[],"depth":2}],"depth":1}"><!-- HEAD_svelte-u9bgzb_END --> <p></p> <div class="items-center shrink-0 min-w-[100px] max-sm:min-w-[50px] justify-end ml-auto flex" style="float: right; margin-left: 10px; display: inline-flex; position: relative; z-index: 10;"><div class="inline-flex rounded-md max-sm:rounded-sm"><button class="inline-flex items-center gap-1 h-7 max-sm:h-7 px-2 max-sm:px-1.5 text-sm font-medium text-gray-800 border border-r-0 rounded-l-md max-sm:rounded-l-sm border-gray-200 bg-white hover:shadow-inner dark:border-gray-850 dark:bg-gray-950 dark:text-gray-200 dark:hover:bg-gray-800" aria-live="polite"><span class="inline-flex items-center justify-center rounded-md p-0.5 max-sm:p-0 hover:text-gray-800 dark:hover:text-gray-200"><svg class="sm:size-3.5 size-3" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg></span> <span>Copy page</span></button> <button class="inline-flex items-center justify-center w-6 max-sm:w-5 h-7 max-sm:h-7 disabled:pointer-events-none text-sm text-gray-500 hover:text-gray-700 dark:hover:text-white rounded-r-md max-sm:rounded-r-sm border border-l transition border-gray-200 bg-white hover:shadow-inner dark:border-gray-850 dark:bg-gray-950 dark:text-gray-200 dark:hover:bg-gray-800" aria-haspopup="menu" aria-expanded="false" aria-label="Open copy menu"><svg class="transition-transform text-gray-400 overflow-visible sm:size-3.5 size-3 rotate-0" width="1em" height="1em" viewBox="0 0 12 7" fill="none" xmlns="http://www.w3.org/2000/svg"><path d="M1 1L6 6L11 1" stroke="currentColor"></path></svg></button></div> </div> <h1 class="relative group"><a id="openenv-wordle-with-grpo-using-trl" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#openenv-wordle-with-grpo-using-trl"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>OpenEnv Wordle with GRPO using TRL</span></h1> <p data-svelte-h="svelte-1fc6auj"><a href="https://colab.research.google.com/github/huggingface/trl/blob/main/examples/notebooks/openenv_wordle_grpo.ipynb" rel="nofollow"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a></p> <p data-svelte-h="svelte-cos6ht"><img src="https://huggingface.co/datasets/trl-lib/documentation-images/resolve/main/trl_banner_dark.png" alt="trl banner"></p> <p data-svelte-h="svelte-10kmmwt">With <a href="https://github.com/huggingface/trl" rel="nofollow"><strong>Transformers Reinforcement Learning (TRL)</strong></a>, you can train a model that learns to <strong>play Wordle</strong>, a word-guessing game, through interaction and reinforcement.</p> <ul data-svelte-h="svelte-1au2mwt"><li><a href="https://github.com/huggingface/trl" rel="nofollow">TRL GitHub Repository</a> — star us to support the project!</li> <li><a href="https://huggingface.co/docs/trl/example_overview" rel="nofollow">Official TRL Examples</a></li> <li><a href="https://huggingface.co/docs/trl/community_tutorials" rel="nofollow">Community Tutorials</a></li> <li><a href="https://github.com/huggingface/OpenEnv" rel="nofollow">OpenEnv</a></li></ul> <p data-svelte-h="svelte-z0zrou">An <strong>agentic environment</strong> is a setting where a model can take actions, observe outcomes, and adjust its behavior based on feedback, similar to how humans learn from trial and error. | |
| In this case, the agent interacts with the <strong>Wordle</strong> environment through the <a href="https://github.com/huggingface/OpenEnv" rel="nofollow"><strong>OpenEnv</strong></a> framework, which standardizes multi-agent and RL-style text environments.</p> <p data-svelte-h="svelte-ecsij6"><a href="https://en.wikipedia.org/wiki/Wordle" rel="nofollow">Wordle</a> is a popular word puzzle where the player must guess a secret five-letter word within six tries. | |
| After each guess, feedback indicates whether each letter is:</p> <ul data-svelte-h="svelte-ld1v3k"><li><strong>GREEN (G)</strong>: Correct and in the right position</li> <li><strong>YELLOW (Y)</strong>: Present but in the wrong position</li> <li><strong>GRAY (X)</strong>: Not in the word</li></ul> <p data-svelte-h="svelte-1ywqk01">This feedback loop makes Wordle a perfect environment for <strong>RL with LLMs</strong>, where the goal is to maximize the probability of guessing the correct word efficiently.</p> <p data-svelte-h="svelte-13ba4lb">We’ll fine-tune a model using <strong>GRPO</strong> (Group Relative Policy Optimization) via TRL. | |
| Using <code>environment_factory</code>, the trainer automatically handles:</p> <ol data-svelte-h="svelte-10sg7gd"><li>Creating environment instances for each rollout.</li> <li>Generating model completions and parsing tool calls.</li> <li>Stepping through the environment with the model’s actions.</li> <li>Collecting rewards and managing the interaction loop.</li></ol> <p data-svelte-h="svelte-152af3i">This means you only need to define the environment class and reward function — the trainer takes care of the rest.</p> <h2 class="relative group"><a id="install-dependencies" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#install-dependencies"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Install dependencies</span></h2> <p data-svelte-h="svelte-mzecib">We’ll start by installing <strong>TRL</strong> (with vLLM support), the <strong>OpenEnv</strong> Wordle environment, and <strong>trackio</strong> for logging.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-bash "><!-- HTML_TAG_START -->pip install -Uq trl[vllm] git+https://huggingface.co/spaces/openenv/wordle trackio<!-- HTML_TAG_END --></pre></div> <h3 class="relative group"><a id="log-in-to-hugging-face" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#log-in-to-hugging-face"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Log in to Hugging Face</span></h3> <p data-svelte-h="svelte-lyetiz">Log in to your <strong>Hugging Face</strong> account to save your fine-tuned model, track your experiment results directly on the Hub or access gated models. You can find your <strong>access token</strong> on your <a href="https://huggingface.co/settings/tokens" rel="nofollow">account settings page</a>.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> huggingface_hub <span class="hljs-keyword">import</span> notebook_login | |
| notebook_login()<!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="define-the-system-prompt" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#define-the-system-prompt"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Define the system prompt</span></h2> <p data-svelte-h="svelte-698zkb">This prompt instructs the model on how to play Wordle. It includes the game rules, feedback format, and importantly, tells the model to use the <code>guess</code> tool to submit guesses. The <code>environment_factory</code> pattern uses tool calling to interact with the environment, so the model needs to know which tool to call.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START -->prompt = <span class="hljs-string">"""You are an expert Wordle solver with deep knowledge of English vocabulary, letter frequency patterns, and optimal guessing strategies. | |
| Follow these rules to play Wordle: | |
| 1. The target is a 5-letter English word | |
| 2. You have 6 attempts to guess the correct word | |
| 3. After each guess, you receive color-coded feedback: | |
| - GREEN (G): Letter is correct and in the correct position | |
| - YELLOW (Y): Letter is in the word but in the wrong position | |
| - GRAY (X): Letter is not in the word at all | |
| 4. All guesses must be valid 5-letter English words | |
| 5. You cannot reuse a word you've already guessed | |
| 6. Use the tool `guess` to make a guess. | |
| """</span><!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="define-the-environment" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#define-the-environment"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Define the environment</span></h2> <p data-svelte-h="svelte-tiz0kv">The <code>WordleEnv</code> class wraps the OpenEnv TextArena Wordle environment into the interface expected by <code>environment_factory</code>.</p> <p data-svelte-h="svelte-129airi">When you pass <code>environment_factory=WordleEnv</code> to the trainer, it will:</p> <ol data-svelte-h="svelte-10q3kdb"><li>Create a new <code>WordleEnv()</code> instance for each rollout episode.</li> <li>Call <code>reset()</code> to start a new game (returns the initial observation or <code>None</code>).</li> <li>Automatically generate model completions, parse tool calls, and invoke the corresponding methods (e.g., <code>guess(...)</code>).</li> <li>Repeat until the environment signals <code>done=True</code> or the max completion length is reached.</li></ol> <p data-svelte-h="svelte-17xvc4o">The environment exposes its public methods as tools. Any public method (other than <code>reset</code>) with a docstring is automatically discovered and exposed as a callable tool. Here, the <code>guess</code> method lets the model submit a Wordle guess and receive feedback.</p> <p data-svelte-h="svelte-153e7gn">For this example, we connect to the hosted environment at <a href="https://huggingface.co/spaces/openenv/wordle" rel="nofollow">openenv/wordle</a>. | |
| For production use, we recommend duplicating the Space to your own account or running it locally via Docker, as the hosted versions have limited concurrency.</p> <p data-svelte-h="svelte-1k082v1">For more information, refer to the <a href="https://huggingface.co/docs/trl/main/en/openenv" rel="nofollow">TRL-OpenEnv documentation</a>.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> textarena_env <span class="hljs-keyword">import</span> TextArenaAction, TextArenaEnv | |
| <span class="hljs-keyword">class</span> <span class="hljs-title class_">WordleEnv</span>: | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self</span>): | |
| self.client = TextArenaEnv(base_url=<span class="hljs-string">"https://openenv-wordle.hf.space"</span>) | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">reset</span>(<span class="hljs-params">self, **kwargs</span>) -> <span class="hljs-literal">None</span> | <span class="hljs-built_in">str</span>: | |
| result = self.client.reset() | |
| <span class="hljs-comment"># The game returns cumulative feedback each turn (new text appended at the end), so</span> | |
| <span class="hljs-comment"># we store the previous full response and slice out only the newly appended part.</span> | |
| self._last_full_feedback = result.observation.messages[<span class="hljs-number">0</span>].content | |
| self.reward = <span class="hljs-number">0.0</span> | |
| self.done = <span class="hljs-literal">False</span> | |
| <span class="hljs-keyword">return</span> self._last_full_feedback | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">guess</span>(<span class="hljs-params">self, guess: <span class="hljs-built_in">str</span></span>) -> <span class="hljs-built_in">str</span>: | |
| <span class="hljs-string">""" | |
| Make a guess in the Wordle environment. | |
| Args: | |
| guess: The guessed word, formatted as '[abcde]' | |
| Returns: | |
| The feedback message from the environment. | |
| """</span> | |
| <span class="hljs-keyword">if</span> self.done: | |
| <span class="hljs-keyword">raise</span> ValueError(<span class="hljs-string">"Game over."</span>) | |
| result = self.client.step(TextArenaAction(message=guess)) | |
| _full_feedback = result.observation.messages[<span class="hljs-number">0</span>].content | |
| <span class="hljs-comment"># Just take the new feedback since the last guess</span> | |
| feedback = _full_feedback[<span class="hljs-built_in">len</span>(self._last_full_feedback):] | |
| self._last_full_feedback = _full_feedback | |
| <span class="hljs-comment"># Penalize invalid moves</span> | |
| <span class="hljs-keyword">if</span> <span class="hljs-string">"You attempted an invalid move"</span> <span class="hljs-keyword">in</span> feedback: | |
| self.reward = <span class="hljs-number">0.0</span> | |
| <span class="hljs-keyword">else</span>: | |
| self.reward = result.reward | |
| self.done = result.done | |
| <span class="hljs-keyword">return</span> feedback<!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="define-the-reward-function" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#define-the-reward-function"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Define the reward function</span></h2> <p data-svelte-h="svelte-bofu78">The reward function receives the list of environment instances after each episode completes. Since the <code>WordleEnv</code> tracks its own reward (updated after each <code>guess</code> call), we simply read it out.</p> <p data-svelte-h="svelte-s5bfmy">This is much simpler than defining multiple reward functions manually — the environment already knows the game outcome.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">def</span> <span class="hljs-title function_">reward_func</span>(<span class="hljs-params">environments, **kwargs</span>) -> <span class="hljs-built_in">list</span>[<span class="hljs-built_in">float</span>]: | |
| <span class="hljs-keyword">return</span> [env.reward <span class="hljs-keyword">for</span> env <span class="hljs-keyword">in</span> environments]<!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="create-the-dataset" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#create-the-dataset"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Create the dataset</span></h2> <p data-svelte-h="svelte-x18exd">We create a dataset with repeated prompts to control the number of training episodes. | |
| Each entry triggers one rollout episode during training. The prompt is formatted as a chat message.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> datasets <span class="hljs-keyword">import</span> Dataset | |
| dataset = Dataset.from_dict({<span class="hljs-string">"prompt"</span>: [[{<span class="hljs-string">"role"</span>: <span class="hljs-string">"user"</span>, <span class="hljs-string">"content"</span>: prompt}] <span class="hljs-keyword">for</span> _ <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(<span class="hljs-number">3000</span>)]})<!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="set-grpo-config" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#set-grpo-config"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Set GRPO Config</span></h2> <p data-svelte-h="svelte-1dfiz93">Next, we define the <strong>GRPOConfig</strong>, which controls all key training parameters. | |
| This configuration specifies how the model interacts with vLLM, manages memory, and logs results.</p> <p data-svelte-h="svelte-1hsds5z">Note the <code>chat_template_kwargs={"enable_thinking": False}</code> parameter — this disables Qwen3’s thinking mode so the model responds directly with tool calls instead of generating internal reasoning tokens first.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> trl <span class="hljs-keyword">import</span> GRPOConfig | |
| model_name = <span class="hljs-string">"Qwen/Qwen3-1.7B"</span> | |
| output_dir = <span class="hljs-string">"wordle-grpo-Qwen3-1.7B"</span> | |
| grpo_config = GRPOConfig( | |
| <span class="hljs-comment"># Training schedule / optimization</span> | |
| num_train_epochs=<span class="hljs-number">1</span>, | |
| learning_rate=<span class="hljs-number">1e-6</span>, | |
| gradient_accumulation_steps=<span class="hljs-number">64</span>, | |
| per_device_train_batch_size=<span class="hljs-number">1</span>, | |
| warmup_steps=<span class="hljs-number">10</span>, | |
| optim=<span class="hljs-string">"adamw_torch"</span>, | |
| max_grad_norm=<span class="hljs-number">1.0</span>, | |
| <span class="hljs-comment"># GRPO configuration</span> | |
| num_generations=<span class="hljs-number">2</span>, | |
| max_completion_length=<span class="hljs-number">1024</span>, | |
| log_completions=<span class="hljs-literal">True</span>, | |
| num_completions_to_print=<span class="hljs-number">2</span>, | |
| chat_template_kwargs={<span class="hljs-string">"enable_thinking"</span>: <span class="hljs-literal">False</span>}, | |
| <span class="hljs-comment"># vLLM configuration</span> | |
| use_vllm=<span class="hljs-literal">True</span>, | |
| vllm_mode=<span class="hljs-string">"colocate"</span>, | |
| vllm_gpu_memory_utilization=<span class="hljs-number">0.15</span>, | |
| vllm_max_model_length=<span class="hljs-number">3072</span>, | |
| <span class="hljs-comment"># Logging / reporting</span> | |
| output_dir=output_dir, | |
| report_to=<span class="hljs-string">"trackio"</span>, | |
| trackio_space_id=output_dir, | |
| logging_steps=<span class="hljs-number">1</span>, | |
| save_steps=<span class="hljs-number">10</span>, | |
| save_total_limit=<span class="hljs-number">1</span>, | |
| <span class="hljs-comment"># Memory optimization</span> | |
| gradient_checkpointing=<span class="hljs-literal">True</span>, | |
| <span class="hljs-comment"># Hub integration</span> | |
| push_to_hub=<span class="hljs-literal">True</span>, | |
| )<!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="create-the-grpotrainer-and-start-training" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#create-the-grpotrainer-and-start-training"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Create the GRPOTrainer and start training</span></h2> <p data-svelte-h="svelte-cp0d1t">Now we initialize the <code>GRPOTrainer</code> with <code>environment_factory=WordleEnv</code>.</p> <p data-svelte-h="svelte-v455eo">This tells the trainer to automatically handle the entire interaction loop:</p> <ul data-svelte-h="svelte-1m4xrmo"><li>It creates a <code>WordleEnv</code> instance for each episode.</li> <li>It generates model completions, parses tool calls (like <code>guess</code>), and steps through the environment.</li> <li>It collects rewards and manages the <code>tool_mask</code> (which tokens are model-generated vs environment-generated) automatically.</li></ul> <p data-svelte-h="svelte-ooa6j6">No need to write a custom <code>rollout_func</code> or manage tokenization manually.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> trl <span class="hljs-keyword">import</span> GRPOTrainer | |
| trainer = GRPOTrainer( | |
| model=model_name, | |
| reward_funcs=reward_func, | |
| train_dataset=dataset, | |
| args=grpo_config, | |
| environment_factory=WordleEnv, | |
| )<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1pw21xs">Show memory stats before training</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> torch | |
| gpu_stats = torch.cuda.get_device_properties(<span class="hljs-number">0</span>) | |
| start_gpu_memory = <span class="hljs-built_in">round</span>(torch.cuda.max_memory_reserved() / <span class="hljs-number">1024</span> / <span class="hljs-number">1024</span> / <span class="hljs-number">1024</span>, <span class="hljs-number">3</span>) | |
| max_memory = <span class="hljs-built_in">round</span>(gpu_stats.total_memory / <span class="hljs-number">1024</span> / <span class="hljs-number">1024</span> / <span class="hljs-number">1024</span>, <span class="hljs-number">3</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f"GPU = <span class="hljs-subst">{gpu_stats.name}</span>. Max memory = <span class="hljs-subst">{max_memory}</span> GB."</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f"<span class="hljs-subst">{start_gpu_memory}</span> GB of memory reserved."</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1b9poti">And train!</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START -->trainer_stats = trainer.train()<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1t6bnxf">Show memory stats after training</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START -->used_memory = <span class="hljs-built_in">round</span>(torch.cuda.max_memory_reserved() / <span class="hljs-number">1024</span> / <span class="hljs-number">1024</span> / <span class="hljs-number">1024</span>, <span class="hljs-number">3</span>) | |
| used_memory_for_training = <span class="hljs-built_in">round</span>(used_memory - start_gpu_memory, <span class="hljs-number">3</span>) | |
| used_percentage = <span class="hljs-built_in">round</span>(used_memory / max_memory * <span class="hljs-number">100</span>, <span class="hljs-number">3</span>) | |
| training_memory_percentage = <span class="hljs-built_in">round</span>(used_memory_for_training / max_memory * <span class="hljs-number">100</span>, <span class="hljs-number">3</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f"<span class="hljs-subst">{trainer_stats.metrics[<span class="hljs-string">'train_runtime'</span>]}</span> seconds used for training."</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f"<span class="hljs-subst">{<span class="hljs-built_in">round</span>(trainer_stats.metrics[<span class="hljs-string">'train_runtime'</span>]/<span class="hljs-number">60</span>, <span class="hljs-number">2</span>)}</span> minutes used for training."</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f"Peak reserved memory = <span class="hljs-subst">{used_memory}</span> GB."</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f"Peak reserved memory for training = <span class="hljs-subst">{used_memory_for_training}</span> GB."</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f"Peak reserved memory % of max memory = <span class="hljs-subst">{used_percentage}</span> %."</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f"Peak reserved memory for training % of max memory = <span class="hljs-subst">{training_memory_percentage}</span> %."</span>)<!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="save-and-push-to-hub" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#save-and-push-to-hub"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Save and push to Hub</span></h2> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START -->trainer.save_model(output_dir) | |
| trainer.push_to_hub()<!-- HTML_TAG_END --></pre></div> <h2 class="relative group"><a id="load-the-fine-tuned-model-and-run-inference" class="header-link block pr-1.5 text-lg no-hover:hidden with-hover:absolute with-hover:p-1.5 with-hover:opacity-0 with-hover:group-hover:opacity-100 with-hover:right-full" href="#load-the-fine-tuned-model-and-run-inference"><span><svg class="" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" aria-hidden="true" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 256 256"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span></a> <span>Load the fine-tuned model and run inference</span></h2> <p data-svelte-h="svelte-5ktyje">Now let’s test our fine-tuned model by loading it and playing a game of Wordle. | |
| We use the same <code>WordleEnv</code> class to interact with the environment, and generate model responses with standard Transformers inference.</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">from</span> transformers <span class="hljs-keyword">import</span> AutoModelForCausalLM, AutoTokenizer | |
| model_name = <span class="hljs-string">"sergiopaniego/wordle-grpo-Qwen3-1.7B"</span> <span class="hljs-comment"># Replace with your HF username or organization</span> | |
| fine_tuned_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=<span class="hljs-string">"float32"</span>, device_map=<span class="hljs-string">"auto"</span>) | |
| tokenizer = AutoTokenizer.from_pretrained(model_name)<!-- HTML_TAG_END --></pre></div> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START --><span class="hljs-keyword">import</span> json | |
| <span class="hljs-keyword">def</span> <span class="hljs-title function_">play_wordle</span>(<span class="hljs-params">model, tokenizer</span>): | |
| env = WordleEnv() | |
| initial_observation = env.reset() | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">"Initial observation:"</span>) | |
| <span class="hljs-built_in">print</span>(initial_observation) | |
| <span class="hljs-built_in">print</span>() | |
| messages = [{<span class="hljs-string">"role"</span>: <span class="hljs-string">"user"</span>, <span class="hljs-string">"content"</span>: prompt}] | |
| <span class="hljs-keyword">if</span> initial_observation: | |
| messages.append({<span class="hljs-string">"role"</span>: <span class="hljs-string">"user"</span>, <span class="hljs-string">"content"</span>: initial_observation}) | |
| <span class="hljs-keyword">for</span> turn <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(<span class="hljs-number">6</span>): | |
| <span class="hljs-keyword">if</span> env.done: | |
| <span class="hljs-keyword">break</span> | |
| prompt_text = tokenizer.apply_chat_template( | |
| messages, | |
| add_generation_prompt=<span class="hljs-literal">True</span>, | |
| tokenize=<span class="hljs-literal">False</span>, | |
| enable_thinking=<span class="hljs-literal">False</span>, | |
| ) | |
| model_inputs = tokenizer([prompt_text], return_tensors=<span class="hljs-string">"pt"</span>).to(model.device) | |
| generated_ids = model.generate(**model_inputs, max_new_tokens=<span class="hljs-number">512</span>) | |
| output_ids = generated_ids[<span class="hljs-number">0</span>][<span class="hljs-built_in">len</span>(model_inputs.input_ids[<span class="hljs-number">0</span>]):] | |
| generated_text = tokenizer.decode(output_ids, skip_special_tokens=<span class="hljs-literal">True</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f"Turn <span class="hljs-subst">{turn + <span class="hljs-number">1</span>}</span> - Model output: <span class="hljs-subst">{generated_text}</span>"</span>) | |
| <span class="hljs-comment"># Try to parse tool call from the generated text</span> | |
| <span class="hljs-keyword">try</span>: | |
| <span class="hljs-comment"># Try to extract a guess from tool call format or bracket format</span> | |
| <span class="hljs-keyword">if</span> <span class="hljs-string">"guess"</span> <span class="hljs-keyword">in</span> generated_text: | |
| <span class="hljs-comment"># Parse JSON tool call</span> | |
| start = generated_text.index(<span class="hljs-string">"{"</span>) | |
| end = generated_text.rindex(<span class="hljs-string">"}"</span>) + <span class="hljs-number">1</span> | |
| args = json.loads(generated_text[start:end]) | |
| <span class="hljs-keyword">if</span> <span class="hljs-string">"arguments"</span> <span class="hljs-keyword">in</span> args: | |
| args = args[<span class="hljs-string">"arguments"</span>] | |
| guess_word = args.get(<span class="hljs-string">"guess"</span>, <span class="hljs-string">""</span>) | |
| <span class="hljs-keyword">else</span>: | |
| <span class="hljs-comment"># Fallback: extract from brackets [word]</span> | |
| <span class="hljs-keyword">import</span> re | |
| <span class="hljs-keyword">match</span> = re.search(<span class="hljs-string">r"\[([a-zA-Z]{5})\]"</span>, generated_text) | |
| guess_word = <span class="hljs-keyword">match</span>.group(<span class="hljs-number">1</span>) <span class="hljs-keyword">if</span> <span class="hljs-keyword">match</span> <span class="hljs-keyword">else</span> generated_text.strip()[:<span class="hljs-number">5</span>] | |
| feedback = env.guess(<span class="hljs-string">f"[<span class="hljs-subst">{guess_word}</span>]"</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f" Guess: <span class="hljs-subst">{guess_word}</span> | Reward: <span class="hljs-subst">{env.reward}</span>"</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f" Feedback: <span class="hljs-subst">{feedback.strip()}</span>"</span>) | |
| <span class="hljs-built_in">print</span>() | |
| messages.append({<span class="hljs-string">"role"</span>: <span class="hljs-string">"assistant"</span>, <span class="hljs-string">"content"</span>: generated_text}) | |
| messages.append({<span class="hljs-string">"role"</span>: <span class="hljs-string">"user"</span>, <span class="hljs-string">"content"</span>: feedback}) | |
| <span class="hljs-keyword">except</span> Exception <span class="hljs-keyword">as</span> e: | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f" Error: <span class="hljs-subst">{e}</span>"</span>) | |
| <span class="hljs-keyword">break</span> | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f"Game finished! Final reward: <span class="hljs-subst">{env.reward}</span>"</span>) | |
| <span class="hljs-built_in">print</span>(<span class="hljs-string">f"Done: <span class="hljs-subst">{env.done}</span>"</span>)<!-- HTML_TAG_END --></pre></div> <p data-svelte-h="svelte-1lrqr39">Let’s play the game!</p> <div class="code-block relative "><div class="absolute top-2.5 right-4"><button class="inline-flex items-center relative text-sm focus:text-green-500 cursor-pointer focus:outline-none transition duration-200 ease-in-out opacity-0 mx-0.5 text-gray-600 " title="code excerpt" type="button"><svg class="" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M28,10V28H10V10H28m0-2H10a2,2,0,0,0-2,2V28a2,2,0,0,0,2,2H28a2,2,0,0,0,2-2V10a2,2,0,0,0-2-2Z" transform="translate(0)"></path><path d="M4,18H2V4A2,2,0,0,1,4,2H18V4H4Z" transform="translate(0)"></path><rect fill="none" width="32" height="32"></rect></svg> <div class="absolute pointer-events-none transition-opacity bg-black text-white py-1 px-2 leading-tight rounded font-normal shadow left-1/2 top-full transform -translate-x-1/2 translate-y-2 opacity-0"><div class="absolute bottom-full left-1/2 transform -translate-x-1/2 w-0 h-0 border-black border-4 border-t-0" style="border-left-color: transparent; border-right-color: transparent; "></div> Copied</div></button></div> <pre class="language-python "><!-- HTML_TAG_START -->play_wordle(fine_tuned_model, tokenizer)<!-- HTML_TAG_END --></pre></div> <a class="!text-gray-400 !no-underline text-sm flex items-center not-prose mt-4" href="https://github.com/huggingface/openenv/blob/main/docs/source/tutorials/wordle-grpo.md" target="_blank"><svg class="mr-1" xmlns="http://www.w3.org/2000/svg" aria-hidden="true" fill="currentColor" focusable="false" role="img" width="1em" height="1em" preserveAspectRatio="xMidYMid meet" viewBox="0 0 32 32"><path d="M31,16l-7,7l-1.41-1.41L28.17,16l-5.58-5.59L24,9l7,7z"></path><path d="M1,16l7-7l1.41,1.41L3.83,16l5.58,5.59L8,23l-7-7z"></path><path d="M12.419,25.484L17.639,6.552l1.932,0.518L14.351,26.002z"></path></svg> <span data-svelte-h="svelte-zjs2n5"><span class="underline">Update</span> on GitHub</span></a> <p></p> | |
| <script> | |
| { | |
| __sveltekit_1qwoa43 = { | |
| assets: "/docs/openenv/pr_749/en", | |
| base: "/docs/openenv/pr_749/en", | |
| env: {} | |
| }; | |
| const element = document.currentScript.parentElement; | |
| const data = [null,null]; | |
| Promise.all([ | |
| import("/docs/openenv/pr_749/en/_app/immutable/entry/start.85477f45.js"), | |
| import("/docs/openenv/pr_749/en/_app/immutable/entry/app.51835dc5.js") | |
| ]).then(([kit, app]) => { | |
| kit.start(app, element, { | |
| node_ids: [0, 66], | |
| data, | |
| form: null, | |
| error: null | |
| }); | |
| }); | |
| } | |
| </script> | |
Xet Storage Details
- Size:
- 63.5 kB
- Xet hash:
- 9ae9aeb8b7b76e2c114eb68228aee9d7d3cec4151b011fa96a0e14cc9fe6f156
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.