Hello I have some problems with implementing the Actor Critic Policy Gradient Algorithm,
When I implement REINFORCE like this, everything is okay:
class Estimator(nn.Module):
    def __init__(self, num_actions):
        super().__init__()
        self.num_actions = num_actions
        self.dense_1 = nn.Linear(4, 32)
        self.out = nn.Linear(32, num_actions)

    def forward(self,x):
        x =self.dense_1(x)
        x = F.softmax(self.out(x))
        return x
    
env = gym.make("CartPole-v0")
estimator = Estimator(2)
estimator.cuda()
opt = optim.Adam(estimator.parameters())
loss = []
running_reward = 0
for i in range(100000): # number episodes
    episode = []
    chosen_actions = []
    rewards = []
    done = False
    state = env.reset()
                  
    while not done:
        probs = estimator(Variable(torch.unsqueeze(torch.from_numpy(state),0).float().cuda())) # calculate the probs of choosing actions
        action = probs.multinomial()
        chosen_actions.append(action)
        next_state, reward, done, _ = env.step(action.data[0,0])
        rewards.append(reward)
        state = next_state
       
    
    R = 0
    for r in rewards[::-1]:
        R = r +  R
        rewards.insert(0, R)
        
    for action, r in zip(chosen_actions, rewards):
        action.reinforce(r)
        
    opt.zero_grad()
    autograd.backward(chosen_actions, [None for _ in chosen_actions])
    opt.step()
    running_reward = running_reward * 0.99 + len(chosen_actions) * 0.01
    if (i+1) % 10 == 0:
        print("Episode: {} Running Reward: {}".format(i+1,round(running_reward,2)))

But if I try to implement the actor critic with the using the Generalized Advantage Estimator 20 the Algorithm fails.
Two things are happening:

The policy is learning to ALWAYS choose one off the actions (the probability is approaching 1)
the loss (and the output) of the state value estimator explode.

I checked the implementation of the policy gradient by pretraining a state value estimator using TD and then plugging it into the code. That works just fine, so I suspect I have made some error implementing the state value updates.
Here is the code:
(I tried a lot of different hyperparamenters like learning rate of both optimizers, the number of state value estimator updates per timestep and also starting the policy gradient algorithm after a short amount of time where the state value predictions are already learning, unfortunatly nothing has worked…)
class Estimator(nn.Module):
    def __init__(self, num_actions):
        super().__init__()
        self.num_actions = num_actions
        self.dense_1 = nn.Linear(4, 32)
        self.out = nn.Linear(32, num_actions)

    def forward(self,x):
        x =self.dense_1(x)
        x = F.softmax(self.out(x))
        return x

class V_Estimator(nn.Module):
    def __init__(self):
        super().__init__()
        self.dense_1 = nn.Linear(4, 32)
        self.out = nn.Linear(32, 1)

    def forward(self,x):
        x = F.relu(self.dense_1(x))
        x = self.out(x)
        return x
    
    
