# On Google Colab, install the scikit-learn version the course was built with.
# On your own computer, install requirements.txt instead; this cell then does nothing.
import subprocess, sys
if "google.colab" in sys.modules:
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "scikit-learn==1.9.0"], check=True)

Lab 6c. Q-learning on a small network

This lab trains a tabular Q-learning agent on a small simulated network modeled as a Markov decision process (MDP). The agent starts on one host and must reach a target host while avoiding a monitored host that triggers a detection penalty. Simulators of this kind are used in network-defense research to study planning agents (Lesson 6.3). Because the network is small, the optimal policy can also be computed exactly by value iteration, so you can check what the agent learned rather than assume it. The lab runs in under a minute.

You will:

1. define the network as an MDP with stochastic action outcomes;
2. compute the optimal policy and the optimal route for reference;
3. train Q-learning with epsilon-greedy exploration and random tie-breaking, over several seeds;
4. compare the learned policy with the optimal one;
5. vary epsilon;
6. reproduce a common bug, an argmax that always returns the first index, and see what it does to an agent whose action values start at zero.

The environment is synthetic. It has eight hosts, a fixed known topology, and fixed transition probabilities. It is meant to show the mechanics of reinforcement learning and the tie-breaking pitfall, not to model a real network.

import os

import numpy as np
import matplotlib.pyplot as plt

1. The network as an MDP

The agent occupies one host at a time, so the state is the current host index. There are eight actions, one per host, read as "attempt the transition to host j". If host j is a neighbor of the current host, the transition succeeds with a probability that depends on the destination host j (P_SUCCESS[j]); on failure the agent stays where it is. Selecting a non-neighbor, or the current host, wastes a step and leaves the state unchanged.

Host 4 is a monitored host. Entering it ends the episode with probability 0.9 and a reward of -1 (the agent is caught and removed). Reaching host 7, the target, gives a reward of +1 and ends the episode. Every other step gives reward 0. The discount factor gamma = 0.95 makes shorter paths worth more, and episodes are cut off after 50 steps.

NAMES = ["web server", "workstation A", "mail server", "file server",
         "monitored host", "workstation B", "relay host", "target host"]
EDGES = [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (3, 6), (5, 6), (4, 7), (6, 7)]
P_SUCCESS = np.array([1.0, 0.8, 0.7, 0.6, 0.9, 0.8, 0.5, 0.7])  # success probability of a transition into host j
N = 8
START, TARGET, MONITORED = 0, 7, 4
P_DETECT = 0.9
GAMMA = 0.95
MAX_STEPS = 50

ADJ = np.zeros((N, N), dtype=bool)
for a, b in EDGES:
    ADJ[a, b] = ADJ[b, a] = True


def transition_model(s, a):
    """Return a list of (probability, next_state, reward, done) for taking action a in state s."""
    if s in (TARGET,):
        return [(1.0, s, 0.0, True)]
    if not ADJ[s, a]:
        return [(1.0, s, 0.0, False)]  # not a neighbor: wasted step
    p = P_SUCCESS[a]
    if a == MONITORED:
        # On a successful move into the monitored host, detection ends the episode with a penalty.
        success = [(p * P_DETECT, a, -1.0, True),
                   (p * (1 - P_DETECT), a, 0.0, False)]
        return success + [(1 - p, s, 0.0, False)]
    if a == TARGET:
        return [(p, a, 1.0, True), (1 - p, s, 0.0, False)]
    return [(p, a, 0.0, False), (1 - p, s, 0.0, False)]


# Precompute the transition table so sampling avoids rebuilding lists in the training loop.
_TABLE = {}
for _s in range(N):
    for _a in range(N):
        outs = transition_model(_s, _a)
        cum = np.cumsum([o[0] for o in outs])
        _TABLE[(_s, _a)] = (cum, [(o[1], o[2], o[3]) for o in outs])


def step(rng, s, a):
    """Sample one transition using the precomputed table."""
    cum, outs = _TABLE[(s, a)]
    u = rng.random()
    i = int(np.searchsorted(cum, u))
    if i >= len(outs):
        i = len(outs) - 1
    return outs[i]

2. The optimal policy by value iteration

