3️⃣ Main DQN Algorithm

Main DQN Algorithm

We now combine all the elements we have designed thus far into the final DQN algorithm. Here, we assume the environment returns three parameters $(s_{new}, r, d)$, a new state $s_{new}$, a reward $r$ and a boolean $d$ indicating whether interaction has terminated yet.

Our Q-value function $Q(s,a)$ is now a network $Q(s,a ; \theta)$ parameterised by weights $\theta$. The key idea, as in Q-learning, is to ensure the Q-value function satisfies the optimal Bellman equation

Q(s,a ; \theta) = \mathbb{E}_{s',r \sim p(\cdot \mid s,a)} \left[r + \gamma \max_{a'} Q(s', a' ;\theta) \right]

which means the expected TD error will be zero (where expectation here is taken over randomly sampled trajectories):

\mathbb{E} \left[ r_{t+1} + \gamma \max_a Q(s_{t+1}, a) - Q(s_t, a_t) \right] = 0

Note that we also have to deal with the case where the episode terminates: i.e. $s_{t+1}$ is a terminal state, and $d_{t+1} = 1$. Since the Q-value of a terminal state is always zero, we can just rewrite the expected TD error expression as:

\mathbb{E} \left[ r_{t+1} + (1 - d_{t+1}) \, \gamma \max_a Q(s_{t+1}, a) - Q(s_t, a_t) \right] = 0

since this makes sure the term is zero whenever $d_{t+1} = 1$. We can now see again why we used terminated not terminated | truncated here - we don't want the agent to learn that its value is always zero just before the episode ends and so there's no point in continuing to perform well!

Since we have an expression which should be zero in expectation for our true Q-value function, and we want the model to learn from a variety of experiences at once, we can sample batches of experiences $B = \{s_{t_i}, a_{t_i}, r_{t_i+1}, d_{t_i+1}, s_{t_i+1}\}_{i=1}^{|B|}$ from the replay buffer, and train against the loss function which equals the squared temporal difference error:

L(\theta) = \frac{1}{|B|} \sum_{i=1}^{|B|} \left( r_{t_i+1} + (1 - d_{t_i+1}) \gamma \max_a Q(s_{t_i+1}, a ; \theta_\text{target}) - Q(s_{t_i}, a_{t_i} ; \theta) \right)^2

Here, $\theta_\text{target}$ is a previous copy of the parameters $\theta$, so we're updating our $s_t$ estimates to catch up with our $s_{t+1}$ estimates (just like in standard Q-learning from earlier!). Every so often, we then update the target $\theta_\text{target} \leftarrow \theta$ as the agent improves its Q-values from experience.

Below is the full DQN algorithm from a paper, for reference. The notation isn't identical to ours (e.g. they use an if/else statement to handle the terminal state case), but the basic algorithm is the same.

DQN Dataclass

Below is a dataclass for training your DQN - read through it to see what you'll be working with (the arguments are grouped by what they control).

The exact breakdown of training is as follows:

