{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Module 5: Train a Wordle Agent with GRPO\n",
    "\n",
    "Fine-tune Qwen3-1.7B to play Wordle using GRPO (Group Relative Policy Optimization) via TRL and OpenEnv.\n",
    "\n",
    "**Time:** ~90 min (training) · **Difficulty:** Advanced · **GPU:** A100 required (Colab Pro or similar)\n",
    "\n",
    "Based on the [TRL OpenEnv Wordle example](https://github.com/huggingface/trl/blob/main/examples/notebooks/openenv_wordle_grpo.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "# Requires GPU (A100 recommended). Run on Colab Pro or similar.\n!pip install -Uq \"trl>=0.17.0\" openenv-core transformers datasets accelerate vllm trackio\n!git clone --depth=1 -q https://github.com/meta-pytorch/OpenEnv.git 2>/dev/null || true\n\nimport sys, os\nrepo = os.path.abspath('OpenEnv')\nfor p in [repo, os.path.join(repo, 'src')]:\n    if p not in sys.path:\n        sys.path.insert(0, p)\nprint(\"Setup complete!\")"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Log in to Hugging Face (required for model access and pushing results)\n",
    "from huggingface_hub import notebook_login\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Initialize the Environment\n",
    "\n",
    "Connect to the TextArena Wordle environment hosted on HF Spaces.\n",
    "\n",
    "> **For production use:** Duplicate the Space to your own account to avoid concurrency limits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from envs.textarena_env import TextArenaEnv\n",
    "\n",
    "textarena_url = 'https://burtenshaw-textarena.hf.space'  # Duplicate this Space for production use!\n",
    "\n",
    "# Verify connection\n",
    "with TextArenaEnv(base_url=textarena_url).sync() as _check:\n",
    "    result = _check.reset()\n",
    "    print(f'Connected to: {textarena_url}')\n",
    "    print(f'Prompt preview: {str(result.observation.prompt)[:100]}...')\n",
    "\n",
    "# Create a persistent sync client for training.\n",
    "# A single WebSocket connection is reused across all rollouts instead of\n",
    "# opening/closing one per episode, which matters at training throughput.\n",
    "env = TextArenaEnv(base_url=textarena_url)\n",
    "sync_env = env.sync()\n",
    "sync_env.connect()\n",
    "print('Persistent training connection established.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Init Model and Tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "model_name = \"Qwen/Qwen3-1.7B\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "print(f\"Model: {model_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Define the System Prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt = \"\"\"\n",
    "You are an expert Wordle solver with deep knowledge of English vocabulary, letter frequency patterns, and optimal guessing strategies.\n",
    "\n",
    "## GAME RULES\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: Letter is correct and in the correct position\n",
    "   - YELLOW: Letter is in the word but in the wrong position\n",
    "   - GRAY: 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",
    "\n",
    "## RESPONSE FORMAT\n",
    "\n",
    "Only respond with your next guess in square brackets, e.g., [crane].\n",
    "\n",
    "## STRATEGIC APPROACH\n",
    "\n",
    "Do not repeat the same guess twice.\n",
    "\n",
    "### Opening Strategy\n",
    "- Start with words rich in common vowels (A, E, I, O, U) and consonants (R, S, T, L, N)\n",
    "- Optimal starters: CRANE, SLATE, STARE, AROSE, IRATE\n",
    "\n",
    "### Mid-Game Strategy\n",
    "- Use confirmed GREEN letters in their correct positions\n",
    "- Place YELLOW letters in different positions than where they appeared\n",
    "- Eliminate GRAY letters from consideration\n",
    "\n",
    "## YOUR GOAL\n",
    "\n",
    "Solve the Wordle in as few guesses as possible.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_user_prompt(prompt_text, messages):\n",
    "    \"\"\"Build a structured prompt from game state and message history.\"\"\"\n",
    "    history = format_history(messages)\n",
    "    prompt_section = prompt_text.strip() if prompt_text.strip() else \"Wordle-v0\"\n",
    "    history_section = history if history else \"[PROMPT] Awaiting first feedback.\"\n",
    "    return (\n",
    "        f\"Game prompt:\\n{prompt_section}\\n\\n\"\n",
    "        f\"Conversation so far:\\n{history_section}\\n\\n\"\n",
    "        \"Reply with your next guess enclosed in square brackets.\"\n",
    "    )\n",
    "\n",
    "\n",
    "def format_history(messages):\n",
    "    \"\"\"Format message history with category tags.\"\"\"\n",
    "    lines = []\n",
    "    for message in messages:\n",
    "        tag = message.category or \"MESSAGE\"\n",
    "        content = message.content.strip()\n",
    "        if content:\n",
    "            lines.append(f\"[{tag}] {content}\")\n",
    "    return \"\\n\".join(lines)\n",
    "\n",
    "\n",
    "def scale_repetition_score(previous_occurrences, max_occurrences):\n",
    "    \"\"\"Scale repetition penalty: 1.0 = novel guess, 0.0 = fully repeated.\"\"\"\n",
    "    if max_occurrences == 0:\n",
    "        return 0.0\n",
    "    return (max_occurrences - previous_occurrences) / max_occurrences\n",
    "\n",
    "\n",
    "print(\"Helper functions defined.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Define the Rollout Function\n",
    "\n",
    "The rollout function plays one full Wordle game per prompt. It's called by `GRPOTrainer` during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from envs.textarena_env.models import TextArenaAction\n",
    "from envs.textarena_env.rewards import extract_feedback_counts, extract_guess, extract_wordle_feedback\n",
    "from trl.experimental.openenv import generate_rollout_completions\n",
    "\n",
    "\n",
    "def rollout_once(trainer, sync_env, tokenizer, dataset_prompt, system_prompt, max_turns):\n",
    "    \"\"\"Execute one full Wordle episode using an already-connected sync client.\"\"\"\n",
    "    result = sync_env.reset()\n",
    "    observation = result.observation\n",
    "\n",
    "    prompt_ids = []\n",
    "    completion_ids = []\n",
    "    logprobs = []\n",
    "    green_scores = []\n",
    "    yellow_scores = []\n",
    "    repetition_scores = []\n",
    "    correct_scores = []\n",
    "    guess_counts = defaultdict(int)\n",
    "\n",
    "    for _turn in range(max_turns):\n",
    "        if result.done:\n",
    "            break\n",
    "\n",
    "        base_prompt = observation.prompt or dataset_prompt\n",
    "        user_prompt = make_user_prompt(base_prompt, observation.messages)\n",
    "        messages = [\n",
    "            {'role': 'system', 'content': system_prompt},\n",
    "            {'role': 'user', 'content': user_prompt},\n",
    "        ]\n",
    "        prompt_text = tokenizer.apply_chat_template(\n",
    "            messages,\n",
    "            add_generation_prompt=True,\n",
    "            tokenize=False,\n",
    "            enable_thinking=False,\n",
    "        )\n",
    "\n",
    "        rollout_outputs = generate_rollout_completions(trainer, [prompt_text])[0]\n",
    "        prompt_ids.extend(rollout_outputs['prompt_ids'])\n",
    "        completion_ids.extend(rollout_outputs['completion_ids'])\n",
    "        logprobs.extend(rollout_outputs['logprobs'])\n",
    "        completion_text = rollout_outputs.get('text') or tokenizer.decode(\n",
    "            rollout_outputs['completion_ids'], skip_special_tokens=True\n",
    "        )\n",
    "\n",
    "        guess = extract_guess(completion_text)\n",
    "        result = sync_env.step(TextArenaAction(message=guess))\n",
    "        observation = result.observation\n",
    "        correct_score = float(result.reward or 0.0)\n",
    "        feedback = extract_wordle_feedback(observation)\n",
    "\n",
    "        previous_occurrences = guess_counts[guess]\n",
    "        repetition_score = max(0.0, 1.0 - previous_occurrences)\n",
    "        guess_counts[guess] += 1\n",
    "\n",
    "        if not feedback:\n",
    "            green_score, yellow_score = 0.0, 0.0\n",
    "        else:\n",
    "            green_count, yellow_count = extract_feedback_counts(feedback)\n",
    "            green_score = green_count / 5.0\n",
    "            yellow_score = yellow_count / 5.0\n",
    "\n",
    "        repetition_scores.append(repetition_score)\n",
    "        green_scores.append(green_score)\n",
    "        yellow_scores.append(yellow_score)\n",
    "        correct_scores.append(correct_score)\n",
    "\n",
    "    return {\n",
    "        'prompt_ids': prompt_ids,\n",
    "        'completion_ids': completion_ids,\n",
    "        'logprobs': logprobs,\n",
    "        'correct_reward': correct_scores[-1] if correct_scores else 0.0,\n",
    "        'green_reward': green_scores[-1] if green_scores else 0.0,\n",
    "        'yellow_reward': yellow_scores[-1] if yellow_scores else 0.0,\n",
    "        'repetition_reward': repetition_scores[-1] if repetition_scores else 0.0,\n",
    "    }\n",
    "\n",
    "\n",
    "def rollout_func(prompts, trainer=None):\n",
    "    \"\"\"Rollout function called by GRPOTrainer. Uses the module-level sync_env.\"\"\"\n",
    "    episode_prompt_ids = []\n",
    "    episode_completion_ids = []\n",
    "    episode_logprobs = []\n",
    "    correctness_rewards = []\n",
    "    green_rewards = []\n",
    "    yellow_rewards = []\n",
    "    repetition_rewards = []\n",
    "\n",
    "    for prompt_text in prompts:\n",
    "        episode = rollout_once(\n",
    "            trainer=trainer,\n",
    "            sync_env=sync_env,     # Persistent connection — no reconnect per episode\n",
    "            tokenizer=tokenizer,\n",
    "            dataset_prompt=prompt_text,\n",
    "            system_prompt=system_prompt,\n",
    "            max_turns=6,\n",
    "        )\n",
    "        episode_prompt_ids.append(episode['prompt_ids'])\n",
    "        episode_completion_ids.append(episode['completion_ids'])\n",
    "        episode_logprobs.append(episode['logprobs'])\n",
    "        correctness_rewards.append(episode['correct_reward'])\n",
    "        green_rewards.append(episode['green_reward'])\n",
    "        yellow_rewards.append(episode['yellow_reward'])\n",
    "        repetition_rewards.append(episode['repetition_reward'])\n",
    "\n",
    "    return {\n",
    "        'prompt_ids': episode_prompt_ids,\n",
    "        'completion_ids': episode_completion_ids,\n",
    "        'logprobs': episode_logprobs,\n",
    "        'correct_reward': correctness_rewards,\n",
    "        'green_reward': green_rewards,\n",
    "        'yellow_reward': yellow_rewards,\n",
    "        'repetition_reward': repetition_rewards,\n",
    "    }\n",
    "\n",
    "\n",
    "print('Rollout functions defined.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Define Reward Functions\n",
    "\n",
    "Four reward signals for richer gradient information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reward_correct(completions, **kwargs):\n",
    "    rewards = kwargs.get(\"correct_reward\")\n",
    "    return [float(r) for r in rewards] if rewards else [0.0] * len(completions)\n",
    "\n",
    "def reward_greens(completions, **kwargs):\n",
    "    rewards = kwargs.get(\"green_reward\")\n",
    "    return [float(r) for r in rewards] if rewards else [0.0] * len(completions)\n",
    "\n",
    "def reward_yellows(completions, **kwargs):\n",
    "    rewards = kwargs.get(\"yellow_reward\")\n",
    "    return [float(r) for r in rewards] if rewards else [0.0] * len(completions)\n",
    "\n",
    "def reward_repetition(completions, **kwargs):\n",
    "    rewards = kwargs.get(\"repetition_reward\")\n",
    "    return [float(r) for r in rewards] if rewards else [0.0] * len(completions)\n",
    "\n",
    "print(\"Reward functions: correct, greens, yellows, repetition\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Create Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "\n",
    "dataset_size = 1000\n",
    "dataset = Dataset.from_dict({\"prompt\": [\"Play Wordle like an expert.\"] * dataset_size})\n",
    "print(f\"Dataset: {len(dataset)} prompts\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Configure GRPO Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import GRPOConfig\n",
    "\n",
    "output_dir = \"wordle-grpo-Qwen3-1.7B\"\n",
    "\n",
    "grpo_config = GRPOConfig(\n",
    "    num_train_epochs=1,\n",
    "    learning_rate=5e-6,\n",
    "    gradient_accumulation_steps=64,\n",
    "    per_device_train_batch_size=1,\n",
    "    warmup_steps=20,\n",
    "    num_generations=2,\n",
    "    max_completion_length=8,\n",
    "    max_prompt_length=1400,\n",
    "    use_vllm=True,\n",
    "    vllm_mode=\"colocate\",\n",
    "    vllm_gpu_memory_utilization=0.1,\n",
    "    output_dir=output_dir,\n",
    "    report_to=\"trackio\",\n",
    "    trackio_space_id=output_dir,\n",
    "    logging_steps=1,\n",
    "    save_steps=10,\n",
    "    gradient_checkpointing=True,\n",
    "    gradient_checkpointing_kwargs={\"use_reentrant\": False},\n",
    "    push_to_hub=True,\n",
    ")\n",
    "\n",
    "print(f\"Output: {output_dir}\")\n",
    "print(f\"vLLM mode: colocate (generation + training on same GPU)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Create Trainer and Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import GRPOTrainer\n",
    "\n",
    "trainer = GRPOTrainer(\n",
    "    model=model_name,\n",
    "    processing_class=tokenizer,\n",
    "    reward_funcs=[\n",
    "        reward_correct,\n",
    "        reward_greens,\n",
    "        reward_yellows,\n",
    "        reward_repetition,\n",
    "    ],\n",
    "    train_dataset=dataset,\n",
    "    args=grpo_config,\n",
    "    rollout_func=rollout_func,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check GPU before training\n",
    "import torch\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",
    "print(f\"GPU: {gpu_stats.name} — {max_memory} GB total, {start_gpu_memory} GB reserved\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train (~90 minutes on A100)\n",
    "trainer_stats = trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Memory stats after training\n",
    "used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)\n",
    "used_for_training = round(used_memory - start_gpu_memory, 3)\n",
    "\n",
    "print(f\"Training time: {round(trainer_stats.metrics['train_runtime']/60, 1)} minutes\")\n",
    "print(f\"Peak memory: {used_memory} GB ({round(used_memory/max_memory*100, 1)}% of {max_memory} GB)\")\n",
    "print(f\"Memory for training: {used_for_training} GB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Save and Push"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Close the persistent environment connection before saving\n",
    "sync_env.close()\n",
    "\n",
    "trainer.save_model(output_dir)\n",
    "trainer.push_to_hub()\n",
    "print(f'Model saved to {output_dir} and pushed to Hub.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Evaluate: Play Wordle with the Trained Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM\n",
    "from envs.textarena_env.models import TextArenaAction\n",
    "from envs.textarena_env.rewards import extract_guess\n",
    "\n",
    "# Load the fine-tuned model (replace with your HF repo id if you pushed)\n",
    "fine_tuned_model = AutoModelForCausalLM.from_pretrained(\n",
    "    output_dir, torch_dtype='auto', device_map='auto'\n",
    ")\n",
    "\n",
    "\n",
    "def play_wordle(sync_env, model, tokenizer, max_turns=6):\n",
    "    \"\"\"Play one Wordle game and print each turn.\"\"\"\n",
    "    result = sync_env.reset()\n",
    "    observation = result.observation\n",
    "    print(f'Prompt: {observation.prompt[:100]}...')\n",
    "\n",
    "    for turn in range(max_turns):\n",
    "        if result.done:\n",
    "            break\n",
    "\n",
    "        user_prompt = make_user_prompt(observation.prompt, observation.messages)\n",
    "        messages = [\n",
    "            {'role': 'system', 'content': system_prompt},\n",
    "            {'role': 'user', 'content': user_prompt},\n",
    "        ]\n",
    "        prompt_text = tokenizer.apply_chat_template(\n",
    "            messages, add_generation_prompt=True,\n",
    "            tokenize=False, enable_thinking=False,\n",
    "        )\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",
    "        guess = extract_guess(generated_text)\n",
    "\n",
    "        print(f'\\nTurn {turn + 1}: {guess}')\n",
    "        result = sync_env.step(TextArenaAction(message=guess))\n",
    "        observation = result.observation\n",
    "        for msg in observation.messages:\n",
    "            print(f'  [{msg.category}] {msg.content}')\n",
    "\n",
    "    print(f'\\nResult: reward={result.reward}, done={result.done}')\n",
    "\n",
    "\n",
    "# Evaluation uses a fresh per-game context (not the training connection)\n",
    "eval_env = TextArenaEnv(base_url=textarena_url)\n",
    "with eval_env.sync() as eval_sync:\n",
    "    play_wordle(eval_sync, fine_tuned_model, tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "What you did:\n",
    "1. Connected to the TextArena Wordle environment via OpenEnv\n",
    "2. Defined a system prompt, rollout function, and 4 reward signals\n",
    "3. Trained Qwen3-1.7B with GRPO for ~90 minutes on an A100\n",
    "4. Evaluated the trained model on live Wordle games\n",
    "\n",
    "The key insight: **OpenEnv makes the environment a plug-in.** Swap Wordle for any other OpenEnv environment — your Module 4 word game, a coding environment, a math problem — and the training pipeline stays the same.\n",
    "\n",
    "### What's next\n",
    "\n",
    "- **Improve the model:** Longer training, larger models, better reward shaping\n",
    "- **Build your own environment:** Use Module 4's pattern, plug it into this pipeline\n",
    "- **Scale up:** See the [Scaling appendix](../README.md#bonus-scaling-openenv) for multi-container deployment\n",
    "- **Explore the Hub:** Browse [openenv environments](https://huggingface.co/collections/openenv/environment-hub) for inspiration"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.0"
  },
  "accelerator": "GPU",
  "gpuClass": "premium"
 },
 "nbformat": 4,
 "nbformat_minor": 4
}