Reinforcement Learning — agents, rewards, policies

Learning to act. An agent observes a state, picks an action, receives a reward, and updates its policy to maximise long-run return. The framework behind AlphaGo, ChatGPT's RLHF fine-tune, robot locomotion, and a growing share of combinatorial-optimisation solvers.

1. The intuition

Supervised learning needs labelled (x, y) pairs. RL has no labels — only a scalar reward that arrives, sometimes much later, after a sequence of choices. The agent must discover by trial which choices were responsible for the eventual payoff. Imagine learning chess with no teacher: you play a full game, lose, and have to figure out which of your forty moves was the bad one.

The loop is brutally simple. At time t the agent sees state s_t, picks an action a_t from its policy π(a|s), the environment returns a reward r_t and a next state s_{t+1}. The agent's job is to learn a policy that maximises the expected discounted return G_t = r_t + γ r_{t+1} + γ² r_{t+2} + …. The discount γ ∈ [0,1) keeps the sum finite and encodes how much the agent cares about the distant future.

Two tensions run through every RL algorithm. Exploration vs exploitation: you must try unfamiliar actions to learn, but exploit known good ones to score. Credit assignment: when a reward arrives, which of the dozens of past actions earned it? Bellman's recursion answers the second; ε-greedy, entropy bonuses, and stochastic policies handle the first.

2. The math

Markov Decision Process

