Learn AI Series (#116) - Mini Project: Training a Game-Playing AI

avatar

Learn AI Series (#116) - Mini Project: Training a Game-Playing AI

variant-a-12-green.png

What will I learn

  • You will learn how to wire an entire RL training pipeline together, from raw environment to a trained agent that actually plays;
  • PPO with vectorized environments -- the boring, reliable workhorse that solves most control problems without drama;
  • GAE for advantage estimation, orthogonal initialization, and the small implementation details that quietly decide whether RL works or silently fails;
  • how to read the diagnostic signals (entropy, clip fraction, value loss) and turn "it's not learning" into a specific, fixable cause;
  • why you must compare against baselines -- random and a hand-crafted heuristic -- before believing your agent learned anything real;
  • and how every piece we've built across this whole reinforcement-learning arc snaps into one working system.

Requirements

  • A working modern computer running macOS, Windows or Ubuntu;
  • An installed Python 3(.11+) distribution with NumPy, PyTorch and Gymnasium;
  • You've followed the RL arc -- especially episodes #107 (DQN), #108-109 (policy gradients and PPO) and #115 (offline RL) -- because today we stop collecting techniques and start building with them.

Difficulty

  • Beginner

Curriculum (of the Learn AI Series):

Learn AI Series (#116) - Mini Project: Training a Game-Playing AI

Solutions to Episode #115 Exercises

Before we build our own agent from the ground up, let's clear last episode's three exercises. They lean on the NaiveOfflineDQN, CQLAgent and DecisionTransformer classes from episode #115, so I'm assuming those are imported and in scope.

Exercise 1: the "naive fails, conservative saves it" experiment, end to end. Collect an offline dataset from a deliberately mediocre behaviour policy, train both agents on that same frozen dataset, then watch the naive agent's Q-values sail off into fantasy while CQL keeps its feet on the ground.

The whole lesson lives in the two printed max-Q numbers -- one climbs toward the moon, the other stays sane:

import gymnasium as gym
import numpy as np
import torch
# Assumes NaiveOfflineDQN and CQLAgent from episode #115.

def behaviour_policy(obs):
    """70% random, 30% a crude 'push toward upright' heuristic. Deliberately mediocre."""
    if np.random.rand() < 0.7:
        return np.random.randint(2)
    angle, ang_vel = obs[2], obs[3]
    return 1 if (angle + 0.5 * ang_vel) > 0 else 0


def collect_offline(n_episodes=400):
    env = gym.make("CartPole-v1")
    S, A, R, S2, D = [], [], [], [], []
    for _ in range(n_episodes):
        obs, _ = env.reset()
        done = False
        while not done:
            a = behaviour_policy(obs)
            nobs, r, term, trunc, _ = env.step(a)
            done = term or trunc
            S.append(obs); A.append(a); R.append(r)
            S2.append(nobs); D.append(float(done))
            obs = nobs
    return (torch.FloatTensor(np.array(S)), torch.LongTensor(A),
            torch.FloatTensor(R), torch.FloatTensor(np.array(S2)),
            torch.FloatTensor(D))


def batches(data, batch_size=128, n_batches=2000):
    S = data[0]
    for _ in range(n_batches):
        idx = torch.randint(0, len(S), (batch_size,))
        yield tuple(t[idx] for t in data)


def rollout_return(policy_fn, n=30):
    env = gym.make("CartPole-v1")
    totals = []
    for _ in range(n):
        obs, _ = env.reset()
        done, ep = False, 0.0
        while not done:
            obs, r, term, trunc, _ = env.step(policy_fn(obs))
            ep += r; done = term or trunc
        totals.append(ep)
    return np.mean(totals)


data = collect_offline()
naive = NaiveOfflineDQN(state_dim=4, n_actions=2)
cql = CQLAgent(state_dim=4, n_actions=2, cql_alpha=1.0)

for i, b in enumerate(batches(data)):
    naive.train_step(b)
    cql.train_step(b)
    if i % 500 == 0:
        with torch.no_grad():
            qn = naive.q_net(data[0]).max().item()
            qc = cql.q_net(data[0]).max().item()
        print(f"step {i:4d}: max-Q naive={qn:9.1f}   cql={qc:9.1f}")

greedy = lambda net: (lambda o: net(torch.FloatTensor(o)).argmax().item())
print("naive rollout return:", round(rollout_return(greedy(naive.q_net)), 1))
print("cql   rollout return:", round(rollout_return(greedy(cql.q_net)), 1))

Watch that max-Q naive column and it tells the whole story of episode #115 in one running number: it inflates round after round, because the max keeps latching onto out-of-distribution actions the network never learned and hallucinated a high value for. CQL's column stays grounded, the pessimism penalty pressing exactly those delusions back down. And the rollout returns follow suit -- the naive agent, happily maximising its own fantasy, plays badly, while CQL sticks to behaviour it has receipts for and comes out ahead. Same frozen data, same network, one penalty term between "learns something" and "believes nonsense".

Exercise 2: turn the cql_alpha knob and find the U-shape. Reuse the helpers from Exercise 1, train a fresh CQL agent per alpha on the same dataset, and plot final return against the conservatism dial:

import numpy as np
import torch
# Reuses collect_offline, batches, rollout_return from Exercise 1.

data = collect_offline()          # ONE frozen dataset, shared across the whole sweep
results = {}
for alpha in [0.0, 0.1, 1.0, 5.0, 20.0]:
    agent = CQLAgent(state_dim=4, n_actions=2, cql_alpha=alpha)
    for b in batches(data, n_batches=2000):
        agent.train_step(b)
    ret = rollout_return(
        lambda o: agent.q_net(torch.FloatTensor(o)).argmax().item())
    results[alpha] = ret
    print(f"cql_alpha={alpha:5.1f} -> return {ret:6.1f}")

best = max(results, key=results.get)
print(f"sweet spot around cql_alpha={best}")

Note that cql_alpha = 0.0 is the naive agent again -- no penalty, pure overestimation -- so the left end of the curve should be grim. Crank alpha too high and the policy turns so timid it just clings to the mediocre behaviour policy and never improves, so the right end sags too. Somewhere in the middle sits the sweet spot, and finding it by hand is a good, humbling taste of why offline RL is fiddly in practice ;-) Nota bene: because CartPole is so forgiving, don't expect a razor-sharp valley -- the shape is the point, not the exact numbers.

Exercise 3: give Decision Transformer the "return dial" test. Train it on a dataset that mixes genuinely good and genuinely bad trajectories, then at inference ask for different returns and measure what you actually get.

The trick is building a dataset with real spread in it -- I do that by sampling from behaviour policies of increasing competence, so the returns range from feeble to decent:

import gymnasium as gym
import numpy as np
import torch
# Assumes DecisionTransformer, train_decision_transformer, generate_action from episode #115.

def mixed_dataset(skills=(0.0, 0.3, 0.6, 0.9), episodes_per=60):
    """Collect trajectories from policies of rising skill (P of using the heuristic)."""
    env = gym.make("CartPole-v1")
    trajectories = []
    for skill in skills:
        for _ in range(episodes_per):
            obs, _ = env.reset()
            states, actions, rewards, done = [], [], [], False
            while not done:
                if np.random.rand() < skill:
                    a = 1 if (obs[2] + 0.5 * obs[3]) > 0 else 0
                else:
                    a = np.random.randint(2)
                states.append(obs)
                actions.append([float(a)])            # 1-dim continuous encoding
                obs, r, term, trunc, _ = env.step(a)
                rewards.append(r); done = term or trunc
            trajectories.append({
                "states": np.array(states, dtype=np.float32),
                "actions": np.array(actions, dtype=np.float32),
                "rewards": np.array(rewards, dtype=np.float32)})
    return {"trajectories": trajectories}


dataset = mixed_dataset()
data_returns = [t["rewards"].sum() for t in dataset["trajectories"]]
print(f"dataset returns: min={min(data_returns):.0f}  max={max(data_returns):.0f}")

dt = DecisionTransformer(state_dim=4, action_dim=1)
train_decision_transformer(dt, dataset, n_epochs=40, context_len=20)


def achieved_return(model, target, n=20):
    env = gym.make("CartPole-v1")
    totals = []
    for _ in range(n):
        obs, _ = env.reset()
        states, actions, times, done, ep = [obs], [[0.0]], [0], False, 0.0
        while not done and len(states) < 500:
            s = torch.FloatTensor(np.array(states)).unsqueeze(0)
            a = torch.FloatTensor(np.array(actions)).unsqueeze(0)
            t = torch.LongTensor(np.array(times)).unsqueeze(0)
            pred = generate_action(model, target, s, a, t)
            act = int(pred.item() > 0.5)              # threshold the 1-dim action
            obs, r, term, trunc, _ = env.step(act)
            ep += r; done = term or trunc
            states.append(obs); actions.append([float(act)]); times.append(len(times))
        totals.append(ep)
    return np.mean(totals)


for target in [50, 100, 200, 500]:                    # 500 is BEYOND anything in the data
    print(f"asked for {target:4d} -> got {achieved_return(dt, target):6.1f}")

If it's working, the achieved return should track the requested one -- ask for more, get more -- right up until you hit the ceiling of what the dataset actually contains. That last target of 500 is deliberately past anything the behaviour policies ever managed, and it's the honest part of the experiment: the model cannot conjure competence that isn't somewhere in its training data, so the achieved number tops out in stead of climbing forever. A model that conditions on returns is only ever as good as the best trajectories it was shown -- a lesson worth carrying into every "just tell it what you want" system you'll ever build.

Right -- fourteen episodes of RL theory, algorithms and individual components. Time to wire the whole thing together.

For this mini project we build a complete training pipeline: environment setup, a PPO agent with vectorized environments for speed, proper logging and visualization, hyperparameter tuning, and honest evaluation against baselines. Not a toy example you run once and forget -- a system you could genuinely adapt to a new environment and iterate on.

The environments? CartPole and LunarLander from Gymnasium. Simple enough to train on a laptop, complex enough to expose every mistake in your pipeline. If your PPO can't solve CartPole in under a minute and LunarLander in under ten, something is wrong -- and (this is the useful part) we'll talk about exactly how to diagnose what.

The environment wrapper

First, we need vectorized environments. Running one environment per step is painfully slow -- you spend most of your wall-clock waiting on Python overhead in stead of learning. So we run N environments in parallel, collecting N transitions per step:

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
from collections import deque
import time

class VectorizedEnv:
    """Run multiple environments in parallel for efficient data collection."""
    def __init__(self, env_name, n_envs=8):
        self.envs = [gym.make(env_name) for _ in range(n_envs)]
        self.n_envs = n_envs

    def reset(self):
        states = []
        for env in self.envs:
            state, _ = env.reset()
            states.append(state)
        return np.array(states)

    def step(self, actions):
        results = []
        for env, action in zip(self.envs, actions):
            state, reward, terminated, truncated, info = env.step(action)
            done = terminated or truncated
            if done:
                state, _ = env.reset()
            results.append((state, reward, done))
        states = np.array([r[0] for r in results])
        rewards = np.array([r[1] for r in results])
        dones = np.array([r[2] for r in results])
        return states, rewards, dones

Simple and synchronous. Gymnasium's AsyncVectorEnv uses multiprocessing for true parallelism, but this version is far easier to debug and quite fast enough for our environments. (When you graduate to something heavier, swap it in -- the interface is deliberately the same.)

The PPO network

We use the shared actor-critic architecture from episode #109. One backbone, two heads -- a policy head that says what to do, a value head that says how good the situation is:

class PPONetwork(nn.Module):
    def __init__(self, state_dim, n_actions, hidden=64):
        super().__init__()
        self.shared = nn.Sequential(
            nn.Linear(state_dim, hidden), nn.Tanh(),
            nn.Linear(hidden, hidden), nn.Tanh(),
        )
        self.policy_head = nn.Linear(hidden, n_actions)
        self.value_head = nn.Linear(hidden, 1)

        # Orthogonal initialization (matters more for RL than you'd think)
        for layer in self.shared:
            if isinstance(layer, nn.Linear):
                nn.init.orthogonal_(layer.weight, gain=np.sqrt(2))
                nn.init.zeros_(layer.bias)
        nn.init.orthogonal_(self.policy_head.weight, gain=0.01)
        nn.init.orthogonal_(self.value_head.weight, gain=1.0)

    def forward(self, x):
        features = self.shared(x)
        logits = self.policy_head(features)
        value = self.value_head(features).squeeze(-1)
        return logits, value

    def get_action(self, state):
        logits, value = self.forward(state)
        dist = torch.distributions.Categorical(logits=logits)
        action = dist.sample()
        return action, dist.log_prob(action), value

That orthogonal initialization with a specific gain per head is one of those details that quietly separates working RL from broken RL. The policy head uses a tiny gain (0.01) so the initial action probabilities come out close to uniform -- you want to explore broadly at the very start, not accidentically commit to some random preference the network happened to be born with. It looks like a throwaway line; it is not.

The rollout buffer with GAE

Now the part that collects data and turns raw rewards into learning signal:

class RolloutBuffer:
    def __init__(self):
        self.states = []
        self.actions = []
        self.log_probs = []
        self.rewards = []
        self.dones = []
        self.values = []

    def store(self, state, action, log_prob, reward, done, value):
        self.states.append(state)
        self.actions.append(action)
        self.log_probs.append(log_prob)
        self.rewards.append(reward)
        self.dones.append(done)
        self.values.append(value)

    def compute_gae(self, last_value, gamma=0.99, lam=0.95):
        """Generalized Advantage Estimation (episode #109)."""
        rewards = np.array(self.rewards)
        dones = np.array(self.dones)
        values = np.array(self.values + [last_value])

        advantages = np.zeros_like(rewards)
        gae = 0
        for t in reversed(range(len(rewards))):
            delta = rewards[t] + gamma * values[t + 1] * (1 - dones[t]) - values[t]
            gae = delta + gamma * lam * (1 - dones[t]) * gae
            advantages[t] = gae

        returns = advantages + values[:-1]
        return advantages, returns

    def get_batches(self, advantages, returns, batch_size=256):
        """Yield mini-batches for PPO updates."""
        states = torch.FloatTensor(np.array(self.states))
        actions = torch.LongTensor(np.array(self.actions))
        old_log_probs = torch.FloatTensor(np.array(self.log_probs))
        advantages = torch.FloatTensor(advantages)
        returns = torch.FloatTensor(returns)

        # Normalize advantages (a small trick with a big stabilizing effect)
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        indices = np.arange(len(states))
        np.random.shuffle(indices)
        for start in range(0, len(states), batch_size):
            idx = indices[start:start + batch_size]
            yield (states[idx], actions[idx], old_log_probs[idx],
                   advantages[idx], returns[idx])

    def clear(self):
        self.__init__()

GAE (Generalized Advantage Estimation) with lambda = 0.95 is the standard, and for good reason. Lower lambda gives you more bias but less variance; higher lambda creeps toward raw Monte Carlo returns (episode #105) with all their noise. For as good as every environment you'll meet early on, 0.95 just works without tuning -- one of the rare free lunches in RL. That advantage-normalization line, by the way, is not optional decoration: it keeps the gradient scale sane batch to batch, and yanking it out is a classic way to make training mysteriously unstable.

The PPO training loop

Here's where everything we've built across the arc comes together:

class PPOTrainer:
    def __init__(self, env_name, n_envs=8, lr=3e-4, gamma=0.99,
                 lam=0.95, clip_eps=0.2, entropy_coef=0.01,
                 value_coef=0.5, max_grad_norm=0.5, n_epochs=4,
                 steps_per_rollout=128, batch_size=256):
        self.env = VectorizedEnv(env_name, n_envs)
        self.n_envs = n_envs
        self.gamma = gamma
        self.lam = lam
        self.clip_eps = clip_eps
        self.entropy_coef = entropy_coef
        self.value_coef = value_coef
        self.max_grad_norm = max_grad_norm
        self.n_epochs = n_epochs
        self.steps_per_rollout = steps_per_rollout
        self.batch_size = batch_size

        sample_env = gym.make(env_name)
        state_dim = sample_env.observation_space.shape[0]
        n_actions = sample_env.action_space.n
        sample_env.close()

        self.network = PPONetwork(state_dim, n_actions)
        self.optimizer = torch.optim.Adam(self.network.parameters(), lr=lr, eps=1e-5)

        # Logging
        self.episode_rewards = deque(maxlen=100)
        self.log = {'rewards': [], 'policy_loss': [], 'value_loss': [],
                    'entropy': [], 'clip_fraction': []}

    def collect_rollout(self, states):
        buffer = RolloutBuffer()
        for _ in range(self.steps_per_rollout):
            state_tensor = torch.FloatTensor(states)
            with torch.no_grad():
                actions, log_probs, values = self.network.get_action(state_tensor)

            next_states, rewards, dones = self.env.step(actions.numpy())
            for i in range(self.n_envs):
                buffer.store(states[i], actions[i].item(), log_probs[i].item(),
                             rewards[i], dones[i], values[i].item())
            states = next_states

        # Bootstrap the last value so the final transition isn't treated as terminal
        with torch.no_grad():
            _, last_values = self.network(torch.FloatTensor(states))
        advantages, returns = buffer.compute_gae(
            last_values.mean().item(), self.gamma, self.lam
        )
        return buffer, advantages, returns, states

    def update(self, buffer, advantages, returns):
        total_policy_loss = 0
        total_value_loss = 0
        total_entropy = 0
        total_clip_frac = 0
        n_updates = 0

        for _ in range(self.n_epochs):
            for batch in buffer.get_batches(advantages, returns, self.batch_size):
                states, actions, old_log_probs, advs, rets = batch

                logits, values = self.network(states)
                dist = torch.distributions.Categorical(logits=logits)
                new_log_probs = dist.log_prob(actions)
                entropy = dist.entropy().mean()

                # Policy loss (the clipped surrogate objective, episode #109)
                ratio = torch.exp(new_log_probs - old_log_probs)
                surr1 = ratio * advs
                surr2 = torch.clamp(ratio, 1 - self.clip_eps,
                                    1 + self.clip_eps) * advs
                policy_loss = -torch.min(surr1, surr2).mean()

                # Value loss
                value_loss = nn.functional.mse_loss(values, rets)

                # Total loss: maximise reward, fit the value function, stay curious
                loss = (policy_loss + self.value_coef * value_loss
                        - self.entropy_coef * entropy)

                self.optimizer.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(self.network.parameters(),
                                         self.max_grad_norm)
                self.optimizer.step()

                # Track metrics (these are your dashboard, not decoration)
                clip_frac = ((ratio - 1).abs() > self.clip_eps).float().mean()
                total_policy_loss += policy_loss.item()
                total_value_loss += value_loss.item()
                total_entropy += entropy.item()
                total_clip_frac += clip_frac.item()
                n_updates += 1

        self.log['policy_loss'].append(total_policy_loss / n_updates)
        self.log['value_loss'].append(total_value_loss / n_updates)
        self.log['entropy'].append(total_entropy / n_updates)
        self.log['clip_fraction'].append(total_clip_frac / n_updates)

    def train(self, total_timesteps=200_000):
        states = self.env.reset()
        steps_done = 0
        iteration = 0

        while steps_done < total_timesteps:
            buffer, advantages, returns, states = self.collect_rollout(states)
            self.update(buffer, advantages, returns)

            steps_done += self.steps_per_rollout * self.n_envs
            iteration += 1

            # Evaluate every 10 iterations
            if iteration % 10 == 0:
                eval_reward = self.evaluate(n_episodes=10)
                self.log['rewards'].append(eval_reward)
                print(f"Step {steps_done:>7d} | "
                      f"Eval reward: {eval_reward:>7.1f} | "
                      f"Policy loss: {self.log['policy_loss'][-1]:.4f} | "
                      f"Entropy: {self.log['entropy'][-1]:.3f}")

    def evaluate(self, n_episodes=10):
        env = gym.make(self.env.envs[0].spec.id)
        rewards = []
        for _ in range(n_episodes):
            state, _ = env.reset()
            total_reward = 0
            done = False
            while not done:
                state_t = torch.FloatTensor(state).unsqueeze(0)
                with torch.no_grad():
                    logits, _ = self.network(state_t)
                action = logits.argmax(dim=1).item()  # greedy at eval
                state, reward, terminated, truncated, _ = env.step(action)
                total_reward += reward
                done = terminated or truncated
            rewards.append(total_reward)
        env.close()
        return np.mean(rewards)

A couple of things to notice, because they trip up nearly everyone the first time. We evaluate with greedy actions (argmax), not sampled ones. During training, exploration is the stochastic policy's job -- during evaluation, you want to see the agent's genuine best behaviour, not a noisy sample of it. And the clip fraction quietly tells you how hard the policy is lurching each update: if it sits consistently above 0.2, your learning rate is probably too hot. Having said that, don't stare at a single number -- watch the trends.

Running the training

# CartPole: should solve (reward >= 475) in ~50k steps
trainer = PPOTrainer('CartPole-v1', n_envs=8, lr=3e-4,
                     steps_per_rollout=128)
trainer.train(total_timesteps=100_000)

# LunarLander: should solve (reward >= 200) in ~500k steps
trainer = PPOTrainer('LunarLander-v3', n_envs=16, lr=2.5e-4,
                     steps_per_rollout=256, n_epochs=4,
                     entropy_coef=0.01)
trainer.train(total_timesteps=1_000_000)

Run the CartPole one first. It solves so fast that if it doesn't, you've got a bug in your pipeline rather than a hard problem -- which makes it the perfect smoke test before you burn ten minutes on LunarLander.

Comparing against baselines

Here's a trap I've watched quite some people fall into: an agent that beats random isn't impressive, because random is a spectacularly low bar. You need real baselines before you're allowed to feel proud:

def evaluate_baseline(env_name, policy_fn, n_episodes=100):
    env = gym.make(env_name)
    rewards = []
    for _ in range(n_episodes):
        state, _ = env.reset()
        total_reward = 0
        done = False
        while not done:
            action = policy_fn(state, env)
            state, reward, terminated, truncated, _ = env.step(action)
            total_reward += reward
            done = terminated or truncated
        rewards.append(total_reward)
    env.close()
    return np.mean(rewards), np.std(rewards)

# Random baseline
random_mean, random_std = evaluate_baseline(
    'LunarLander-v3',
    lambda s, e: e.action_space.sample()
)
print(f"Random:    {random_mean:.1f} +/- {random_std:.1f}")

# Heuristic baseline for LunarLander
def lunar_heuristic(state, env):
    """Simple heuristic: fire main engine when falling, side when tilted."""
    x, y, vx, vy, angle, angular_vel, left_leg, right_leg = state
    if left_leg or right_leg:
        return 0  # do nothing when landed
    if abs(angle) > 0.2:
        return 1 if angle > 0 else 3  # fire left/right
    if vy < -0.5:
        return 2  # fire main engine
    return 0

heuristic_mean, heuristic_std = evaluate_baseline(
    'LunarLander-v3', lunar_heuristic
)
print(f"Heuristic: {heuristic_mean:.1f} +/- {heuristic_std:.1f}")

# Trained PPO agent
ppo_mean = trainer.evaluate(n_episodes=100)
print(f"PPO:       {ppo_mean:.1f}")

Typical results: random scores around -200, the hand-rolled heuristic around 50-100, and a well-trained PPO agent around 200-280. The rule to burn into your head: if your agent can't beat the heuristic, your training has a bug -- because a dozen lines of hand-written if-statements just embarrassed a neural network, and that only happens when something's broken.

Diagnosing training failures

RL training fails silently more often than any other flavour of ML. There's no exception, no NaN, no red text -- just a reward curve that stubbornly refuses to climb. So here are the failure modes I reach for first, and the specific signal each one shows:

Entropy collapses to zero early. The policy commits to one action before it has explored enough, and now it's stuck. Fix: raise entropy_coef to 0.02-0.05, or drop the learning rate so it stops rushing to conclusions.

Reward flatlines. The agent isn't learning at all. Check the plumbing: are rewards actually flowing through? Is the advantage computation right? Print the raw rewards straight out of a rollout. A classic bug is resetting the environment but quietly dropping the terminal reward on the floor.

Reward spikes then crashes. The policy changed too aggressively and fell off a cliff it had just climbed. Fix: reduce the learning rate, widen clip_eps a touch, or cut n_epochs. Glance at the clip fraction -- if it's parked above 0.3, your updates are simply too large.

Value loss explodes. The value function can't keep pace with the reward signal. Fix: bump value_coef, lower the learning rate, or clip how much V is allowed to move per update.

Nota bene: the reason these are worth memorising is that in RL the symptom (flat reward) is almost never the cause. The diagnostic signals above are how you turn "it's not working" into an actual bug report.

Hyperparameter sensitivity

RL is notoriously twitchy about hyperparameters -- far more than the supervised learning we did earlier in the series. Here's what actually moves the needle, roughly in order of how much it matters:

# The hyperparameters that matter most, in order:
configs = {
    'learning_rate':       [1e-4, 2.5e-4, 3e-4, 5e-4],  # most critical
    'n_envs':              [4, 8, 16, 32],               # more = smoother gradients
    'steps_per_rollout':   [64, 128, 256, 512],          # longer = less bias
    'n_epochs':            [3, 4, 6, 10],                # more = sample-efficient but risky
    'clip_eps':            [0.1, 0.2, 0.3],              # 0.2 is the standard
    'entropy_coef':        [0.0, 0.01, 0.02],            # subtle but important
    'gamma':               [0.99, 0.995, 0.999],         # longer horizon
}

Learning rate and the number of parallel environments have by far the biggest effect. With 8 environments and a learning rate of 3e-4 you'll solve most simple control problems without a fuss. LunarLander specifically prefers a slightly gentler learning rate (2.5e-4) and more environments (16), because its reward signal is noisier and a bigger batch of parallel experience smooths that noise out.

Saving and loading

def save_checkpoint(trainer, path):
    torch.save({
        'network_state': trainer.network.state_dict(),
        'optimizer_state': trainer.optimizer.state_dict(),
        'log': trainer.log,
    }, path)

def load_checkpoint(trainer, path):
    checkpoint = torch.load(path, weights_only=True)
    trainer.network.load_state_dict(checkpoint['network_state'])
    trainer.optimizer.load_state_dict(checkpoint['optimizer_state'])
    trainer.log = checkpoint['log']

Save checkpoints during training, not just at the end. RL reward curves are noisy enough that the best policy isn't always the last one you trained -- sometimes the agent peaks at 300k steps and then wanders off. So save every time you hit a new best evaluation reward, and you'll never lose the good version to a bad final stretch.

What we actually built

This little project ties together a genuinely large chunk of the series: the RL framework and MDP formalization (#102), PPO's clipped surrogate objective and GAE for advantages (#109), neural network architecture and initialization (#38, #40), the training-loop pattern from supervised learning bent to fit RL (#42-43), and evaluation methodology (#13, adapted for RL's stochastic nature). Fourteen episodes of scattered pieces, one working system.

The codebase is deliberately bare-bones. Production RL libraries like Stable-Baselines3 or CleanRL handle a pile of edge cases we skipped: properly parallel vectorized environments with shared memory, observation and reward normalization, learning-rate annealing, and plenty more. But -- and this is the whole point of building it yourself first -- every scrap of complexity in those libraries exists to solve a problem you now understand from the inside, in stead of a magic flag you flip and pray over.

What to remember from this one

  • Vectorized environments collect N transitions per step, making training dramatically more sample-efficient per second of wall-clock;
  • orthogonal initialization with a per-head gain (0.01 on the policy head) stops the network committing to random action preferences before it has explored;
  • GAE with lambda = 0.95 balances bias and variance in advantage estimation, and almost never needs tuning;
  • evaluate with greedy actions but train with stochastic sampling -- exploration is for training, exploitation is for evaluation;
  • watch entropy, clip fraction and value loss as diagnostic signals: entropy collapse, a high clip fraction, or an exploding value loss each point to a specific, different fix;
  • learning rate and the number of parallel environments are the two most impactful hyperparameters, full stop;
  • always compare against baselines -- random and a hand-crafted heuristic -- because beating random proves nothing, while beating a heuristic proves your agent learned something real.

And with that, the reinforcement-learning arc closes -- from the very first "what even is a reward?" back in episode #102 all the way to an agent you trained yourself and watched beat a heuristic on the scoreboard. Next time we step back from algorithms entirely and start asking the questions that decide whether any of this survives contact with reality: how you actually design a machine-learning system that has to run, and keep running, outside a notebook. Different muscle, same series ;-)

De groeten!

@scipio

stem stemsocial steemstem python programming
0
0
0.000
0 comments