Reinforcement Learning Topics in AI Engineering From Scratch: DQN, PPO, and RLHF Explained

TLDR: The "09-reinforcement-learning" phase of the rohitg00/ai-engineering-from-scratch repository provides self-contained Python implementations covering the full reinforcement learning spectrum—from Markov Decision Processes and Dynamic Programming to Deep Q-Networks, Proximal Policy Optimization, and Reinforcement Learning from Human Feedback—using pure Python without external ML frameworks.

The rohitg00/ai-engineering-from-scratch repository offers a comprehensive educational resource for understanding reinforcement learning topics through hands-on coding. Each lesson resides in phases/09-reinforcement-learning/ as an executable main.py file that demonstrates algorithmic implementations using only standard Python data structures and NumPy-style operations, stripping away framework abstraction to reveal the underlying mathematical mechanics.

From Markov Decision Processes to Deep Function Approximation

The curriculum follows a logical progression from classical tabular methods to modern deep reinforcement learning. Early lessons establish the theoretical foundation with Markov Decision Processes (MDPs) and exact solution methods, while later modules introduce neural function approximation and alignment techniques.

MDP Foundations and Exact Methods

The first four lessons establish the mathematical framework:

• MDP Formalism: Lesson 01 in phases/09-reinforcement-learning/01-mdps-states-actions-rewards/code/main.py defines the grid-world environment with tuple-based state representations and deterministic transition dynamics.
• Dynamic Programming: Lesson 02 implements value iteration and policy iteration using tabular Bellman updates on the same 4×4 grid.
• Monte-Carlo Methods: Lesson 03 demonstrates first-visit MC evaluation through episode-wise return averaging.
• Temporal-Difference Control: Lesson 04 covers both on-policy SARSA and off-policy Q-learning with ε-greedy exploration and tabular Q-tables.

Deep Q-Network (DQN): Neural Function Approximation with Experience Replay

The DQN implementation in phases/09-reinforcement-learning/05-dqn/code/main.py demonstrates how to stabilize learning when replacing lookup tables with neural networks.

Architecture and Target Networks

The network architecture consists of a two-layer MLP initialized via init_net() with weight matrices W1, b1, W2, and b2. The input uses one-hot encoding of the 4×4 grid state (state_features), processed through a ReLU hidden layer to produce linear Q-value outputs for each of the four actions.

def epsilon_greedy(net, state, rng, epsilon):
    if rng.random() < epsilon:                     # explore

        return rng.randrange(len(ACTIONS))
    q, _ = forward(net, state_features(state))     # exploit

    return max(range(len(ACTIONS)), key=lambda i: q[i])

The implementation employs a target network (target = clone(online)) that synchronizes with the online network every sync_every steps. This stabilizes temporal-difference targets by decoupling the action selection from the value estimation.

Experience Replay and Training

The train_step() function samples random minibatches from a replay buffer with fixed capacity, breaking correlations in sequential data:

# Inside the main loop:

if len(buffer) >= batch:
    mb = rng.sample(buffer, batch)                # sample replay batch

    train_step(online, target, mb, gamma, lr)     # TD-error back-prop

After 400 episodes, the learned Q-values approximate the optimal policy, with the greedy policy visualizing directional arrows toward the terminal state.

Proximal Policy Optimization (PPO): Clipped Surrogate Objectives

The PPO implementation in phases/09-reinforcement-learning/08-ppo/code/main.py demonstrates modern policy gradient methods with stabilized updates.

Policy Architecture and Rollout Collection

Unlike value-based methods, PPO maintains a policy network (linear softmax weights theta) and a value network (linear regressor weights w). The collect_rollout() function gathers experience from n_envs=8 parallel environments, storing states, actions, rewards, and the old policy log-probability (log_pi_old).

Generalized Advantage Estimation and Clipping

The gae() function computes advantages using Generalized Advantage Estimation with λ-returns. The ppo_update() function implements the clipped surrogate objective:

ratio = math.exp(logp - rec["log_pi_old"])
clipped = (adv > 0 and ratio > 1 + eps) or (adv < 0 and ratio < 1 - eps)
if not clipped:
    pg_scale = ratio * adv
else:
    pg_scale = 0.0

# policy gradient

theta[action][j] += lr_a * pg_scale * grad_logpi * x[j]

The clipping mechanism restricts the probability ratio to [1-ε, 1+ε] (typically ε=0.2), preventing overly aggressive policy updates that could destabilize training. After 60 updates, the policy converges to the optimal grid-world solution with an average return of approximately -6.0.

Reinforcement Learning from Human Feedback (RLHF): Reward Modeling and KL Regularization