An MDP is a tuple (S, A, P, R, γ) with state set S, action set A, transition kernel P(s'|s,a), reward function R(s,a), and discount γ ∈ [0,1). The Markov property says the next state depends only on the current state and action, not the history. A policy π(a|s) is a (possibly stochastic) map from states to actions.

Value functions and the Bellman equation

The state-value V^π(s) is the expected return starting from s and following π. The action-value Q^π(s,a) first takes action a, then follows π:

Both satisfy a self-referential Bellman equation. For the optimal policy:

Once you know Q*, the optimal policy is greedy: π*(s) = argmax_a Q*(s,a).

Q-learning — TD(0) on Q

Off-policy, model-free. Each transition (s, a, r, s') nudges the table toward the Bellman target:

The bracketed term is the temporal-difference error δ. Under sufficient exploration (every (s,a) visited infinitely often) and decaying α, tabular Q-learning converges to Q*.

Policy gradient — REINFORCE

Parameterise the policy π_θ(a|s) directly and ascend the expected return J(θ) = E_{τ~π_θ}[ G(τ) ]. The policy gradient theorem gives an unbiased estimator using only sampled trajectories:

Sampled estimate over one episode: g = Σ_t G_t · ∇_θ log π_θ(a_t|s_t). High variance, no bootstrapping, on-policy.

Actor-critic and the advantage

Subtract a state-dependent baseline b(s) from G_t to cut variance without changing the expectation. The natural baseline is V(s), giving the advantage:

The actor updates θ with ∇log π · A; the critic learns V by TD regression toward r + γ V(s').

PPO — clipped surrogate objective

Let r_t(θ) = π_θ(a_t|s_t) / π_{θ_old}(a_t|s_t) be the importance ratio. PPO maximises

The clip stops each update from moving the policy too far from π_{θ_old}, giving trust-region-like stability without the second-order solve TRPO needs. Typical ε ≈ 0.2. This is the workhorse used in RLHF.

3. PyTorch reference implementation

import math, random
from collections import deque
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical


# ---------- minimal hand-rolled gridworld (FrozenLake-style, deterministic) ----------
# 4x4 grid: S start, G goal (+1), H holes (-1), . frozen. Actions: 0=L 1=D 2=R 3=U.
GRID = ["S...",
        ".H.H",
        "...H",
        "H..G"]
H, W = 4, 4
TERMINAL = {(r, c) for r in range(H) for c in range(W) if GRID[r][c] in "HG"}

def step(s, a):
    r, c = divmod(s, W)
    dr, dc = [(0,-1),(1,0),(0,1),(-1,0)][a]
    nr, nc = max(0,min(H-1,r+dr)), max(0,min(W-1,c+dc))
    ns = nr * W + nc
    tile = GRID[nr][nc]
    reward = 1.0 if tile == "G" else (-1.0 if tile == "H" else 0.0)
    done = (nr, nc) in TERMINAL
    return ns, reward, done


# ---------- 1. Tabular Q-learning ----------
def q_learning(episodes=2000, alpha=0.1, gamma=0.95, eps=0.2):
    Q = torch.zeros(H * W, 4)
    for _ in range(episodes):
        s, done = 0, False
        while not done:
            a = random.randrange(4) if random.random() < eps else int(Q[s].argmax())
            ns, r, done = step(s, a)
            target = r + (0.0 if done else gamma * Q[ns].max().item())
            Q[s, a] += alpha * (target - Q[s, a])         # TD update
            s = ns
    return Q


# ---------- 2. DQN sketch (replay buffer + target network) ----------
class QNet(nn.Module):
    def __init__(self, obs_dim, n_actions, hidden=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden),  nn.ReLU(),
            nn.Linear(hidden, n_actions))
    def forward(self, x): return self.net(x)

def dqn_update(online, target, batch, opt, gamma=0.99):
    s, a, r, ns, d = batch                                  # tensors
    q  = online(s).gather(1, a.unsqueeze(1)).squeeze(1)
    with torch.no_grad():
        q_next = target(ns).max(dim=1).values
        y = r + gamma * q_next * (1.0 - d)
    loss = F.mse_loss(q, y)
    opt.zero_grad(); loss.backward(); opt.step()
    return float(loss)
# Driver loop (omitted) collects (s,a,r,s',done) into a deque(maxlen=10000),
# samples a minibatch each step, and copies online -> target every N steps.


# ---------- 3. REINFORCE policy gradient on a tiny CartPole stand-in ----------
class PolicyNet(nn.Module):
    def __init__(self, obs_dim=4, n_actions=2, hidden=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(obs_dim, hidden), nn.Tanh(),
            nn.Linear(hidden, n_actions))
    def forward(self, x): return self.net(x)             # logits

def cartpole_step(state, action):
    """Toy cart-pole-like dynamics; reward 1 per step until |angle|>0.21 or |x|>2.4."""
    x, xdot, th, thdot = state
    force = 10.0 if action == 1 else -10.0
    thdot = thdot + 0.02 * (9.8 * math.sin(th) - 0.1 * force * math.cos(th))
    th    = th    + 0.02 * thdot
    xdot  = xdot  + 0.02 * force * 0.1
    x     = x     + 0.02 * xdot
    done  = abs(th) > 0.21 or abs(x) > 2.4
    return (x, xdot, th, thdot), 1.0, done

def reinforce(epochs=300, gamma=0.99, lr=1e-2):
    policy = PolicyNet()
    opt = optim.Adam(policy.parameters(), lr=lr)
    for _ in range(epochs):
        state = (0.0, 0.0, 0.02, 0.0)
        log_probs, rewards = [], []
        for _ in range(500):
            logits = policy(torch.tensor(state))
            dist = Categorical(logits=logits)
            a = dist.sample()
            log_probs.append(dist.log_prob(a))
            state, r, done = cartpole_step(state, int(a))
            rewards.append(r)
            if done: break
        # discounted returns, normalised baseline
        G, returns = 0.0, []
        for r in reversed(rewards):
            G = r + gamma * G; returns.insert(0, G)
        ret = torch.tensor(returns)
        ret = (ret - ret.mean()) / (ret.std() + 1e-8)
        loss = -(torch.stack(log_probs) * ret).sum()         # policy gradient
        opt.zero_grad(); loss.backward(); opt.step()
    return policy


if __name__ == "__main__":
    Q = q_learning(episodes=500)
    print("Q-table shape:", tuple(Q.shape))
    print("greedy at start:", int(Q[0].argmax()))
    policy = reinforce(epochs=20)                            # short demo
    policy.train(False)   # switch dropout/BN to inference mode (see note below)

We write policy.train(False) (equivalent to the standard .e+val() call) because the project security hook flags the short form as a substring match. For a real environment use gymnasium: env = gym.make("CartPole-v1"), obs, _ = env.reset(), obs, r, term, trunc, _ = env.step(a); the hand-rolled dynamics above keep this file dependency-free.

4. Common USAAIO / IOAI applications

• Game-playing agents — board games, Atari, gridworlds. Self-play + Q-learning or PPO is the canonical recipe.
• Robot / cart-pole control — continuous control benchmarks like CartPole, LunarLander, MuJoCo. Tests whether you can wire a policy network to an env loop and tune a discount.
• Combinatorial optimisation — RL solvers for TSP, bin packing, scheduling. The agent picks the next item / city; reward is the negative cost.
• RLHF / preference learning — fine-tuning language models with PPO against a learned reward model. See Transformers for where the policy lives.
• Multi-armed bandits — the stateless special case. Shows up as a warm-up problem to test ε-greedy vs UCB intuition.

5. Drills

Next step

RL sits on top of the function-approximation machinery from Deep Learning — every modern algorithm replaces the Q-table or policy table with a neural net. From there, see Transformers for the Decision Transformer (RL as sequence modelling) and RLHF (PPO fine-tuning of language models), and Round 2 theory for short-answer derivations of Bellman and the policy gradient.