diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 07b8a75f1..6a835c033 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -48,7 +48,7 @@ - local: tutorials/end-to-end-walkthrough title: Train a Reasoning Model - local: tutorials/wordle-grpo - title: Play Wordle with GRPO + title: Train a Wordle Agent - local: tutorials/rl-training-2048 title: Play 2048 with GRPO - local: tutorials/sft-warmup diff --git a/docs/source/environments/textarena.md b/docs/source/environments/textarena.md index f8e5a288f..d84b0998e 100644 --- a/docs/source/environments/textarena.md +++ b/docs/source/environments/textarena.md @@ -11,6 +11,7 @@ An OpenEnv wrapper for [TextArena](https://github.com/textarena/textarena) game > Generic wrapper for any [TextArena](https://www.textarena.ai/docs/overview) game inside OpenEnv. This module exposes the TextArena `Env` interface through the standard HTTP server/client APIs used by other OpenEnv environments, enabling quick experimentation with the full suite of word, reasoning, and multi-agent games. Try Wordle on the [`openenv/wordle`](https://huggingface.co/spaces/openenv/wordle) Space, or connect a client to `https://openenv-wordle.hf.space`. +To train a model on it, follow the [Wordle tutorial](https://huggingface.co/docs/openenv/tutorials/wordle-grpo): an SFT warmup, then GRPO with TRL. ## Quick Start diff --git a/docs/source/tutorials/index.md b/docs/source/tutorials/index.md index 8a5050e77..2ca6f5913 100644 --- a/docs/source/tutorials/index.md +++ b/docs/source/tutorials/index.md @@ -3,7 +3,7 @@ Choose a learning path from the sidebar: - **Basics:** start with [Hello World](openenv-tutorial), build [Your First Environment](../guides/first-environment) and [deploy it](../getting_started/environment-builder). -- **Training:** [the ways to train and the supported frameworks](../guides/training), then train a [reasoning model](end-to-end-walkthrough), play [Wordle](wordle-grpo) or [2048](rl-training-2048) with GRPO, or [collect rollouts for SFT](sft-warmup). Each framework's own examples (ART, Miles, Oumi, SkyRL, torchforge, TRL, Unsloth, …) are listed in [Integrations](../guides/training#integrations). +- **Training:** [the ways to train and the supported frameworks](../guides/training), then train a [reasoning model](end-to-end-walkthrough), a [Wordle agent](wordle-grpo) with SFT and GRPO, play [2048](rl-training-2048) with GRPO, or [collect rollouts for SFT](sft-warmup). Each framework's own examples (ART, Miles, Oumi, SkyRL, torchforge, TRL, Unsloth, …) are listed in [Integrations](../guides/training#integrations). - **Harnesses:** [pick a path](harnesses): train with the trainer's own loop on [BrowserGym](browsergym-harness) (white-box), train real agents through [Harbor](harbor-harness) (black-box), or [evaluate Claude Code inside an environment](claude-code-harness). - **Evals:** follow [Evaluating with Environments](evaluation-inspect). @@ -33,7 +33,7 @@ Already familiar with the basics? These tutorials cover specific workflows in de | [Train a Reasoning Model](end-to-end-walkthrough) | The full pipeline: connect to `reasoning_gym`, wire it into TRL via `environment_factory`, fine-tune with GRPO, and push the checkpoint to the Hub. | Yes | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/OpenEnv/blob/main/examples/end_to_end_walkthrough.ipynb) | | [MCP Environments](mcp-environment) | Consume and build MCP-backed environments: list and call tools through `step()`, register Python functions as tools with FastMCP. | No | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/OpenEnv/blob/main/examples/mcp_environment.ipynb) | | [Rubrics](rubrics) | Compose reward functions from reusable pieces using `Gate`, `WeightedSum`, `LLMJudge`, and `TrajectoryRubric`. | No | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/OpenEnv/blob/main/examples/rubrics.ipynb) | -| [Play Wordle with GRPO](wordle-grpo) | Train an agent to play Wordle using GRPO via TRL's `environment_factory`. | Yes | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/trl/blob/main/examples/grpo_wordle/grpo_wordle.ipynb) | +| [Train a Wordle Agent](wordle-grpo) | Take Qwen3-1.7B from 0 to 19 wins out of 30: generate SFT data with a solver, run a short SFT warmup, then GRPO with TRL's `environment_factory`. Every step with its hardware, time and cost. | Yes | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/openenv/OpenEnv/blob/main/examples/wordle/wordle.ipynb) | | [Play 2048 with GRPO](rl-training-2048) | Train a language model to play 2048 on the OpenSpiel environment with TRL's `GRPOTrainer` and `environment_factory`. Unsloth's own [2048 notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/OpenEnv_gpt_oss_(20B)_Reinforcement_Learning_2048_Game.ipynb) trains `gpt-oss-20b` to write a 2048 strategy. | Yes | — | | [Evaluating with Environments](evaluation-inspect) | Wrap an OpenEnv environment in an Inspect AI `Task`, run it via `InspectAIHarness`, and get a structured `EvalResult`. | No | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/OpenEnv/blob/main/examples/evaluation_inspect.ipynb) | | [White-Box: Train a Web Agent on BrowserGym](browsergym-harness) | Train a vision-language model on BrowserGym web tasks with TRL's `GRPOTrainer` and `environment_factory`, so TRL runs the multi-turn tool loop. | Yes | — | diff --git a/docs/source/tutorials/wordle-grpo.md b/docs/source/tutorials/wordle-grpo.md index a23dc6312..ede82f463 100644 --- a/docs/source/tutorials/wordle-grpo.md +++ b/docs/source/tutorials/wordle-grpo.md @@ -1,359 +1,145 @@ -# Play Wordle with GRPO +# Train a Wordle Agent: SFT, then GRPO -[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/trl/blob/main/examples/grpo_wordle/grpo_wordle.ipynb) +This tutorial trains [Qwen3-1.7B](https://huggingface.co/Qwen/Qwen3-1.7B) to play [Wordle](https://en.wikipedia.org/wiki/Wordle) against the OpenEnv TextArena environment. The model starts out unable to win a single game. A short supervised warmup teaches it the game, and in our run GRPO then doubled its wins and removed most of its invalid guesses. Every step runs on one GPU, and the whole pipeline costs about $6 on [Hugging Face Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs). -![trl banner](https://huggingface.co/datasets/trl-lib/documentation-images/resolve/main/trl_banner_dark.png) +![Base model, SFT and SFT + GRPO playing the same 30 words](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/openenv/wordle-progress.gif) +| Model | Wins on 30 words | Guesses per win | Invalid guesses | +|---|---|---|---| +| [Qwen3-1.7B](https://huggingface.co/Qwen/Qwen3-1.7B) | 0/30 | - | 24 | +| + SFT ([model](https://huggingface.co/openenv/Qwen3-1.7B-Wordle-SFT)) | 10/30 | 3.2 | 41 | +| + SFT + GRPO ([model](https://huggingface.co/openenv/Qwen3-1.7B-Wordle-SFT-GRPO)) | 20/30 | 4.1 | 3 | -With [**Transformers Reinforcement Learning (TRL)**](https://github.com/huggingface/trl), you can train a model that learns to **play Wordle**, a word-guessing game, through interaction and reinforcement. +These numbers come from our run, with one training run per step, evaluated greedily on the same 30 words ([`secrets30.txt`](https://github.com/openenv/OpenEnv/blob/main/examples/wordle/secrets30.txt)). Thirty words is a small sample and training is not deterministic, so a new run will land a few wins away. Two SFT runs with the same recipe won 10 and 13 games, for example. Expect the same ordering (base, SFT, then SFT + GRPO clearly ahead) rather than the exact numbers. -- [TRL GitHub Repository](https://github.com/huggingface/trl) -- star us to support the project! -- [Official TRL Examples](https://huggingface.co/docs/trl/example_overview) -- [Community Tutorials](https://huggingface.co/docs/trl/community_tutorials) -- [OpenEnv](https://github.com/huggingface/OpenEnv) +The GRPO run used an earlier version of the invalid-move penalty, which missed the invalid move that ends a game, and capped rollouts at six tool rounds, including invalid warnings. The current `grpo.py` counts terminal invalid moves and allows up to 12 tool rounds and 12 submissions, while the environment still ends games after six legal guesses. The reported results have not been regenerated with this corrected budget. +The 30 words come from the 156 nouns TextArena draws its secrets from, the same list used for the SFT games and during GRPO. So the table measures how well each model plays the words of this environment. We did not test words outside that list. -An **agentic environment** 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 **Wordle** environment through the [**OpenEnv**](https://github.com/huggingface/OpenEnv) framework, which standardizes multi-agent and RL-style text environments. +Watch all three play any of the 30 words in the [Wordle Arena](https://huggingface.co/spaces/openenv/wordle-arena) Space, and see the training curves on [Trackio](https://huggingface.co/spaces/openenv/wordle-runs). The scripts are in [`examples/wordle`](https://github.com/openenv/OpenEnv/tree/main/examples/wordle). -[Wordle](https://en.wikipedia.org/wiki/Wordle) 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: -- **GREEN (G)**: Correct and in the right position -- **YELLOW (Y)**: Present but in the wrong position -- **GRAY (X)**: Not in the word +## What you need -This feedback loop makes Wordle a perfect environment for **RL with LLMs**, where the goal is to maximize the probability of guessing the correct word efficiently. +| Step | Script | Hardware | Time | Cost | +|---|---|---|---|---| +| 1. Generate SFT data | `generate_sft_data.py` | Laptop CPU | 10 s | $0 | +| 2. SFT warmup | `sft.py` | 1x A100 80GB (`a100-large`) | 7 min | $0.30 | +| 3. GRPO | `grpo.py` | 1x A100 80GB (`a100-large`) | 2 h 15 min | $5.60 | +| 4. Evaluate and make the GIF | `evaluate.py`, `make_gif.py` | Laptop (Apple silicon) or any GPU | 20 min for three models | $0 | +Times and costs are from our runs, at the `a100-large` price on Hugging Face Jobs ($2.50/h), and include the time to install packages and load the model. You need a Hugging Face account with a payment method for Jobs, and [uv](https://docs.astral.sh/uv/). The scripts pin their dependencies (`trl==1.16.0`, `openenv==0.10.0`, vLLM 0.31), so `uv run` and `hf jobs uv run` install the same versions used here. -We'll fine-tune a model using **GRPO** (Group Relative Policy Optimization) via TRL. -Using `environment_factory`, the trainer automatically handles: -1. Creating environment instances for each rollout. -2. Generating model completions and parsing tool calls. -3. Stepping through the environment with the model's actions. -4. Collecting rewards and managing the interaction loop. +## 1. Play the environment -This means you only need to define the environment class and reward function -- the trainer takes care of the rest. - - -## Install dependencies - -We'll start by installing **TRL** (with vLLM support), the **OpenEnv** Wordle environment, and **trackio** for logging. - -```bash -pip install -Uq trl[vllm] git+https://huggingface.co/spaces/openenv/wordle trackio -``` - -### Log in to Hugging Face - -Log in to your **Hugging Face** account to save your fine-tuned model, track your experiment results directly on the Hub or access gated models. You can find your **access token** on your [account settings page](https://huggingface.co/settings/tokens). +The Wordle environment runs as a Space, [`openenv/wordle`](https://huggingface.co/spaces/openenv/wordle). Open its playground to play a game by hand, or connect from Python: ```python -from huggingface_hub import notebook_login +from textarena_env import TextArenaAction, TextArenaEnv -notebook_login() +with TextArenaEnv(base_url="https://openenv-wordle.hf.space").sync() as env: + result = env.reset() + print(result.observation.messages[0].content) + result = env.step(TextArenaAction(message="[crane]")) + print(result.observation.messages[0].content) ``` -## Define the system prompt - -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 `guess` tool to submit guesses. The `environment_factory` pattern uses tool calling to interact with the environment, so the model needs to know which tool to call. - -```python -prompt = """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. -""" -``` +Install the client with `pip install "openenv-textarena @ git+https://github.com/openenv/OpenEnv.git@v0.10.0#subdirectory=envs/textarena_env"`. -## Define the environment +Each guess returns the feedback for every letter: `G` (right letter, right spot), `Y` (in the word, wrong spot) and `X` (not in the word). A repeated word or a word that isn't English is an invalid move. The first one is a warning, and a second one in the same turn ends the game. A win ends with `Congratulations! You guessed the word correctly!`, and a loss shows the secret word. -The `WordleEnv` class wraps the OpenEnv TextArena Wordle environment into the interface expected by `environment_factory`. +The environment's own reward is `1.0` for a win. A lost game scores (greens + half the yellows) / 5 for its last guess. A game that never reaches its end scores `0`. -When you pass `environment_factory=WordleEnv` to the trainer, it will: -1. Create a new `WordleEnv()` instance for each rollout episode. -2. Call `reset()` to start a new game (returns the initial observation or `None`). -3. Automatically generate model completions, parse tool calls, and invoke the corresponding methods (e.g., `guess(...)`). -4. Repeat until the environment signals `done=True` or the max completion length is reached. +## 2. Why start with SFT -The environment exposes its public methods as tools. Any public method (other than `reset`) with a docstring is automatically discovered and exposed as a callable tool. Here, the `guess` method lets the model submit a Wordle guess and receive feedback. +Training Qwen3-1.7B with GRPO directly does not work well, for two reasons we saw in practice: -For this example, we connect to the hosted environment at [openenv/wordle](https://huggingface.co/spaces/openenv/wordle). -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. +- **The base model almost never wins.** It always opens with `apple` and soon repeats a word, so it wins around 2-3% of games while sampling. GRPO learns from the difference between good and bad games in a group, and with almost no wins there is little to learn from. A 25-step run from the base model stayed at a 2-5% win rate. +- **The environment reward can be gamed.** A game only scores once it ends, and a win pays the same after 1 guess or 6. With Qwen3's long thinking, the easiest way to score more was to think less and use all 6 guesses. In one run, reasoning dropped from 7,000 to 2,300 tokens in 13 steps while the reward fell below where it started. -For more information, refer to the [TRL-OpenEnv documentation](https://huggingface.co/docs/trl/main/en/openenv). +So the recipe starts with a short SFT on well-played games, which teaches the format and a reasonable strategy. GRPO then trains against the environment, with a reward that pays for winning fast, and fixes what the SFT model still gets wrong. -```python -from textarena_env import TextArenaAction, TextArenaEnv +## 3. Generate the SFT data +`generate_sft_data.py` plays 1,000 games with a solver, against the same environment class the Space runs (in-process, so it takes seconds). Each turn the solver keeps the words that fit all the feedback so far and picks the guess that splits them best, and writes a short reasoning before each guess: -class WordleEnv: - def __init__(self): - self.client = TextArenaEnv(base_url="https://openenv-wordle.hf.space") - - def reset(self, **kwargs) -> None | str: - result = self.client.reset() - # The game returns cumulative feedback each turn (new text appended at the end), so - # we store the previous full response and slice out only the newly appended part. - self._last_full_feedback = result.observation.messages[0].content - self.reward = 0.0 - self.done = False - return self._last_full_feedback - - def guess(self, guess: str) -> str: - """ - Make a guess in the Wordle environment. - - Args: - guess: The guessed word, formatted as '[abcde]' - - Returns: - The feedback message from the environment. - """ - if self.done: - raise ValueError("Game over.") - result = self.client.step(TextArenaAction(message=guess)) - _full_feedback = result.observation.messages[0].content - # Just take the new feedback since the last guess - feedback = _full_feedback[len(self._last_full_feedback):] - self._last_full_feedback = _full_feedback - # Penalize invalid moves - if "You attempted an invalid move" in feedback: - self.reward = 0.0 - else: - self.reward = result.reward - self.done = result.done - return feedback +```text +ARISE gave X X X Y Y: A gray, R gray, I gray, S yellow, E yellow. +In the word but misplaced: E (not position 5), S (not position 4). Ruled out: A, I, R. +Words that fit: SLEEP, SHEEP, SMELL, STEEL, SHELF, and more. +SHEEP fits the clues and its letters split the remaining options best, so I'll guess it. ``` -## Define the reward function - -The reward function receives the list of environment instances after each episode completes. Since the `WordleEnv` tracks its own reward (updated after each `guess` call), we simply read it out. - -This is much simpler than defining multiple reward functions manually -- the environment already knows the game outcome. +Each game is stored as a tool-calling conversation: the user prompt, then one assistant turn per guess (the reasoning as `reasoning_content` and a call to the `guess` tool), each followed by the environment feedback. This is the same format TRL renders during GRPO, so the SFT model plays exactly the way RL will train it. -```python -def reward_func(environments, **kwargs) -> list[float]: - return [env.reward for env in environments] +```bash +uv run generate_sft_data.py --games 1000 --push-to-hub /wordle-tool-calls-sft ``` -## Create the dataset +The dataset used here is [`openenv/wordle-tool-calls-sft`](https://huggingface.co/datasets/openenv/wordle-tool-calls-sft). TextArena draws the secret from 156 common nouns, so the solver wins every game, in 2.7 guesses on average. -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. +## 4. SFT warmup -```python -from datasets import Dataset +`sft.py` fine-tunes Qwen3-1.7B on those games for one epoch (63 steps), training only the assistant turns: -dataset = Dataset.from_dict({"prompt": [[{"role": "user", "content": prompt}] for _ in range(3000)]}) +```bash +hf jobs uv run --flavor a100-large --timeout 1h --secrets HF_TOKEN sft.py \ + --dataset /wordle-tool-calls-sft --hub-model-id /Qwen3-1.7B-Wordle-SFT ``` -## Set GRPO Config +In our run the loss went from 4.2 to 0.09. Keep the trainable weights in fp32, with bf16 mixed-precision compute. When we loaded the weights in bf16 for training, the model came out clearly worse (final loss 0.21), because bf16 weights lose the small updates. The script sets `model_init_kwargs={"dtype": "float32"}` explicitly (also the default in TRL 1.16.0) and keeps `bf16=True` for compute. It casts the weights to bf16 after training, just before the final upload. -Next, we define the **GRPOConfig**, which controls all key training parameters. -This configuration specifies how the model interacts with vLLM, manages memory, and logs results. +## 5. GRPO -Note the `chat_template_kwargs={"enable_thinking": False}` parameter -- this disables Qwen3's thinking mode so the model responds directly with tool calls instead of generating internal reasoning tokens first. +`grpo.py` starts from the SFT model and trains it against the environment with TRL's `GRPOTrainer` and `environment_factory`. The environment class exposes one tool, `guess`, and keeps what the reward needs: ```python -from trl import GRPOConfig - -model_name = "Qwen/Qwen3-1.7B" -output_dir = "wordle-grpo-Qwen3-1.7B" - -grpo_config = GRPOConfig( - # Training schedule / optimization - num_train_epochs=1, - learning_rate=1e-6, - gradient_accumulation_steps=64, - per_device_train_batch_size=1, - warmup_steps=10, - optim="adamw_torch", - max_grad_norm=1.0, - - # GRPO configuration - num_generations=2, - max_completion_length=1024, - log_completions=True, - num_completions_to_print=2, - chat_template_kwargs={"enable_thinking": False}, - - # vLLM configuration - use_vllm=True, - vllm_mode="colocate", - vllm_gpu_memory_utilization=0.15, - vllm_max_model_length=3072, - - # Logging / reporting - output_dir=output_dir, - report_to="trackio", - trackio_space_id=output_dir, - logging_steps=1, - save_steps=10, - save_total_limit=1, - - # Memory optimization - gradient_checkpointing=True, - - # Hub integration - push_to_hub=True, -) -``` +def reward_win(environments, **kwargs): + return [1.0 + 0.5 * (1 - env.num_guesses / 6) if env.won else 0.0 for env in environments] -## Create the `GRPOTrainer` and start training +def reward_partial(environments, **kwargs): + return [0.0 if env.won else 0.1 * env.best_greens for env in environments] -Now we initialize the `GRPOTrainer` with `environment_factory=WordleEnv`. - -This tells the trainer to automatically handle the entire interaction loop: -- It creates a `WordleEnv` instance for each episode. -- It generates model completions, parses tool calls (like `guess`), and steps through the environment. -- It collects rewards and manages the `tool_mask` (which tokens are model-generated vs environment-generated) automatically. - -No need to write a custom `rollout_func` or manage tokenization manually. - -```python -from trl import GRPOTrainer - -trainer = GRPOTrainer( - model=model_name, - reward_funcs=reward_func, - train_dataset=dataset, - args=grpo_config, - environment_factory=WordleEnv, -) +def reward_invalid(environments, **kwargs): + return [-0.1 * env.invalid_moves for env in environments] ``` -Show memory stats before training +- A win pays more the faster it comes (1.42 in one guess, 1.0 in six). +- A loss pays for the best guess of the game, whether or not the game reached its end. +- Each invalid move costs 0.1. -```python -import torch +TRL logs the three terms separately, so `rewards/reward_win/mean` follows the wins (each one weighted by how fast it came). -gpu_stats = torch.cuda.get_device_properties(0) -start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3) -max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) +Each step plays 4 prompts x 16 games (64 games) with thinking on, and runs 150 steps at a constant learning rate of 1e-6 (in an earlier run with a linear decay to 0, the reward was still rising when the learning rate ran out): -print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.") -print(f"{start_gpu_memory} GB of memory reserved.") +```bash +hf jobs uv run --flavor a100-large --timeout 4h --secrets HF_TOKEN grpo.py \ + --model /Qwen3-1.7B-Wordle-SFT --hub-model-id /Qwen3-1.7B-Wordle-SFT-GRPO \ + --env-url https://-wordle.hf.space ``` -And train! +Every game is one environment session, so the environment must accept 64 at once. [Duplicate the `openenv/wordle` Space](https://huggingface.co/spaces/openenv/wordle?duplicate=true) and set the `MAX_CONCURRENT_ENVS=64` variable on your copy, rather than training against the shared Space. A private copy works too, since the script sends your token. -```python -trainer_stats = trainer.train() -``` +![Reward during GRPO](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/openenv/wordle-grpo-reward.png) -Show memory stats after training +In our run, the reward (averaged over 10 steps) went from about 0.37 to about 0.58, with a lot of noise from step to step. Most of the gain came from more wins and fewer invalid moves, while the reasoning stayed around 500 tokens per game and entropy stayed stable. A new run will give different numbers, but it should show the same trend. -```python -used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3) -used_memory_for_training = round(used_memory - start_gpu_memory, 3) -used_percentage = round(used_memory / max_memory * 100, 3) -training_memory_percentage = round(used_memory_for_training / max_memory * 100, 3) - -print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.") -print(f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.") -print(f"Peak reserved memory = {used_memory} GB.") -print(f"Peak reserved memory for training = {used_memory_for_training} GB.") -print(f"Peak reserved memory % of max memory = {used_percentage} %.") -print(f"Peak reserved memory for training % of max memory = {training_memory_percentage} %.") -``` +## 6. Evaluate and make the GIF -## Save and push to Hub +`evaluate.py` plays a model greedily and saves every game. With `--secrets`, it runs the environment in-process and sets the secret word itself, so every model plays the same words: -```python -trainer.save_model(output_dir) -trainer.push_to_hub() +```bash +uv run evaluate.py --model /Qwen3-1.7B-Wordle-SFT-GRPO --thinking --max-new-tokens 512 \ + --secrets "$(cat secrets30.txt)" --out grpo.jsonl +uv run make_gif.py base.jsonl sft.jsonl grpo.jsonl --labels "Qwen3-1.7B" "+ SFT" "+ SFT + GRPO" --out progress.gif ``` -## Load the fine-tuned model and run inference - -Now let's test our fine-tuned model by loading it and playing a game of Wordle. -We use the same `WordleEnv` class to interact with the environment, and generate model responses with standard Transformers inference. - -```python -from transformers import AutoModelForCausalLM, AutoTokenizer - -model_name = "sergiopaniego/wordle-grpo-Qwen3-1.7B" # Replace with your HF username or organization +The base model runs without thinking (`--max-new-tokens 256`). With thinking it writes thousands of tokens per guess, which is slow on a laptop and, in our tests, didn't help it win. -fine_tuned_model = AutoModelForCausalLM.from_pretrained(model_name, dtype="float32", device_map="auto") -tokenizer = AutoTokenizer.from_pretrained(model_name) -``` +## Run it in Colab -```python -import json - - -def play_wordle(model, tokenizer): - env = WordleEnv() - initial_observation = env.reset() - - print("Initial observation:") - print(initial_observation) - print() - - messages = [{"role": "user", "content": prompt}] - if initial_observation: - messages.append({"role": "user", "content": initial_observation}) - - for turn in range(6): - if env.done: - break - - prompt_text = tokenizer.apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=False, - enable_thinking=False, - ) - model_inputs = tokenizer([prompt_text], return_tensors="pt").to(model.device) - generated_ids = model.generate(**model_inputs, max_new_tokens=512) - output_ids = generated_ids[0][len(model_inputs.input_ids[0]):] - generated_text = tokenizer.decode(output_ids, skip_special_tokens=True) - - print(f"Turn {turn + 1} - Model output: {generated_text}") - - # Try to parse tool call from the generated text - try: - # Try to extract a guess from tool call format or bracket format - if "guess" in generated_text: - # Parse JSON tool call - start = generated_text.index("{") - end = generated_text.rindex("}") + 1 - args = json.loads(generated_text[start:end]) - if "arguments" in args: - args = args["arguments"] - guess_word = args.get("guess", "") - else: - # Fallback: extract from brackets [word] - import re - match = re.search(r"\[([a-zA-Z]{5})\]", generated_text) - guess_word = match.group(1) if match else generated_text.strip()[:5] - - feedback = env.guess(f"[{guess_word}]") - print(f" Guess: {guess_word} | Reward: {env.reward}") - print(f" Feedback: {feedback.strip()}") - print() - - messages.append({"role": "assistant", "content": generated_text}) - messages.append({"role": "user", "content": feedback}) - except Exception as e: - print(f" Error: {e}") - break - - print(f"Game finished! Final reward: {env.reward}") - print(f"Done: {env.done}") -``` +The [notebook](https://colab.research.google.com/github/openenv/OpenEnv/blob/main/examples/wordle/wordle.ipynb) runs the same pipeline. On a free T4 it loads the published models and plays them against each other, and launches SFT and GRPO on Hugging Face Jobs. On an A100 (Colab Pro) it can train in the notebook itself. -Let's play the game! +## Going further -```python -play_wordle(fine_tuned_model, tokenizer) -``` +- The [SFT warmup tutorial](sft-warmup) shows the same idea with rollouts collected from a stronger model. +- Training against a larger model (Qwen3-4B), more games per step or a richer SFT reasoning (listing every remaining candidate) are the next things to try for a higher win rate. diff --git a/envs/textarena_env/README.md b/envs/textarena_env/README.md index c830b8952..681ce7db3 100644 --- a/envs/textarena_env/README.md +++ b/envs/textarena_env/README.md @@ -23,6 +23,7 @@ An OpenEnv wrapper for [TextArena](https://github.com/textarena/textarena) game > Generic wrapper for any [TextArena](https://www.textarena.ai/docs/overview) game inside OpenEnv. This module exposes the TextArena `Env` interface through the standard HTTP server/client APIs used by other OpenEnv environments, enabling quick experimentation with the full suite of word, reasoning, and multi-agent games. Try Wordle on the [`openenv/wordle`](https://huggingface.co/spaces/openenv/wordle) Space, or connect a client to `https://openenv-wordle.hf.space`. +To train a model on it, follow the [Wordle tutorial](https://huggingface.co/docs/openenv/tutorials/wordle-grpo): an SFT warmup, then GRPO with TRL. ## Quick Start diff --git a/examples/wordle/README.md b/examples/wordle/README.md new file mode 100644 index 000000000..21f4f634a --- /dev/null +++ b/examples/wordle/README.md @@ -0,0 +1,25 @@ +# Wordle: SFT, then GRPO + +Train Qwen3-1.7B to play Wordle against the [`openenv/wordle`](https://huggingface.co/spaces/openenv/wordle) environment. The [tutorial](https://huggingface.co/docs/openenv/tutorials/wordle-grpo) walks through each step, with results, hardware and costs. + +| Script | What it does | Where it runs | +|---|---|---| +| `generate_sft_data.py` | A solver plays 1,000 games and writes them as tool-calling conversations | Laptop, 10 s | +| `sft.py` | SFT warmup on those games | Hugging Face Jobs, `a100-large` | +| `grpo.py` | GRPO against the environment, starting from the SFT model | Hugging Face Jobs, `a100-large` | +| `evaluate.py` | Plays a model and saves every game (`--secrets` fixes the words, so models play the same ones) | Laptop or any GPU | +| `make_gif.py` | Turns saved games into a side-by-side GIF | Laptop | +| `wordle.ipynb` | The same pipeline in a notebook | Colab | + +```bash +uv run generate_sft_data.py --games 1000 --push-to-hub /wordle-tool-calls-sft +hf jobs uv run --flavor a100-large --timeout 1h --secrets HF_TOKEN sft.py \ + --dataset /wordle-tool-calls-sft --hub-model-id /Qwen3-1.7B-Wordle-SFT +hf jobs uv run --flavor a100-large --timeout 4h --secrets HF_TOKEN grpo.py \ + --model /Qwen3-1.7B-Wordle-SFT --hub-model-id /Qwen3-1.7B-Wordle-SFT-GRPO \ + --env-url https://-wordle.hf.space +uv run evaluate.py --model /Qwen3-1.7B-Wordle-SFT-GRPO --thinking --max-new-tokens 512 \ + --secrets "$(cat secrets30.txt)" --out grpo.jsonl +``` + +GRPO plays 64 games at once, one environment session each. Duplicate the `openenv/wordle` Space and set `MAX_CONCURRENT_ENVS=64` on your copy before training. diff --git a/examples/wordle/evaluate.py b/examples/wordle/evaluate.py new file mode 100644 index 000000000..147ee15ac --- /dev/null +++ b/examples/wordle/evaluate.py @@ -0,0 +1,221 @@ +# /// script +# dependencies = [ +# "transformers", +# "torch", +# "openenv==0.10.0", +# "openenv-textarena @ git+https://github.com/openenv/OpenEnv.git@v0.10.0#subdirectory=envs/textarena_env", +# ] +# /// +""" +Play Wordle with a model and save every game, to compare models or make a GIF with `make_gif.py`. + +The model plays as in `grpo.py`: the same prompt and the same `guess` tool. Pass `--thinking` for models trained with +thinking on. With `--secrets`, the environment runs in-process and the secret words are fixed, so every model plays the +same games. Without it, games are played against `--env-url`. + +```sh +uv run evaluate.py --model Qwen/Qwen3-1.7B --max-new-tokens 256 --secrets "$(cat secrets30.txt)" --out base.jsonl +uv run evaluate.py --model openenv/Qwen3-1.7B-Wordle-SFT-GRPO --thinking --max-new-tokens 512 \ + --secrets "$(cat secrets30.txt)" --out grpo.jsonl +``` +""" + +import argparse +import json +import re +from contextlib import nullcontext +from types import SimpleNamespace + +import torch +from huggingface_hub import get_token +from textarena_env import TextArenaAction, TextArenaEnv +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.utils import get_json_schema + + +PROMPT = """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. +""" + + +def guess(guess: str) -> str: + """ + Make a guess in the Wordle environment. + + Args: + guess: The guessed word, formatted as '[abcde]' + + Returns: + The feedback message from the environment. + """ + + +TOOLS = [get_json_schema(guess)] +TOOL_CALL = re.compile(r"\s*(\{.*?\})\s*", re.S) +THINK = re.compile(r"(.*?)", re.S) +FEEDBACK = re.compile(r"Feedback:\n([A-Z ]+)\n([GYX ]+)") + + +def play(model, tokenizer, env, args): + result = env.reset() + full = result.observation.messages[0].content + # Same as TRL: the reset text is appended to the prompt with no separator. + messages = [{"role": "user", "content": PROMPT + full}] + rows, reward, done = [], 0.0, False + legal_guesses = 0 + # Allow one invalid warning per legal guess, but bound nonterminating servers. + while not done and legal_guesses < 6 and len(rows) < 12: + text = tokenizer.apply_chat_template( + messages, + tools=TOOLS, + add_generation_prompt=True, + tokenize=False, + enable_thinking=args.thinking, + ) + inputs = tokenizer(text, return_tensors="pt").to(model.device) + sampling = ( + {"do_sample": True, "temperature": args.temperature} + if args.temperature > 0 + else {"do_sample": False} + ) + output = model.generate( + **inputs, max_new_tokens=args.max_new_tokens, **sampling + ) + reply = tokenizer.decode( + output[0][inputs.input_ids.shape[1] :], skip_special_tokens=True + ) + call = TOOL_CALL.search(reply) + if call is None: + break # the model answered without calling the tool + word = json.loads(call.group(1)).get("arguments", {}).get("guess", "") + result = env.step(TextArenaAction(message=word)) + new_full = result.observation.messages[0].content + feedback, full = new_full[len(full) :], new_full + invalid = "invalid move" in feedback.lower() + legal_guesses += not invalid + reward, done = ( + (0.0 if invalid else result.reward or 0.0), + result.done, + ) + match = FEEDBACK.search(feedback) + thinking = THINK.search(reply) + reasoning = thinking.group(1).strip() if thinking else "" + rows.append( + { + "guess": word, + "letters": match.group(1).replace(" ", "") + if match + else word.strip("[]").upper(), + "colors": match.group(2).replace(" ", "") + if match + else "G" * 5 + if reward >= 1.0 + else None, + "reasoning": reasoning, + } + ) + messages += [ + { + "role": "assistant", + "content": "", + "reasoning_content": reasoning, + "tool_calls": [ + { + "type": "function", + "function": {"name": "guess", "arguments": {"guess": word}}, + } + ], + }, + {"role": "tool", "name": "guess", "content": feedback}, + ] + secret = env.state().raw_state.get("game_state", {}).get("secret_word") + return {"rows": rows, "reward": reward, "won": reward >= 1.0, "secret": secret} + + +class LocalWordle: + """The Wordle server class run in-process, with the secret word set by us, so every model plays the same words.""" + + def __init__(self, secrets): + from textarena_env.server.environment import TextArenaEnvironment + + self.env, self.secrets = TextArenaEnvironment("Wordle-v0"), iter(secrets) + + def reset(self): + observation = self.env.reset() + game = self.env._ta_env + while hasattr(game, "env"): + game = game.env + game.state.game_state["secret_word"] = next(self.secrets) + return SimpleNamespace(observation=observation) + + def step(self, action): + observation = self.env.step(action) + return SimpleNamespace( + observation=observation, reward=observation.reward, done=observation.done + ) + + def state(self): + return self.env.state + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", default="Qwen/Qwen3-1.7B") + parser.add_argument("--env-url", default="https://openenv-wordle.hf.space") + parser.add_argument("--games", type=int, default=20) + parser.add_argument("--max-new-tokens", type=int, default=256) + parser.add_argument("--thinking", action="store_true") + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--out", required=True) + parser.add_argument( + "--secrets", + help="comma-separated words: play them in-process instead of against --env-url", + ) + args = parser.parse_args() + + tokenizer = AutoTokenizer.from_pretrained(args.model) + device = ( + "cuda" + if torch.cuda.is_available() + else "mps" + if torch.backends.mps.is_available() + else "cpu" + ) + # T4 GPUs (free Colab) have no bf16 support. + dtype = ( + torch.float16 + if device == "cuda" and not torch.cuda.is_bf16_supported() + else torch.bfloat16 + ) + model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype).to(device) + # A private Space needs your token. A public one ignores it. + headers = {"Authorization": f"Bearer {get_token()}"} if get_token() else None + + if args.secrets: + secrets = args.secrets.split(",") + args.games = len(secrets) + env_context = nullcontext(LocalWordle(secrets)) + else: + env_context = TextArenaEnv(base_url=args.env_url, headers=headers).sync() + with env_context as env, open(args.out, "w") as f: + for i in range(args.games): + game = play(model, tokenizer, env, args) + f.write(json.dumps({"model": args.model, **game}) + "\n") + print( + f"game {i + 1}: {'won' if game['won'] else 'lost'} in {len(game['rows'])} guesses, reward {game['reward']}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/wordle/generate_sft_data.py b/examples/wordle/generate_sft_data.py new file mode 100644 index 000000000..9fca4f860 --- /dev/null +++ b/examples/wordle/generate_sft_data.py @@ -0,0 +1,268 @@ +# /// script +# dependencies = [ +# "openenv==0.10.0", +# "openenv-textarena @ git+https://github.com/openenv/OpenEnv.git@v0.10.0#subdirectory=envs/textarena_env", +# "transformers", +# "datasets", +# ] +# /// +""" +Generate Wordle games for a short SFT warmup before GRPO. + +A solver plays real games against the TextArena Wordle environment (run in-process) and writes each one as a +tool-calling conversation: a short reasoning and a `guess` call per turn, then the environment feedback. The format +matches what TRL's `GRPOTrainer` renders with `environment_factory`, so the model plays the same way during RL. + +```sh +uv run generate_sft_data.py --games 1000 --push-to-hub /wordle-tool-calls-sft +``` +""" + +import argparse +import json +import random +from collections import Counter + +from datasets import Dataset +from textarena_env import TextArenaAction +from textarena_env.server.environment import TextArenaEnvironment +from transformers.utils import get_json_schema + + +PROMPT = """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. +""" + +OPENERS = [ + "crane", + "slate", + "raise", + "trace", + "stare", + "crate", + "arise", + "least", + "alert", + "roast", + "saint", + "stale", + "snare", + "trade", +] +NAMES = {"G": "green", "Y": "yellow", "X": "gray"} +ORDINALS = {1: "1st", 2: "2nd", 3: "3rd", 4: "4th", 5: "5th"} + + +def guess(guess: str) -> str: + """ + Make a guess in the Wordle environment. + + Args: + guess: The guessed word, formatted as '[abcde]' + + Returns: + The feedback message from the environment. + """ + + +def score(word, secret): + feedback, rest = ["X"] * 5, list(secret) + for i in range(5): + if word[i] == secret[i]: + feedback[i], rest[i] = "G", None + for i in range(5): + if feedback[i] == "X" and word[i] in rest: + feedback[i] = "Y" + rest[rest.index(word[i])] = None + return "".join(feedback) + + +def pick(candidates, rng): + # The guess that leaves the fewest candidates on average, at random among near-ties. + def cost(word): + return sum( + n * n + for n in Counter(score(word, secret) for secret in candidates).values() + ) + + ranked = sorted(rng.sample(candidates, min(len(candidates), 150)), key=cost) + return rng.choice( + [word for word in ranked[:3] if cost(word) <= cost(ranked[0]) + 2] + ) + + +def reasoning(history, candidates, choice, rng): + if not history: + letters = ", ".join(c.upper() for c in choice) + vowels = " and ".join(c.upper() for c in choice if c in "aeiou") + return rng.choice( + [ + f"No feedback yet, so I want an opener with common letters. {choice.upper()} tests {letters}, including the vowels {vowels}.", + f"First guess. I'll start with {choice.upper()}: five distinct, frequent letters ({letters}) that should narrow the word down quickly.", + f"Nothing to go on yet. A good opener covers frequent letters and vowels, so I'll try {choice.upper()}.", + ] + ) + word, feedback = history[-1] + lines = [ + f"{word.upper()} gave {' '.join(feedback)}: " + + ", ".join(f"{c.upper()} {NAMES[f]}" for c, f in zip(word, feedback)) + + "." + ] + greens, misplaced, found = ["_"] * 5, {}, set() + for word, feedback in history: + for i, (c, f) in enumerate(zip(word, feedback)): + if f == "G": + greens[i] = c.upper() + if f in "GY": + found.add(c) + if f == "Y": + misplaced.setdefault(c, set()).add(i + 1) + grays = { + c + for word, feedback in history + for c, f in zip(word, feedback) + if f == "X" and c not in found + } + known = [] + if greens != ["_"] * 5: + known.append(f"Pattern so far: {' '.join(greens)}.") + if misplaced: + known.append( + "In the word but misplaced: " + + ", ".join( + f"{c.upper()} (not {'/'.join(ORDINALS[p] for p in sorted(ps))})" + for c, ps in sorted(misplaced.items()) + ) + + "." + ) + if grays: + known.append(f"Ruled out: {', '.join(sorted(c.upper() for c in grays))}.") + lines.append(" ".join(known)) + if len(candidates) == 1: + lines.append( + rng.choice( + [ + f"Only {choice.upper()} fits all of that, so it should be the answer.", + f"The one word that fits every clue is {choice.upper()}.", + ] + ) + ) + elif len(candidates) == 2: + other = next(w for w in candidates if w != choice) + lines.append( + f"Two words fit: {choice.upper()} and {other.upper()}. I'll try {choice.upper()} first." + ) + else: + shown = [choice] + [w for w in candidates if w != choice][: rng.randint(2, 4)] + rng.shuffle(shown) + lines.append( + f"Words that fit: {', '.join(w.upper() for w in shown)}" + + (", and a few more." if len(candidates) > len(shown) else ".") + ) + lines.append( + rng.choice( + [ + f"{choice.upper()} fits the clues and its letters split the remaining options best, so I'll guess it.", + f"I'll go with {choice.upper()}: it is consistent with everything and tests the letters that separate these words.", + f"{choice.upper()} is consistent with all feedback and tells the others apart, so it's my next guess.", + ] + ) + ) + return "\n".join(lines) + + +def play(env, rng): + full = env.reset().messages[0].content + game = env._ta_env + while hasattr(game, "env"): + game = game.env + secret = game.state.game_state["secret_word"] + # TRL appends the reset text to the prompt with no separator. + messages = [{"role": "user", "content": PROMPT + full}] + candidates, history, observation = list(game.word_list), [], None + while observation is None or not observation.done: + choice = rng.choice(OPENERS) if not history else pick(candidates, rng) + call = [ + { + "type": "function", + "function": {"name": "guess", "arguments": {"guess": f"[{choice}]"}}, + } + ] + messages.append( + { + "role": "assistant", + "content": "", + "reasoning_content": reasoning(history, candidates, choice, rng), + "tool_calls": call, + } + ) + observation = env.step(TextArenaAction(message=f"[{choice}]")) + new_full = observation.messages[0].content + messages.append( + {"role": "tool", "name": "guess", "content": new_full[len(full) :]} + ) + full = new_full + history.append((choice, score(choice, secret))) + candidates = [ + w for w in candidates if score(choice, w) == history[-1][1] and w != choice + ] + word = secret.upper() + closing = rng.choice( + [ + f"The game says [{secret}] was correct, so the game is over.", + f"{word} was the word: the game congratulated me and there are no guesses left to make.", + f"Correct on guess {len(history)}. Nothing more to do.", + ] + ) + messages.append( + { + "role": "assistant", + "content": f"Solved: the word is {word}.", + "reasoning_content": closing, + } + ) + return { + "messages": messages, + "tools": json.dumps([get_json_schema(guess)]), + }, observation.reward >= 1.0 + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--games", type=int, default=1000) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--out", default="wordle_sft.jsonl") + parser.add_argument( + "--push-to-hub", + help="dataset repo to push to, e.g. /wordle-tool-calls-sft", + ) + args = parser.parse_args() + + random.seed( + args.seed + ) # TextArena draws the secret word from the global random state + rng, env, rows = random.Random(args.seed), TextArenaEnvironment("Wordle-v0"), [] + while len(rows) < args.games: + row, won = play(env, rng) + if won: + rows.append(row) + with open(args.out, "w") as f: + f.writelines(json.dumps(row) + "\n" for row in rows) + print(f"wrote {len(rows)} games to {args.out}") + if args.push_to_hub: + Dataset.from_list(rows).push_to_hub(args.push_to_hub) + + +if __name__ == "__main__": + main() diff --git a/examples/wordle/grpo.py b/examples/wordle/grpo.py new file mode 100644 index 000000000..3b48fff61 --- /dev/null +++ b/examples/wordle/grpo.py @@ -0,0 +1,177 @@ +# /// script +# dependencies = [ +# "trl[vllm]==1.16.0", +# "trackio", +# "openenv==0.10.0", +# "openenv-textarena @ git+https://github.com/openenv/OpenEnv.git@v0.10.0#subdirectory=envs/textarena_env", +# ] +# /// +""" +GRPO on the OpenEnv Wordle environment, starting from the SFT warmup model, with thinking on. + +```sh +hf jobs uv run --flavor a100-large --timeout 4h --secrets HF_TOKEN grpo.py \ + --model /Qwen3-1.7B-Wordle-SFT --hub-model-id /Qwen3-1.7B-Wordle-SFT-GRPO +``` +""" + +import argparse +import re + +import torch +from datasets import Dataset +from huggingface_hub import get_token +from textarena_env import TextArenaAction, TextArenaEnv +from trl import GRPOConfig, GRPOTrainer + + +PROMPT = """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. +""" + +FEEDBACK = re.compile(r"Feedback:\n([A-Z ]+)\n([GYX ]+)") + + +# The environment pays a win the same at any speed and pays partial credit only at the end of a game, so a policy can +# learn to think less and spend every guess. These terms reward winning fast, give partial credit from the best guess +# whether or not the game finished, and penalize invalid moves. Each one is logged separately. +def reward_win(environments, **kwargs) -> list[float]: + return [ + 1.0 + 0.5 * (1 - env.num_guesses / 6) if env.won else 0.0 + for env in environments + ] + + +def reward_partial(environments, **kwargs) -> list[float]: + return [0.0 if env.won else 0.1 * env.best_greens for env in environments] + + +def reward_invalid(environments, **kwargs) -> list[float]: + return [-0.1 * env.invalid_moves for env in environments] + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True) + parser.add_argument("--hub-model-id", required=True) + parser.add_argument("--env-url", default="https://openenv-wordle.hf.space") + parser.add_argument("--max-steps", type=int, default=150) + parser.add_argument("--trackio-space-id", default=None) + args = parser.parse_args() + env_url = args.env_url + # A private Space needs your token. A public one ignores it. + headers = {"Authorization": f"Bearer {get_token()}"} if get_token() else None + + class WordleEnv: + def __init__(self): + self.client = TextArenaEnv(base_url=env_url, headers=headers) + + def reset(self, **kwargs) -> str | None: + result = self.client.reset() + # The game returns cumulative feedback each turn, so keep the full text and return only the new part. + self._last_full_feedback = result.observation.messages[0].content + self.done = False + self.won = False + self.num_guesses = 0 + self.best_greens = 0 + self.invalid_moves = 0 + return self._last_full_feedback + + def guess(self, guess: str) -> str: + """ + Make a guess in the Wordle environment. + + Args: + guess: The guessed word, formatted as '[abcde]' + + Returns: + The feedback message from the environment. + """ + if self.done or self.num_guesses + self.invalid_moves >= 12: + raise ValueError("Game over.") + result = self.client.step(TextArenaAction(message=guess)) + full_feedback = result.observation.messages[0].content + feedback = full_feedback[len(self._last_full_feedback) :] + self._last_full_feedback = full_feedback + if "invalid move" in feedback.lower(): + self.invalid_moves += 1 + else: + self.num_guesses += 1 + match = FEEDBACK.search(feedback) + if match: + self.best_greens = max(self.best_greens, match.group(2).count("G")) + self.won = result.done and result.reward >= 1.0 + self.done = result.done + return feedback + + # 4 prompts x 16 games per step. Each game is one environment session, so the Space must allow 64 at once. + prompts_per_step, num_generations, per_device_batch_size = 4, 16, 8 + dataset = Dataset.from_dict( + { + "prompt": [[{"role": "user", "content": PROMPT}]] + * (args.max_steps * prompts_per_step) + } + ) + + trainer = GRPOTrainer( + model=args.model, + reward_funcs=[reward_win, reward_partial, reward_invalid], + train_dataset=dataset, + environment_factory=WordleEnv, + args=GRPOConfig( + output_dir=args.hub_model_id.split("/")[-1], + max_steps=args.max_steps, + learning_rate=1e-6, + # Constant: with a linear decay to 0 the reward was still rising when the learning rate ran out. + lr_scheduler_type="constant", + num_generations=num_generations, + per_device_train_batch_size=per_device_batch_size, + gradient_accumulation_steps=prompts_per_step + * num_generations + // per_device_batch_size, + gradient_checkpointing=True, + bf16=True, + chat_template_kwargs={"enable_thinking": True}, + max_completion_length=4096, + mask_truncated_completions=True, + # Multi-turn rollouts: clip per-token ratios instead of masking whole sequences. + vllm_importance_sampling_mode="token_truncate", + # Invalid warnings use tool rounds, but not legal guesses. + # The adapter caps submissions at 12; the environment still allows six legal guesses. + max_tool_calling_iterations=12, + use_vllm=True, + vllm_mode="colocate", + vllm_gpu_memory_utilization=0.35, + report_to="trackio", + trackio_space_id=args.trackio_space_id, + project="openenv-wordle", + run_name="grpo", + logging_steps=1, + log_completions=True, + num_completions_to_print=1, + save_steps=25, + save_total_limit=2, + push_to_hub=True, + hub_model_id=args.hub_model_id, + hub_private_repo=False, + ), + ) + trainer.train() + # Share the weights in bf16, as the SFT model. + trainer.model.to(torch.bfloat16) + trainer.push_to_hub() + + +if __name__ == "__main__": + main() diff --git a/examples/wordle/make_gif.py b/examples/wordle/make_gif.py new file mode 100644 index 000000000..7eccd6df2 --- /dev/null +++ b/examples/wordle/make_gif.py @@ -0,0 +1,146 @@ +# /// script +# dependencies = ["pillow"] +# /// +""" +Turn games saved by `evaluate.py` into a side-by-side GIF, one column per model. + +```sh +uv run make_gif.py base.jsonl sft.jsonl grpo.jsonl --labels "Qwen3-1.7B" "+ SFT" "+ SFT + GRPO" --out progress.gif +``` +""" + +import argparse +import json + +from PIL import Image, ImageDraw, ImageFont + + +TILE, GAP, PAD = 52, 6, 24 +COLORS = {"G": (106, 170, 100), "Y": (201, 180, 88), "X": (120, 124, 126)} +EMPTY, BORDER, TEXT, BG = ( + (255, 255, 255), + (211, 214, 218), + (255, 255, 255), + (250, 250, 250), +) + + +def font(size): + for name in ("DejaVuSans-Bold.ttf", "Arial Bold.ttf", "Helvetica.ttc"): + try: + return ImageFont.truetype(name, size) + except OSError: + pass + return ImageFont.load_default(size) + + +def board(draw, x0, y0, game, rows_shown): + big = font(28) + # Rejected submissions have no colors and do not consume a Wordle row. + rows = [row for row in game["rows"][:rows_shown] if row["colors"]] + for r in range(6): + row = rows[r] if r < len(rows) else None + for c in range(5): + x, y = x0 + c * (TILE + GAP), y0 + r * (TILE + GAP) + colors = row and row["colors"] + fill = COLORS[colors[c]] if colors else EMPTY + draw.rectangle( + [x, y, x + TILE, y + TILE], + fill=fill, + outline=None if colors else BORDER, + width=2, + ) + if row and c < len(row["letters"]): + draw.text( + (x + TILE / 2, y + TILE / 2), + row["letters"][c], + font=big, + anchor="mm", + fill=TEXT if colors else (60, 60, 60), + ) + + +def frames(columns, labels): + width = 5 * TILE + 4 * GAP + height = 6 * TILE + 5 * GAP + W = PAD + len(columns) * (width + PAD) + H = 3 * PAD + 40 + height + 70 + games = min(len(c) for c in columns) + out = [] + for g in range(games): + wins = [sum(x["won"] for x in column[:g]) for column in columns] + # One frame per submission, including invalids; preserve legacy pacing. + submissions = max(6, *(len(column[g]["rows"]) for column in columns)) + for shown in range(1, submissions + 1): + img = Image.new("RGB", (W, H), BG) + draw = ImageDraw.Draw(img) + for i, (column, label) in enumerate(zip(columns, labels)): + game, x0 = column[g], PAD + i * (width + PAD) + done = shown >= len(game["rows"]) + won_so_far = wins[i] + (done and game["won"]) + draw.text( + (x0 + width / 2, PAD), + label, + font=font(20), + anchor="mt", + fill=(30, 30, 30), + ) + draw.text( + (x0 + width / 2, PAD + 26), + f"wins {won_so_far}/{g + 1}", + font=font(15), + anchor="mt", + fill=(90, 90, 90), + ) + board(draw, x0, 2 * PAD + 40, game, shown) + visible = game["rows"][:shown] + if visible and not visible[-1]["colors"]: + # Keep rejected input out of the board and verdict area. + letters = visible[-1]["letters"] + letters = letters[:13] + "..." if len(letters) > 16 else letters + draw.text( + (x0 + width / 2, 2 * PAD + 40 + height + 14), + f"invalid: {letters}", + font=font(15), + anchor="mm", + fill=(180, 60, 60), + ) + if done: + verdict = ( + "solved" + if game["won"] + else f"missed: {(game['secret'] or '?').upper()}" + ) + draw.text( + (x0 + width / 2, H - PAD - 10), + verdict, + font=font(18), + anchor="mm", + fill=COLORS["G"] if game["won"] else (180, 60, 60), + ) + out.append(img) + out += [out[-1]] * 3 # hold the final board + return out + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "games", nargs="+", help="jsonl files written by wordle_eval.py, one per model" + ) + parser.add_argument("--labels", nargs="+") + parser.add_argument("--out", default="progress.gif") + args = parser.parse_args() + + columns = [] + for path in args.games: + with open(path) as games_file: + columns.append([json.loads(line) for line in games_file]) + labels = args.labels or [c[0]["model"].split("/")[-1] for c in columns] + imgs = frames(columns, labels) + imgs[0].save(args.out, save_all=True, append_images=imgs[1:], duration=450, loop=0) + print(f"wrote {args.out} ({len(imgs)} frames)") + + +if __name__ == "__main__": + main() diff --git a/examples/wordle/secrets30.txt b/examples/wordle/secrets30.txt new file mode 100644 index 000000000..d0a25438e --- /dev/null +++ b/examples/wordle/secrets30.txt @@ -0,0 +1 @@ +snake,front,dress,chain,verse,level,plane,steam,sugar,sheep,thing,crime,paper,cause,slope,watch,angle,nerve,stamp,light,scale,value,price,awake,chalk,steel,offer,stick,smile,jewel diff --git a/examples/wordle/sft.py b/examples/wordle/sft.py new file mode 100644 index 000000000..29c9cf3c0 --- /dev/null +++ b/examples/wordle/sft.py @@ -0,0 +1,66 @@ +# /// script +# dependencies = [ +# "trl==1.16.0", +# "trackio", +# ] +# /// +""" +Short SFT warmup on Wordle tool-calling games, before GRPO. + +```sh +hf jobs uv run --flavor a100-large --timeout 1h --secrets HF_TOKEN sft.py \ + --dataset /wordle-tool-calls-sft --hub-model-id /Qwen3-1.7B-Wordle-SFT +``` +""" + +import argparse + +import torch +from datasets import load_dataset +from trl import SFTConfig, SFTTrainer + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", default="Qwen/Qwen3-1.7B") + parser.add_argument("--dataset", required=True) + parser.add_argument("--hub-model-id", required=True) + parser.add_argument("--trackio-space-id", default=None) + args = parser.parse_args() + + trainer = SFTTrainer( + model=args.model, + train_dataset=load_dataset(args.dataset, split="train"), + args=SFTConfig( + output_dir=args.hub_model_id.split("/")[-1], + # Train every assistant turn: reasoning, tool calls and the closing turn. + assistant_only_loss=True, + max_length=2048, + num_train_epochs=1, + per_device_train_batch_size=8, + gradient_accumulation_steps=2, + learning_rate=1e-5, + lr_scheduler_type="cosine", + warmup_steps=3, + # Keep small weight updates in fp32 while computing in bf16. + model_init_kwargs={"dtype": "float32"}, + bf16=True, + gradient_checkpointing=True, + logging_steps=1, + report_to="trackio", + trackio_space_id=args.trackio_space_id, + project="openenv-wordle", + run_name="sft", + push_to_hub=True, + hub_model_id=args.hub_model_id, + hub_private_repo=False, + ), + ) + trainer.train() + # Cast only after training, to share the weights in bf16. + trainer.model.to(torch.bfloat16) + trainer.push_to_hub(dataset_name=args.dataset) + + +if __name__ == "__main__": + main() diff --git a/examples/wordle/wordle.ipynb b/examples/wordle/wordle.ipynb new file mode 100644 index 000000000..a26f24667 --- /dev/null +++ b/examples/wordle/wordle.ipynb @@ -0,0 +1,216 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Train a Wordle agent with OpenEnv: SFT, then GRPO\n", + "\n", + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/openenv/OpenEnv/blob/main/examples/wordle/wordle.ipynb)\n", + "\n", + "This notebook follows the [OpenEnv Wordle tutorial](https://huggingface.co/docs/openenv/tutorials/wordle-grpo). Qwen3-1.7B can't win a game of Wordle. A short SFT warmup teaches it the game, and GRPO against the [`openenv/wordle`](https://huggingface.co/spaces/openenv/wordle) environment cuts its mistakes.\n", + "\n", + "- **On a free T4** you can play the environment, generate the SFT data, launch SFT and GRPO on [Hugging Face Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs) and compare the published models.\n", + "- **On an A100 (Colab Pro)** you can also run SFT and GRPO in the notebook itself." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "!pip install -q \"openenv==0.10.0\" \"openenv-textarena @ git+https://github.com/openenv/OpenEnv.git@v0.10.0#subdirectory=envs/textarena_env\" \"trl==1.16.0\" trackio" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Log in to Hugging Face to push datasets and models and to launch Jobs. You need a token with write access." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from huggingface_hub import notebook_login\n", + "\n", + "notebook_login()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Play the environment\n", + "\n", + "The environment runs as a Space. Each guess returns the feedback for every letter: `G` (right letter, right spot), `Y` (in the word, wrong spot) and `X` (not in the word)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from textarena_env import TextArenaAction, TextArenaEnv\n", + "\n", + "with TextArenaEnv(base_url=\"https://openenv-wordle.hf.space\").sync() as env:\n", + " result = env.reset()\n", + " print(result.observation.messages[0].content)\n", + " result = env.step(TextArenaAction(message=\"[crane]\"))\n", + " print(result.observation.messages[0].content[-200:])" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2. Generate the SFT data\n", + "\n", + "A solver plays 1,000 games against the same environment class the Space runs (in-process, so it takes seconds) and writes each one as a tool-calling conversation, with a short reasoning before every guess. Enter your Hugging Face username or organization below. Set `env_url` to the URL of your duplicated Wordle Space before starting GRPO." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import subprocess\n", + "import sys\n", + "\n", + "hub_namespace = input(\"Hugging Face username or organization: \").strip()\n", + "if not hub_namespace or any(c not in \"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_\" for c in hub_namespace):\n", + " raise ValueError(\"Enter a valid Hub username or organization\")\n", + "env_url = input(\"URL of your duplicated Wordle Space: \").strip()\n", + "if not env_url.startswith(\"https://\") or any(c.isspace() for c in env_url) or any(c in env_url for c in \"<>\"):\n", + " raise ValueError(\"Enter your Space's HTTPS URL\")\n", + "dataset_id = f\"{hub_namespace}/wordle-tool-calls-sft\"\n", + "sft_model_id = f\"{hub_namespace}/Qwen3-1.7B-Wordle-SFT\"\n", + "grpo_model_id = f\"{hub_namespace}/Qwen3-1.7B-Wordle-SFT-GRPO\"\n", + "\n", + "!wget -q https://raw.githubusercontent.com/openenv/OpenEnv/main/examples/wordle/generate_sft_data.py https://raw.githubusercontent.com/openenv/OpenEnv/main/examples/wordle/sft.py https://raw.githubusercontent.com/openenv/OpenEnv/main/examples/wordle/grpo.py https://raw.githubusercontent.com/openenv/OpenEnv/main/examples/wordle/evaluate.py https://raw.githubusercontent.com/openenv/OpenEnv/main/examples/wordle/make_gif.py https://raw.githubusercontent.com/openenv/OpenEnv/main/examples/wordle/secrets30.txt\n", + "subprocess.run([sys.executable, \"generate_sft_data.py\", \"--games\", \"1000\", \"--push-to-hub\", dataset_id], check=True)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 3. SFT and GRPO on Hugging Face Jobs\n", + "\n", + "Both steps run on one A100 80GB (`a100-large`, $2.50/h). The SFT warmup takes 7 min and GRPO 2 h 15 min, about $6 in total.\n", + "\n", + "GRPO plays 64 games at once, one environment session each. [Duplicate the `openenv/wordle` Space](https://huggingface.co/spaces/openenv/wordle?duplicate=true), set the `MAX_CONCURRENT_ENVS=64` variable on your copy and pass its URL as `--env-url`." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "subprocess.run([\n", + " \"hf\", \"jobs\", \"uv\", \"run\", \"--flavor\", \"a100-large\", \"--timeout\", \"1h\",\n", + " \"--secrets\", \"HF_TOKEN\", \"sft.py\",\n", + " \"--dataset\", dataset_id, \"--hub-model-id\", sft_model_id,\n", + "], check=True)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Once the SFT job finishes, start GRPO from the SFT model. Follow the run with `hf jobs logs ` and on Trackio." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "subprocess.run([\n", + " \"hf\", \"jobs\", \"uv\", \"run\", \"--flavor\", \"a100-large\", \"--timeout\", \"4h\",\n", + " \"--secrets\", \"HF_TOKEN\", \"grpo.py\",\n", + " \"--model\", sft_model_id, \"--hub-model-id\", grpo_model_id,\n", + " \"--env-url\", env_url,\n", + "], check=True)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Or train in the notebook (A100)\n", + "\n", + "With an A100 runtime, run the same scripts here instead. GRPO needs vLLM." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# !pip install -q \"trl[vllm]==1.16.0\"\n", + "# subprocess.run([sys.executable, \"sft.py\", \"--dataset\", dataset_id, \"--hub-model-id\", sft_model_id], check=True)\n", + "# subprocess.run([sys.executable, \"grpo.py\", \"--model\", sft_model_id, \"--hub-model-id\", grpo_model_id, \"--env-url\", env_url], check=True)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 4. Compare the models\n", + "\n", + "`evaluate.py` plays a model greedily and saves every game. `--secrets` runs the environment in-process with fixed secret words, so every model plays the same ones. This uses the published models, swap in yours once they're trained. The cells below play the first 10 of the 30 words. On a laptop with Apple silicon, all 30 take 3 to 10 minutes per model. The evaluator allows six legal guesses (up to 12 submissions including invalid warnings) and stops when the environment ends the game. The tutorial’s published scores used the earlier six-submission cap and have not been recomputed with this budget." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "!python evaluate.py --model Qwen/Qwen3-1.7B --max-new-tokens 256 --secrets \"$(head -c 59 secrets30.txt)\" --out base.jsonl\n", + "!python evaluate.py --model openenv/Qwen3-1.7B-Wordle-SFT --thinking --max-new-tokens 512 --secrets \"$(head -c 59 secrets30.txt)\" --out sft.jsonl\n", + "!python evaluate.py --model openenv/Qwen3-1.7B-Wordle-SFT-GRPO --thinking --max-new-tokens 512 --secrets \"$(head -c 59 secrets30.txt)\" --out grpo.jsonl" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "!pip install -q pillow\n", + "!python make_gif.py base.jsonl sft.jsonl grpo.jsonl --labels \"Qwen3-1.7B\" \"+ SFT\" \"+ SFT + GRPO\" --out progress.gif\n", + "\n", + "from IPython.display import Image\n", + "\n", + "Image(\"progress.gif\")" + ], + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "gpuType": "T4", + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/tests/scripts/test_wordle_evaluate.py b/tests/scripts/test_wordle_evaluate.py new file mode 100644 index 000000000..4673e0646 --- /dev/null +++ b/tests/scripts/test_wordle_evaluate.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: BSD-3-Clause +"""Exercise the evaluation loop without loading model weights or a server.""" + +import ast +import json +import re +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +ROOT = Path(__file__).resolve().parents[2] + + +def load_play(): + tree = ast.parse((ROOT / "examples/wordle/evaluate.py").read_text()) + names = {"PROMPT", "TOOL_CALL", "THINK", "FEEDBACK"} + nodes = [ + node + for node in tree.body + if (isinstance(node, ast.FunctionDef) and node.name == "play") + or ( + isinstance(node, ast.Assign) + and any(isinstance(t, ast.Name) and t.id in names for t in node.targets) + ) + ] + namespace = dict(json=json, re=re, TextArenaAction=SimpleNamespace, TOOLS=[]) + exec( + compile(ast.Module(body=nodes, type_ignores=[]), "evaluate.py", "exec"), + namespace, + ) + return namespace["play"] + + +class Inputs(dict): + input_ids = SimpleNamespace(shape=(1, 0)) + + def to(self, device): + return self + + +class Tokenizer: + def apply_chat_template(self, *args, **kwargs): + return "prompt" + + def __call__(self, *args, **kwargs): + return Inputs() + + def decode(self, *args, **kwargs): + return '{"arguments": {"guess": "[apple]"}}' + + +class Environment: + def __init__(self, invalid_turns=(), terminal_at=None, win_at=None): + self.invalid_turns = invalid_turns + self.terminal_at = terminal_at + self.win_at = win_at + self.calls = 0 + self.full = "" + + def reset(self): + return SimpleNamespace( + observation=SimpleNamespace(messages=[SimpleNamespace(content=self.full)]) + ) + + def step(self, action): + self.calls += 1 + invalid = self.calls in self.invalid_turns + self.full += ( + "Invalid Move: rejected" if invalid else "Feedback:\nA P P L E\nX X X X X" + ) + return SimpleNamespace( + observation=self.reset().observation, + reward=1.0 if self.calls == self.win_at else 0.0, + done=self.calls == self.terminal_at, + ) + + def state(self): + return SimpleNamespace(raw_state={"game_state": {"secret_word": "apple"}}) + + +@pytest.mark.parametrize( + "invalid_turns,terminal_at,win_at,expected_calls,won", + [ + ({1}, 7, 7, 7, True), + ({1, 3, 5, 7, 9, 11}, 12, 12, 12, True), + (set(), None, None, 6, False), + (set(range(1, 20)), None, None, 12, False), + ({1, 2}, 2, 2, 2, False), + ], +) +def test_guess_budget(invalid_turns, terminal_at, win_at, expected_calls, won): + env = Environment(invalid_turns, terminal_at, win_at) + model = SimpleNamespace(device="cpu", generate=lambda **kwargs: [[]]) + args = SimpleNamespace(thinking=False, temperature=0, max_new_tokens=16) + result = load_play()(model, Tokenizer(), env, args) + assert env.calls == expected_calls + assert len(result["rows"]) == expected_calls + assert result["won"] is won + + +def test_notebook_commands_use_configured_arguments(monkeypatch): + import subprocess + + notebook = json.loads((ROOT / "examples/wordle/wordle.ipynb").read_text()) + answers = iter(["example-org", "https://example-org-wordle.hf.space"]) + monkeypatch.setattr("builtins.input", lambda prompt: next(answers)) + calls = [] + monkeypatch.setattr( + subprocess, "run", lambda argv, **kwargs: calls.append((argv, kwargs)) + ) + namespace = {} + for cell in notebook["cells"]: + source = "".join(cell["source"]) + if cell["cell_type"] == "code" and "subprocess.run" in source: + # Skip the download magic; all publication/job calls are intercepted. + source = "\n".join( + line for line in source.splitlines() if not line.startswith("!") + ) + exec(compile(source, "wordle.ipynb", "exec"), namespace) + assert len(calls) == 3 + for argv, kwargs in calls: + assert kwargs == {"check": True} + assert not any("" in arg for arg in argv) + assert calls[0][0][-1] == "example-org/wordle-tool-calls-sft" + assert calls[1][0][-1] == "example-org/Qwen3-1.7B-Wordle-SFT" + assert calls[2][0][-1] == "https://example-org-wordle.hf.space" + + +@pytest.mark.parametrize( + "answers", [[""], ["example-org", "https://-wordle.hf.space"]] +) +def test_notebook_rejects_unconfigured_placeholders(monkeypatch, answers): + import subprocess + + notebook = json.loads((ROOT / "examples/wordle/wordle.ipynb").read_text()) + replies = iter(answers) + monkeypatch.setattr("builtins.input", lambda prompt: next(replies)) + monkeypatch.setattr( + subprocess, "run", lambda *a, **kw: pytest.fail("Command ran before validation") + ) + source = next( + "".join(c["source"]) + for c in notebook["cells"] + if "hub_namespace = input" in "".join(c["source"]) + ) + source = "\n".join(line for line in source.splitlines() if not line.startswith("!")) + with pytest.raises(ValueError): + exec(compile(source, "wordle.ipynb", "exec"), {}) diff --git a/tests/scripts/test_wordle_gif.py b/tests/scripts/test_wordle_gif.py new file mode 100644 index 000000000..11560cacc --- /dev/null +++ b/tests/scripts/test_wordle_gif.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: BSD-3-Clause +"""Render synthetic evaluation traces without inference or a server.""" + +import importlib.util +from pathlib import Path + +import pytest + + +pytest.importorskip("PIL") +ROOT = Path(__file__).resolve().parents[2] + + +@pytest.fixture +def renderer(monkeypatch): + spec = importlib.util.spec_from_file_location( + "wordle_gif", ROOT / "examples/wordle/make_gif.py" + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + texts = [] + original = module.ImageDraw.Draw + + def record(image): + draw = original(image) + text = draw.text + labels = [] + texts.append(labels) + + def capture(xy, value, **kwargs): + labels.append(value) + return text(xy, value, **kwargs) + + draw.text = capture + return draw + + monkeypatch.setattr(module.ImageDraw, "Draw", record) + return module, texts + + +def game(rows, won=False): + return {"rows": rows, "won": won, "secret": "apple"} + + +def row(colors="XXXXX"): + return {"letters": "APPLE", "colors": colors} + + +@pytest.mark.parametrize("alternating", [False, True]) +def test_late_win_uses_six_legal_rows(renderer, alternating): + module, texts = renderer + rows = ( + [item for _ in range(6) for item in (row(None), row())] + if alternating + else [row(None)] + [row() for _ in range(6)] + ) + rows[-1] = row("GGGGG") + images = module.frames([[game(rows, True)]], ["model"]) + assert "solved" in texts[-1] + assert "wins 1/1" in texts[-1] + assert all("wins 0/1" in frame for frame in texts[:-1]) + assert len(images) == len(rows) + 3 + # The winning sixth legal guess is visible, after the rejected submission(s). + x = module.PAD + 4 + y = 2 * module.PAD + 40 + 5 * (module.TILE + module.GAP) + 4 + assert images[-1].getpixel((x, y)) == module.COLORS["G"] + assert "invalid: APPLE" in texts[0] + + +def test_invalid_only_trace_finishes_and_leaves_board_empty(renderer): + module, texts = renderer + images = module.frames([[game([row(None)] * 12)]], ["model"]) + assert len(images) == 15 + assert "missed: APPLE" in texts[-1] + assert "wins 0/1" in texts[-1] + assert "invalid: APPLE" in texts[-1] + assert images[-1].getpixel((module.PAD + 4, 2 * module.PAD + 44)) == module.EMPTY + + +def test_legacy_short_games_and_running_wins(renderer): + module, texts = renderer + early = game([row("GGGGG")], True) + late = game([row()] * 5 + [row("GGGGG")], True) + module.frames([[early, late], [late, game([row()] * 6)]], ["a", "b"]) + assert texts[0].count("solved") == 1 + assert "wins 1/1" in texts[0] and "wins 0/1" in texts[0] + assert texts[5].count("wins 1/1") == 2 + assert texts[6].count("wins 1/2") == 2 + assert "wins 2/2" in texts[-1] and "wins 1/2" in texts[-1] + assert "solved" in texts[-1] and "missed: APPLE" in texts[-1] + + +def test_empty_trace_has_verdict(renderer): + module, texts = renderer + images = module.frames([[game([])]], ["model"]) + assert images + assert "missed: APPLE" in texts[-1] diff --git a/tests/scripts/test_wordle_training.py b/tests/scripts/test_wordle_training.py new file mode 100644 index 000000000..f3403f941 --- /dev/null +++ b/tests/scripts/test_wordle_training.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: BSD-3-Clause +"""Run the training entry points with small fakes, without model or network access.""" + +import runpy +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +ROOT = Path(__file__).resolve().parents[2] + + +def run_training(monkeypatch, script, trainer, client=None): + monkeypatch.setattr( + sys, + "argv", + [script, "--model", "example/model", "--hub-model-id", "example/output"] + + (["--dataset", "example/games"] if script == "sft.py" else []), + ) + modules = { + "torch": SimpleNamespace(bfloat16="bfloat16", float32="float32"), + "datasets": SimpleNamespace( + load_dataset=lambda name, split: [(name, split)], + Dataset=SimpleNamespace(from_dict=lambda data: data), + ), + "trl": SimpleNamespace( + SFTConfig=SimpleNamespace, + SFTTrainer=trainer, + GRPOConfig=SimpleNamespace, + GRPOTrainer=trainer, + ), + "huggingface_hub": SimpleNamespace(get_token=lambda: None), + "textarena_env": SimpleNamespace( + TextArenaEnv=lambda **kwargs: client, TextArenaAction=SimpleNamespace + ), + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + runpy.run_path(str(ROOT / "examples/wordle" / script), run_name="__main__") + + +def test_sft_explicit_fp32_weights_and_bf16_compute_then_upload(monkeypatch): + events = [] + + class Trainer: + def __init__(self, *, model, train_dataset, args): + assert model == "example/model" + assert train_dataset == [("example/games", "train")] + assert args.model_init_kwargs == {"dtype": "float32"} + assert args.bf16 is True + self.model = SimpleNamespace( + to=lambda dtype: events.append(("cast", dtype)) + ) + + def train(self): + events.append("train") + + def push_to_hub(self, **kwargs): + events.append(("upload", kwargs)) + + run_training(monkeypatch, "sft.py", Trainer) + assert events == [ + "train", + ("cast", "bfloat16"), + ("upload", {"dataset_name": "example/games"}), + ] + + +class WordleClient: + """Only legal guesses enter history; six legal guesses or repeated invalids end play.""" + + def __init__(self, invalid_turns, *, win=False, terminal_invalid=False): + self.invalid_turns = invalid_turns + self.win = win + self.terminal_invalid = terminal_invalid + self.calls = self.legal = self.consecutive_invalids = 0 + self.full = "Wordle: six legal guesses." + + def result(self, *, done=False, reward=0.0): + return SimpleNamespace( + observation=SimpleNamespace(messages=[SimpleNamespace(content=self.full)]), + done=done, + reward=reward, + ) + + def reset(self): + return self.result() + + def step(self, action): + self.calls += 1 + if self.calls in self.invalid_turns: + self.consecutive_invalids += 1 + self.full += "\nInvalid move: rejected." + return self.result( + done=self.terminal_invalid and self.consecutive_invalids == 2 + ) + self.consecutive_invalids = 0 + self.legal += 1 + self.full += "\nFeedback:\nA P P L E\nG X X X X" + return self.result( + done=self.legal == 6, reward=float(self.win and self.legal == 6) + ) + + +@pytest.mark.parametrize("calls_per_round", [1, 2]) +@pytest.mark.parametrize( + "invalid_turns,win,terminal_invalid,expected_calls,expected_legal", + [ + ({1}, True, False, 7, 6), + ({1, 3, 5, 7, 9, 11}, True, False, 12, 6), + (set(), False, False, 6, 6), + ({1, 2}, False, True, 2, 0), + (set(range(1, 20)), False, False, 12, 0), + ], +) +def test_grpo_rollout_budget( + monkeypatch, + invalid_turns, + win, + terminal_invalid, + expected_calls, + expected_legal, + calls_per_round, +): + client = WordleClient(invalid_turns, win=win, terminal_invalid=terminal_invalid) + + class Trainer: + def __init__(self, *, environment_factory, args, **kwargs): + self.env = environment_factory() + self.args = args + self.model = SimpleNamespace(to=lambda dtype: None) + + def train(self): + self.env.reset() + # TRL 1.16.0 counts tool rounds, even for invalid results. It can keep + # calling tools after done; the adapter must reject those calls. + for _ in range(self.args.max_tool_calling_iterations): + for _ in range(calls_per_round): + try: + self.env.guess("[apple]") + except ValueError as error: + assert str(error) == "Game over." + assert client.calls == expected_calls + assert client.legal == self.env.num_guesses == expected_legal + assert self.env.invalid_moves == expected_calls - expected_legal + assert self.env.won is win + assert self.args.max_tool_calling_iterations == 12 + assert self.args.bf16 is True + + def push_to_hub(self): + pass + + run_training(monkeypatch, "grpo.py", Trainer, client)