{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Module 2: Policy Competition on OpenSpiel\n",
    "\n",
    "Build 4 policies, compete them on Catch, then switch to another game with the same code.\n",
    "\n",
    "**Time:** ~20 min · **Difficulty:** Beginner · **GPU:** Not required"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "!pip install -q openenv-core\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": "markdown",
   "metadata": {},
   "source": [
    "## 1. Connect to Catch\n",
    "\n",
    "Catch: a ball falls from the top of a 10×5 grid. Move your paddle to catch it.\n",
    "\n",
    "- Actions: `0` = LEFT, `1` = STAY, `2` = RIGHT\n",
    "- Reward: `+1` if caught, `0` if missed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from envs.openspiel_env import OpenSpielEnv\n",
    "from envs.openspiel_env.models import OpenSpielAction, OpenSpielObservation\n",
    "import random\n",
    "\n",
    "CATCH_URL = 'https://openenv-openspiel-catch.hf.space'\n",
    "\n",
    "# Quick sanity check\n",
    "with OpenSpielEnv(base_url=CATCH_URL).sync() as env:\n",
    "    result = env.reset()\n",
    "    print(f'Legal actions: {result.observation.legal_actions}')\n",
    "    print(f'Info state shape: {len(result.observation.info_state)} values')\n",
    "    print(f'Game phase: {result.observation.game_phase}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Define Four Policies\n",
    "\n",
    "Each policy takes an `OpenSpielObservation` and returns an action ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RandomPolicy:\n",
    "    \"\"\"Pure random — baseline.\"\"\"\n",
    "    name = \"Random\"\n",
    "\n",
    "    def select_action(self, obs: OpenSpielObservation) -> int:\n",
    "        return random.choice(obs.legal_actions)\n",
    "\n",
    "\n",
    "class AlwaysStayPolicy:\n",
    "    \"\"\"Never moves — hopes ball lands on paddle.\"\"\"\n",
    "    name = \"Always Stay\"\n",
    "\n",
    "    def select_action(self, obs: OpenSpielObservation) -> int:\n",
    "        return 1  # STAY\n",
    "\n",
    "\n",
    "class SmartPolicy:\n",
    "    \"\"\"Moves paddle toward ball — optimal for Catch.\"\"\"\n",
    "    name = \"Smart Heuristic\"\n",
    "\n",
    "    def select_action(self, obs: OpenSpielObservation) -> int:\n",
    "        info_state = obs.info_state\n",
    "        grid_width = 5\n",
    "\n",
    "        # Find ball column (first 1.0 in the flattened grid)\n",
    "        ball_col = None\n",
    "        for idx, val in enumerate(info_state):\n",
    "            if abs(val - 1.0) < 0.01:\n",
    "                ball_col = idx % grid_width\n",
    "                break\n",
    "\n",
    "        # Paddle is in the last row\n",
    "        last_row = info_state[-grid_width:]\n",
    "        paddle_col = last_row.index(1.0)\n",
    "\n",
    "        if ball_col is not None:\n",
    "            if paddle_col < ball_col:\n",
    "                return 2  # RIGHT\n",
    "            elif paddle_col > ball_col:\n",
    "                return 0  # LEFT\n",
    "        return 1  # STAY\n",
    "\n",
    "\n",
    "class LearningPolicy:\n",
    "    \"\"\"Epsilon-greedy — starts random, learns to be smart.\"\"\"\n",
    "    name = \"Epsilon-Greedy\"\n",
    "\n",
    "    def __init__(self):\n",
    "        self.steps = 0\n",
    "        self._smart = SmartPolicy()\n",
    "\n",
    "    def select_action(self, obs: OpenSpielObservation) -> int:\n",
    "        self.steps += 1\n",
    "        epsilon = max(0.1, 1.0 - self.steps / 100)\n",
    "        if random.random() < epsilon:\n",
    "            return random.choice(obs.legal_actions)\n",
    "        return self._smart.select_action(obs)\n",
    "\n",
    "\n",
    "print(\"Policies defined: Random, Always Stay, Smart Heuristic, Epsilon-Greedy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Run a Single Episode\n",
    "\n",
    "Helper to play one full game and return whether the ball was caught."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import here so run_episode is self-contained even if cell[3] is skipped\n",
    "from envs.openspiel_env.models import OpenSpielAction\n",
    "\n",
    "def run_episode(env, policy, verbose=False):\n",
    "    \"\"\"Play one episode. Returns 1 if caught, 0 if missed.\"\"\"\n",
    "    result = env.reset()\n",
    "    step = 0\n",
    "\n",
    "    while not result.done:\n",
    "        action_id = policy.select_action(result.observation)\n",
    "        if verbose:\n",
    "            name = {0: 'LEFT', 1: 'STAY', 2: 'RIGHT'}.get(action_id, str(action_id))\n",
    "            print(f'  Step {step}: {name}')\n",
    "        result = env.step(OpenSpielAction(action_id=action_id, game_name='catch'))\n",
    "        step += 1\n",
    "\n",
    "    caught = 1 if result.reward and result.reward > 0 else 0\n",
    "    if verbose:\n",
    "        status = 'Caught!' if caught else 'Missed'\n",
    "        print(f'  Result: {status} (reward={result.reward})')\n",
    "    return caught\n",
    "\n",
    "\n",
    "# Demo: one verbose episode with SmartPolicy\n",
    "with OpenSpielEnv(base_url=CATCH_URL).sync() as env:\n",
    "    print('Smart policy — single episode:')\n",
    "    run_episode(env, SmartPolicy(), verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Policy Competition\n",
    "\n",
    "Run 50 episodes per policy and compare success rates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_EPISODES = 50\n",
    "\n",
    "policies = [\n",
    "    RandomPolicy(),\n",
    "    AlwaysStayPolicy(),\n",
    "    SmartPolicy(),\n",
    "    LearningPolicy(),\n",
    "]\n",
    "\n",
    "results = {}\n",
    "\n",
    "with OpenSpielEnv(base_url=CATCH_URL).sync() as env:\n",
    "    for policy in policies:\n",
    "        wins = sum(run_episode(env, policy) for _ in range(NUM_EPISODES))\n",
    "        rate = wins / NUM_EPISODES * 100\n",
    "        results[policy.name] = rate\n",
    "        print(f\"{policy.name:20s} — {rate:5.1f}% ({wins}/{NUM_EPISODES})\")\n",
    "\n",
    "print(\"\\n--- Results ---\")\n",
    "for name, rate in sorted(results.items(), key=lambda x: -x[1]):\n",
    "    bar = \"█\" * int(rate / 2)\n",
    "    print(f\"{name:20s} [{bar:<50}] {rate:.1f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Expected results:\n",
    "- **Random**: ~20% (pure luck)\n",
    "- **Always Stay**: ~20% (terrible strategy)\n",
    "- **Smart Heuristic**: ~100% (optimal)\n",
    "- **Epsilon-Greedy**: ~80-90% (improves over episodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Switch Games\n",
    "\n",
    "The same `OpenSpielEnv` client works for all 6 OpenSpiel games. Let's try Tic-Tac-Toe — the observation format is identical, only the game logic changes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TTT_URL = \"https://openenv-openspiel-tictactoe.hf.space\"\n",
    "\n",
    "with OpenSpielEnv(base_url=TTT_URL).sync() as env:\n",
    "    result = env.reset()\n",
    "    print(f\"Game: Tic-Tac-Toe\")\n",
    "    print(f\"Legal actions: {result.observation.legal_actions}\")\n",
    "    print(f\"Info state: {result.observation.info_state}\")\n",
    "    print(f\"Current player: {result.observation.current_player_id}\")\n",
    "    print()\n",
    "\n",
    "    # Play randomly until game ends\n",
    "    step = 0\n",
    "    while not result.done:\n",
    "        action_id = random.choice(result.observation.legal_actions)\n",
    "        result = env.step(OpenSpielAction(action_id=action_id, game_name=\"tic_tac_toe\"))\n",
    "        step += 1\n",
    "        print(f\"Step {step}: action={action_id}, reward={result.reward}, done={result.done}\")\n",
    "\n",
    "    print(f\"\\nGame over! Final reward: {result.reward}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same client class. Same observation type. Different game. That's the OpenEnv promise.\n",
    "\n",
    "## Summary\n",
    "\n",
    "- Built 4 policies with increasing sophistication\n",
    "- Ran a 50-episode competition on Catch\n",
    "- Switched to Tic-Tac-Toe with zero code changes to the client\n",
    "\n",
    "All policies work with `OpenSpielObservation` — you read `info_state`, `legal_actions`, and `done`. The game logic is on the server. Your code is on the client.\n",
    "\n",
    "**Next:** [Module 3](../module-3/README.md) — Deploying environments to HF Spaces."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}