{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3",
   "language": "python"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "name": "reinforcementlearning.ipynb"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# reinforcement learning \u2014 Python demo\n\nNumerical companion to the entry [reinforcement learning](https://dictionaryofml.org/terms/reinforcementlearning.html) of the [Dictionary of Applied Machine Learning](https://dictionaryofml.org/): it recomputes what the entry states and prints one line per check.\n\nNumerical companion to the glossary entry 'reinforcementlearning' (reinforcement learning (RL)). The gridworld is the one of the entry 'gridworld': a 4 x 4 lattice of cells (row, col) with row 0 at the bottom, a wall at cell (2, 2), a goal at cell (3, 3), and the four moves up, down, left, right. A move that would leave the grid or enter the wall leaves the agent in place. Every move yields reward -1, except the move that enters the goal, which yields +1 and ends the episode (the goal is absorbing). No discounting (gamma = 1): every episode ends after finitely many moves under every policy considered here.\n\nRequires NumPy and Matplotlib only, and uses fixed seeds, so the printed numbers reproduce exactly. Generated from [`pythondemos/reinforcementlearning.py`](https://dictionaryofml.org/terms/reinforcementlearning.py); CC BY 4.0."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# Notebook shim: the script resolves output paths relative to __file__,\n# which a notebook kernel does not define; everything lands in the\n# working directory instead.\nimport os\n__file__ = os.path.join(os.getcwd(), \"reinforcementlearning.py\")\nos.makedirs(\"pythondemos\", exist_ok=True)"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "\"\"\"Reinforcement learning on a 4 x 4 gridworld: the value function of a\npolicy is the fixed point of its Bellman operator, the greedy policy is\noptimal, and Q-learning finds the same policy from interaction alone.\n\nPurpose\n-------\nNumerical companion to the glossary entry 'reinforcementlearning'\n(reinforcement learning (RL)).  The gridworld is the one of the entry\n'gridworld': a 4 x 4 lattice of cells (row, col) with row 0 at the\nbottom, a wall at cell (2, 2), a goal at cell (3, 3), and the four moves\nup, down, left, right.  A move that would leave the grid or enter the\nwall leaves the agent in place.  Every move yields reward -1, except the\nmove that enters the goal, which yields +1 and ends the episode (the\ngoal is absorbing).  No discounting (gamma = 1): every episode ends after\nfinitely many moves under every policy considered here.\n\nThe demo checks the entry's claims: (1) the state-value function of the\nequiprobable random policy is the fixed point of the policy's Bellman\noperator, reached by iterating it from zero (iterative policy\nevaluation); (2) the greedy policy with respect to that value function\nalready reaches the goal on a shortest path from every cell, and its\nvalue function is the optimal one, 2 - d(s) with d(s) the shortest path\nlength to the goal; (3) Q-learning with random moves for exploration, which\nnever uses the transition function and learns only from the rewards of\nthe actions it tried, arrives at the same optimal policy, and its return\nper episode rises from that of a random walk to the optimum -- delayed\nreward at work, since the +1 arrives only at the last move.\n\nDeterministic: fixed seed for the exploration noise.  Self-contained:\nnumpy + matplotlib only.\n\nBlocks\n------\n[B-gridworld] Build the MDP: 15 states (16 cells minus the wall), 4\n              actions, deterministic transition function, reward -1 per\n              move and +1 for entering the goal; check the boundary and\n              wall rules.\n[B-evaluate]  Iterative policy evaluation of the random policy: apply the\n              Bellman operator of the policy from v = 0 until the change\n              is below 1e-10; check the result is a fixed point, that\n              the iteration converged, and that the values are negative\n              and decrease with the distance to the goal.\n[B-improve]   Greedy policy with respect to the random policy's value\n              function: check it reaches the goal from every cell in the\n              shortest number of moves, and that its value function is\n              the optimal one, 2 - d(s).\n[B-qlearning] Q-learning (exploration probability 0.1, update weight 0.5) over 400\n              episodes from random start cells, learning from the\n              rewards of the actions tried only; check that the greedy\n              policy of the learned Q equals the optimal policy and that\n              the average return of the last 50 episodes exceeds that of\n              the first 50.\n[B-plot]      Preview: the gridworld with the optimal values and policy\n              arrows, and the return per episode of Q-learning.\n\nOutputs\n-------\nreinforcementlearning_values.csv   : row, col, v_random, v_optimal, action\n                                     (optimal move) per non-wall cell\nreinforcementlearning_returns.csv  : episode, return, running mean over\n                                     20 episodes (Q-learning)\nreinforcementlearning.png          : preview (checking only)\n\"\"\"\n\nimport numpy as np\nimport matplotlib\n\nmatplotlib.use(\"Agg\")\nimport matplotlib.pyplot as plt\n\nfrom pathlib import Path\n\nOUT_DIR = Path(__file__).parent\n\nreport = []\n\n\ndef check(name, ok):\n    report.append((name, bool(ok)))\n    print(f\"  [{'ok' if ok else 'FAIL'}] {name}\")"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-gridworld]** Build the MDP: 15 states (16 cells minus the wall), 4 actions, deterministic transition function, reward -1 per move and +1 for entering the goal; check the boundary and wall rules."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "H, W = 4, 4\nWALL, GOAL = (2, 2), (3, 3)\nACTIONS = {\"up\": (1, 0), \"down\": (-1, 0), \"left\": (0, -1), \"right\": (0, 1)}\nSTATES = [(r, c) for r in range(H) for c in range(W) if (r, c) != WALL]\n\n\ndef step(s, a):\n    \"\"\"Next state and reward of a deterministic move; the goal is absorbing.\"\"\"\n    if s == GOAL:\n        return s, 0.0\n    dr, dc = ACTIONS[a]\n    n = (s[0] + dr, s[1] + dc)\n    if not (0 <= n[0] < H and 0 <= n[1] < W) or n == WALL:\n        n = s\n    return n, (1.0 if n == GOAL else -1.0)\n\n\ncheck(\"[B-gridworld] 15 states and 4 actions\", len(STATES) == 15 and len(ACTIONS) == 4)\ncheck(\"[B-gridworld] a move into the boundary or the wall leaves the agent in place\",\n      step((0, 0), \"down\")[0] == (0, 0) and step((1, 2), \"up\")[0] == (1, 2))\ncheck(\"[B-gridworld] entering the goal yields +1, any other move -1\",\n      step((3, 2), \"right\") == (GOAL, 1.0) and step((0, 0), \"up\") == ((1, 0), -1.0))"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-evaluate]** Iterative policy evaluation of the random policy: apply the Bellman operator of the policy from v = 0 until the change is below 1e-10; check the result is a fixed point, that the iteration converged, and that the values are negative and decrease with the distance to the goal."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "def bellman_policy(v, policy):\n    \"\"\"One application of the Bellman operator of a policy (gamma = 1):\n    (T v)(s) = sum_a pi(a|s) [r(s,a) + v(s')].\"\"\"\n    out = {}\n    for s in STATES:\n        if s == GOAL:\n            out[s] = 0.0\n            continue\n        out[s] = sum(p * (step(s, a)[1] + v[step(s, a)[0]])\n                     for a, p in policy[s].items())\n    return out\n\n\nrandom_policy = {s: {a: 0.25 for a in ACTIONS} for s in STATES}\nv = {s: 0.0 for s in STATES}\nsweeps, deltas = 0, []\nwhile True:\n    v_new = bellman_policy(v, random_policy)\n    delta = max(abs(v_new[s] - v[s]) for s in STATES)\n    deltas.append(delta); v = v_new; sweeps += 1\n    if delta < 1e-10 or sweeps > 10_000:\n        break\nv_random = v\nprint(f\"  policy evaluation: {sweeps} sweeps; v_random(0,0) = {v_random[(0, 0)]:.2f}, \"\n      f\"v_random(3,2) = {v_random[(3, 2)]:.2f}\")\ncheck(\"[B-evaluate] the iteration converged (change below 1e-10)\", deltas[-1] < 1e-10)\ncheck(\"[B-evaluate] the result is a fixed point of the policy's Bellman operator\",\n      max(abs(bellman_policy(v_random, random_policy)[s] - v_random[s]) for s in STATES) < 1e-8)\n\n\ndef shortest(s):\n    \"\"\"Shortest number of moves from s to the goal (breadth-first search).\"\"\"\n    from collections import deque\n    seen, q = {s: 0}, deque([s])\n    while q:\n        u = q.popleft()\n        if u == GOAL:\n            return seen[u]\n        for a in ACTIONS:\n            n = step(u, a)[0]\n            if n not in seen:\n                seen[n] = seen[u] + 1; q.append(n)\n    return None\n\n\nd = {s: shortest(s) for s in STATES}\ncheck(\"[B-evaluate] the random policy's values are negative away from the goal \"\n      \"and decrease with the distance to it\",\n      all(v_random[s] < 0 for s in STATES if s != GOAL)\n      and all(v_random[s] >= v_random[t] for s in STATES for t in STATES\n              if d[s] < d[t]))"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-improve]** Greedy policy with respect to the random policy's value function: check it reaches the goal from every cell in the shortest number of moves, and that its value function is the optimal one, 2 - d(s)."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "def greedy(vfun):\n    pol = {}\n    for s in STATES:\n        if s == GOAL:\n            pol[s] = None; continue\n        pol[s] = max(ACTIONS, key=lambda a: step(s, a)[1] + vfun[step(s, a)[0]])\n    return pol\n\n\ngreedy_policy = greedy(v_random)\n\n\ndef rollout(pol, s, limit=100):\n    \"\"\"Return and number of moves when following pol from s.\"\"\"\n    G, n = 0.0, 0\n    while s != GOAL and n < limit:\n        s, r = step(s, pol[s]); G += r; n += 1\n    return G, n\n\n\ncheck(\"[B-improve] the greedy policy reaches the goal on a shortest path from every cell\",\n      all(rollout(greedy_policy, s)[1] == d[s] for s in STATES if s != GOAL))\nv_greedy = {s: rollout(greedy_policy, s)[0] for s in STATES}\nv_optimal = {s: (0.0 if s == GOAL else 2.0 - d[s]) for s in STATES}\ncheck(\"[B-improve] its value function is the optimal one, 2 - d(s)\",\n      all(abs(v_greedy[s] - v_optimal[s]) < 1e-12 for s in STATES))\ncheck(\"[B-improve] the optimal values are a fixed point of the optimal Bellman \"\n      \"operator\", all(abs(max(step(s, a)[1] + v_optimal[step(s, a)[0]] for a in ACTIONS)\n                           - v_optimal[s]) < 1e-12 for s in STATES if s != GOAL))"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-qlearning]** Q-learning (exploration probability 0.1, update weight 0.5) over 400 episodes from random start cells, learning from the rewards of the actions tried only; check that the greedy policy of the learned Q equals the optimal policy and that the average return of the last 50 episodes exceeds that of the first 50."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "rng = np.random.default_rng(0)\nEPS, ALPHA_Q, EPISODES = 0.1, 0.5, 400\nQ = {s: {a: 0.0 for a in ACTIONS} for s in STATES}\nstarts = [s for s in STATES if s != GOAL]\nreturns = []\nfor ep in range(EPISODES):\n    s = starts[rng.integers(len(starts))]\n    G, n = 0.0, 0\n    while s != GOAL and n < 200:\n        if rng.random() < EPS:\n            a = list(ACTIONS)[rng.integers(4)]\n        else:\n            a = max(ACTIONS, key=lambda x: Q[s][x])\n        n_s, r = step(s, a)\n        target = r + (0.0 if n_s == GOAL else max(Q[n_s].values()))\n        Q[s][a] += ALPHA_Q * (target - Q[s][a])\n        G += r; s = n_s; n += 1\n    returns.append(G)\nreturns = np.array(returns)\nq_policy = {s: (None if s == GOAL else max(ACTIONS, key=lambda a: Q[s][a])) for s in STATES}\ncheck(\"[B-qlearning] the greedy policy of the learned Q reaches the goal on a \"\n      \"shortest path from every cell\",\n      all(rollout(q_policy, s)[1] == d[s] for s in STATES if s != GOAL))\nfirst, last = returns[:50].mean(), returns[-50:].mean()\nprint(f\"  Q-learning: mean return first 50 episodes {first:.1f}, last 50 episodes {last:.1f}\"\n      f\" (optimum averaged over start cells {np.mean([v_optimal[s] for s in starts]):.1f})\")\ncheck(\"[B-qlearning] the average return of the last 50 episodes exceeds that of \"\n      \"the first 50\", last > first)\n\n# ---- CSVs\nwith open(OUT_DIR / \"reinforcementlearning_values.csv\", \"w\") as f:\n    f.write(\"row,col,v_random,v_optimal,action\\n\")\n    for (r, c) in STATES:\n        f.write(f\"{r},{c},{v_random[(r, c)]:.2f},{v_optimal[(r, c)]:.0f},\"\n                f\"{greedy_policy[(r, c)] or 'goal'}\\n\")\nrunning = np.convolve(returns, np.ones(20) / 20, mode=\"full\")[:EPISODES]\nrunning[:19] = np.nan\nwith open(OUT_DIR / \"reinforcementlearning_returns.csv\", \"w\") as f:\n    f.write(\"episode,ret,mean20\\n\")\n    for i, (g, m) in enumerate(zip(returns, running), 1):\n        f.write(f\"{i},{g:.0f},{'nan' if np.isnan(m) else f'{m:.2f}'}\\n\")"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-plot]** Preview: the gridworld with the optimal values and policy arrows, and the return per episode of Q-learning."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4.2))\nARROW = {\"up\": (0, 0.3), \"down\": (0, -0.3), \"left\": (-0.3, 0), \"right\": (0.3, 0)}\nfor r in range(H):\n    for c in range(W):\n        if (r, c) == WALL:\n            ax1.add_patch(plt.Rectangle((c, r), 1, 1, color=\"0.6\"))\n            continue\n        if (r, c) == GOAL:\n            ax1.add_patch(plt.Rectangle((c, r), 1, 1, color=\"#c8e6c9\"))\n            ax1.text(c + 0.5, r + 0.5, \"goal\", ha=\"center\", va=\"center\", fontsize=9)\n            continue\n        ax1.text(c + 0.5, r + 0.78, f\"{v_random[(r, c)]:.1f}\", ha=\"center\",\n                 va=\"center\", fontsize=7, color=\"0.35\")\n        ax1.text(c + 0.5, r + 0.22, f\"{v_optimal[(r, c)]:.0f}\", ha=\"center\",\n                 va=\"center\", fontsize=9)\n        dx, dy = ARROW[greedy_policy[(r, c)]]\n        ax1.annotate(\"\", xy=(c + 0.5 + dx, r + 0.5 + dy), xytext=(c + 0.5, r + 0.5),\n                     arrowprops=dict(arrowstyle=\"->\", lw=1.2))\nax1.set_xlim(0, W); ax1.set_ylim(0, H); ax1.set_xticks(range(W + 1)); ax1.set_yticks(range(H + 1))\nax1.grid(True, color=\"black\", lw=0.6); ax1.set_aspect(\"equal\")\nax1.set_xlabel(\"column\"); ax1.set_ylabel(\"row\")\nax1.set_title(\"Values of the random policy (small) and optimal values with policy\")\nax2.plot(np.arange(1, EPISODES + 1), returns, \".\", ms=3, color=\"0.6\", label=\"return of the episode\")\nax2.plot(np.arange(1, EPISODES + 1), running, \"k-\", lw=1.5, label=\"mean over 20 episodes\")\nax2.axhline(np.mean([v_optimal[s] for s in starts]), ls=\"--\", color=\"black\",\n            lw=1, label=\"optimal (mean over start cells)\")\nax2.set_xlabel(\"episode\"); ax2.set_ylabel(\"return\")\nax2.set_title(\"Q-learning: return per episode\")\nax2.legend(frameon=False, fontsize=8)\nfig.tight_layout()\nfig.savefig(OUT_DIR / \"reinforcementlearning.png\", dpi=110)\n\nn_ok = sum(ok for _, ok in report)\nprint(f\"\\n{n_ok}/{len(report)} checks pass\")\nif n_ok != len(report):\n    raise SystemExit(1)"
  }
 ]
}