estimator = Estimator(2)
estimator.cuda()
v_estimator = V_Estimator()
v_estimator.cuda()
opt = optim.Adam(estimator.parameters(), lr=0.0001)
v_opt = optim.Adam(v_estimator.parameters(), lr=0.0001)
env = gym.make("CartPole-v0")
mse = nn.MSELoss()
buffer = ReplayBuffer(100000)
running_reward = 0
for i in range(10000): # number episodes
    episode_len = 0
    done = False
    state = env.reset()
    

    while not done:
        episode_len += 1
        state_python = state
        state = Variable(torch.unsqueeze(torch.from_numpy(state),0).float().cuda())
        probs = estimator(state) 
        #print(probs.data.cpu().numpy()) # one of the action probabilites just approaches 1
        action = probs.multinomial()

        action_python = action.data[0,0]
        v_estimate_curr = v_estimator(state)
        #v_estimate_curr = v_estimate(state)
        #print(v_estimate_curr)
        next_state, reward, done, _ = env.step(action_python)
        v_estimate_next = v_estimator(Variable(torch.unsqueeze(torch.from_numpy(next_state),0).float().cuda()))
        #v_estimate_next = v_estimate(Variable(torch.unsqueeze(torch.from_numpy(next_state),0).float().cuda()))
        #print(v_estimate_next)
        
        td_error = reward + v_estimate_next - v_estimate_curr
        
        
        buffer.add(state_python, action_python, reward, done, next_state)
        state = next_state
        
        #refit v-estimator
        average_state_value_loss = 0
        state_value_updates = 30
        for j in range(state_value_updates):
            s_batch, a_batch, r_batch, d_batch, s2_batch  = buffer.sample_batch(128)
            #print("s_batch shape: {}".format(s_batch.shape))
            targ = v_estimator(Variable(torch.from_numpy(s2_batch)).float().cuda())
            #print("targ shape: {}".format(targ.data.cpu().numpy().shape))
            torch_rew_batch = Variable(torch.unsqueeze(torch.from_numpy(r_batch).float().cuda(),-1))
            #print("torch_rew_batch shape: {}".format(torch_rew_batch.data.cpu().numpy().shape))
            targ = targ + torch_rew_batch
            targ = targ.detach()
            targ.requires_grad = False
            #print("targ shape: {}".format(targ.data.cpu().numpy().shape))
            out = v_estimator(Variable(torch.from_numpy(s_batch)).float().cuda())
            #print("out shape: {}".format(out.data.cpu().numpy().shape))
            v_loss = mse(out, targ)
            average_state_value_loss += v_loss.data[0] / state_value_updates
            
            v_opt.zero_grad()
            v_loss.backward()
            v_opt.step()

        # update policy gradient
        #if i > 100: # starting after 100 episodes to give the state value nn some time to learn
        opt.zero_grad()
        action.reinforce(td_error.data)
        action.backward()
        opt.step()
    running_reward = running_reward * 0.9 + episode_len * 0.1
    print("current episode: " + str(i)+ " - running reward: " + str(round(running_reward,2)) + " - average state value estimator loss: {}".format(average_state_value_loss))


I looked at the implementation in the pytorch examples repo but they do things a little differently (like sharing policy parameters)
If anybody has any idea on how to fix the error I would greatly appreciate it
Johannes
PS: for reproducebility execute this first, then both code samples run:
import torch
from torch.autograd import Variable
import gym
import torch.nn as nn
import torch.nn.functional as F
import torch
from torch.autograd import Variable
from torch import autograd
from torch import optim

from collections import deque
import random
import numpy as np

class ReplayBuffer(object):

    def __init__(self, buffer_size):
        self.buffer_size = buffer_size
        self.count = 0
        self.buffer = deque()

    def add(self, s, a, r, d, s2):
        experience = (s, a, r, d, s2)
        if self.count < self.buffer_size:
            self.buffer.append(experience)
            self.count += 1
        else:
            self.buffer.popleft()
            self.buffer.append(experience)

    def size(self):
        return self.count

    def sample_batch(self, batch_size):
        '''
        batch_size specifies the number of experiences to add
        to the batch. If the replay buffer has less than batch_size
        elements, simply return all of the elements within the buffer.
        '''

        if self.count < batch_size:
            batch = random.sample(self.buffer, self.count)
        else:
            batch = random.sample(self.buffer, batch_size)

        s_batch = np.array([np.array(_[0]) for _ in batch])
        a_batch = np.array([_[1] for _ in batch])
        r_batch = np.array([_[2] for _ in batch])
        d_batch = np.array([_[3] for _ in batch])
        s2_batch = np.array([np.array(_[4]) for _ in batch])

        return s_batch, a_batch, r_batch, d_batch, s2_batch

    def clear(self):
        self.buffer.clear()
        self.count = 0