The agent takes total_timesteps environment steps during the training loop, summed over all of our num_envs parallel environments: each call to agent.play_step advances every environment once, so each call adds num_envs to the environment step counter.
The first buffer_size environment steps are used to fill the replay buffer (we don't update gradients until the buffer is full).
After this point, we perform one optimization step (a gradient update of the Q-network) every steps_per_update calls to agent.play_step, i.e. every steps_per_update * num_envs environment steps. We also copy the weights from our Q-network to our target network every trains_per_target_update optimization steps.

It's important to keep these two kinds of steps distinct: with our default arguments, we take 300k environment steps but only ~18k optimization steps.

This is shown in the diagram below (the actual numbers aren't representative of the values in our dataclass, they're just to make sure the diagram is understandable - obviously the scale is very different in our actual training).

For example, in the code below we change total_timesteps, and this also changes the total number of optimization steps, total_training_steps (which is computed in the __post_init__ method of our dataclass, as a function of total_timesteps).

@dataclass
class DQNArgs:
    # Basic / global
    seed: int = 1
    env_id: str = "CartPole-gpu"
    num_envs: int = 16 # number of parallel environments to run
    device: str = device

    # Wandb / logging
    use_wandb: bool = False
    wandb_project_name: str = "DQNCartPole"
    wandb_entity: str | None = None
    video_log_freq: int | None = None # number of optimization steps between wandb video upload
    steps_per_live_video: int | None = 5_000 # number of optimization steps between live video display
    overwrite_video: bool = True # older live videos are overwritten by newer ones if True. Otherwise, they are plotted one after the other.

    # Duration of different phases / buffer memory settings
    total_timesteps: int = 300_000 # total environment steps (summed over all envs), including buffer filling
    steps_per_update: int = 1 # calls to `agent.play_step` (each = num_envs environment steps) per optimization step
    trains_per_target_update: int = 50 # optimization steps between target network updates
    buffer_size: int = 10_000 # replay buffer capacity (environment steps used to fill it before training)

    # Optimization hparams
    batch_size: int = 512
    learning_rate: float = 5e-3

    # RL-specific
    gamma: float = 0.99
    exploration_fraction: float = 0.1
    start_e: float = 1.0
    end_e: float = 0.02

    def __post_init__(self):
        assert self.total_timesteps - self.buffer_size >= self.steps_per_update * self.num_envs
        # Every call to `agent.play_step` advances all `num_envs` environments by one step; we take one optimization
        # step per `steps_per_update` such calls, so this is the total number of optimization steps
        self.total_training_steps = (self.total_timesteps - self.buffer_size) // (self.steps_per_update * self.num_envs)
        self.device = t.device(self.device)

Exercise - fill in the agent class

Difficulty: 🔴🔴🔴🔴⚪

Importance: 🔵🔵🔵🔵⚪

You should spend up to 25-45 minutes on this exercise.

You should now fill in the methods for the DQNAgent class below. This is a class which is designed to handle taking steps in the environment (with an epsilon greedy policy), and updating the buffer.

play_step should be somewhat similar to the demo code you saw earlier, which sampled a batch of experiences to add to the buffer. It should:
- Get actions (using self.get_actions rather than randomly sampling like we did in the demo code before)
- Step our environment with these actions
- Add the new experiences to the buffer: remembering that for environments which just finished an episode, next_obs is already the reset observation, and the true final observation is in infos["final_observation"].
- Set your new observation as self.obs, ready for the next step
get_actions should do the following:
- Set self.epsilon according to the linear schedule function & the current global step counter
- Sample actions according to the epsilon-greedy policy (i.e. using your epsilon_greedy_policy function), and return them

A small note on code practices here - the implementation below was designed to follow separation of concerns (SoC), a design principle used in software engineering. The DQNAgent class is only responsible for interacting with the environment; it doesn't do anything like create the Q-network or buffer on initialization. This is further reflected in the fact that we don't pass in args to our DQN agent, but instead pass in all the relevant variables separately (if we were forced to pass in args, this would be a sign that the DQN agent class might be doing too much work!).

class DQNAgent:
    """Base Agent class handling the interaction with the environment."""

    def __init__(
        self,
        envs: CartPole,
        buffer: ReplayBuffer,
        q_network: QNetwork,
        start_e: float,
        end_e: float,
        exploration_fraction: float,
        total_timesteps: int,
        rng: t.Generator,
    ):
        self.envs = envs
        self.buffer = buffer
        self.q_network = q_network
        self.start_e = start_e
        self.end_e = end_e
        self.exploration_fraction = exploration_fraction
        self.total_timesteps = total_timesteps
        self.rng = rng

        self.step = 0  # Tracking number of environment steps taken (summed across all environments)
        self.obs, _ = self.envs.reset()  # Need a starting observation
        self.epsilon = start_e  # Starting value (will be updated in `get_actions`)

    def play_step(self) -> dict:
        """
        Carries out a single interaction step between agent & environment, and adds results to the
        replay buffer.

        Returns `infos` (a dict containing info we will log).
        """
        raise NotImplementedError()

        self.step += self.envs.num_envs
        return infos

    def get_actions(self, obs: Float[Tensor, " num_envs *obs_shape"]) -> Int[Tensor, " num_envs"]:
        """
        Samples actions according to the epsilon-greedy policy using the linear schedule for epsilon.
        """
        raise NotImplementedError()


tests.test_agent(DQNAgent)

Solution

class DQNAgent:
    """Base Agent class handling the interaction with the environment."""

    def __init__(
        self,
        envs: CartPole,
        buffer: ReplayBuffer,
        q_network: QNetwork,
        start_e: float,
        end_e: float,
        exploration_fraction: float,
        total_timesteps: int,
        rng: t.Generator,
    ):
        self.envs = envs
        self.buffer = buffer
        self.q_network = q_network
        self.start_e = start_e
        self.end_e = end_e
        self.exploration_fraction = exploration_fraction
        self.total_timesteps = total_timesteps
        self.rng = rng

        self.step = 0  # Tracking number of environment steps taken (summed across all environments)
        self.obs, _ = self.envs.reset()  # Need a starting observation
        self.epsilon = start_e  # Starting value (will be updated in `get_actions`)

    def play_step(self) -> dict:
        """
        Carries out a single interaction step between agent & environment, and adds results to the
        replay buffer.

        Returns `infos` (a dict containing info we will log).
        """
        actions = self.get_actions(self.obs)
        next_obs, rewards, terminated, truncated, infos = self.envs.step(actions)

        # Get `true_next_obs` by finding all environments where we terminated & replacing `next_obs`
        # with the actual terminal states (the env has already auto-reset those environments)
        done = terminated | truncated
        true_next_obs = next_obs.clone()
        true_next_obs[done] = infos["final_observation"][done]

        self.buffer.add(self.obs, actions, rewards, terminated, true_next_obs)
        self.obs = next_obs

        self.step += self.envs.num_envs
        return infos

    def get_actions(self, obs: Float[Tensor, " num_envs *obs_shape"]) -> Int[Tensor, " num_envs"]:
        """
        Samples actions according to the epsilon-greedy policy using the linear schedule for epsilon.
        """
        self.epsilon = linear_schedule(
            self.step, self.start_e, self.end_e, self.exploration_fraction, self.total_timesteps
        )
        actions = epsilon_greedy_policy(self.envs, self.q_network, self.rng, obs, self.epsilon)
        assert actions.shape == (self.envs.num_envs,)
        return actions

Before we move on to the big exercise of today (completing the DQNTrainer class), we'll briefly discuss logging to Weights and Biases in RL, plus some general advice on what kinds of variables you should be logging.

Logging to `wandb` in RL

In previous exercises in this chapter, we've just trained the agent, and then plotted the reward per episode after training. For small toy examples that train in a few seconds this is fine, but for longer runs we'd like to watch the run live and make sure the agent is doing something interesting (especially if we were planning to run the model overnight). Luckily, Weights and Biases has got us covered! When you run your experiments, you'll be able to view not only live plots of the loss and average reward per episode while the agent is training - you can also log and view animations, which visualise your agent's progress in real time! The code below will handle all logging.

Sadly, effective logging & debugging in RL isn't just about watching videos, since in the vast majority of cases where your algorithm has a bug, the agent will just fail to learn anything useful and the videos won't be informative. Debugging RL requires knowing what variables to log and how to interpret the results you're getting, which requires some understanding of the underlying theory! This is part of the reason why we've spent so much time discussing the theory behind DQN and other RL algorithms, rather than just giving you a black box to train.

As an example of how logged variables can be misleading and hard to interpret, consider our TD loss function in DQN. This loss function just reflects how close together the Q-network's estimates are to the experiences currently sampled from the replay buffer, which might not adequately represent what the world actually looks like. This means that once the agent starts to learn something and do better at the problem, it's expected for the loss to increase. For example, maybe the Q-network initially learned some state was bad, because an agent that reached them was just flapping around randomly and died shortly after. But now it's getting evidence that the same state is good, now that the agent that reached the state has a better idea what to do next. A higher loss is thus actually a good sign that something is happening (the agent hasn't stagnated), but it's not clear if it's learning anything useful without also checking how the total reward per episode has changed. Key point - just looking at one variable can be misleading, we need to log multiple variables and derive a picture of what's happening from taking all of them into account!

