Reinforcement Learning Foundations

The reward was sparse, the episodes were long, and my discount factor suggested I should have studied supervised learning instead.

Tensor, Reward-Starved AI Agent

1. The Reinforcement Learning Framework

Supervised learning needs labeled data: input X, correct output Y. But what if there is no single "correct" answer, only better and worse ones? What if the learner must try something, observe the consequence, and improve over time? That is reinforcement learning.

Think of training a dog. You say "sit." The dog tries something (maybe it lies down). You give a treat only when it actually sits. Over many repetitions, the dog learns which action earns the treat. This is exactly the RL loop: an agent (the dog) interacts with an environment (you and the room), takes actions (sitting, lying down), and receives rewards (treats or nothing).

Let us define each piece of the RL framework precisely:

Figure 0.4.2: The agent-environment interaction loop. At each step the agent observes a state, takes an action, and receives a reward.

2. Policies: The Agent's Strategy

A policy is the agent's strategy: a rule that maps each state to an action. It answers the question, "Given what I see right now, what should I do?"

Policies come in two flavors. A deterministic policy always picks the same action for a given state (if the dog sees a hand signal, it always sits). A stochastic policy assigns probabilities to each possible action, then samples from that distribution.

Formally, a stochastic policy is written as π(a | s), the probability of choosing action a when in state s. For an LLM, this is exactly the softmax output: the probability the model assigns to each token given the context so far. The entire goal of RL training is to adjust the parameters of π so that high-reward actions become more probable.

3. Value Functions and the Bellman Equation

Rewards tell us how good a single step was. But we need a way to evaluate long-term prospects. Is this state a good place to be? Is this action a wise choice? Value functions answer these questions.

State-Value Function V(s)

What: V(s) estimates the total future reward the agent expects to accumulate starting from state s and following its current policy onward. Why it matters: It lets the agent judge whether its current situation is promising or dire. How it works: Think of it as a "mood meter." A high V(s) means "things are going well from here"; a low V(s) means "trouble ahead."

Action-Value Function Q(s, a)

What: Q(s, a) estimates the total future reward if the agent takes action a in state s and then follows its policy. Why it matters: It lets the agent compare actions and pick the best one. How it works: If V(s) is the mood meter, Q(s, a) is like asking, "If I take this specific action right now, will things go better or worse than average?"

The Bellman Equation (Intuition)

The Bellman equation expresses a simple but powerful recursive idea: the value of a state equals the immediate reward plus the (discounted) value of the next state. In plain English: "How good is it to be here? Well, how much reward do I get right now, plus how good is the place I end up?"

