NiiL

DQN · Study

English translation

Execution conditions

Run the baseline Queue example with NumPy. Compare policies on seed=100…104, which was not used for training.

PyTorch is required. This is a short synthetic experiment; its configuration does not guarantee learning performance.

Public environment activities use gymnasium and stable-baselines3. gymnasium[classic-control] is required only for visualization.

If the required package is missing, run %pip install control==0.10.2 gymnasium==1.3.0 stable-baselines3==2.9.2 in a Colab or Jupyter code cell.

Topics covered

As the state space grows, verify whether replacing a Q-table with a neural network actually improves the policy.

Function approximation

DQN replaces the Q-table with a neural network. Replay reuses stored experiences, and the target network separates updates to the parameters used to form targets.

Read a public implementation

Map the buffer, learning start, target updates, input normalization, and exploration schedule as well as the policy name. Returns from different environments do not share the same performance unit.

Worked example

Compare state representations in a Q-table and a neural network on the same Queue problem, and evaluate the learned policy against rule-based policies.

"""Synthetic service queue. An action chooses capacity, including its energy cost."""
import numpy as np

class Queue:
    def reset(self, seed=0):
        self.rng = np.random.default_rng(seed)
        self.q, self.t = 0, 0
        return self.q

    def step(self, action):
        capacity = (0, 2, 4)[int(action)]
        arriving = int(self.rng.poisson(2))
        raw = max(0, self.q - capacity) + arriving
        dropped = max(0, raw - 20)
        self.q = min(20, raw)
        self.t += 1
        reward = -(self.q + .2 * capacity**2 + 10 * dropped)
        return self.q, reward, False, self.t >= 100, {'queue': self.q, 'capacity': capacity, 'dropped': dropped}

def evaluate(policy, seeds=range(100, 105)):
    rows = []
    for seed in seeds:
        env = Queue(); q = env.reset(seed)
        total = waiting = energy = dropped = 0
        while True:
            q, r, terminated, truncated, info = env.step(policy(q))
            total += r; waiting += q; energy += info['capacity']**2; dropped += info['dropped']
            if terminated or truncated: break
        rows.append([total, waiting / env.t, energy / env.t, dropped])
    return np.array(rows)

policies = {'constant': lambda q: 1, 'threshold': lambda q: 0 if q == 0 else (1 if q < 5 else 2)}
for name, policy in policies.items():
    print(name, evaluate(policy).mean(axis=0))

# Requires the Queue and evaluate definitions in the preceding cell.
import random
from collections import deque
import torch
from torch import nn

def train_dqn(seed=0, variant='dqn', steps=1800):
    torch.manual_seed(seed); rng = random.Random(seed)
    class Network(nn.Module):
        def __init__(self):
            super().__init__()
            self.shared = nn.Sequential(nn.Linear(1, 32), nn.ReLU())
            self.values = nn.Linear(32, 1 if variant == 'dueling' else 3)
            self.advantage = nn.Linear(32, 3)
        def forward(self, x):
            h = self.shared(x); v = self.values(h)
            if variant != 'dueling': return v
            a = self.advantage(h)
            return v + a - a.mean(dim=-1, keepdim=True)
    online = Network(); target = Network(); target.load_state_dict(online.state_dict())
    optimizer = torch.optim.Adam(online.parameters(), lr=.001)
    buffer = deque(maxlen=3000); priorities = deque(maxlen=3000)
    env = Queue(); q = env.reset(seed); losses = []; sample_counts = [0, 0]
    def tensor(states): return torch.tensor(states, dtype=torch.float32).reshape(-1, 1) / 20
    for step in range(steps):
        epsilon = max(.1, 1 - step / 1000)
        action = rng.randrange(3) if rng.random() < epsilon else int(online(tensor([q])).argmax())
        nq, reward, terminated, truncated, _ = env.step(action)
        buffer.append((q, action, reward / 20, nq, terminated))
        priorities.append(max(priorities, default=1.0))
        q = env.reset(seed + step + 1) if terminated or truncated else nq
        if len(buffer) < 64: continue
        probabilities = np.array(priorities, dtype=float)**.6
        probabilities /= probabilities.sum()
        if variant == 'per':
            # Sampling with replacement; importance weighting uses the same probabilities.
            ids = rng.choices(range(len(buffer)), weights=probabilities, k=32)
            weights = (len(buffer) * probabilities[ids])**(-.4)
            weights /= weights.max()
        else:
            ids = rng.sample(range(len(buffer)), 32); weights = np.ones(32)
        bq, ba, br, bn, bt = zip(*(buffer[i] for i in ids))
        predicted = online(tensor(bq)).gather(1, torch.tensor(ba).reshape(-1, 1)).flatten()
        with torch.no_grad():
            if variant == 'double':
                selected = online(tensor(bn)).argmax(dim=1, keepdim=True)
                next_value = target(tensor(bn)).gather(1, selected).flatten()
            else: next_value = target(tensor(bn)).max(dim=1).values
            expected = torch.tensor(br, dtype=torch.float32) + .95 * (~torch.tensor(bt)) * next_value
        td = expected - predicted
        loss = (torch.tensor(weights, dtype=torch.float32) * td.square()).mean()
        optimizer.zero_grad(); loss.backward(); nn.utils.clip_grad_norm_(online.parameters(), 5); optimizer.step()
        for i, delta in zip(ids, td.detach().numpy()): priorities[i] = abs(float(delta)) + 1e-5
        losses.append(float(loss.detach()))
        if step % 100 == 0: target.load_state_dict(online.state_dict())
    def policy(q):
        with torch.no_grad(): return int(online(tensor([q])).argmax())
    return policy, losses

policy,losses=train_dqn(0,variant='dqn')
print(evaluate(policy).mean(0))

constant [-921.2      8.052    4.       3.6  ]
threshold [-387.4     2.77    5.52    0.  ]
[-437.32     2.114   11.296    0.   ]

Official documentation and data

Reference example

View more examples