Reinforcement learning

Problem

Supervised learning needs the right answer for every input. Many problems only say how well things went, and only afterwards: a game is won or lost after hundreds of moves, a robot reaches the target or falls over, a chatbot's reply is rated after it is written. Nobody can label which move was right. The learner has to try actions, see what reward they lead to, and work out which earlier choices deserve the credit, all while the data it sees depends on the choices it makes. Reinforcement learning (RL) is the set of methods for that: learn a policy, a rule for picking actions, that maximizes the total reward collected over time.

Intuition

An agent sees a state, picks an action, and the environment answers with a reward and the next state. Two ideas carry almost every method. First, a state is worth the reward you can expect from it onwards, and that worth can be written recursively: the reward now plus the (discounted) worth of where you land. Learning those worths turns a delayed reward into an immediate signal for every step. Second, if a sampled action turned out better than expected, make it more likely; if worse, less likely. Value-based methods (Q-learning, DQN) follow the first idea; policy-gradient methods (REINFORCE, PPO) follow the second; actor–critic methods use the first to make the second less noisy.

The data an RL agent learns from is produced by its own actions, so changing the policy changes the data.

Mechanics

MDPs, returns and discounting

The standard model is a Markov decision process (MDP): states s, actions a, a transition distribution p(s′, r | s, a) and a discount factor γ ∈ [0, 1). "Markov" means the next state and reward depend only on the current state and action, not on how the agent got there. A policy π(a | s) gives the probability of each action in each state. From time t, the return is the discounted sum of future rewards:

G_t = r_t+1 + γ r_t+2 + γ² r_t+3 + … = r_t+1 + γ G_t+1

Discounting keeps the sum finite in endless tasks and prefers sooner rewards to later ones; a reward k steps away counts γ^k as much, so γ = 0.9 looks roughly 1/(1 − γ) = 10 steps ahead. The goal is a policy that maximizes the expected return J(π) = E[G₀].

Value functions and the Bellman equation

The state value V^π(s) is the expected return from s when following π; the action value Q^π(s, a) is the expected return after taking a in s and following π from then on. Because G_t = r + γ G_t+1, both satisfy a one-step recursion, the Bellman equation:

Q^π(s, a) = E[ r + γ V^π(s′) ]     V^π(s) = Σ_a π(a | s) Q^π(s, a)

The best possible values satisfy the Bellman optimality equation, with a max in place of the average over the policy:

Q*(s, a) = E[ r + γ max_a′ Q*(s′, a′) ]     π*(s) = argmax_a Q*(s, a)

Knowing Q* is enough to act optimally: pick the action with the highest value. When the transition probabilities are known, value iteration finds V* by applying the right-hand side as an update, over and over, to every state; with γ < 1 each sweep shrinks the error by at least a factor γ, so it converges to the unique fixed point. The worked example runs it on the chain below.

The values fall by more than γ per state away from the goal because each move can slip and cost a step.

Q-learning and DQN

Without a model of the environment, the expectation in the Bellman equation cannot be computed, but it can be sampled. Q-learning keeps a table of estimates and, after each transition (s, a, r, s′), nudges Q(s, a) toward a one-sample Bellman target:

Q(s, a) ← Q(s, a) + α [ r + γ max_a′ Q(s′, a′) − Q(s, a) ]

The bracket is the temporal-difference (TD) error: how much better or worse the step went than the current estimate predicted. The target uses the max over next actions whatever action the agent will actually take next, so Q-learning is off-policy: it learns the greedy policy's values while behaving differently, typically ε-greedy (a random action with probability ε, otherwise the best one), which keeps it exploring. With every state–action pair visited often enough and a suitably decaying α, the table converges to Q*.

