Double backprop error in Catan RL network

rosh_m: <p>My friend and I are trying to code a catan RL agent and are running into this backprop problem with the PPO agent policy backward function saying “RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.” even though we have retain_graph=True in the code. Some of the code is chatbot generated as a starting point to simplify the coding process, and we’re running all of this in a Google Colab environment. The library used for training and running games, Catanatron, was made from bcollazo on Github, and we tested an existing model to confirm that there’s no issues with running Catanatron in Google Colab.  Any help would be greatly appreciated. The following is the code.</p>
<pre><code class="lang-auto">class PolicyNetwork(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(PolicyNetwork, self).__init__()
        self.fc1 = nn.Linear(input_dim, 128)
        self.fc2 = nn.Linear(128, 128)
        self.fc3 = nn.Linear(128, output_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

class ValueNetwork(nn.Module):
    def __init__(self, input_dim):
        super(ValueNetwork, self).__init__()
        self.fc1 = nn.Linear(input_dim, 128)
        self.fc2 = nn.Linear(128, 128)
        self.fc3 = nn.Linear(128, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

class PPOAgent:
    def __init__(self, env, policy_lr=3e-4, value_lr=1e-3, gamma=0.99, clip_epsilon=0.2, k_epochs=4):
        self.env = env
        self.policy_net = PolicyNetwork(env.observation_space.shape[0], env.action_space.n)
        self.value_net = ValueNetwork(env.observation_space.shape[0])
        self.policy_optimizer = optim.Adam(self.policy_net.parameters(), lr=policy_lr)
        self.value_optimizer = optim.Adam(self.value_net.parameters(), lr=value_lr)
        self.gamma = gamma
        self.clip_epsilon = clip_epsilon
        self.k_epochs = k_epochs

    def select_action(self, state):
        state = torch.tensor(state, dtype=torch.float32)
        logits = (self.policy_net)(state)
        action_probs = torch.softmax(logits, dim=-1)
        dist = Categorical(action_probs)
        action = dist.sample()
        return action.item(), dist.log_prob(action)

    def compute_returns(self, rewards, dones, next_value):
        returns = []
        G = next_value
        for reward, done in zip(reversed(rewards), reversed(dones)):
            G = reward + self.gamma * G * (1 - done)
            returns.insert(0, G)
        return returns

    def train(self, num_episodes=1000, max_steps=200):
        for episode in range(num_episodes):
            state, valid_actions = self.env.reset()
            log_probs = []
            values = []
            rewards = []
            dones = []
            counter = 0
            done = False

            for step in range(max_steps):
              if done:
                break

              action, log_prob = self.select_action(state)
              value = self.value_net(torch.tensor(state, dtype=torch.float32))
              next_state, reward, markov_decision_process, done, available_actions = self.env.step(action)

              log_probs.append(log_prob)
              values.append(value)
              rewards.append(reward)
              dones.append(done)

              state = next_state


            next_value = self.value_net(torch.tensor(state, dtype=torch.float32)).item()
            returns = self.compute_returns(rewards, dones, next_value)

            log_probs = torch.stack(log_probs)
            values = torch.stack(values).squeeze()
            returns = torch.tensor(returns, dtype=torch.float32)
            advantages = returns - values

            for something in range(self.k_epochs):
                new_log_probs = []
                state, _ = self.env.reset()
                done = False
                counter = 1
                for step in range(len(log_probs)):
                  if done:
                    break

                  action, new_log_prob = self.select_action(state)
                  new_log_probs.append(new_log_prob)
                  next_state, _, _, done, _ = self.env.step(action)
                  counter+= 1

                new_log_probs = torch.stack(new_log_probs)

                if len(new_log_probs) &gt; len(log_probs):
                    new_log_probs = new_log_probs[:len(log_probs)]
                elif len(new_log_probs) &lt; len(log_probs):
                    log_probs = log_probs[:len(new_log_probs)]
                    advantages = advantages[:len(new_log_probs)]

                ratios = torch.exp(new_log_probs - log_probs.detach())
                surr1 = ratios * advantages
                surr2 = torch.clamp(ratios, 1 - self.clip_epsilon, 1 + self.clip_epsilon) * advantages
                policy_loss = -torch.min(surr1, surr2).mean()

                value_loss = nn.MSELoss()(values, returns)

                self.policy_optimizer.zero_grad()
                policy_loss.backward(retain_graph = True)
                self.policy_optimizer.step()

                self.value_optimizer.zero_grad()
                value_loss.backward()
                self.value_optimizer.step()

            if episode % 50 == 0:
                print(f"Episode {episode}, Policy Loss: {policy_loss.item()}, Value Loss: {value_loss.item()}")

env = gym.make("catanatron_gym:catanatron-v1")
agent = PPOAgent(env)
agent.train()
</code></pre>