Reinforcement Learning

From returns and Bellman backups through Q-learning, policy gradients, GAE and PPO to the reward-model, KL and DPO losses used to align LLMs.

Reinforcement Learning

Why it matters

Reinforcement learning (RL) trains an agent by trial and error to maximize reward instead of imitating labels. It powers game-playing agents, robotics, recommendation and, most visibly today, the post-training of LLMs (RLHF, DPO, GRPO). Interviews test both layers: the tabular basics (Bellman equations, Q-learning vs SARSA) and the policy-gradient machinery behind PPO and RLHF.

Core ideas

An agent in state s picks action a from its policy π(a|s), receives reward r and lands in s'. It wants to maximize the discounted return:

G_t = r_t + γ·r_{t+1} + γ²·r_{t+2} + ...  =  r_t + γ·G_{t+1}        (0 ≤ γ ≤ 1)

Value functions predict that return: V(s) from a state, Q(s, a) after taking a. They satisfy the Bellman equations:

V(s)    = E[ r + γ·V(s') ]                      expectation: evaluate a policy
Q*(s,a) = E[ r + γ·max_a' Q*(s', a') ]          optimality: find the best policy

Most algorithms solve these equations, exactly or from samples:

Family	Learns	Update target	Examples
Dynamic programming	V with a known model	full expectation over s'	value iteration
TD, value-based	V or Q from samples	r + γ·(estimate at s')	TD(0), SARSA, Q-learning, DQN
Policy gradient	the policy π_θ directly	return or advantage × ∇ log π	REINFORCE, PPO
Actor-critic	π_θ and a critic V	advantage A = Q - V, e.g. from GAE	A2C, PPO with GAE

Key distinctions:

On-policy vs off-policy. SARSA bootstraps from the next action it actually takes (on-policy). Q-learning bootstraps from the greedy action whatever it does next (off-policy).
Bias vs variance. Monte Carlo returns are unbiased but noisy. TD targets use the critic's estimate: biased, but low variance. n-step returns and GAE's λ interpolate between the two.
Exploration. ε-greedy acts randomly with probability ε so that better actions can still be found.
Policy gradient. ∇J = E[ A_t · ∇ log π(a_t|s_t) ]. PPO clips the ratio π_θ / π_old to [1-ε, 1+ε] so that several gradient steps on one batch stay safe.

RL for language models

For an LLM, the state is the prompt plus the tokens so far, an action is the next token, and the reward usually comes once per response.

Reward model: learn a scalar score from human preference pairs with the Bradley–Terry loss -log σ(r_chosen - r_rejected).
KL penalty: keep π_θ close to the frozen reference model, reward - β·KL(π_θ ‖ π_ref), so it doesn't reward-hack.
PPO-RLHF: the clipped surrogate, with the KL folded into the reward.
DPO: skip the reward model and the RL loop: -log σ(β·(Δ_chosen - Δ_rejected)) on policy-vs-reference log-ratios.

Worked micro-example

Discounted return of rewards [1, 1, 1, 1] with γ = 0.9:

G_0 = 1 + 0.9 + 0.81 + 0.729 = 3.439

One Q-learning step with Q(s, a) = 1, r = 5, Q(s', ·) = [3, 4], γ = 0.9, α = 0.5:

target  = 5 + 0.9 · max(3, 4) = 8.6
δ       = 8.6 - 1 = 7.6
Q(s, a) = 1 + 0.5 · 7.6 = 4.8

SARSA with next action a' = 0 uses Q(s', 0) = 3 instead: target 7.7, new value 4.35.

Common pitfalls

Terminal states: a terminal next state has value 0, so don't bootstrap through done. A truncated (time-limit) episode is different: it still bootstraps from V(s_T).
Off-by-one in returns: the first reward has discount γ⁰; values for GAE has length T + 1.
In-place updates: a synchronous value-iteration sweep reads the old V only. Copy Q-tables instead of mutating the caller's array.
Wrong sign: policy gradients and PPO objectives are maximized; the training loss is their negative.
Unstable log-sigmoid: log(sigmoid(x)) underflows and log(p + 1e-10) caps the loss. Use -logaddexp(0, -x) (NumPy) or F.logsigmoid (PyTorch).
Double-counting KL: in PPO-RLHF, put the KL either in the reward or in the loss, never both.

Interview angle

Expect to derive the Q-learning and SARSA updates, explain why Q-learning overestimates (the max of noisy estimates; Double DQN fixes it) and contrast on- and off-policy learning. For policy gradients: derive ∇ log π for a softmax, say why a baseline reduces variance without adding bias, and walk through how PPO's min and clip behave for positive and negative advantages. For RLHF: why the KL penalty is there, how DPO's implicit reward β·log(π_θ/π_ref) arises, and what changes in GRPO. Complexity: a value-iteration sweep is O(|S|²·|A|); a TD update is O(1) and GAE is O(T) per trajectory.

Problems in this section

Discounted Return: vectorized G_0 with a discount vector.
Epsilon-Greedy Policy: explore or exploit, with deterministic tie-breaking.
Bellman Expectation (State Value): a one-state expected backup.
Q-Learning Update: off-policy TD control with a max over next actions.
SARSA Update: on-policy TD control with the action actually taken.
TD(0) Value Prediction: bootstrapped value learning with terminal handling.
Value Iteration Step: a synchronous Bellman optimality sweep over a known MDP.
N-Step Return: n real rewards plus a bootstrap, with truncation.
REINFORCE Policy Gradient: the Monte Carlo policy gradient for a softmax policy.
Advantage Estimation (GAE): the backward GAE(γ, λ) recursion.
PPO Clipped Objective: the clipped surrogate.
RLHF Reward Model Loss: the Bradley–Terry pairwise loss, computed stably.
Direct Preference Optimization (DPO): preference learning without a reward model.
KL Divergence Penalty (RLHF): the sampled k1 KL estimate.
PPO-RLHF Objective: the clipped surrogate with a KL-shaped reward.

Next, Post-training & Efficiency builds GRPO (PPO without a critic, with a k3 KL term) and int8 quantization. Softmax and cross-entropy, which policy gradients rely on, are practiced in Neural Networks.

Reinforcement Learning