Tables do not scale to images or large state spaces. DQN replaces the table with a neural network Q_θ(s, ·) that outputs one value per action, and trains it by regression on the TD target. Naively this is unstable, because consecutive samples are correlated and the target moves with every update. The DQN paper learned to play Atari games from pixels with two fixes: an experience replay buffer of past transitions sampled at random in minibatches, and targets computed from parameters held fixed while each update is made. The later Nature version made this a separate target network, copied from the online network every C steps, and clipped the error term, which amounts to a Huber loss (PyTorch's nn.HuberLoss).

Policy gradients: REINFORCE, derived

Value-based methods need a max over actions, which is awkward for continuous actions, and the policy they produce is deterministic. Policy-gradient methods instead parameterize the policy directly, π_θ(a | s), for example a softmax over a network's logits, and do gradient ascent on J(θ). The difficulty is that θ changes which actions are sampled, and you cannot backpropagate through a random sample. The log-derivative trick gets around it. For one action drawn from π_θ with reward R(a):

∇_θ E_a∼π[R(a)] = Σ_a R(a) ∇_θπ_θ(a) = Σ_a π_θ(a) R(a) ∇_θ log π_θ(a) = E_a∼π[ R(a) ∇_θ log π_θ(a) ]

using ∇π = π ∇log π. The right-hand side is an expectation, so a sampled action gives an unbiased estimate of the gradient, and it needs only the reward's value, not its derivative. For a whole trajectory, the probability is a product of policy terms and environment terms; the environment terms do not depend on θ and drop out of ∇log, leaving

∇_θ J(θ) = E[ Σ_t ∇_θ log π_θ(a_t | s_t) · G_t ]

where each action is credited only with the return that follows it (rewards before a_t cannot depend on it, and their terms average to zero). REINFORCE estimates this from sampled episodes and takes a gradient step. For a softmax policy, ∇_logits log π(a) is one-hot(a) − π: raise the chosen action's logit, lower the rest in proportion to their probability. In code, the surrogate loss −log π(a | s) · G has exactly this gradient, which is how PyTorch's distributions documentation writes REINFORCE with Categorical.log_prob.

The estimate is unbiased but very noisy. Subtracting a baseline b(s) from G_t keeps it unbiased, because E[∇log π] = Σ ∇π = ∇ Σ π = ∇1 = 0, and can cut the variance a lot (worked example below).

Actor–critic

The best common baseline is the state's value. An actor–critic method learns a critic V_φ(s) alongside the actor π_θ, and weights each log-probability by an advantage, how much better the action did than the state's average. The simplest estimate is the TD error:

A(s, a) ≈ δ = r + γ V_φ(s′) − V_φ(s)     actor: ∇θ ∝ δ ∇_θ log π_θ(a | s)     critic: minimize δ²

Using V(s′) instead of the full sampled return lowers variance further but adds bias wherever the critic is wrong; methods such as generalized advantage estimation blend one-step and multi-step estimates to trade the two off. The critic is trained by regression, like DQN's network, and the actor by the policy gradient. Under the true values the advantages average to zero in every state, which is what makes them a clean "better or worse than usual" signal.

PPO's clipped objective

Policy gradients are on-policy: each batch of experience is valid for the policy that collected it, and one gradient step uses it up. Taking several steps on the same batch, or one large step, can move the policy so far that performance collapses. PPO (proximal policy optimization) collects a batch with the current policy π_old, then runs several epochs of minibatch updates on it, maximizing a surrogate that stops rewarding moves too far from π_old. With the probability ratio ρ = π_θ(a | s) / π_old(a | s):

L^CLIP(θ) = E[ min( ρ A, clip(ρ, 1 − ε, 1 + ε) A ) ]

with ε = 0.2 in the paper. For a good action (A > 0) the objective grows with ρ only up to 1 + ε, then goes flat, so there is no gradient pushing that action's probability further. For a bad action (A < 0) it stops at 1 − ε. The min makes the clip one-sided: if an update has already moved the policy the wrong way (made a bad action more likely, or a good one less likely), the unclipped term is the smaller one and its full gradient still applies. In practice the loss adds a value loss for the critic and an entropy bonus, which keeps the policy from collapsing onto one action too early. PPO is the RL algorithm behind RLHF (see preference optimization).

Green: L^CLIP for one sample; dashed grey: the unclipped ρA. Dotted verticals mark 1 − ε and 1 + ε.

The PyTorch API

PyTorch has no RL training loop of its own (libraries such as TorchRL build on it). What it provides are the pieces: distributions to sample actions and score them, autograd for the surrogate losses, and the regression losses for critics and Q-networks.

Worked example

Implementation

import torch
from torch.distributions import Categorical

def discounted(rewards, gamma=0.99):
    returns, g = [], 0.0
    for r in reversed(rewards):
        g = r + gamma * g
        returns.append(g)
    return torch.tensor(returns[::-1])

# REINFORCE with a mean baseline; env is any environment with reset() and step().
state, done, log_probs, rewards = env.reset(), False, [], []
while not done:
    dist = Categorical(logits=policy(state))
    action = dist.sample()
    log_probs.append(dist.log_prob(action))
    state, reward, done = env.step(action.item())
    rewards.append(reward)
G = discounted(rewards)
loss = -(torch.stack(log_probs) * (G - G.mean())).sum()
loss.backward(); optimizer.step(); optimizer.zero_grad()

# PPO: several epochs of minibatches from one batch collected with the old policy.
def ppo_loss(policy, critic, states, actions, old_log_probs, returns, advantages,
             clip_eps=0.2, value_coef=0.5, entropy_coef=0.01):
    dist = Categorical(logits=policy(states))
    ratio = torch.exp(dist.log_prob(actions) - old_log_probs)
    clipped = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps)
    policy_loss = -torch.min(ratio * advantages, clipped * advantages).mean()
    value_loss = (critic(states).squeeze(-1) - returns).pow(2).mean()
    return policy_loss + value_coef * value_loss - entropy_coef * dist.entropy().mean()

