4️⃣ Beyond CartPole & Bonus

Beyond CartPole

If things go well and your agent masters CartPole, the next harder challenge is MountainCar-v0: gpu_env.py ships a batched GPU version of it too, "MountainCar-gpu"(it has 3 discrete actions and a 2-dimensional observation, soQNetwork` needs no changes). It's a much harder exploration problem than CartPole - the reward is -1 every step until the car reaches the flag, so random exploration almost never sees a success. Feel free to Google for appropriate hyperparameters - in a real RL problem you would have to do hyperparameter search using the techniques we learned on a previous day, because bad hyperparameters in RL often completely fail to learn even if the algorithm is perfectly correct. There are many more exciting environments to play in, but generally they're going to require more compute and more optimization than we have time for today. If you finish the main material, some we recommend are:

Minimalistic Gridworld Environments - a fast gridworld environment for experiments with sparse rewards and natural language instruction.
microRTS - a small real-time strategy game suitable for experimentation.
Megastep - RL environment that runs fully on the GPU (fast!)
Procgen - A family of 16 procedurally generated gym environments to measure the ability for an agent to generalize. Optimized to run quickly on the CPU.
Atari - although you might want to wait until the PPO section to try this on DQN, because we'll be going through some guided exercises implementing Atari with PPO there!

Some (very unpolished) code for setting up Atari with DQN

This is based on a hybrid of the agent/critic network setup for Atari, and the DQN implementation in this notebook. I've achieved decent performance in 40 mins training this, but not as good as we get when we do PPO on Atari in the PPO section, so I think this is somewhat underoptimized - if anyone finds improvements then feel free to make a PR!

It uses the same vectorised, GPU-tensor Atari environments as the PPO material (rl_utils.AtariEnvs, built on EnvPool). Observations are the standard stack of 4 grayscale 84x84 frames, already scaled to $[0, 1]$, so the only thing you need to swap is the Q-network - the MLP below is replaced by the Nature-DQN convolutional net. Since DQNTrainer looks up the name QNetwork when it builds its networks, redefining the class here is all it takes. Two things to know: EnvPool runs with episodic life and reward clipping (so the progress bar's "episodes" are lives and rewards are $\pm 1$), and it resets an environment on the step after the one it ended on, so one terminal-observation-to-reset-observation transition per life ends up in the replay buffer - harmless here, but worth knowing about.

def layer_init(layer: nn.Linear, std=np.sqrt(2), bias_const=0.0):
    t.nn.init.orthogonal_(layer.weight, std)
    t.nn.init.constant_(layer.bias, bias_const)
    return layer


class QNetwork(nn.Module):
    """The Nature-DQN convolutional Q-network, for (4, 84, 84) stacked-frame observations in [0, 1]."""

    layers: nn.Sequential

    def __init__(self, obs_shape: tuple[int, ...], num_actions: int, hidden_sizes: list[int] = [120, 84]):
        super().__init__()
        assert len(obs_shape) == 3, "Expecting (frames, height, width) observations - use the MLP QNetwork for CartPole"
        assert obs_shape[-1] % 8 == 4
        L_after_convolutions = (obs_shape[-1] // 8) - 3
        in_features = 64 * L_after_convolutions * L_after_convolutions

        self.layers = nn.Sequential(
            layer_init(nn.Conv2d(obs_shape[0], 32, 8, stride=4, padding=0)),
            nn.ReLU(),
            layer_init(nn.Conv2d(32, 64, 4, stride=2, padding=0)),
            nn.ReLU(),
            layer_init(nn.Conv2d(64, 64, 3, stride=1, padding=0)),
            nn.ReLU(),
            nn.Flatten(),
            layer_init(nn.Linear(in_features, 512)),
            nn.ReLU(),
            layer_init(nn.Linear(512, num_actions), std=0.01),
        )

    def forward(self, x: Tensor) -> Tensor:
        return self.layers(x)


# Register the environment: `ENVS` entries take (num_envs, seed, device); EnvPool picks its own device, so we ignore it
ENVS["ALE/Breakout-v5"] = lambda num_envs, seed, device: AtariEnvs("ALE/Breakout-v5", num_envs=num_envs, seed=seed)

args = DQNArgs(
    env_id="ALE/Breakout-v5",
    num_envs=4,
    buffer_size=1_000,  # 1000 x (4, 84, 84) float32 observations, twice (obs and next_obs) ~ 0.2 GB on the GPU
    batch_size=32,
    end_e=0.01,
    learning_rate=1e-4,
    total_timesteps=20_000,
    steps_per_update=5,
    steps_per_live_video=None,  # the 4x4 CartPole grid video doesn't apply here
    use_wandb=True,
    wandb_project_name="DQNAtari",
)
trainer = DQNTrainer(args)
trainer.train()

Bonus

Target Network

Why have the target network? Modify the DQN code above, but this time use the same network for both the target and the Q-value network, rather than updating the target every so often.

Compare the performance of this against using the target network.

Shrink the Brain

Can DQN still learn to solve CartPole with a Q-network with fewer parameters? Could we get away with three-quarters or even half as many parameters? Try comparing the resulting training curves with a shrunken version of the Q-network. What about the same number of parameters, but with more/less layers, and less/more parameters per layer?

Dueling DQN

Implement dueling DQN according to the paper and compare its performance.