Python Reference

A practical programming guide for the reinforcement learning labs

Python Reference

Note

This page is a practical Python reference for The Mathematics of Reinforcement Learning. It summarizes the coding patterns used throughout the computer labs.

Purpose

The reinforcement-learning labs use Python to make the mathematics concrete. This reference collects common patterns that appear across the notebooks:

• arrays and tables,
• random simulation,
• policies,
• value functions,
• Q-functions,
• training loops,
• plotting,
• diagnostics,
• reproducibility,
• debugging.

The goal is not to teach all of Python. The goal is to provide enough Python to read, run, modify, and extend the labs.

Repository

The book repository is:

https://github.com/wanghemath/Book-MathRL

The labs are stored in:

labs/

A typical lab notebook is:

labs/chapter-01-lab.ipynb

A typical Colab URL is:

https://colab.research.google.com/github/wanghemath/Book-MathRL/blob/main/labs/chapter-01-lab.ipynb

Standard imports

Most labs use the following imports:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import random

These libraries are enough for most tabular and small-scale reinforcement-learning experiments.

Library	Main use
numpy	arrays, random numbers, vectorized computation
pandas	tables and experiment summaries
matplotlib	plots and visualizations
random	simple Python random choices

Reproducibility

Set random seeds so that experiments can be repeated.

def set_seed(seed=0):
    np.random.seed(seed)
    random.seed(seed)

set_seed(42)

In reinforcement learning, results may change across random seeds because of:

• random initialization,
• random actions,
• stochastic transitions,
• random datasets,
• minibatch sampling.

For final projects, it is often better to run several seeds and report averages.

Basic NumPy arrays

Create arrays

x = np.array([1, 2, 3])
zeros = np.zeros(5)
ones = np.ones(5)
matrix = np.zeros((3, 4))

Shape

Q = np.zeros((10, 4))
print(Q.shape)

Output:

(10, 4)

For a tabular Q-function:

Q.shape == (number_of_states, number_of_actions)

Indexing

s = 3
a = 2

Q[s, a] = 1.5
value = Q[s, a]

Row operations

best_action = np.argmax(Q[s])
best_value = np.max(Q[s])

These appear frequently in greedy policy extraction.

Random choices

Random integer

a = np.random.randint(4)

This samples an integer from:

0, 1, 2, 3

Random action from probabilities

probs = np.array([0.1, 0.2, 0.6, 0.1])
a = np.random.choice(4, p=probs)

Shuffle data

np.random.shuffle(transitions)

This is useful in offline reinforcement learning and fitted Q iteration.

Tables with pandas

Pandas is useful for displaying trajectories, summaries, and experiment results.

rows = []

rows.append({
    "episode": 0,
    "return": 1.2,
    "length": 15,
    "success": True
})

df = pd.DataFrame(rows)
df

Summary table

summary = pd.DataFrame([
    {"method": "Q-learning", "success rate": 0.82, "mean return": 0.51},
    {"method": "Dyna-Q", "success rate": 0.91, "mean return": 0.63},
])

summary

Group by

df.groupby("method")["return"].mean()

This is useful when comparing random seeds or algorithms.

Plotting with matplotlib

Basic line plot

plt.figure(figsize=(7, 4))
plt.plot(returns)
plt.xlabel("episode")
plt.ylabel("return")
plt.title("Learning curve")
plt.grid(True, alpha=0.3)
plt.show()

Plot several curves

plt.figure(figsize=(7, 4))
plt.plot(q_learning_returns, label="Q-learning")
plt.plot(dyna_returns, label="Dyna-Q")
plt.xlabel("episode")
plt.ylabel("return")
plt.title("Comparison")
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()

Bar chart

methods = ["Q-learning", "Dyna-Q", "Offline FQI"]
success_rates = [0.76, 0.88, 0.65]

