{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OpenEnv Wordle with GRPO using TRL\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/trl/blob/main/examples/notebooks/openenv_wordle_grpo.ipynb)\n",
    "\n",
    "![trl banner](https://huggingface.co/datasets/trl-lib/documentation-images/resolve/main/trl_banner_dark.png)\n",
    "\n",
    "\n",
    "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.\n",
    "\n",
    "- [TRL GitHub Repository](https://github.com/huggingface/trl) -- star us to support the project!\n",
    "- [Official TRL Examples](https://huggingface.co/docs/trl/example_overview)\n",
    "- [Community Tutorials](https://huggingface.co/docs/trl/community_tutorials)\n",
    "- [OpenEnv](https://github.com/huggingface/OpenEnv)\n",
    "\n",
    "\n",
    "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.\n",
    "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.\n",
    "\n",
    "[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.\n",
    "After each guess, feedback indicates whether each letter is:\n",
    "- **GREEN (G)**: Correct and in the right position\n",
    "- **YELLOW (Y)**: Present but in the wrong position\n",
    "- **GRAY (X)**: Not in the word\n",
    "\n",
    "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.\n",
    "\n",
    "\n",
    "We'll fine-tune a model using **GRPO** (Group Relative Policy Optimization) via TRL.\n",
    "Using `environment_factory`, the trainer automatically handles:\n",
    "1. Creating environment instances for each rollout.\n",
    "2. Generating model completions and parsing tool calls.\n",
    "3. Stepping through the environment with the model's actions.\n",
    "4. Collecting rewards and managing the interaction loop.\n",
    "\n",
    "This means you only need to define the environment class and reward function -- the trainer takes care of the rest.\n",
    "\n",
    "\n",
    "## Install dependencies\n",
    "\n",
    "We'll start by installing **TRL** (with vLLM support), the **OpenEnv** Wordle environment, and **trackio** for logging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -Uq trl[vllm] git+https://huggingface.co/spaces/openenv/wordle trackio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Log in to Hugging Face\n",
    "\n",
    "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)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the system prompt\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"You are an expert Wordle solver with deep knowledge of English vocabulary, letter frequency patterns, and optimal guessing strategies.\n",
    "\n",
    "Follow these rules to play Wordle:\n",
    "\n",
    "1. The target is a 5-letter English word\n",
    "2. You have 6 attempts to guess the correct word\n",
    "3. After each guess, you receive color-coded feedback:\n",
    "   - GREEN (G): Letter is correct and in the correct position\n",
    "   - YELLOW (Y): Letter is in the word but in the wrong position\n",
    "   - GRAY (X): Letter is not in the word at all\n",
    "4. All guesses must be valid 5-letter English words\n",
    "5. You cannot reuse a word you've already guessed\n",
    "6. Use the tool `guess` to make a guess.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the environment\n",
    "\n",
    "The `WordleEnv` class wraps the OpenEnv TextArena Wordle environment into the interface expected by `environment_factory`.\n",
    "\n",
    "When you pass `environment_factory=WordleEnv` to the trainer, it will:\n",
    "1. Create a new `WordleEnv()` instance for each rollout episode.\n",
    "2. Call `reset()` to start a new game (returns the initial observation or `None`).\n",
    "3. Automatically generate model completions, parse tool calls, and invoke the corresponding methods (e.g., `guess(...)`).\n",
    "4. Repeat until the environment signals `done=True` or the max completion length is reached.\n",
    "\n",
    "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.\n",
    "\n",
    "For this example, we connect to the hosted environment at [openenv/wordle](https://huggingface.co/spaces/openenv/wordle).\n",
    "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.\n",
    "\n",
    "For more information, refer to the [TRL-OpenEnv documentation](https://huggingface.co/docs/trl/main/en/openenv)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from textarena_env import TextArenaAction, TextArenaEnv\n",
    "\n",
    "\n",
    "class WordleEnv:\n",
    "    def __init__(self):\n",
    "        self.client = TextArenaEnv(base_url=\"https://openenv-wordle.hf.space\")\n",
    "\n",
    "    def reset(self, **kwargs) -> None | str:\n",
    "        result = self.client.reset()\n",
    "        # The game returns cumulative feedback each turn (new text appended at the end), so\n",
    "        # we store the previous full response and slice out only the newly appended part.\n",
    "        self._last_full_feedback = result.observation.messages[0].content\n",
    "        self.reward = 0.0\n",
    "        self.done = False\n",
    "        return self._last_full_feedback\n",
    "\n",
    "    def guess(self, guess: str) -> str:\n",
    "        \"\"\"\n",
    "        Make a guess in the Wordle environment.\n",
    "\n",
    "        Args:\n",
    "            guess: The guessed word, formatted as '[abcde]'\n",
    "\n",
    "        Returns:\n",
    "            The feedback message from the environment.\n",
    "        \"\"\"\n",
    "        if self.done:\n",
    "            raise ValueError(\"Game over.\")\n",
    "        result = self.client.step(TextArenaAction(message=guess))\n",
    "        _full_feedback = result.observation.messages[0].content\n",
    "        # Just take the new feedback since the last guess\n",
    "        feedback = _full_feedback[len(self._last_full_feedback):]\n",
    "        self._last_full_feedback = _full_feedback\n",
    "        # Penalize invalid moves\n",
    "        if \"You attempted an invalid move\" in feedback:\n",
    "            self.reward = 0.0\n",
    "        else:\n",
    "            self.reward = result.reward\n",
    "        self.done = result.done\n",
    "        return feedback"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the reward function\n",
    "\n",
    "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.\n",
    "\n",
    "This is much simpler than defining multiple reward functions manually -- the environment already knows the game outcome."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reward_func(environments, **kwargs) -> list[float]:\n",
    "    return [env.reward for env in environments]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the dataset\n",
    "\n",
    "We create a dataset with repeated prompts to control the number of training episodes.\n",
    "Each entry triggers one rollout episode during training. The prompt is formatted as a chat message."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "\n",
    "dataset = Dataset.from_dict({\"prompt\": [[{\"role\": \"user\", \"content\": prompt}] for _ in range(3000)]})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set GRPO Config\n",
    "\n",
    "Next, we define the **GRPOConfig**, which controls all key training parameters.\n",
    "This configuration specifies how the model interacts with vLLM, manages memory, and logs results.\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import GRPOConfig\n",
    "\n",
    "model_name = \"Qwen/Qwen3-1.7B\"\n",
    "output_dir = \"wordle-grpo-Qwen3-1.7B\"\n",
    "\n",
    "grpo_config = GRPOConfig(\n",
    "    # Training schedule / optimization\n",
    "    num_train_epochs=1,\n",
    "    learning_rate=1e-6,\n",
    "    gradient_accumulation_steps=64,\n",
    "    per_device_train_batch_size=1,\n",
    "    warmup_steps=10,\n",
    "    optim=\"adamw_torch\",\n",
    "    max_grad_norm=1.0,\n",
    "\n",
    "    # GRPO configuration\n",
    "    num_generations=2,\n",
    "    max_completion_length=1024,\n",
    "    log_completions=True,\n",
    "    num_completions_to_print=2,\n",
    "    chat_template_kwargs={\"enable_thinking\": False},\n",
    "\n",
    "    # vLLM configuration\n",
    "    use_vllm=True,\n",
    "    vllm_mode=\"colocate\",\n",
    "    vllm_gpu_memory_utilization=0.15,\n",
    "    vllm_max_model_length=3072,\n",
    "\n",
    "    # Logging / reporting\n",
    "    output_dir=output_dir,\n",
    "    report_to=\"trackio\",\n",
    "    trackio_space_id=output_dir,\n",
    "    logging_steps=1,\n",
    "    save_steps=10,\n",
    "    save_total_limit=1,\n",
    "\n",
    "    # Memory optimization\n",
    "    gradient_checkpointing=True,\n",
    "\n",
    "    # Hub integration\n",
    "    push_to_hub=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the `GRPOTrainer` and start training\n",
    "\n",
    "Now we initialize the `GRPOTrainer` with `environment_factory=WordleEnv`.\n",
    "\n",
    "This tells the trainer to automatically handle the entire interaction loop:\n",
    "- It creates a `WordleEnv` instance for each episode.\n",
    "- It generates model completions, parses tool calls (like `guess`), and steps through the environment.\n",
    "- It collects rewards and manages the `tool_mask` (which tokens are model-generated vs environment-generated) automatically.\n",
    "\n",
    "No need to write a custom `rollout_func` or manage tokenization manually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import GRPOTrainer\n",
    "\n",
    "trainer = GRPOTrainer(\n",
    "    model=model_name,\n",
    "    reward_funcs=reward_func,\n",
    "    train_dataset=dataset,\n",
    "    args=grpo_config,\n",
    "    environment_factory=WordleEnv,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Show memory stats before training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "gpu_stats = torch.cuda.get_device_properties(0)\n",
    "start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)\n",
    "max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)\n",
    "\n",
    "print(f\"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.\")\n",
    "print(f\"{start_gpu_memory} GB of memory reserved.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And train!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer_stats = trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Show memory stats after training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)\n",
    "used_memory_for_training = round(used_memory - start_gpu_memory, 3)\n",
    "used_percentage = round(used_memory / max_memory * 100, 3)\n",
    "training_memory_percentage = round(used_memory_for_training / max_memory * 100, 3)\n",
    "\n",
    "print(f\"{trainer_stats.metrics['train_runtime']} seconds used for training.\")\n",
    "print(f\"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training.\")\n",
    "print(f\"Peak reserved memory = {used_memory} GB.\")\n",
    "print(f\"Peak reserved memory for training = {used_memory_for_training} GB.\")\n",
    "print(f\"Peak reserved memory % of max memory = {used_percentage} %.\")\n",
    "print(f\"Peak reserved memory for training % of max memory = {training_memory_percentage} %.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save and push to Hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_model(output_dir)\n",
    "trainer.push_to_hub()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the fine-tuned model and run inference\n",
    "\n",
    "Now let's test our fine-tuned model by loading it and playing a game of Wordle.\n",
    "We use the same `WordleEnv` class to interact with the environment, and generate model responses with standard Transformers inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "model_name = \"sergiopaniego/wordle-grpo-Qwen3-1.7B\"  # Replace with your HF username or organization\n",
    "\n",
    "fine_tuned_model = AutoModelForCausalLM.from_pretrained(model_name, dtype=\"float32\", device_map=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "\n",
    "def play_wordle(model, tokenizer):\n",
    "    env = WordleEnv()\n",
    "    initial_observation = env.reset()\n",
    "\n",
    "    print(\"Initial observation:\")\n",
    "    print(initial_observation)\n",
    "    print()\n",
    "\n",
    "    messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "    if initial_observation:\n",
    "        messages.append({\"role\": \"user\", \"content\": initial_observation})\n",
    "\n",
    "    for turn in range(6):\n",
    "        if env.done:\n",
    "            break\n",
    "\n",
    "        prompt_text = tokenizer.apply_chat_template(\n",
    "            messages,\n",
    "            add_generation_prompt=True,\n",
    "            tokenize=False,\n",
    "            enable_thinking=False,\n",
    "        )\n",
    "        model_inputs = tokenizer([prompt_text], return_tensors=\"pt\").to(model.device)\n",
    "        generated_ids = model.generate(**model_inputs, max_new_tokens=512)\n",
    "        output_ids = generated_ids[0][len(model_inputs.input_ids[0]):]\n",
    "        generated_text = tokenizer.decode(output_ids, skip_special_tokens=True)\n",
    "\n",
    "        print(f\"Turn {turn + 1} - Model output: {generated_text}\")\n",
    "\n",
    "        # Try to parse tool call from the generated text\n",
    "        try:\n",
    "            # Try to extract a guess from tool call format or bracket format\n",
    "            if \"guess\" in generated_text:\n",
    "                # Parse JSON tool call\n",
    "                start = generated_text.index(\"{\")\n",
    "                end = generated_text.rindex(\"}\") + 1\n",
    "                args = json.loads(generated_text[start:end])\n",
    "                if \"arguments\" in args:\n",
    "                    args = args[\"arguments\"]\n",
    "                guess_word = args.get(\"guess\", \"\")\n",
    "            else:\n",
    "                # Fallback: extract from brackets [word]\n",
    "                import re\n",
    "                match = re.search(r\"\\[([a-zA-Z]{5})\\]\", generated_text)\n",
    "                guess_word = match.group(1) if match else generated_text.strip()[:5]\n",
    "\n",
    "            feedback = env.guess(f\"[{guess_word}]\")\n",
    "            print(f\"         Guess: {guess_word} | Reward: {env.reward}\")\n",
    "            print(f\"         Feedback: {feedback.strip()}\")\n",
    "            print()\n",
    "\n",
    "            messages.append({\"role\": \"assistant\", \"content\": generated_text})\n",
    "            messages.append({\"role\": \"user\", \"content\": feedback})\n",
    "        except Exception as e:\n",
    "            print(f\"         Error: {e}\")\n",
    "            break\n",
    "\n",
    "    print(f\"Game finished! Final reward: {env.reward}\")\n",
    "    print(f\"Done: {env.done}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's play the game!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "play_wordle(fine_tuned_model, tokenizer)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
