{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OpenEnv Sudoku 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_sudoku_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 Sudoku**, a logic-based number puzzle, 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 **Sudoku** environment through the [**OpenEnv**](https://github.com/huggingface/OpenEnv) framework, which standardizes multi-agent and RL-style text environments.\n",
    "\n",
    "[Sudoku](https://en.wikipedia.org/wiki/Sudoku) is a classic logic-based puzzle where the objective is to fill a **9x9 grid** so that each **row**, **column**, and **3x3 subgrid** contains all digits from **1 to 9** exactly once. This structured yet challenging setup makes Sudoku an excellent benchmark for reasoning and decision-making tasks.\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 functions -- 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** Sudoku 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/sudoku 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 Sudoku. It includes the game rules, board reading format, strategic approaches, and importantly, tells the model to use the `place` tool to submit moves. 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": [
    "system_prompt = \"\"\"You are an expert Sudoku player with deep knowledge of logical deduction strategies and number placement techniques.\n",
    "\n",
    "## GAME RULES\n",
    "\n",
    "1. The puzzle is a 9x9 grid divided into nine 3x3 subgrids (boxes)\n",
    "2. Some cells are pre-filled with numbers 1-9\n",
    "3. You must fill in the empty cells (shown as '.') with numbers 1-9\n",
    "4. Each row must contain numbers 1-9 without repetition\n",
    "5. Each column must contain numbers 1-9 without repetition\n",
    "6. Each 3x3 subgrid must contain numbers 1-9 without repetition\n",
    "7. You cannot overwrite pre-filled cells\n",
    "8. Invalid moves result in penalties (-1 reward)\n",
    "\n",
    "## HOW TO PLAY\n",
    "\n",
    "Use the `place` tool to make a move. The tool takes three arguments:\n",
    "- `row`: Row number (1-9)\n",
    "- `col`: Column number (1-9)\n",
    "- `number`: The digit to place (1-9)\n",
    "\n",
    "## STRATEGIC APPROACH\n",
    "\n",
    "Do not repeat the same move twice.\n",
    "\n",
    "### Basic Strategies\n",
    "- **Naked Singles**: If a cell has only one possible candidate, fill it in immediately.\n",
    "- **Hidden Singles**: If a number can only go in one cell within a row, column, or box, place it there.\n",
    "- **Scanning**: Look at each row, column, and box to find where specific numbers can go.\n",
    "\n",
    "### Solving Process\n",
    "1. Start by scanning the entire grid to identify easy fills (cells with few candidates)\n",
    "2. Look for rows, columns, or boxes with many numbers already placed\n",
    "3. Fill all naked singles first\n",
    "4. Then look for hidden singles in each row, column, and box\n",
    "\n",
    "### Common Pitfalls to Avoid\n",
    "- Don't guess randomly - Sudoku is pure logic\n",
    "- Don't overlook any constraint (row, column, or box)\n",
    "- Don't try to overwrite pre-filled cells\n",
    "- Don't place invalid numbers (must be 1-9)\n",
    "- Don't use invalid coordinates (must be 1-9)\n",
    "- Don't repeat a move that was already made\n",
    "\n",
    "## BOARD READING\n",
    "\n",
    "The board is displayed as a 9x9 grid:\n",
    "- Numbers 1-9 are pre-filled or already placed\n",
    "- Empty cells are shown as '.'\n",
    "- Rows are labeled R1-R9 (top to bottom)\n",
    "- Columns are labeled C1-C9 (left to right)\n",
    "\n",
    "## IMPORTANT CONSTRAINTS\n",
    "\n",
    "- Coordinates are 1-indexed (1-9 for both row and column)\n",
    "- Numbers must be 1-9\n",
    "- One move per response\n",
    "- Must be a valid move (no rule violations)\n",
    "- Never repeat a previous move\n",
    "\"\"\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the environment\n",
    "\n",
    "The `SudokuEnv` class wraps the OpenEnv TextArena Sudoku environment into the interface expected by `environment_factory`.\n",
    "\n",
    "When you pass `environment_factory=SudokuEnv` to the trainer, it will:\n",
    "1. Create a new `SudokuEnv()` instance for each rollout episode.\n",
    "2. Call `reset()` to start a new game (returns the initial board state).\n",
    "3. Automatically generate model completions, parse tool calls, and invoke the `place(row, col, number)` method.\n",
    "4. Repeat until the environment signals `done=True` or the max completion length is reached.\n",
    "\n",
    "The environment tracks multiple reward signals as properties:\n",
    "- **correct_reward**: Did the puzzle get solved?\n",
    "- **valid_move_reward**: Average rate of valid moves.\n",
    "- **empty_cell_reward**: Did the model target empty cells?\n",
    "- **repetition_reward**: Penalty for repeating moves.\n",
    "- **progress_reward**: How many cells were filled (normalized 0-1).\n",
    "\n",
    "For this example, we connect to the hosted environment at [openenv/sudoku](https://huggingface.co/spaces/openenv/sudoku).\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": [
    "# @title SudokuEnv class (click to expand)\n",
    "from collections import defaultdict\n",
    "\n",
    "from textarena_env import TextArenaAction, TextArenaEnv\n",
    "\n",
    "\n",
    "def _is_valid_board_state(board_str: str) -> bool:\n",
    "    return \"R1\" in board_str and \"R9\" in board_str and \"|\" in board_str\n",
    "\n",
    "\n",
    "def _parse_board(board_str: str) -> list[list[int]]:\n",
    "    grid = [[0] * 9 for _ in range(9)]\n",
    "    if not _is_valid_board_state(board_str):\n",
    "        return grid\n",
    "    for line in board_str.split(\"\\n\"):\n",
    "        line_stripped = line.strip()\n",
    "        if line_stripped and line_stripped[0] == \"R\" and len(line_stripped) > 1 and line_stripped[1].isdigit():\n",
    "            row = int(line_stripped[1]) - 1\n",
    "            cell_part = line_stripped[2:]\n",
    "            col = 0\n",
    "            for char in cell_part:\n",
    "                if char == \".\":\n",
    "                    grid[row][col] = 0\n",
    "                    col += 1\n",
    "                elif char.isdigit():\n",
    "                    grid[row][col] = int(char)\n",
    "                    col += 1\n",
    "    return grid\n",
    "\n",
    "\n",
    "def _count_filled_cells(board_str: str) -> int:\n",
    "    if not _is_valid_board_state(board_str):\n",
    "        return 0\n",
    "    grid = _parse_board(board_str)\n",
    "    return sum(1 for row in grid for cell in row if cell != 0)\n",
    "\n",
    "\n",
    "def _get_valid_numbers(grid: list[list[int]], row: int, col: int) -> set[int]:\n",
    "    if grid[row][col] != 0:\n",
    "        return set()\n",
    "    used = set()\n",
    "    for c in range(9):\n",
    "        if grid[row][c] != 0:\n",
    "            used.add(grid[row][c])\n",
    "    for r in range(9):\n",
    "        if grid[r][col] != 0:\n",
    "            used.add(grid[r][col])\n",
    "    box_row, box_col = 3 * (row // 3), 3 * (col // 3)\n",
    "    for r in range(box_row, box_row + 3):\n",
    "        for c in range(box_col, box_col + 3):\n",
    "            if grid[r][c] != 0:\n",
    "                used.add(grid[r][c])\n",
    "    return set(range(1, 10)) - used\n",
    "\n",
    "\n",
    "def _extract_empty_cells_with_candidates(board_str: str, sort_by_difficulty: bool = True):\n",
    "    grid = _parse_board(board_str)\n",
    "    cells_with_candidates = []\n",
    "    for row in range(9):\n",
    "        for col in range(9):\n",
    "            if grid[row][col] == 0:\n",
    "                candidates = _get_valid_numbers(grid, row, col)\n",
    "                cells_with_candidates.append((row + 1, col + 1, candidates))\n",
    "    if sort_by_difficulty:\n",
    "        cells_with_candidates.sort(key=lambda x: len(x[2]))\n",
    "    return cells_with_candidates\n",
    "\n",
    "\n",
    "def _extract_empty_cells(board_str: str) -> list[tuple[int, int]]:\n",
    "    empty_cells = []\n",
    "    if not _is_valid_board_state(board_str):\n",
    "        return empty_cells\n",
    "    for line in board_str.split(\"\\n\"):\n",
    "        line_stripped = line.strip()\n",
    "        if line_stripped and line_stripped[0] == \"R\" and len(line_stripped) > 1 and line_stripped[1].isdigit():\n",
    "            row = int(line_stripped[1])\n",
    "            cell_part = line_stripped[2:]\n",
    "            col = 0\n",
    "            for char in cell_part:\n",
    "                if char == \".\":\n",
    "                    col += 1\n",
    "                    empty_cells.append((row, col))\n",
    "                elif char.isdigit():\n",
    "                    col += 1\n",
    "    return empty_cells\n",
    "\n",
    "\n",
    "def _extract_board_only(text: str) -> str:\n",
    "    if not text:\n",
    "        return \"\"\n",
    "    lines = text.split(\"\\n\")\n",
    "    board_lines = []\n",
    "    in_board = False\n",
    "    for line in lines:\n",
    "        stripped = line.strip()\n",
    "        if stripped.startswith(\"C1\") or (\n",
    "            stripped and stripped[0] == \"R\" and len(stripped) > 1 and stripped[1].isdigit()\n",
    "        ):\n",
    "            in_board = True\n",
    "        if in_board and (stripped.startswith(\"-\") or stripped.startswith(\"R\") or stripped.startswith(\"C1\")):\n",
    "            board_lines.append(line)\n",
    "        elif (\n",
    "            in_board\n",
    "            and stripped\n",
    "            and not stripped.startswith(\"-\")\n",
    "            and not (stripped[0] == \"R\" and len(stripped) > 1 and stripped[1].isdigit())\n",
    "        ):\n",
    "            break\n",
    "    return \"\\n\".join(board_lines) if board_lines else \"\"\n",
    "\n",
    "\n",
    "def _make_hints(board, successful_moves, failed_moves, difficulty=\"easy\"):\n",
    "    \"\"\"Generate hint text for the model.\"\"\"\n",
    "    parts = []\n",
    "    all_tried = successful_moves + failed_moves\n",
    "    if all_tried:\n",
    "        parts.append(f\"\\nMOVES ALREADY TRIED (do not repeat): {', '.join(all_tried)}\")\n",
    "    if not board:\n",
    "        return \"\\n\".join(parts)\n",
    "    if difficulty == \"easy\":\n",
    "        cells = _extract_empty_cells_with_candidates(board, sort_by_difficulty=True)\n",
    "        if cells:\n",
    "            guaranteed = []\n",
    "            other = []\n",
    "            for r, c, candidates in cells[:10]:\n",
    "                if len(candidates) == 1:\n",
    "                    guaranteed.append(f\"[{r} {c} {list(candidates)[0]}]\")\n",
    "                elif len(candidates) <= 3:\n",
    "                    nums = \",\".join(str(n) for n in sorted(candidates))\n",
    "                    other.append(f\"({r},{c})->{nums}\")\n",
    "            if guaranteed:\n",
    "                parts.append(f\"\\nGUARANTEED MOVES: {', '.join(guaranteed[:5])}\")\n",
    "            if other:\n",
    "                parts.append(f\"Other options: {' | '.join(other[:5])}\")\n",
    "    elif difficulty == \"medium\":\n",
    "        cells = _extract_empty_cells_with_candidates(board, sort_by_difficulty=False)\n",
    "        if cells:\n",
    "            cell_hints = []\n",
    "            for r, c, candidates in cells[:10]:\n",
    "                nums = \",\".join(str(n) for n in sorted(candidates))\n",
    "                cell_hints.append(f\"({r},{c})->{nums}\")\n",
    "            parts.append(f\"\\nEmpty cells: {' | '.join(cell_hints)}\")\n",
    "    return \"\\n\".join(parts)\n",
    "\n",
    "\n",
    "class SudokuEnv:\n",
    "    def __init__(self):\n",
    "        self.client = TextArenaEnv(base_url=\"https://openenv-sudoku.hf.space\")\n",
    "        self.difficulty = \"easy\"\n",
    "        self.max_turns = 100\n",
    "        self._turn = 0\n",
    "        self._move_counts = defaultdict(int)\n",
    "        self._successful_moves = []\n",
    "        self._failed_moves = []\n",
    "        self._valid_move_scores = []\n",
    "        self._empty_cell_scores = []\n",
    "        self._correct_scores = []\n",
    "        self._repetition_scores = []\n",
    "        self._last_board_state = \"\"\n",
    "        self._last_full_content = \"\"\n",
    "        self._initial_filled = 0\n",
    "        self._max_filled = 0\n",
    "        self.done = False\n",
    "\n",
    "    def reset(self, **kwargs) -> str | None:\n",
    "        result = self.client.reset()\n",
    "        observation = result.observation\n",
    "        self.done = False\n",
    "        self._turn = 0\n",
    "        self._move_counts = defaultdict(int)\n",
    "        self._successful_moves = []\n",
    "        self._failed_moves = []\n",
    "        self._valid_move_scores = []\n",
    "        self._empty_cell_scores = []\n",
    "        self._correct_scores = []\n",
    "        self._repetition_scores = []\n",
    "        self._last_board_state = \"\"\n",
    "        self._initial_filled = 0\n",
    "        self._max_filled = 0\n",
    "\n",
    "        # Store full message content for diffing (messages are cumulative)\n",
    "        self._last_full_content = observation.messages[0].content if observation.messages else \"\"\n",
    "\n",
    "        for message in observation.messages:\n",
    "            if message.content and _is_valid_board_state(message.content):\n",
    "                self._last_board_state = message.content\n",
    "                self._initial_filled = _count_filled_cells(self._last_board_state)\n",
    "                self._max_filled = self._initial_filled\n",
    "                break\n",
    "\n",
    "        board = _extract_board_only(self._last_board_state) if self._last_board_state else \"No board available.\"\n",
    "        hints = _make_hints(self._last_board_state, [], [], self.difficulty)\n",
    "        return f\"Step 0. Progress: 0 cells filled.\\n\\nBoard:\\n{board}{hints}\"\n",
    "\n",
    "    def place(self, row: int, col: int, number: int) -> str:\n",
    "        \"\"\"Place a number on the Sudoku board.\n",
    "\n",
    "        Args:\n",
    "            row: Row number (1-9).\n",
    "            col: Column number (1-9).\n",
    "            number: Number to place (1-9).\n",
    "\n",
    "        Returns:\n",
    "            The result of the move and updated board state.\n",
    "        \"\"\"\n",
    "        if self.done:\n",
    "            return \"Game is over. No more moves allowed.\"\n",
    "\n",
    "        self._turn += 1\n",
    "        move = f\"[{row} {col} {number}]\"\n",
    "\n",
    "        # Step environment\n",
    "        result = self.client.step(TextArenaAction(message=move))\n",
    "        observation = result.observation\n",
    "        correct_score = float(result.reward or 0.0)\n",
    "        self.done = result.done\n",
    "\n",
    "        # Only check the NEW content for feedback (messages are cumulative)\n",
    "        full_content = observation.messages[0].content if observation.messages else \"\"\n",
    "        new_content = full_content[len(self._last_full_content):]\n",
    "        self._last_full_content = full_content\n",
    "\n",
    "        new_content_lower = new_content.lower()\n",
    "        env_says_invalid = any(\n",
    "            kw in new_content_lower for kw in [\"invalid\", \"error\", \"cannot\", \"already\", \"violation\", \"lost\"]\n",
    "        )\n",
    "        got_warning = \"please resubmit\" in new_content_lower or \"avoid penalties\" in new_content_lower\n",
    "\n",
    "        # Also verify against our own board state: placing on a non-empty cell is always invalid\n",
    "        if self._last_board_state:\n",
    "            empty_cells = _extract_empty_cells(self._last_board_state)\n",
    "            targets_empty = (row, col) in empty_cells\n",
    "        else:\n",
    "            empty_cells = []\n",
    "            targets_empty = True  # Can't verify, assume valid\n",
    "\n",
    "        is_valid = not env_says_invalid and targets_empty\n",
    "\n",
    "        # Empty cell score: did the model target an empty cell?\n",
    "        empty_cell_score = 1.0 if targets_empty else -1.0\n",
    "\n",
    "        # Repetition tracking\n",
    "        is_new_move = self._move_counts[move] == 0\n",
    "        repetition_count = self._move_counts[move]\n",
    "        self._move_counts[move] += 1\n",
    "        repetition_score = -min(2 ** (repetition_count - 1), 10.0) if repetition_count > 0 else 0.0\n",
    "\n",
    "        # Valid move score\n",
    "        if is_valid and is_new_move:\n",
    "            valid_move_score = 1.0\n",
    "            self._successful_moves.append(move)\n",
    "        elif got_warning:\n",
    "            valid_move_score = -0.5\n",
    "            self._failed_moves.append(move)\n",
    "        else:\n",
    "            valid_move_score = 0.0\n",
    "\n",
    "        # Update board state from new content\n",
    "        if is_valid and _is_valid_board_state(new_content):\n",
    "            self._last_board_state = new_content\n",
    "            current_filled = _count_filled_cells(self._last_board_state)\n",
    "            if current_filled > self._max_filled:\n",
    "                self._max_filled = current_filled\n",
    "\n",
    "        self._valid_move_scores.append(valid_move_score)\n",
    "        self._empty_cell_scores.append(empty_cell_score)\n",
    "        self._correct_scores.append(correct_score)\n",
    "        self._repetition_scores.append(repetition_score)\n",
    "\n",
    "        # Enforce max turns\n",
    "        if self._turn >= self.max_turns:\n",
    "            self.done = True\n",
    "\n",
    "        # Build response\n",
    "        board = _extract_board_only(self._last_board_state) if self._last_board_state else \"No board available.\"\n",
    "        status = \"valid\" if is_valid else \"invalid\"\n",
    "        cells_filled = len(self._successful_moves)\n",
    "        progress = f\"Step {self._turn}. Progress: {cells_filled} cells filled.\"\n",
    "        hints = _make_hints(self._last_board_state, self._successful_moves, self._failed_moves, self.difficulty)\n",
    "\n",
    "        if self.done:\n",
    "            return f\"Move {move}: {status}. Game over.\\n{progress}\\n\\nFinal board:\\n{board}\"\n",
    "        return f\"Move {move}: {status}\\n{progress}\\n\\nBoard:\\n{board}{hints}\"\n",
    "\n",
    "    # ── Reward properties ──\n",
    "\n",
    "    @property\n",
    "    def correct_reward(self) -> float:\n",
    "        return self._correct_scores[-1] if self._correct_scores else 0.0\n",
    "\n",
    "    @property\n",
    "    def valid_move_reward(self) -> float:\n",
    "        return sum(self._valid_move_scores) / len(self._valid_move_scores) if self._valid_move_scores else 0.0\n",
    "\n",
    "    @property\n",
    "    def empty_cell_reward(self) -> float:\n",
    "        return sum(self._empty_cell_scores) / len(self._empty_cell_scores) if self._empty_cell_scores else 0.0\n",
    "\n",
    "    @property\n",
    "    def repetition_reward(self) -> float:\n",
    "        return sum(self._repetition_scores) / len(self._repetition_scores) if self._repetition_scores else 0.0\n",
    "\n",
    "    @property\n",
    "    def progress_reward(self) -> float:\n",
    "        remaining_to_fill = 81 - self._initial_filled\n",
    "        if remaining_to_fill > 0:\n",
    "            return (self._max_filled - self._initial_filled) / remaining_to_fill\n",
    "        return 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the reward functions\n",
    "\n",
    "The reward functions receive the list of environment instances after each episode completes. Since the `SudokuEnv` tracks multiple reward signals as properties, we simply read them out.\n",
    "\n",
    "Each reward function captures a different aspect of play quality:\n",
    "- **empty_cell_reward**: Did the model target empty cells (vs. trying to overwrite filled ones)?\n",
    "- **valid_move_reward**: Were the moves accepted by the environment?\n",
    "- **repetition_reward**: Penalty for repeating the same move.\n",
    "- **progress_reward**: How much of the puzzle was filled (0 to 1).\n",
    "- **correct_reward**: Did the model solve the puzzle completely?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reward_empty_cell(environments, **kwargs) -> list[float]:\n",
    "    \"\"\"Reward for targeting empty cells.\"\"\"\n",
    "    return [env.empty_cell_reward for env in environments]\n",
    "\n",
    "\n",
    "def reward_valid_moves(environments, **kwargs) -> list[float]:\n",
    "    \"\"\"Reward for making valid moves.\"\"\"\n",
    "    return [env.valid_move_reward for env in environments]\n",
    "\n",
    "\n",
    "def reward_repetition(environments, **kwargs) -> list[float]:\n",
    "    \"\"\"Penalty for repeating moves.\"\"\"\n",
    "    return [env.repetition_reward for env in environments]\n",
    "\n",
    "\n",
    "def reward_progress(environments, **kwargs) -> list[float]:\n",
    "    \"\"\"Reward for filling more cells in the board.\"\"\"\n",
    "    return [env.progress_reward for env in environments]\n",
    "\n",
    "\n",
    "def reward_correct(environments, **kwargs) -> list[float]:\n",
    "    \"\"\"Reward for solving the puzzle.\"\"\"\n",
    "    return [env.correct_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 with the system prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "\n",
    "dataset_size = 3000\n",
    "dataset = Dataset.from_dict({\n",
    "    \"prompt\": [[\n",
    "        {\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": \"Play Sudoku like an expert.\"},\n",
    "    ]] * dataset_size\n",
    "})"
   ]
  },
  {
   "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 = \"sudoku-grpo-Qwen3-1.7B\"\n",
    "\n",
    "grpo_config = GRPOConfig(\n",
    "    # Training schedule / optimization\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",
    "    optim=\"adamw_torch\",\n",
    "    max_grad_norm=1.0,\n",
    "\n",
    "    # GRPO configuration\n",
    "    num_generations=2,\n",
    "    max_completion_length=16384,\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",
    "\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",
    "    # Hub integration\n",
    "    push_to_hub=True,\n",
    "\n",
    "    # Sampling\n",
    "    temperature=0.8,\n",
    "    top_k=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the `GRPOTrainer` and start training\n",
    "\n",
    "Now we initialize the `GRPOTrainer` with `environment_factory=SudokuEnv`.\n",
    "\n",
    "This tells the trainer to automatically handle the entire interaction loop:\n",
    "- It creates a `SudokuEnv` instance for each episode.\n",
    "- It generates model completions, parses tool calls (like `place`), 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=[\n",
    "        reward_empty_cell,\n",
    "        reward_valid_moves,\n",
    "        reward_repetition,\n",
    "        reward_progress,\n",
    "        reward_correct,\n",
    "    ],\n",
    "    train_dataset=dataset,\n",
    "    args=grpo_config,\n",
    "    environment_factory=SudokuEnv,\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 Sudoku.\n",
    "We use the same `SudokuEnv` 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/sudoku-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_sudoku(model, tokenizer):\n",
    "    env = SudokuEnv()\n",
    "    initial_observation = env.reset()\n",
    "\n",
    "    print(\"Initial observation:\")\n",
    "    print(initial_observation)\n",
    "    print()\n",
    "\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system_prompt},\n",
    "        {\"role\": \"user\", \"content\": \"Play Sudoku like an expert.\"},\n",
    "    ]\n",
    "    if initial_observation:\n",
    "        messages.append({\"role\": \"user\", \"content\": initial_observation})\n",
    "\n",
    "    for turn in range(20):  # Play up to 20 turns\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:\n",
    "            # Try to parse tool call arguments\n",
    "            if \"place\" in generated_text:\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",
    "                row = int(args.get(\"row\", 0))\n",
    "                col = int(args.get(\"col\", 0))\n",
    "                number = int(args.get(\"number\", 0))\n",
    "            else:\n",
    "                # Fallback: extract [row col number] pattern\n",
    "                import re\n",
    "                match = re.search(r\"\\[(\\d)\\s+(\\d)\\s+(\\d)\\]\", generated_text)\n",
    "                if match:\n",
    "                    row, col, number = int(match.group(1)), int(match.group(2)), int(match.group(3))\n",
    "                else:\n",
    "                    print(\"         Could not parse move.\")\n",
    "                    break\n",
    "\n",
    "            feedback = env.place(row, col, number)\n",
    "            print(f\"         Move: [{row} {col} {number}]\")\n",
    "            print(f\"         Progress reward: {env.progress_reward:.2f}\")\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! Correct reward: {env.correct_reward}\")\n",
    "    print(f\"Progress: {env.progress_reward:.2f}\")\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_sudoku(fine_tuned_model, tokenizer)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