plt.figure(figsize=(7, 4))
plt.bar(methods, success_rates)
plt.ylabel("success rate")
plt.title("Policy comparison")
plt.grid(True, axis="y", alpha=0.3)
plt.show()

Heatmap

values = np.arange(16).reshape(4, 4)

plt.figure(figsize=(5, 4))
plt.imshow(values)
plt.colorbar(label="value")
plt.title("Value heatmap")
plt.show()

Heatmaps are useful for value functions and state-visitation counts.

Rolling averages

RL learning curves are noisy. Rolling averages make trends easier to see.

def rolling_mean(x, window=20):
    x = np.asarray(x, dtype=float)
    if len(x) < window:
        return x
    return np.convolve(x, np.ones(window) / window, mode="valid")

Use it like this:

plt.figure(figsize=(7, 4))
plt.plot(rolling_mean(returns, window=20))
plt.xlabel("episode")
plt.ylabel("20-episode rolling return")
plt.title("Smoothed learning curve")
plt.grid(True, alpha=0.3)
plt.show()

Grid-world state indexing

Many labs use grid-world environments.

A two-dimensional position:

(row, col)

is converted to a one-dimensional state index:

state = row * n_cols + col

The inverse operation is:

row, col = divmod(state, n_cols)

Example:

rows = 4
cols = 5

def to_state(row, col):
    return row * cols + col

def to_pos(state):
    return divmod(state, cols)

s = to_state(2, 3)
print(s)

print(to_pos(s))

Minimal environment interface

Most labs use small custom environments with this interface:

s = env.reset()
sp, r, done, info = env.step(a)

The common variables are:

Variable	Meaning
s	current state
a	action
sp	next state
r	reward
done	whether the episode ended
info	optional diagnostic dictionary

A minimal episode loop:

s = env.reset()
done = False
total_reward = 0.0

while not done:
    a = np.random.randint(env.n_actions)
    sp, r, done, info = env.step(a)

    total_reward += r
    s = sp

Policy representations

A policy tells the agent how to choose actions.

Deterministic policy

A deterministic policy can be stored as an integer array:

policy = np.zeros(env.n_states, dtype=int)
policy[s] = 1

Then action selection is:

a = int(policy[s])

Stochastic policy

A stochastic policy can be stored as a matrix:

policy_probs = np.ones((env.n_states, env.n_actions)) / env.n_actions

Then action selection is:

a = np.random.choice(env.n_actions, p=policy_probs[s])

Greedy policy from Q

policy = np.argmax(Q, axis=1)

For terminal states, it is common to set the action to zero:

for s in terminal_states:
    policy[s] = 0

Epsilon-greedy action selection

Epsilon-greedy exploration is used in many labs.

def epsilon_greedy(Q, state, epsilon):
    if np.random.rand() < epsilon:
        return np.random.randint(Q.shape[1])
    return int(np.argmax(Q[state]))

Interpretation:

• with probability epsilon, explore randomly;
• with probability 1 - epsilon, exploit the current best action.

A decaying epsilon schedule:

epsilon = eps_end + (eps_start - eps_end) * np.exp(-episode / decay)

Example:

eps_start = 1.0
eps_end = 0.05
decay = 300

epsilon = eps_end + (eps_start - eps_end) * np.exp(-episode / decay)

Value functions

A tabular value function is a vector:

V = np.zeros(env.n_states)

The value of state s is:

V[s]

A Bellman-style update:

V[s] += alpha * (target - V[s])

For TD(0), the target is:

target = r + gamma * V[sp]

if the next state is nonterminal.

If the transition ended the episode:

target = r

Q-functions

A tabular Q-function is a matrix:

Q = np.zeros((env.n_states, env.n_actions))

The value of state-action pair (s,a) is:

Q[s, a]

The greedy action is:

a_star = np.argmax(Q[s])

The greedy value is:

max_q = np.max(Q[s])

Q-learning update

The Q-learning target is:

target = r
if not done:
    target += gamma * np.max(Q[sp])