# DQN: regress Q(s, a) on a target computed by a frozen copy of the network.
q = q_net(s).gather(1, a.unsqueeze(1)).squeeze(1)
with torch.no_grad():
    target = r + gamma * (1 - done) * target_net(s_next).max(dim=1).values
loss = torch.nn.HuberLoss()(q, target)

The ratio is computed from log-probabilities, which is stabler than dividing probabilities. old_log_probs and advantages are computed once when the batch is collected and must carry no gradient; advantages are usually normalized to zero mean and unit variance per batch. The coefficients shown are common defaults rather than values from the pinned source.

Tradeoffs

• Value-based versus policy gradient. Q-learning and DQN are off-policy and reuse old data through replay, so they are sample-efficient, but they need a max over actions (discrete actions) and can diverge with function approximation. Policy gradients handle continuous actions and stochastic policies and optimize the objective directly, but are on-policy and noisy.
• Bias versus variance. Monte Carlo returns are unbiased but high-variance; bootstrapping from a learned value (TD targets, critics) cuts variance and adds bias wherever the value is wrong.
• Exploration versus exploitation. ε-greedy and entropy bonuses keep trying new actions at the cost of reward now; too little exploration converges to a poor policy, too much never uses what it learned.
• Discount. A larger γ values the long run but makes returns noisier and credit assignment harder; a smaller γ learns faster and is short-sighted.
• PPO's ε and epochs. A smaller clip range or fewer epochs is safer but squeezes less learning from each batch.
• Reward design. The agent optimizes exactly the reward given; a reward that is easy to game gets gamed, which is the reward-hacking problem in RLHF.

Connections

References

• Sutton and Barto. Reinforcement Learning: An Introduction, 2nd edition. MIT Press, 2018 (MDPs, Bellman equations, Q-learning, policy gradients).
• Williams. Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning. Machine Learning 8, 1992 (REINFORCE).
• Mnih, Kavukcuoglu, Silver, Graves, Antonoglou, Wierstra and Riedmiller. Playing Atari with Deep Reinforcement Learning. 2013. arXiv:1312.5602 (DQN).
• Schulman, Wolski, Dhariwal, Radford and Klimov. Proximal Policy Optimization Algorithms. 2017. arXiv:1707.06347.
• PyTorch v2.15.0: torch/distributions/categorical.py, torch/distributions/__init__.py.