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 steps in the environment during the training loop. These are environment steps, summed over all of our num_envs parallel environments: each call to agent.play_step advances every environment once, so we advance the total timesteps counter by num_envs on each step.
The first buffer_size of these steps are used to fill the replay buffer (we don't update gradients until the buffer is full).
After this point, we perform an optimizer step every steps_per_train calls to agent.play_step (i.e. every steps_per_train * num_envs environment steps). We also copy the weights from our Q-network to our target network every trains_per_target_update steps of our Q-network.

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 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 training steps between wandb video upload
    steps_per_live_video: int | None = 5_000 # number of training 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
    steps_per_update: int = 1
    trains_per_target_update: int = 50
    buffer_size: int = 10_000

    # 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
        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 steps taken (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 steps taken (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 training steps (where we sample from the buffer) & adding to the buffer (adding args.steps_per_update * num_envs new experiences 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() to take n steps in the environment, which adds the results to the replay buffer
- It's used to fill the buffer before training starts, and before each training step to add new experiences to the buffer
training_step
- This performs an update 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 steps (i.e. whenever the step argument 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 the same format gym's RecordEpisodeStatistics wrapper uses, produced for our GPU env by track_episode_stats in gpu_env.py)
  - 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 the documentation page for an explanation.
- 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 training steps
- Mean TD loss, Q values, and the epsilon hyperparameter are all useful to log, but log them every 10 or so training steps rather than every step. A wandb.log call costs about as much as a whole gradient step on this tiny network, so logging on every 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):
        """
        Takes n steps with the agent, 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 training step.

        Args:
            step (int): The number of training steps taken (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):
    """
    Takes n steps with the agent, adding to the replay buffer (and logging any results). 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 training step. The `step` argument is used to track the
    number of training steps taken.
    """
    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):
    """
    Takes n steps with the agent, 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 training step. The `step` argument is used to track the
    number of training steps taken.
    """
    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)
    expected_value = t.tensor(expected_value_for_probes[probe_idx - 1]).to(device)
    t.testing.assert_close(value, expected_value, atol=tolerances[probe_idx - 1], 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.