DQN cartpole agent from pytorch's tutorial not learning

niccolot: <p>I tried to implement the pytorch’s DQN tutorial <a href="https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html" rel="noopener nofollow ugc">https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html</a> as an agent class for tidyness and in order to implement n-step learning. I wrote the following code rearranging the tutorial:</p>
<pre><code class="lang-auto">class Agent:
  def __init__(self, exp_buffer, policy_net, target_net, device):
    self.exp_buffer = exp_buffer
    self.policy_net = policy_net
    self.target_net = target_net
    self.optimizer = optim.AdamW(self.policy_net.parameters(), lr=1e-4, amsgrad=True)
    self.device = device
    self.steps_done = 0
    self.eps_start = 0.9
    self.eps_end = 0.05
    self.eps_decay = 1000
    self.gamma = 0.99
    self.tau = 0.005


  def _eps_decay(self):
    return self.eps_end + (self.eps_start - self.eps_end) * \
        math.exp(-1. * self.steps_done / self.eps_decay)


  def _select_action(self, state):

    eps_threshold = self._eps_decay()

    if random.random() &gt; eps_threshold:
      with torch.no_grad():
        # .view(1,1) in order for it to work with the .gather() method
        # used to extract state action values
        return self.policy_net(state).max(1)[1].view(1,1)
    else:
      return torch.tensor([[env.action_space.sample()]], device=self.device, dtype=torch.long)


  def sync_nets(self):
    target_net_state_dict = self.target_net.state_dict()
    policy_net_state_dict = self.policy_net.state_dict()

    for key in policy_net_state_dict:
      target_net_state_dict[key] = policy_net_state_dict[key]*self.tau + target_net_state_dict[key]*(1-self.tau)

    self.target_net.load_state_dict(target_net_state_dict)


  def explore(self, env, state, n_steps=1):
    
    state_init = state
    state_init = torch.tensor(state_init, dtype=torch.float32, device=self.device).unsqueeze(0)

    total_reward = 0.0
    
    for step in range(n_steps):
      state = torch.tensor(state, dtype=torch.float32, device=self.device).unsqueeze(0)
      action = self._select_action(state)
      observation, reward, terminated, truncated, _ = env.step(action.item())

      if terminated:
        next_state = None
      else:
        next_state = torch.tensor(observation, dtype=torch.float32, device=self.device).unsqueeze(0)

      reward *= self.gamma**step

      total_reward += reward
      state = next_state

    done = terminated or truncated
    self.steps_done += 1
    total_reward = torch.tensor([total_reward], device=self.device)
    self.exp_buffer.push(state_init, action, next_state, total_reward)

    return done


  def optimize_model(self, batch_size, n_steps=1):

    if len(self.exp_buffer) &lt; batch_size:
      return

    transitions = self.exp_buffer.sample(batch_size)
    batch = Transition(*zip(*transitions))
    
    non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,
                                          batch.next_state)), device=self.device, dtype=torch.bool)
    non_final_next_states = torch.cat([s for s in batch.next_state
                                                if s is not None])


    state_batch = torch.cat(batch.state)
    action_batch = torch.cat(batch.action)
    reward_batch = torch.cat(batch.reward)
    state_action_values = self.policy_net(state_batch).gather(1, action_batch)

    next_state_values = torch.zeros(batch_size, device=device)
    with torch.no_grad():
        next_state_values[non_final_mask] = self.target_net(non_final_next_states).max(1)[0]

    expected_state_action_values = (next_state_values * self.gamma**n_steps) + reward_batch

    criterion = nn.SmoothL1Loss()
    loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1))

    self.optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_value_(self.policy_net.parameters(), 100)
    self.optimizer.step()

episode_durations = []
seed = 1423
n_actions = env.action_space.n
state, info = env.reset()
n_observations = len(state)
torch.manual_seed(seed)
policy_net = DQN(n_observations, n_actions).to(device)
target_net = copy.deepcopy(policy_net)
target_net.to(device)

exp_buffer = ReplayMemory(10000)

if torch.cuda.is_available:
    num_episodes = 600
else:
    num_episodes = 50

agent = Agent(exp_buffer, policy_net, target_net, device)

for i_episode in range(num_episodes):
  state, info = env.reset()

  for t in count():
    
    done = agent.explore(env, state)
    agent.optimize_model(batch_size=128)
    agent.sync_nets()

    if done:
      episode_durations.append(t + 1)
      plot_durations()
      break

print('Complete')
plot_durations(show_result=True)
plt.plot(episode_durations)
plt.show()
</code></pre>
<p>I omitted imports and the definitions of the DQN, buffer and plot function for brevity as they are the same as in the tutorial.</p>
<p>In the <code>.explore()</code> method is implemented the n step dqn but in this case is used as simple 1-step and it should be the same as in the tutorial.</p>
<p>When i run the code the reward is stuck around 10-15 and it doesn’t grow (when obviously running the code as in the tutorial works well)</p>