Because the model is known and small, we can solve for the optimal action values exactly. Value iteration applies the Bellman optimality update until the values stop changing. This gives a reference policy and the true value of the start state, which the learning agent should approach.

def value_iteration(tol=1e-10):
    Q = np.zeros((N, N))
    while True:
        Q_new = np.zeros((N, N))
        for s in range(N):
            for a in range(N):
                q = 0.0
                for p, s2, r, done in transition_model(s, a):
                    q += p * (r + (0.0 if done else GAMMA * Q[s2].max()))
                Q_new[s, a] = q
        if np.abs(Q_new - Q).max() < tol:
            return Q_new
        Q = Q_new


Q_star = value_iteration()
V_start = Q_star[START].max()
optimal_action = int(np.argmax(Q_star[START]))
print(f"optimal value of the start state: {V_start:.4f}")
print(f"first optimal action from {NAMES[START]}: go to {NAMES[optimal_action]}")

optimal value of the start state: 0.7730
first optimal action from web server: go to mail server

The optimal path avoids the monitored host. We read it off by following the greedy action of the exact solution from the start host until the target is reached.

def greedy_path(Q):
    s, path, seen = START, [START], set()
    while s != TARGET and s not in seen:
        seen.add(s)
        a = int(np.argmax(Q[s]))
        if not ADJ[s, a] or a == s:
            break
        s = a
        path.append(s)
    return path


opt_path = greedy_path(Q_star)
print("optimal route:", " -> ".join(NAMES[i] for i in opt_path))

optimal route: web server -> mail server -> workstation B -> relay host -> target host

3. Q-learning with random tie-breaking

The agent does not know the model. It learns action values from experience with the update