The RLHF lesson in phases/09-reinforcement-learning/09-reward-modeling-rlhf/code/main.py implements a two-stage alignment pipeline combining preference learning with constrained policy optimization.

Pairwise Reward Model Training

The train_rm() function learns a linear reward model (weights w) over a vocabulary of tokens using pairwise Bradley-Terry logistic regression. The model trains on synthetic "good" versus "bad" token pairs to predict human preferences:

# 1️⃣ train a pairwise reward model

w = train_rm(n_pairs=600, rng=random.Random(42))

PPO with KL Penalty

The rlhf_loop() function integrates the reward model into PPO with a KL-divergence regularizer that keeps the optimized policy close to a reference (initial) policy:

reward = rm_score - beta * kl_term               # RM + KL penalty

...
ratio = math.exp(logp - log_pi_old)
clipped = (adv > 0 and ratio > 1 + eps) or (adv < 0 and ratio < 1 - eps)
if not clipped:
    for i in range(len(VOCAB)):
        grad = (1.0 if i == token else 0.0) - probs[i]
        theta[p_idx][i] += lr * ratio * adv * grad

The beta parameter controls the trade-off between maximizing the learned reward and maintaining proximity to the reference policy. Higher values (e.g., 1.0) enforce conservative updates, while lower values (e.g., 0.01) allow greater policy drift in pursuit of higher rewards.

Additional Reinforcement Learning Topics

Beyond the three flagship algorithms, the repository covers several advanced paradigms:

Policy Gradient Foundations

Lesson 06 (06-policy-gradients-reinforce/code/main.py) implements REINFORCE with stochastic softmax policies and Monte-Carlo returns. Lesson 07 (07-actor-critic-a2c-a3c/code/main.py) introduces Actor-Critic methods, using bootstrapped advantage estimates to reduce variance while maintaining the policy gradient approach.

Multi-Agent and Transfer Learning

• Multi-Agent RL: Lesson 10 (10-multi-agent-rl/code/main.py) demonstrates independent Q-learners sharing a joint grid-world environment, where each agent maintains separate Q-tables while learning simultaneously.
• Sim-to-Real Transfer: Lesson 11 (11-sim-to-real-transfer/code/main.py) applies domain randomization by training on randomized grid configurations before evaluating on the canonical environment.
• Game Applications: Lesson 12 (12-rl-for-games/code/main.py) wraps custom games (e.g., Tic-Tac-Toe) into the RL scaffolding with shaped reward functions.

Summary

The rohitg00/ai-engineering-from-scratch reinforcement learning curriculum provides executable implementations covering:

• Foundational theory: MDPs, value iteration, and first-visit Monte-Carlo methods in pure Python
• Deep Q-Networks: Two-layer MLPs with target networks and experience replay buffers for stable function approximation
• Proximal Policy Optimization: Clipped surrogate objectives with GAE and multi-environment rollouts for policy optimization
• RLHF alignment: Pairwise reward modeling followed by KL-regularized PPO to align policies with preference data
• Advanced applications: Multi-agent systems, sim-to-real transfer, and game-specific RL implementations

Each lesson maintains consistency through a shared 4×4 grid-world environment and avoids external dependencies, ensuring the focus remains on algorithmic principles rather than framework abstractions.

Frequently Asked Questions

What is the difference between the DQN and PPO implementations in this repository?

The DQN implementation (phases/09-reinforcement-learning/05-dqn/code/main.py) uses a value-based approach with a two-layer neural network to approximate Q-values, employing experience replay and target networks to stabilize off-policy learning. The PPO implementation (phases/09-reinforcement-learning/08-ppo/code/main.py) uses a policy-gradient approach with separate policy and value networks, optimizing via clipped surrogate objectives and generalized advantage estimation for on-policy updates.

How does the RLHF implementation handle the reward model training?

The RLHF lesson (phases/09-reinforcement-learning/09-reward-modeling-rlhf/code/main.py) implements a pairwise Bradley-Terry model through the train_rm() function, which learns linear weights from synthetic "good/bad" token pairs. This reward model then replaces the environment reward during the PPO phase, with a KL penalty (beta) ensuring the policy remains close to a reference distribution.

Can these implementations be run without PyTorch or TensorFlow?

Yes. All implementations in the 09-reinforcement-learning phase use only standard Python libraries and basic arithmetic operations. Neural networks are implemented manually using weight matrices (W1, b1, W2, b2) and forward functions, while gradients are computed analytically or through finite differences, eliminating dependencies on external ML frameworks.

What environment is used consistently across all lessons?

All lessons share a 4×4 grid-world environment with deterministic transitions, defined in the respective main.py files. This consistency allows direct comparison between tabular methods (lessons 01-04) and deep approaches (lessons 05-09), as the optimal policy and value functions remain constant across the curriculum.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}