The update is:

td_error = target - Q[s, a]
Q[s, a] += alpha * td_error

As a function:

def q_learning_update(Q, s, a, r, sp, done, alpha=0.1, gamma=0.95):
    target = r
    if not done:
        target += gamma * np.max(Q[sp])

    td_error = target - Q[s, a]
    Q[s, a] += alpha * td_error

    return td_error

Mathematically:

\[ Q(s,a) \leftarrow Q(s,a) + \alpha \left[ r+\gamma\max_b Q(s',b)-Q(s,a) \right]. \]

SARSA update

SARSA uses the next action actually chosen by the behavior policy.

target = r
if not done:
    target += gamma * Q[sp, ap]

td_error = target - Q[s, a]
Q[s, a] += alpha * td_error

Mathematically:

\[ Q(s,a) \leftarrow Q(s,a) + \alpha \left[ r+\gamma Q(s',a')-Q(s,a) \right]. \]

Expected SARSA update

Expected SARSA uses the expected next Q-value under the policy.

expected_next = np.sum(policy_probs[sp] * Q[sp])

target = r
if not done:
    target += gamma * expected_next

Q[s, a] += alpha * (target - Q[s, a])

Monte Carlo returns

The discounted return from time t is:

\[ G_t = R_{t+1} + \gamma R_{t+2} + \gamma^2 R_{t+3}+\cdots. \]

Code:

def discounted_returns(rewards, gamma=0.95):
    returns = np.zeros(len(rewards))
    G = 0.0

    for t in reversed(range(len(rewards))):
        G = rewards[t] + gamma * G
        returns[t] = G

    return returns

Example:

rewards = [0, 0, 1]
discounted_returns(rewards, gamma=0.9)

Policy gradient basics

For a softmax policy:

def softmax(logits):
    z = logits - np.max(logits)
    exp_z = np.exp(z)
    return exp_z / np.sum(exp_z)

If theta[s, a] stores logits for each state and action:

probs = softmax(theta[s])
a = np.random.choice(env.n_actions, p=probs)

The REINFORCE update is based on:

\[ \nabla_\theta J(\theta) = E[ G_t\nabla_\theta\log\pi_\theta(A_t\mid S_t) ]. \]

For a tabular softmax policy, the gradient of the log probability is:

grad = -probs
grad[a] += 1.0

Then update:

theta[s] += alpha * G * grad

Advantage estimates

An advantage compares an action to the state baseline:

\[ A(s,a)=Q(s,a)-V(s). \]

In actor-critic methods, the TD error often plays the role of advantage:

\[ \delta_t = R_{t+1} + \gamma V(S_{t+1}) - V(S_t). \]

Code:

td_error = r + gamma * V[sp] * (1 - done) - V[s]

Then:

V[s] += alpha_v * td_error
theta[s] += alpha_pi * td_error * grad_log_prob

Linear function approximation

A linear value function has the form:

\[ V(s;w)=\phi(s)^T w. \]

Code:

def value(phi_s, w):
    return phi_s @ w

Prediction:

v_s = features[s] @ w

Update:

target = r + gamma * (features[sp] @ w)
error = target - features[s] @ w
w += alpha * error * features[s]

Feature matrices

A feature matrix stores one row per state:

Phi = np.zeros((n_states, n_features))

The feature vector of state s is:

phi_s = Phi[s]

Linear prediction for all states:

V = Phi @ w

Least-squares solution:

w_hat = np.linalg.solve(Phi.T @ Phi, Phi.T @ y)

Safer version with ridge regularization:

ridge = 1e-6
w_hat = np.linalg.solve(Phi.T @ Phi + ridge * np.eye(Phi.shape[1]), Phi.T @ y)

Replay buffers

Deep Q-learning uses replay buffers to reuse transitions.

A simple replay buffer:

class ReplayBuffer:
    def __init__(self, capacity=10000):
        self.capacity = capacity
        self.data = []

    def add(self, transition):
        if len(self.data) >= self.capacity:
            self.data.pop(0)
        self.data.append(transition)

    def sample(self, batch_size):
        indices = np.random.choice(len(self.data), size=batch_size, replace=False)
        return [self.data[i] for i in indices]

    def __len__(self):
        return len(self.data)

Each transition usually has the form:

(s, a, r, sp, done)

Target networks

A target network is a delayed copy of the current network. In small NumPy examples, this may be represented by copying parameters:

target_weights = current_weights.copy()

Update every fixed number of steps:

if step % target_update_frequency == 0:
    target_weights = current_weights.copy()

The idea is to stabilize bootstrapping targets.

Dyna-Q model dictionary

Dyna-Q stores a learned model:

model = {}
seen_pairs = []

After a real transition:

if (s, a) not in model:
    seen_pairs.append((s, a))

model[(s, a)] = (r, sp, done)

Planning update:

sim_s, sim_a = random.choice(seen_pairs)
sim_r, sim_sp, sim_done = model[(sim_s, sim_a)]

q_learning_update(Q, sim_s, sim_a, sim_r, sim_sp, sim_done, alpha, gamma)

Offline datasets

Offline RL uses a fixed dataset:

rows = []

rows.append({
    "episode": ep,
    "t": t,
    "state": s,
    "action": a,
    "reward": r,
    "next_state": sp,
    "done": done
})

data = pd.DataFrame(rows)

A transition array:

transitions = data[["state", "action", "reward", "next_state", "done"]].to_numpy()

State-action counts:

def state_action_counts(data, n_states, n_actions):
    counts = np.zeros((n_states, n_actions), dtype=int)

    for _, row in data.iterrows():
        s = int(row["state"])
        a = int(row["action"])
        counts[s, a] += 1

    return counts

Coverage matters in offline RL.

Behavior cloning

Behavior cloning learns from expert state-action pairs:

expert_data = pd.DataFrame([
    {"state": 0, "action": 1},
    {"state": 1, "action": 1},
    {"state": 2, "action": 0},
])

Tabular behavior cloning:

def train_tabular_behavior_cloning(data, n_states, n_actions):
    counts = np.zeros((n_states, n_actions), dtype=int)

    for _, row in data.iterrows():
        counts[int(row["state"]), int(row["action"])] += 1

    policy = np.argmax(counts, axis=1)
    return policy, counts

For unseen states, use a fallback policy:

if np.sum(counts[s]) == 0:
    policy[s] = fallback_policy[s]

Preference data for RLHF

Preference data often has the form:

rows.append({
    "prompt": x,
    "winner": a_w,
    "loser": a_l
})

Bradley–Terry preference probability:

\[ P(y_w \succ y_l\mid x) = \sigma(r_\theta(x,y_w)-r_\theta(x,y_l)). \]

Sigmoid function:

def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-z))

Preference loss:

loss = -np.mean(np.log(sigmoid(reward_winner - reward_loser) + 1e-12))

KL divergence

For two discrete policies pi and pi_ref:

def mean_kl(pi, pi_ref):
    return np.mean(np.sum(pi * (np.log(pi + 1e-12) - np.log(pi_ref + 1e-12)), axis=1))

KL regularization appears in RLHF and trust-region policy methods.

Entropy

Policy entropy measures randomness:

def policy_entropy(pi):
    return np.mean(-np.sum(pi * np.log(pi + 1e-12), axis=1))

Entropy is high when the policy is random and low when the policy is nearly deterministic.

Evaluation loops

Training and evaluation should be separated.

During training, the agent may explore:

a = epsilon_greedy(Q, s, epsilon)

During evaluation, use a greedy policy:

a = int(np.argmax(Q[s]))

A standard evaluation function:

def evaluate_policy(env, policy, episodes=500, gamma=0.95, seed=0):
    set_seed(seed)

    returns = []
    lengths = []
    successes = 0

    for ep in range(episodes):
        s = env.reset()
        done = False
        G = 0.0
        discount = 1.0
        steps = 0

        while not done:
            a = int(policy[s])
            sp, r, done, info = env.step(a)

            G += discount * r
            discount *= gamma

            s = sp
            steps += 1

        returns.append(G)
        lengths.append(steps)

    return {
        "mean discounted return": float(np.mean(returns)),
        "std discounted return": float(np.std(returns)),
        "mean episode length": float(np.mean(lengths)),
    }

Modify this function for success rate, trap rate, or other task-specific metrics.

Training-loop pattern

A standard RL training loop looks like this:

returns = []
lengths = []

for episode in range(n_episodes):
    s = env.reset()
    done = False
    G = 0.0
    discount = 1.0
    steps = 0

    while not done:
        a = epsilon_greedy(Q, s, epsilon)
        sp, r, done, info = env.step(a)

        q_learning_update(Q, s, a, r, sp, done, alpha, gamma)

        G += discount * r
        discount *= gamma
        steps += 1

        s = sp

    returns.append(G)
    lengths.append(steps)

This pattern appears throughout the labs.

Common diagnostics

TD error

td_errors.append(abs(td_error))

Large TD errors may indicate rapid learning, instability, or poor value estimates.

State visitation

state_counts = np.zeros(env.n_states, dtype=int)
state_counts[s] += 1

Useful for exploration and offline coverage.

Action counts

action_counts = np.zeros(env.n_actions, dtype=int)
action_counts[a] += 1

Useful for bandits and policy diagnostics.

Success rate

successes.append(final_state == env.goal)

Plot:

plt.plot(rolling_mean(successes, window=50))

Policy disagreement

Compare a learned policy with a benchmark policy:

disagreement = np.mean(policy != reference_policy)

For grid-worlds, exclude walls and terminal states when appropriate.

Numerical stability tips

Avoid log of zero

Use a small constant:

np.log(probs + 1e-12)

Stable softmax

Do not write:

np.exp(logits) / np.sum(np.exp(logits))

Use:

z = logits - np.max(logits)
probs = np.exp(z) / np.sum(np.exp(z))

Avoid singular linear systems

Use ridge regularization:

A_reg = A + 1e-6 * np.eye(A.shape[0])
x = np.linalg.solve(A_reg, b)

Clip probabilities if needed

p = np.clip(p, 1e-8, 1.0)

Common errors and fixes

Error: index out of bounds

Cause: using a state or action outside the valid range.

Check:

print(s, env.n_states)
print(a, env.n_actions)

Valid ranges:

0 <= s < env.n_states
0 <= a < env.n_actions

Error: probabilities do not sum to 1

Cause: numerical or construction error in a probability vector.

Fix:

probs = probs / np.sum(probs)

Check:

print(np.sum(probs))

Error: NaN values

Check for:

• too large learning rate,
• division by zero,
• log of zero,
• unstable exponentials.

Debug:

np.isnan(Q).any()
np.max(np.abs(Q))

Error: policy never reaches the goal

Possible causes:

• insufficient exploration,
• too few episodes,
• reward too sparse,
• learning rate too small or too large,
• bug in environment dynamics,
• terminal states not handled correctly.

Error: notebook results differ from another run

RL is stochastic. Set seeds and compare averages over several seeds.

Suggested code organization

For projects, organize code into sections:

1. imports and random seed,
2. environment class,
3. helper functions,
4. algorithm implementation,
5. training,
6. evaluation,
7. visualization,
8. comparison table,
9. discussion.

Example:

# 1. Imports
# 2. Environment
# 3. Helpers
# 4. Algorithms
# 5. Experiments
# 6. Evaluation
# 7. Plots

Style guidelines

Use clear names:

state
action
next_state
reward
done

or compact RL names:

s
a
sp
r
done

Avoid mixing conventions in the same notebook.

Use functions for repeated logic:

def train_q_learning(...):
    ...

def evaluate_policy(...):
    ...

Store histories in dictionaries:

history = {
    "returns": np.array(returns),
    "lengths": np.array(lengths),
    "successes": np.array(successes)
}

Minimal Q-learning example

This is a compact template for many tabular control labs.

Q = np.zeros((env.n_states, env.n_actions))
returns = []

for episode in range(500):
    epsilon = 0.05 + 0.95 * np.exp(-episode / 200)

    s = env.reset()
    done = False
    G = 0.0
    discount = 1.0

    while not done:
        a = epsilon_greedy(Q, s, epsilon)
        sp, r, done, info = env.step(a)

        target = r
        if not done:
            target += gamma * np.max(Q[sp])

        Q[s, a] += alpha * (target - Q[s, a])

        G += discount * r
        discount *= gamma
        s = sp

    returns.append(G)

policy = np.argmax(Q, axis=1)

Minimal REINFORCE example

theta = np.zeros((env.n_states, env.n_actions))

for episode in range(500):
    states = []
    actions = []
    rewards = []

    s = env.reset()
    done = False

    while not done:
        probs = softmax(theta[s])
        a = np.random.choice(env.n_actions, p=probs)

        sp, r, done, info = env.step(a)

        states.append(s)
        actions.append(a)
        rewards.append(r)

        s = sp

    returns = discounted_returns(rewards, gamma)

    for s, a, G in zip(states, actions, returns):
        probs = softmax(theta[s])
        grad = -probs
        grad[a] += 1.0
        theta[s] += alpha * G * grad

Minimal Dyna-Q example

Q = np.zeros((env.n_states, env.n_actions))
model = {}
seen_pairs = []

for episode in range(500):
    s = env.reset()
    done = False

    while not done:
        a = epsilon_greedy(Q, s, epsilon)
        sp, r, done, info = env.step(a)

        q_learning_update(Q, s, a, r, sp, done, alpha, gamma)

        if (s, a) not in model:
            seen_pairs.append((s, a))

        model[(s, a)] = (r, sp, done)

        for _ in range(planning_steps):
            sim_s, sim_a = random.choice(seen_pairs)
            sim_r, sim_sp, sim_done = model[(sim_s, sim_a)]
            q_learning_update(Q, sim_s, sim_a, sim_r, sim_sp, sim_done, alpha, gamma)

        s = sp

Minimal offline fitted Q iteration example

Q = np.zeros((env.n_states, env.n_actions))
transitions = data[["state", "action", "reward", "next_state", "done"]].to_numpy()

for sweep in range(200):
    np.random.shuffle(transitions)

    for s, a, r, sp, done in transitions:
        s = int(s)
        a = int(a)
        sp = int(sp)
        done = bool(done)

        target = float(r)
        if not done:
            target += gamma * np.max(Q[sp])

        Q[s, a] += alpha * (target - Q[s, a])

Final project coding checklist

Before submitting a project notebook:

• The notebook runs from top to bottom.
• All imports are at the top.
• Random seeds are set.
• Environment parameters are clearly shown.
• Algorithms are implemented as functions.
• Training and evaluation are separated.
• Evaluation metrics are summarized in a table.
• Learning curves are plotted.
• Policies or values are visualized when possible.
• Main conclusions are written in markdown cells.
• Hyperparameters are reported.
• No unused large outputs clutter the notebook.

Final advice

When debugging reinforcement-learning code, always check the loop:

state -> action -> reward -> next state -> update -> next state becomes current state

Most bugs come from one of the following:

• the state is not updated,
• terminal states are mishandled,
• the wrong next value is used,
• exploration is accidentally turned off,
• arrays have the wrong shape,
• rewards do not match the intended environment,
• evaluation still uses exploratory actions.

When in doubt, print one short trajectory and inspect it carefully.