$V(s) = E[ r + \gamma \cdot V(s') ]$

This recursion is the engine behind nearly all RL algorithms. The discount factor γ (gamma, between 0 and 1) controls how much the agent values future rewards relative to immediate ones. When γ is close to 1, the agent is patient and plans ahead. When γ is close to 0, it is greedy and focuses on immediate rewards.

4. Policy Gradients: Learning by Trial and Feedback

Now we arrive at the core question: how does the agent improve its policy? The answer is the policy gradient theorem, and the intuition is remarkably straightforward.

Imagine the dog tries ten different behaviors. Three of them earned treats. The training strategy is simple: make those three behaviors more likely in the future. Behaviors that earned nothing (or a scolding) become less likely. That is the entire idea behind policy gradients.

More precisely: the agent samples actions from its policy, observes the rewards, and then adjusts its parameters (the neural network weights) to increase the probability of actions that led to high rewards and decrease the probability of actions that led to low rewards. The magnitude of each adjustment is proportional to the reward received.

# A 4x4 GridWorld environment: the agent navigates from (0,0) to the goal.
# Each step returns (next_state, reward, done) just like OpenAI Gym.
import numpy as np

class SimpleGridWorld:
 """A 4x4 grid where an agent must reach the goal.
 State: (row, col) position. Actions: 0=up, 1=right, 2=down, 3=left.
 Reward: +1 for reaching the goal, -0.01 per step (to encourage speed).
 LLM analogy: each 'step' is like generating one token, and the final
 reward scores the complete trajectory (the full response)."""

 def __init__(self, size=4):
 self.size = size
 self.goal = (size - 1, size - 1)
 self.reset()

 def reset(self):
 self.pos = (0, 0)
 return self.pos

 def step(self, action):
 r, c = self.pos
 if action == 0: r = max(r - 1, 0) # up
 elif action == 1: c = min(c + 1, self.size - 1) # right
 elif action == 2: r = min(r + 1, self.size - 1) # down
 elif action == 3: c = max(c - 1, 0) # left
 self.pos = (r, c)

 done = (self.pos == self.goal)
 reward = 1.0 if done else -0.01
 return self.pos, reward, done

# Run one episode with a random policy
env = SimpleGridWorld()
state = env.reset()
total_reward = 0
for step in range(100):
 action = np.random.randint(4) # random stochastic policy
 state, reward, done = env.step(action)
 total_reward += reward
 if done:
 break

print(f"Episode finished in {step + 1} steps, total reward: {total_reward:.2f}")

# REINFORCE algorithm: a policy network outputs action probabilities,
# and the loss weights each log-probability by its discounted return.
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical

class PolicyNetwork(nn.Module):
 """A small neural network that maps states to action probabilities.
 In an LLM, this role is played by the entire transformer: it maps
 the token sequence (state) to a distribution over the next token (action)."""

 def __init__(self, state_dim, n_actions, hidden=64):
 super().__init__()
 self.net = nn.Sequential(
 nn.Linear(state_dim, hidden),
 nn.ReLU(),
 nn.Linear(hidden, n_actions),
 nn.Softmax(dim=-1)
 )

 def forward(self, state):
 return self.net(state)

# Simplified REINFORCE training loop
policy = PolicyNetwork(state_dim=2, n_actions=4)
optimizer = optim.Adam(policy.parameters(), lr=1e-3)

def reinforce_episode(env, policy):
 """Collect one episode and update the policy.
 Core idea: increase probability of actions that led to positive reward,
 decrease probability of actions that led to negative reward."""
 state = env.reset()
 log_probs = []
 rewards = []

 # 1. Roll out an episode using the current policy
 for _ in range(200):
 state_tensor = torch.FloatTensor(state)
 probs = policy(state_tensor)
 dist = Categorical(probs)
 action = dist.sample() # stochastic action selection
 log_probs.append(dist.log_prob(action))

 state, reward, done = env.step(action.item())
 rewards.append(reward)
 if done:
 break

 # 2. Compute discounted returns (what was the total reward from each step?)
 returns = []
 G = 0
 for r in reversed(rewards):
 G = r + 0.99 * G
 returns.insert(0, G)
 returns = torch.FloatTensor(returns)
 returns = (returns - returns.mean()) / (returns.std() + 1e-8)

 # 3. Policy gradient: nudge probabilities toward high-reward actions
 loss = -sum(lp * R for lp, R in zip(log_probs, returns))
 optimizer.zero_grad()
 loss.backward()
 optimizer.step()
 return sum(rewards)

# PPO sketch: collect trajectories, compute advantage estimates,
# and update the policy with a clipped surrogate objective.
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import random

# Simple environment: agent must learn to pick action 2 (out of 0,1,2,3)
# Reward = +1 if action == 2, else -0.1
class SimpleEnv:
 def reset(self):
 return [random.random(), random.random()] # random 2D state
 def step(self, action):
 reward = 1.0 if action == 2 else -0.1
 done = True # single-step episodes for clarity
 return self.reset(), reward, done

policy = nn.Sequential(nn.Linear(2, 32), nn.ReLU(), nn.Linear(32, 4), nn.Softmax(dim=-1))
optimizer = optim.Adam(policy.parameters(), lr=3e-3)
env = SimpleEnv()

for episode in range(1, 1001):
 state = torch.FloatTensor(env.reset())
 probs = policy(state)
 dist = Categorical(probs)
 action = dist.sample()
 log_prob = dist.log_prob(action)
 _, reward, _ = env.step(action.item())

 loss = -log_prob * reward
 optimizer.zero_grad()
 loss.backward()
 optimizer.step()

 if episode % 200 == 0:
 # Evaluate: check how often the policy picks action 2
 correct = sum(1 for _ in range(100)
 if policy(torch.FloatTensor(env.reset())).argmax().item() == 2)
 print(f"Episode {episode:4d} | Last reward: {reward:+.1f} | Action 2 rate: {correct}%")

5. PPO: Stable Policy Updates

Proximal Policy Optimization (PPO) solves a critical problem with basic policy gradients: if a single update changes the policy too drastically, performance can collapse and never recover. PPO prevents this by clipping the update, ensuring each step is small and safe.

Here is the intuition. Imagine you are adjusting a recipe. You taste the soup, decide it needs more salt, and add some. With vanilla policy gradients, nothing stops you from dumping in the entire salt shaker. PPO is the rule that says: "Never add more than one pinch at a time." You can always taste again and add another pinch, but you cannot ruin the soup with a single reckless change.

Technically, PPO computes a ratio between the new policy's probability of an action and the old policy's probability. If this ratio drifts too far from 1.0 (typically beyond a range of 0.8 to 1.2), the gradient is clipped to zero. The policy still improves, but it cannot make a catastrophically large jump in a single step.

6. From RL to LLM Training: The Complete Picture

We have built up the vocabulary piece by piece. Now let us assemble the full picture of how reinforcement learning powers LLM alignment. The mapping is direct and concrete:

Figure 0.4.4: How RL concepts map to LLM training with RLHF. The LLM is the policy, tokens are actions, and the reward model provides the training signal.

The RLHF training pipeline works as follows:

1. Supervised fine-tuning (SFT): First, the LLM is fine-tuned on high-quality human demonstrations. This gives the policy a good starting point.
2. Reward model training: Human annotators rank multiple model outputs for the same prompt. A separate neural network (the reward model) is trained to predict these human preferences.
3. RL optimization with PPO: The LLM generates responses to prompts. The reward model scores each response. PPO uses these scores to update the LLM's weights, nudging the probability distribution toward responses that score highly.

A critical detail: during the PPO phase, a KL (Kullback-Leibler) penalty prevents the RL-trained model from drifting too far from the SFT model. Without this constraint, the model might find degenerate responses that exploit the reward model (for example, generating repetitive flattery that scores highly but is useless). The KL penalty keeps the model close to its pre-trained language abilities.

7. Putting It All Together: RL to LLM Alignment Pipeline

Let us trace the full journey from RL foundations to LLM training. An LLM begins as a pretrained model that can generate fluent text but may produce harmful, incorrect, or unhelpful responses. We want to steer it toward better behavior.

We frame this as an RL problem. The LLM is the agent (the policy). Each prompt starts a new episode. At each time step, the model picks a token (an action) based on its current context (the state). When the response is complete, the reward model assigns a score (the reward). PPO adjusts the model's weights so that high-scoring response patterns become more probable, while the clipping mechanism and KL penalty ensure the model does not lose its fluency or degenerate into reward hacking.

The dog analogy still holds at this scale. The dog does not understand English; it learns by trying actions and receiving treats. The LLM does not "understand" human values; it learns which outputs are valued by observing reward signals. The elegance of RL is that this simple loop of action, feedback, and adjustment can produce remarkably sophisticated behavior.

Exercises for Further Practice

1. Modify the grid world reward. Change the step penalty from -0.01 to -0.1 in the SimpleGridWorld class. Run several episodes with a random policy. How does the harsher penalty affect total reward? What would this mean if the grid world were an LLM generating a long response?
2. Experiment with the discount factor. In the REINFORCE code (Code Example 2), change the discount factor gamma from 0.99 to 0.5 and then to 0.999. Observe how this affects learning speed. Relate your observations to the Bellman equation: how does a lower gamma change the agent's planning horizon?
3. Compare random vs. trained policy. Using Code Example 3 (the complete REINFORCE), run 100 evaluation episodes with the untrained policy (before training) and 100 with the trained policy (after 1000 episodes). Compare the average action-2 selection rate. This contrast illustrates the core value of RL: learning from experience.

What Comes Next

With ML fundamentals, deep learning, PyTorch, and reinforcement learning now in your toolkit, you have completed Chapter 00. In Chapter 01: Foundations of NLP and Text Representation, we shift from general machine learning to the specific domain of language. You will learn how text is converted into numbers (tokenization and embeddings), how classical NLP techniques work, and why representing words as vectors was one of the most transformative ideas in the field. These representations are the raw material that transformers (Chapter 04) will learn to process.

pip install matplotlib numpy

import matplotlib.pyplot as plt
import numpy as np

milestones = {
 1957: "Perceptron",
 1986: "Backpropagation",
 1997: "LSTM",
 2012: "AlexNet",
 2013: "Word2Vec",
 2014: "GANs",
 2017: "Transformer",
 2018: "BERT",
 2020: "GPT-3",
 2022: "InstructGPT\n(RLHF)",
 2023: "GPT-4",
}

years = list(milestones.keys())
labels = list(milestones.values())

fig, ax = plt.subplots(figsize=(14, 3))
ax.scatter(years, [0] * len(years), s=80, zorder=5)
for i, (yr, lbl) in enumerate(zip(years, labels)):
 offset = 0.4 if i % 2 == 0 else -0.6
 ax.annotate(f"{lbl}\n({yr})", (yr, 0), (yr, offset),
 fontsize=8, ha="center", va="bottom" if offset > 0 else "top",
 arrowprops=dict(arrowstyle="-", color="gray", lw=0.8))

ax.axhline(0, color="black", lw=0.8)
ax.set_xlim(1955, 2025)
ax.set_ylim(-1.2, 1.2)
ax.axis("off")
ax.set_title("Key ML and AI Milestones", fontsize=13, fontweight="bold")
plt.tight_layout()
plt.savefig("ml_timeline.png", dpi=150, bbox_inches="tight")
plt.show()
print("Timeline saved to ml_timeline.png")

import numpy as np

class GridWorld:
 def __init__(self, size=4):
 self.size = size
 self.state = (0, 0)
 self.goal = (size - 1, size - 1)
 self.actions = [(0, 1), (0, -1), (1, 0), (-1, 0)] # right, left, down, up

 def reset(self):
 self.state = (0, 0)
 return self.state

 def step(self, action_idx):
 dr, dc = self.actions[action_idx]
 r = max(0, min(self.size - 1, self.state[0] + dr))
 c = max(0, min(self.size - 1, self.state[1] + dc))
 self.state = (r, c)
 done = self.state == self.goal
 reward = 1.0 if done else -0.01
 return self.state, reward, done

env = GridWorld()
print(f"Start: {env.reset()}, Goal: {env.goal}")

Q = np.zeros((4, 4, 4)) # states (4x4) x actions (4)
alpha, gamma, epsilon = 0.1, 0.99, 0.1
episodes = 500

for ep in range(episodes):
 state = env.reset()
 for _ in range(100): # max steps
 # Epsilon-greedy action selection
 if np.random.random() < epsilon:
 action = np.random.randint(4)
 else:
 action = np.argmax(Q[state[0], state[1]])

 next_state, reward, done = env.step(action)
 # Q-learning update
 best_next = np.max(Q[next_state[0], next_state[1]])
 Q[state[0], state[1], action] += alpha * (
 reward + gamma * best_next - Q[state[0], state[1], action]
 )
 state = next_state
 if done:
 break

# Show learned policy
action_names = ["right", "left", "down", "up"]
print("Learned policy:")
for r in range(4):
 row = []
 for c in range(4):
 row.append(action_names[np.argmax(Q[r, c])])
 print(f" Row {r}: {row}")

v = np.max(Q, axis=2) # Value function: max Q at each state

fig, ax = plt.subplots(figsize=(5, 5))
im = ax.imshow(v, cmap="YlOrRd", interpolation="nearest")
for r in range(4):
 for c in range(4):
 ax.text(c, r, f"{v[r, c]:.2f}", ha="center", va="center", fontsize=11)
ax.set_title("Learned Value Function V(s) = max Q(s,a)")
ax.set_xlabel("Column")
ax.set_ylabel("Row")
plt.colorbar(im, ax=ax, shrink=0.8)
plt.tight_layout()
plt.savefig("gridworld_values.png", dpi=150)
plt.show()