Q(s, a) <- Q(s, a) + alpha * (r + gamma * max_a' Q(s', a') - Q(s, a)),

where alpha = 0.5 is the learning rate, r is the reward and s' is the next state. When the transition ends the episode (the agent is detected or reaches the target), the term gamma * max_a' Q(s', a') is replaced by 0, as in the code below. Actions are chosen epsilon-greedily. Early on, many action values are equal (all zero), so the greedy choice is a tie. We break ties uniformly at random. The helper below returns a random maximizer rather than the first one.

def argmax_random(rng, row):
    """Index of a maximal entry, ties broken uniformly at random."""
    m = row.max()
    candidates = np.flatnonzero(row == m)
    return int(candidates[rng.integers(len(candidates))])


def q_learning(seed, episodes=3000, alpha=0.5, epsilon=0.1, tie="random"):
    rng = np.random.default_rng(seed)
    Q = np.zeros((N, N))
    returns = np.zeros(episodes)
    for ep in range(episodes):
        s = START
        total, disc = 0.0, 1.0
        for _ in range(MAX_STEPS):
            if rng.random() < epsilon:
                a = int(rng.integers(N))
            elif tie == "random":
                a = argmax_random(rng, Q[s])
            else:
                a = int(np.argmax(Q[s]))  # always the first maximal index
            s2, r, done = step(rng, s, a)
            target = r + (0.0 if done else GAMMA * Q[s2].max())
            Q[s, a] += alpha * (target - Q[s, a])
            total += disc * r
            disc *= GAMMA
            s = s2
            if done:
                break
        returns[ep] = total
    return Q, returns


SEEDS = [0, 1, 2, 3, 4]
runs = [q_learning(s) for s in SEEDS]
Q_seeds = [q for q, _ in runs]
curves = np.array([ret for _, ret in runs])
learned_value = float(np.mean([q[START].max() for q in Q_seeds]))
seeds_above = int(sum(q[START].max() > V_start for q in Q_seeds))
print(f"learned value of the start state (mean of {len(SEEDS)} seeds): {learned_value:.4f}")
print(f"exact optimal value: {V_start:.4f}")
print(f"seeds whose learned start value exceeds the exact optimum: {seeds_above}/{len(SEEDS)}")

learned value of the start state (mean of 5 seeds): 0.8177
exact optimal value: 0.7730
seeds whose learned start value exceeds the exact optimum: 5/5

The learned start value exceeds the exact optimum. The learning rate stays at 0.5, so each action value keeps moving toward the most recent sampled targets instead of settling at its expectation, and the start value is read out as the largest of eight such noisy estimates whose true values are close together. A maximum over noisy estimates is biased upward, which Sutton and Barto call maximization bias. The per-action values, averaged over seeds, are exported with the results so you can compare them with the exact values.

4. What did it learn?

For each seed we follow the learned greedy action from the start host and compare the resulting route with the optimal one. We report the most common learned route and the number of seeds whose route equals the optimal route.

from collections import Counter
seed_paths = [tuple(greedy_path(q)) for q in Q_seeds]
learned_path = list(Counter(seed_paths).most_common(1)[0][0])  # the most common route across seeds
print("most common learned route:", " -> ".join(NAMES[i] for i in learned_path))
print("optimal route:            ", " -> ".join(NAMES[i] for i in opt_path))
matches = sum(1 for p in seed_paths if list(p) == opt_path)
print(f"seeds whose learned route equals the optimal route: {matches}/{len(SEEDS)}")

# The smoothed learning curve, averaged over seeds.
def smooth(x, k=100):
    c = np.cumsum(np.insert(x, 0, 0))
    return (c[k:] - c[:-k]) / k

mean_curve = curves.mean(axis=0)
smoothed = smooth(mean_curve, 100)
print(f"mean return over the last 200 episodes: {mean_curve[-200:].mean():.4f}")

most common learned route: web server -> mail server -> workstation B -> relay host -> target host
optimal route:             web server -> mail server -> workstation B -> relay host -> target host
seeds whose learned route equals the optimal route: 3/5
mean return over the last 200 episodes: 0.7167

5. The effect of epsilon

With epsilon = 0 the only exploration comes from breaking ties at random while action values are still equal, which is enough on this small MDP. With epsilon too large the agent explores so much that it rarely follows the good path and its average return stays low even though its values may be correct. We sweep a few values, using the first three seeds to keep the run short, and report the mean return over the final 200 episodes. Because of the smaller seed set, the value at epsilon = 0.1 differs slightly from the five-seed mean printed in Section 4.

eps_values = [0.0, 0.05, 0.1, 0.3, 0.6]
eps_seeds = SEEDS[:3]
eps_final = {}
for eps in eps_values:
    if eps == 0.1:
        rets = [ret[-200:].mean() for _, ret in runs[:len(eps_seeds)]]  # already trained at epsilon=0.1
    else:
        rets = [q_learning(s, epsilon=eps)[1][-200:].mean() for s in eps_seeds]
    eps_final[eps] = float(np.mean(rets))
    print(f"epsilon={eps:<4} mean return over last 200 episodes: {eps_final[eps]:.4f}")

epsilon=0.0  mean return over last 200 episodes: 0.7750

epsilon=0.05 mean return over last 200 episodes: 0.7310
epsilon=0.1  mean return over last 200 episodes: 0.7194

epsilon=0.3  mean return over last 200 episodes: 0.5901

epsilon=0.6  mean return over last 200 episodes: 0.2854

6. The tie-breaking pitfall

Now we repeat the training with tie="first", so the greedy branch uses np.argmax, which always returns the lowest index among tied actions. At the start every action value is zero, so the greedy action is always action 0. From the start host, action 0 targets the web server itself and wastes a step; from hosts 1 and 2 it moves the agent back to the web server. Only the epsilon branch ever tries anything else. With a small epsilon and a reward that appears only after at least three moves (four on either safe route), this agent rarely reaches the target. We count how often it does, and compare its learned start value and route with the random-tie-breaking agent.

first_runs = [q_learning(s, tie="first") for s in SEEDS]
Q_first_seeds = [q for q, _ in first_runs]
Q_first = np.mean(Q_first_seeds, axis=0)
first_curve = np.array([ret for _, ret in first_runs]).mean(axis=0)
first_value = float(Q_first[START].max())
first_path = greedy_path(Q_first)
first_matches = sum(1 for q in Q_first_seeds if greedy_path(q) == opt_path)
first_successes = int(sum(int((ret > 0).sum()) for _, ret in first_runs))
first_episodes = int(sum(len(ret) for _, ret in first_runs))
print(f"first-index episodes that reached the target: {first_successes} of {first_episodes}")
print(f"first-index tie-breaking, learned start value: {first_value:.4f} (optimal {V_start:.4f})")
print("first-index route:", " -> ".join(NAMES[i] for i in first_path))
print(f"first-index seeds matching optimal route: {first_matches}/{len(SEEDS)}")
print(f"first-index mean return over last 200 episodes: {first_curve[-200:].mean():.4f}")
print(f"random-tie   mean return over last 200 episodes: {mean_curve[-200:].mean():.4f}")

first-index episodes that reached the target: 5 of 15000
first-index tie-breaking, learned start value: 0.0000 (optimal 0.7730)
first-index route: web server
first-index seeds matching optimal route: 0/5
first-index mean return over last 200 episodes: -0.0020
random-tie   mean return over last 200 episodes: 0.7167

The cell above shows the failure described in Lesson 6.3. With deterministic first-index tie-breaking and zero-initialized values, the greedy action out of a tie is fixed, so the agent relies on epsilon alone to discover the path. It reaches the target in only a few of its episodes, too rarely for the reward to propagate back to the start host, so its start value stays at zero. Random tie-breaking removes the bias. Initializing the values optimistically would also work, because then every untried action looks attractive and gets tried.

Try this

• Raise the detection penalty for the monitored host from -1 to -5 and see whether the optimal route changes.
• Set the reward for reaching the target to +10 and predict, before running, whether the random-tie agent learns the route in fewer episodes or whether its learning curve is only rescaled (divide the new curve by 10 and compare it with the original).
• Lower alpha to 0.1 and predict, before running, whether the learning curve rises faster or slower.

Results summary

These are the numbers quoted in the lesson text.

results = {
    "n_hosts": N,
    "gamma": GAMMA,
    "optimal_value_start": float(V_start),
    "learned_value_start": learned_value,
    "seeds_value_above_optimal": seeds_above,
    "q_star_via_host1": float(Q_star[START, 1]),
    "optimal_route_len": len(opt_path),
    "n_seeds": len(SEEDS),
    "seeds_matching_optimal": int(matches),
    "final_return_random_tie": float(mean_curve[-200:].mean()),
    "eps0_final_return": eps_final[0.0],
    "eps_best": float(min(eps_values, key=lambda e: -eps_final[e])),
    "eps_best_return": float(max(eps_final.values())),
    "eps06_final_return": eps_final[0.6],
    "first_tie_value_start": first_value,
    "first_tie_final_return": float(first_curve[-200:].mean()),
    "first_tie_seeds_matching": int(first_matches),
    "first_tie_success_episodes": first_successes,
    "first_tie_total_episodes": first_episodes,
    "n_eps_seeds": len(eps_seeds),
    "p_detect": P_DETECT,
}
for k, v in results.items():
    print(f"{k:<26} {v}")

export = {
    "names": NAMES,
    "edges": EDGES,
    "monitored": MONITORED,
    "start": START,
    "target": TARGET,
    "p_success": P_SUCCESS.tolist(),
    "opt_path": opt_path,
    "learned_path": learned_path,
    "smoothed_random": smoothed.tolist(),
    "smoothed_first": smooth(first_curve, 100).tolist(),
    "eps_values": eps_values,
    "eps_final": [eps_final[e] for e in eps_values],
    "q_star_start": Q_star[START].tolist(),
    "q_learned_start": np.mean([q[START] for q in Q_seeds], axis=0).tolist(),
}

if os.environ.get("COURSE_RESULTS"):
    import json
    with open(os.environ["COURSE_RESULTS"], "w") as f:
        json.dump({"results": results, "data": export}, f)

n_hosts                    8
gamma                      0.95
optimal_value_start        0.7729839586532148
learned_value_start        0.81772302715159
seeds_value_above_optimal  5
q_star_via_host1           0.7657706290702986
optimal_route_len          5
n_seeds                    5
seeds_matching_optimal     3
final_return_random_tie    0.7166791053117731
eps0_final_return          0.7749838014256806
eps_best                   0.0
eps_best_return            0.7749838014256806
eps06_final_return         0.28537201996609246
first_tie_value_start      0.0
first_tie_final_return     -0.001986307959549617
first_tie_seeds_matching   0
first_tie_success_episodes 5
first_tie_total_episodes   15000
n_eps_seeds                3
p_detect                   0.9