Some useful variables to log during DQN training are:

TD loss, i.e. the actual loss you're backpropagating through. This should start off high and decrease pretty quickly, but may not be monotonic (i.e. temporary spikes in loss aren't necessarily a bad thing)
SPS (steps per second), i.e. the total number of agent steps divided by the total time. This helps us debug when the environment steps are a bottleneck (won't be the case in a simple environment like this one, but might matter more when we move to more complex environments)
Q-values, i.e. the predicted Q-values from the Q-network. Can you guess how these should behave?

Question - what do you think the Q values will do when the agent moves closer to solving the cartpole environment?

Initially they should be near zero, thanks to the randomly initialized model weights. As our episode length get closer to 500 (i.e. we can essentially solve the environment), they should tend to the limit of the total possible time-discounted reward available, which is the geometric sum $1 + \gamma + \gamma^2 + \cdots$ (since we get 1 reward for every second we stand up, and as previously discussed, the way we handle dones in the formula above doesn't assume a truncated environment causes future rewards to be terminated). The limit of this sum is $\frac{1 - \gamma^{500}}{1 - \gamma} \approx \frac{1}{1-\gamma}$, which for our default value $\gamma = 0.99$ is approximately 100. Intuitively, the agent isn't "planning" far enough into the future to care about the full 500 timesteps, only the next 100. But this is good enough to ensure sensible behaviour (agents that tend to optimise well over the next 100 steps end up doing the same thing as agents with $\gamma \to 1$.)

Note, the Q values won't increase smoothly, they'll spike up immediately after we copy over the weights from our Q-network to our target network. This is because each time we copy over weights, our gradient changes and the Q-network rapidly "catches up" to this new target network, causing the Q values to change rapidly. However, our copying over of weights will be frequent enough that these jumps will be relatively small, and so the curve should still appear smooth.

Exercise - write DQN training loop

Difficulty: 🔴🔴🔴🔴🔴

Importance: 🔵🔵🔵🔵🔵

You should spend up to 30-60 minutes on this exercise.

Now we'll create a new class DQNTrainer, which will handle the full training loop. We've filled in the __init__ for you, which defines all the things you need (the networks, optimizer, replay buffer, and the agent). We've also filled in train for you, which performs the main training loop: it optionally initializes Weights & Biases, fills the buffer using prepopulate_replay_buffer, then alternates between optimization steps (where we sample from the buffer and update the Q-network) & adding to the buffer (adding args.steps_per_update * num_envs new experiences, i.e. environment steps, each time).

You should fill in the remaining 2 methods. First you should get the basic no-logging version working, then once you're running without error (even if maybe you're not learning anything useful) you should move onto logging as this will help you debug.

add_to_replay_buffer
- This calls self.agent.play_step() n times (i.e. n * num_envs environment steps), which adds the results to the replay buffer
- It's used to fill the buffer before training starts, and before each optimization step to add new experiences to the buffer
training_step
- This performs one optimization step from a batch of experiences from the buffer, sampled using self.buffer.sample with batch size self.args.batch_size
- An update step involves:
  - Getting the predicted Q-values $Q(s_{t_i}, a_{t_i} ; \theta)$ from the Q-network
  - Getting the max target Q-values $\max_a Q(s_{t_i+1}, a ; \theta_\text{target})$ from the target network (remember to use inference mode - we're not training the target network!)
  - Computing the TD loss $L(\theta)$ using the formula we gave earlier (we've also copied it below, for convenience)
  - Performing an update step with this loss
- You should also copy weights from the Q-network to the target network every args.trains_per_target_update optimization steps (i.e. whenever the step argument, which counts optimization steps, is a multiple of this). The load_state_dict method might be useful here

For convenience, here's the full TD loss formula again:

L(\theta) = \frac{1}{|B|} \sum_{i=1}^{|B|} \left( r_{t_i+1} + (1 - d_{t_i+1}) \gamma \max_a Q(s_{t_i+1}, a ; \theta_\text{target}) - Q(s_{t_i}, a_{t_i} ; \theta) \right)^2

When you get to logging, there are 2 types of data you can log:

Data for terminated episodes, during buffer filling
- Terminated episode data can be found in the infos dict returned by the agent.play_step method. If at least one environment finished an episode on that step, then infos["final_info"][0]["episode"] will be a dict containing the length l and reward r of the finished episode(s), averaged over the environments that finished on that step (this is produced for our GPU env by track_episode_stats in gpu_env.py, and mirrors the layout gymnasium 0.29's vector environments used with the RecordEpisodeStatistics wrapper)
  - We've given you a helper function get_episode_data_from_infos which gives you a dict of the episode length & reward, or None if no envs terminated. See gpu_env.py for its implementation, and gymnasium's RecordEpisodeStatistics documentation for the wrapper this format is based on (note that gymnasium 1.x vector environments lay these stats out differently, as infos["episode"] with an infos["_episode"] mask).
- add_to_replay_buffer should just return this dict for the last episode that finished: the training loop we've given you logs it to the progress bar and to wandb (along with the SPS, steps per second) - at most twice a second
Data during optimization steps
- Mean TD loss, Q values, and the epsilon hyperparameter are all useful to log, but log them every 10 or so optimization steps rather than every one. A wandb.log call costs about as much as a whole gradient step on this tiny network, so logging on every optimization step would more than double the run time!

Don't be discouraged if your code takes a while to work - it's normal for debugging RL to take longer than you would expect. Add asserts or your own tests, implement an appropriate probe environment, try anything in the Andy Jones post that sounds promising, and try to notice confusion. Reinforcement Learning is often so tricky as even if the algorithm has bugs, the agent might still learn something useful regardless (albeit maybe not as well), or even if everything is correct, the agent might just fail to learn anything useful (like how DQN failed to do anything on Montezuma's Revenge.)

Since the environment is already known to be one DQN can solve, and we've already provided hyperparameters that work for this environment, hopefully that's isolated a lot of the problems one would usually have with solving real world problems with RL.

def get_episode_data_from_infos(infos: dict) -> dict[str, int | float] | None:
    """
    Helper function: returns a dict of episode data (length, reward and duration, averaged over the
    environments that finished an episode on this step), if at least one
    terminated.
    """
    for final_info in infos.get("final_info", []):
        if final_info is not None and "episode" in final_info:
            return {
                "episode_length": final_info["episode"]["l"].item(),
                "episode_reward": final_info["episode"]["r"].item(),
                "episode_duration": final_info["episode"]["t"].item(),
            }


class DQNTrainer:
    def __init__(self, args: DQNArgs):
        set_global_seeds(args.seed)
        self.args = args
        self.rng = t.Generator(device=args.device).manual_seed(args.seed)
        self.run_name = f"{args.env_id}__{args.wandb_project_name}__seed{args.seed}__{time.strftime('%Y%m%d-%H%M%S')}"

        self.envs = make_envs(args.env_id, args.num_envs, args.seed, args.device)

        # Define some basic variables from our environment (note, we assume a single discrete action space)
        num_envs = self.envs.num_envs
        action_shape = self.envs.single_action_space.shape
        num_actions = self.envs.single_action_space.n
        obs_shape = self.envs.single_observation_space.shape
        assert action_shape == ()

        # Create our replay buffer (on the same device as the environment)
        self.buffer = ReplayBuffer(num_envs, obs_shape, action_shape, args.buffer_size, args.seed, args.device)

        # Create our networks & optimizer (target network should be initialized with a copy of the Q-network's weights)
        self.q_network = QNetwork(obs_shape, num_actions).to(args.device)
        self.target_network = QNetwork(obs_shape, num_actions).to(args.device)
        self.target_network.load_state_dict(self.q_network.state_dict())
        self.optimizer = t.optim.AdamW(self.q_network.parameters(), lr=args.learning_rate)

        # Create our agent
        self.agent = DQNAgent(
            self.envs,
            self.buffer,
            self.q_network,
            args.start_e,
            args.end_e,
            args.exploration_fraction,
            args.total_timesteps,
            self.rng,
        )

    def add_to_replay_buffer(self, n: int, verbose: bool = False):
        """
        Calls `agent.play_step` n times (n * num_envs environment steps), adding to the replay buffer.
        Should return a dict of data from the last terminated episode, if any (the training loop takes
        care of logging it).

        Optional argument `verbose`: if True, we can use a progress bar (useful to check how long
        the initial buffer filling is taking).
        """
        raise NotImplementedError()

    def prepopulate_replay_buffer(self):
        """
        Called to fill the replay buffer before training starts.
        """
        n_steps_to_fill_buffer = self.args.buffer_size // self.args.num_envs
        self.add_to_replay_buffer(n_steps_to_fill_buffer, verbose=True)

    def training_step(self, step: int) -> None:
        """
        Samples once from the replay buffer, and takes a single optimization step.

        Args:
            step (int): The number of optimization steps taken so far (used for logging, and for deciding
            when to update the target network)
        """
        raise NotImplementedError()

    def train(self) -> None:
        if self.args.use_wandb:
            wandb.init(
                project=self.args.wandb_project_name,
                entity=self.args.wandb_entity,
                name=self.run_name,
            )

        self.prepopulate_replay_buffer()

        pbar = tqdm(range(self.args.total_training_steps))
        last_logged_time, last_logged_step = time.time(), self.agent.step  # so we don't log too often

        greedy = lambda obs: self.q_network(obs).argmax(-1)
        fresh_envs = lambda seed: CartPole(num_envs=16, seed=seed, device=self.args.device)
        can_render = self.args.env_id == "CartPole-gpu"
        live_videos = can_render and self.args.steps_per_live_video is not None
        wandb_videos = can_render and self.args.use_wandb and self.args.video_log_freq is not None
        video_slot = LiveVideo(overwrite=self.args.overwrite_video)  # one display slot per run (see rl_utils.LiveVideo)

        for step in pbar:
            data = self.add_to_replay_buffer(self.args.steps_per_update)
            if data is not None:
                if time.time() - last_logged_time > 0.5:
                    now = time.time()
                    if self.args.use_wandb:
                        sps = (self.agent.step - last_logged_step) / (now - last_logged_time)
                        wandb.log({**data, "SPS": sps}, step=self.agent.step)
                    last_logged_time, last_logged_step = now, self.agent.step
                    pbar.set_postfix(**data)

            self.training_step(step)

            live = live_videos and step % self.args.steps_per_live_video == 0
            to_wandb = wandb_videos and step % self.args.video_log_freq == 0
            if live or to_wandb:
                log_greedy_rollout_video(fresh_envs(step), greedy, self.envs.draw, step=self.agent.step, live=live, slot=video_slot)


        self.envs.close()
        if self.args.use_wandb:
            wandb.finish()

Solution (simple, no logging)

def add_to_replay_buffer(self, n: int, verbose: bool = False):
    """
    Calls `agent.play_step` n times (n * num_envs environment steps), adding to the replay buffer. Should return a
    dict of data from the last terminated episode, if any.

    Optional argument `verbose`: if True, we can use a progress bar (useful to check how long the initial buffer
    filling is taking).
    """
    data = None

    for step in tqdm(range(n), disable=not verbose, desc="Adding to replay buffer"):
        infos = self.agent.play_step()
        new_data = get_episode_data_from_infos(infos)
        data = new_data if new_data is not None else data  # keep the *last* terminated episode

    return data

def prepopulate_replay_buffer(self):
    """
    Called to fill the replay buffer before training starts.
    """
    n_steps_to_fill_buffer = self.args.buffer_size // self.args.num_envs
    self.add_to_replay_buffer(n_steps_to_fill_buffer, verbose=True)

def training_step(self, step: int) -> None:
    """
    Samples once from the replay buffer, and takes a single optimization step. The `step` argument is the number of
    optimization steps taken so far.
    """
    data = self.buffer.sample(self.args.batch_size)  # o_t, a_t, r_{t+1}, d_{t+1}, o_{t+1}

    with t.inference_mode():
        target_max = self.target_network(data.next_obs).max(-1).values
    predicted_q_vals = self.q_network(data.obs).gather(-1, data.actions.unsqueeze(-1)).squeeze(-1)

    td_error = data.rewards + self.args.gamma * target_max * (1 - data.terminated.float()) - predicted_q_vals
    loss = td_error.pow(2).mean()
    loss.backward()
    self.optimizer.step()
    self.optimizer.zero_grad()

    if step % self.args.trains_per_target_update == 0:
        self.target_network.load_state_dict(self.q_network.state_dict())

Solution (full logging)

def add_to_replay_buffer(self, n: int, verbose: bool = False):
    """
    Calls `agent.play_step` n times (n * num_envs environment steps), adding to the replay buffer. Should return a
    dict of data from the last terminated episode, if any (the training loop takes care of logging it).

    Optional argument `verbose`: if True, we can use a progress bar (useful to check how long the initial buffer
    filling is taking).
    """
    data = None

    for step in tqdm(range(n), disable=not verbose, desc="Adding to replay buffer"):
        infos = self.agent.play_step()

        # Get data from environments, if some environment did actually terminate
        new_data = get_episode_data_from_infos(infos)
        if new_data is not None:
            data = new_data  # makes sure we return a non-empty dict at the end, if some episode terminates

    return data

def prepopulate_replay_buffer(self):
    """
    Called to fill the replay buffer before training starts.
    """
    n_steps_to_fill_buffer = self.args.buffer_size // self.args.num_envs
    self.add_to_replay_buffer(n_steps_to_fill_buffer, verbose=True)

def training_step(self, step: int) -> None:
    """
    Samples once from the replay buffer, and takes a single optimization step. The `step` argument is the number of
    optimization steps taken so far.
    """
    data = self.buffer.sample(self.args.batch_size)  # o_t, a_t, r_{t+1}, d_{t+1}, o_{t+1}

    with t.inference_mode():
        target_max = self.target_network(data.next_obs).max(-1).values
    predicted_q_vals = self.q_network(data.obs).gather(-1, data.actions.unsqueeze(-1)).squeeze(-1)

    td_error = data.rewards + self.args.gamma * target_max * (1 - data.terminated.float()) - predicted_q_vals
    loss = td_error.pow(2).mean()
    loss.backward()
    self.optimizer.step()
    self.optimizer.zero_grad()

    if step % self.args.trains_per_target_update == 0:
        self.target_network.load_state_dict(self.q_network.state_dict())

    # Log every 10th step
    if self.args.use_wandb and step % 10 == 0:
        wandb.log(
            {"td_loss": loss, "q_values": predicted_q_vals.mean().item(), "epsilon": self.agent.epsilon},
            step=self.agent.step,
        )

Here's some boilerplate code to test out your various probes, which you should make sure you're passing before testing on Cartpole.

def test_probe(probe_idx: int):
    """
    Tests a probe environment by training a network on it & verifying that the value functions are
    in the expected range.
    """
    # Train our network on this probe env.
    args = DQNArgs(
        env_id=f"Probe{probe_idx}-v0",
        wandb_project_name=f"test-probe-{probe_idx}",
        total_timesteps=3000 if probe_idx <= 2 else 5000,
        learning_rate=0.001,
        buffer_size=500,
        use_wandb=False,
        trains_per_target_update=20,
        video_log_freq=None,
        num_envs=1,
        steps_per_update=10,
        exploration_fraction=0.2,
        end_e=0.1,
    )
    trainer = DQNTrainer(args)
    trainer.train()

    # Get the correct set of observations, and corresponding values we expect
    obs_for_probes = [[[0.0]], [[-1.0], [+1.0]], [[0.0], [1.0]], [[0.0]], [[0.0], [1.0]]]
    expected_value_for_probes = [
        [[1.0]],
        [[-1.0], [+1.0]],
        [[args.gamma], [1.0]],
        [[-1.0, 1.0]],
        [[1.0, -1.0], [-1.0, 1.0]],
    ]
    tolerances = [5e-4, 5e-4, 5e-4, 5e-4, 1e-3]
    obs = t.tensor(obs_for_probes[probe_idx - 1]).to(device)

    # Calculate the actual value, and verify it
    value = trainer.q_network(obs)
    # A bf16 network (e.g. under autocast) can't represent these targets to better than its eps
    atol = max(tolerances[probe_idx - 1], t.finfo(value.dtype).eps)
    expected_value = t.tensor(expected_value_for_probes[probe_idx - 1]).to(device)
    t.testing.assert_close(value.float(), expected_value, atol=atol, rtol=0)
    print("Probe tests passed!\n")


for probe_idx in range(1, 6):
    test_probe(probe_idx)

Once you've passed the tests for all 5 probe environments, you should test your model on Cartpole. We recommend you start by not using wandb until you can get it running without error, because this will improve your feedback loops (however if you've passed all probe environments then there's a good chance this code will just work for you).

You may find due to the small size of this network, and the few number of parallel environments, that it may actually run faster on the CPU rather than the GPU. Try both and see.

args = DQNArgs(use_wandb=True)
trainer = DQNTrainer(args)
trainer.train()

Catastrophic forgetting

Note - you might see performance frequently drop off after it's achieved the maximum for a while, before eventually recovering again and repeating the cycle. Here's an example CartPole run using the solution code:

This is a well-known RL phenomenon called catastrophic forgetting. It happens when the replay buffer mostly contains successful experiences, and the model forgets how to adapt or recover from bad states. One way to fix this is to change your buffer to keep 10% of experiences from previous epochs, and 90% of experiences from the current phase. Can you implement this?

When we cover PPO, we'll also introduce reward shaping, which is another way this kind of behaviour can be mitigated.

Mechanism of catastrophic forgetting

Now that you have a working DQNTrainer, we can try to quantify catastrophic forgetting. The cell below re-runs training with the default settings and, every 500 gradient steps, records five things:

how long the greedy policy survives on fresh environments;
what fraction of the replay buffer's transitions are terminal (i.e. whether the data still contains any failures, which if we keep training on, the agent will forget what failure looks like)
the action gap $\max_a Q(s,a) - \min_a Q(s,a)$ averaged over a fixed batch of states
and the value $\max_a Q(s, a)$ the network assigns to a fixed set of doomed states: states from which every action ends the episode on the next step. By definition, doomed states $s$ satisfy $Q_{\pi}(s, a) = 0$ for any policy $\pi$ and any action $a$, as there is no reward, and no next state to bootstrap from.
the TD loss on a large sample from the buffer, the quantity training_step minimises.

Both fixed sets are sampled from the buffer at the start, while it is still full of random-policy crashes.

# Re-run training with the default settings, but every 500 gradient steps record four things:
#   1. how long the greedy policy survives on 16 fresh environments (no exploration)
#   2. how many of the transitions in the replay buffer are terminal, i.e. does the data still contain failures
#   3. the action gap  max_a Q(s,a) - min_a Q(s,a)  averaged over a fixed batch of states from the random-policy phase
#   4. max_a Q(s, a) on a fixed set of doomed states, from which every action terminates next step (correct value 0)
#   5. the TD loss, recomputed on a large sample from the buffer (the same quantity training_step minimises)
from plotly.subplots import make_subplots
from rl_utils import greedy_rollout

# toggle this back to true to run this
RUN_ACTION_GAP_DEMO = False

if RUN_ACTION_GAP_DEMO:
    args = DQNArgs(use_wandb=False, steps_per_live_video=None)
    trainer = DQNTrainer(args)
    trainer.prepopulate_replay_buffer()
    probe_obs = trainer.buffer.sample(2048).obs  # sampled now, while the buffer is still full of random-policy failures
    pre_terminal_obs = trainer.buffer.obs[trainer.buffer.terminated]  # states whose taken action ended the episode
    # Check these are doomed under BOTH actions, by writing each state into a batched env and stepping it with each action.
    check_env = CartPole(num_envs=len(pre_terminal_obs), device=args.device)
    doomed = t.ones(len(pre_terminal_obs), dtype=t.bool, device=args.device)
    for a in (0, 1):
        check_env.state = pre_terminal_obs.clone()
        doomed &= check_env.step(t.full((len(pre_terminal_obs),), a, device=args.device))[2]  # terminated flag
    print(f"{doomed.sum().item()} of {len(doomed)} pre-termination states are doomed under both actions")
    doomed_obs = pre_terminal_obs[doomed]
    greedy = lambda obs: trainer.q_network(obs).argmax(-1)

    def td_loss(n_samples: int = 2048) -> float:
        d = trainer.buffer.sample(n_samples)
        with t.inference_mode():
            target = d.rewards + args.gamma * trainer.target_network(d.next_obs).max(-1).values * (1 - d.terminated.float())
            predicted = trainer.q_network(d.obs).gather(-1, d.actions.unsqueeze(-1)).squeeze(-1)
        return (target - predicted).pow(2).mean().item()

    log = []
    pbar = tqdm(range(args.total_training_steps))
    for step in pbar:
        trainer.add_to_replay_buffer(args.steps_per_update)
        trainer.training_step(step)
        if (step + 1) % 500 == 0:
            _, _, lifespan = greedy_rollout(CartPole(num_envs=16, seed=step, device=args.device), greedy)
            with t.inference_mode():
                q = trainer.q_network(probe_obs)
                q_doomed = trainer.q_network(doomed_obs).max(-1).values
            record = {
                "G" : lifespan.mean().item(),
                "n_term" : trainer.buffer.terminated.sum().item(),
                "act_gap" : (q.max(-1).values - q.min(-1).values).mean().item(),
                "q_doom" : q_doomed.mean().item(),
                "loss" : td_loss(),
            }
            log.append([step + 1, *record.values()])
            pbar.set_postfix({k: f"{v:.3g}" for k, v in record.items()})
else:
    print(f"{RUN_ACTION_GAP_DEMO=} Skipping action gap demo.")

if RUN_ACTION_GAP_DEMO:
    steps, greedy_lifespan, n_terminal, action_gap, q_doomed, td_loss_curve = zip(*log)
    fig = make_subplots(rows=5, cols=1, shared_xaxes=True, vertical_spacing=0.06, subplot_titles=[
        "Greedy episode length on 16 fresh envs",
        "Terminal transitions in the replay buffer (log scale; 0 shown at the bottom)",
        "Action gap  max_a Q(s,a) - min_a Q(s,a)  on a fixed batch of states (log scale)",
        "max_a Q(s, a) on fixed doomed states  (correct value 0)",
        "TD loss on 2048 sampled transitions (log scale)",
    ])
    n_terminal_plot = [max(n, 0.5) for n in n_terminal]  # a log axis can't show 0, so exact zeros sit at 0.5 with a tick labelled "0"
    for row, y in enumerate([greedy_lifespan, n_terminal_plot, action_gap, q_doomed, td_loss_curve], start=1):
        fig.add_scatter(x=steps, y=y, mode="lines+markers", showlegend=False, row=row, col=1)
    ticks = [0.5, 1, 10, 100, 1000]
    fig.update_yaxes(type="log", tickvals=ticks, ticktext=["0"] + [str(v) for v in ticks[1:]], range=[-0.4, 3], row=2, col=1)
    fig.update_yaxes(type="log", row=3, col=1)
    fig.update_yaxes(type="log", row=5, col=1)
    fig.update_layout(height=1100, width=800, xaxis5_title="gradient steps", margin=dict(t=40, b=40, l=40, r=20))
    fig.show()

Panel by panel:

The buffer runs out of failures. 10,000 transitions across 16 envs is the last ~625 steps of play. Once every env survives all 500 steps, the buffer holds no terminal transitions (panel 2).
Q stops distinguishing actions. With no failures in the data, the TD targets say every action is fine, so $Q(s, \text{left})$ and $Q(s, \text{right})$ converge and the action gap decays (panel 3).
Q forgets what doom looks like. Doomed states are worth 0, but with nothing to anchor them their value drifts towards $1/(1-\gamma) = 100$, the value of never falling (panel 4). This begins while failures are merely rare, the impact to the loss is minimal, so they can be ignored.
The loss can't see it. Throughout the drift the TD loss is near zero (panel 5): with no failures in the buffer, $Q$ is perfectly consistent with the data it is trained on. Forgetting is invisible to the training signal.
Collapse and recovery. As the action gap approaches zero, the argmax policy becomes a coinflip, the policy becomes random, and performance tanks (panel 1). This causes failures to re-enter the buffer, the loss spikes, and within a few hundred steps panels 3 and 4 snap back.

Panels 2-4 predict the actual collapse of the policy well in advance; the loss (panel 5) does not. This motivates keeping old experiences, even if your current policy would never generate them, as it reminds the Q-network not to forget.

If you let this run for an extened period (increase total_timesteps) you will see this cycle repeat.