# What is the problem in my implementation of actor critic?

I have been implementing both REINFORCE with baseline and actor-critic to solve "cartpole-v1".

As a reminder, here is the presentation of the algorithms in Sutton and Barto's book (http://incompleteideas.net/book/RLbook2020.pdf):

Despite the codes being super similar, REINFORCE with baseline works well and actor-critic does not. I tried adding some entropy term without success.

Here is the code of REINFORCE:

#%%
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import gym

from tqdm.auto import trange

gamma = 0.99

class ActorCritic(nn.Module):
def __init__(self):
super().__init__()
self.hidden = nn.Linear(4, 32)
self.prob = nn.Linear(32, 2)
self.critic = nn.Linear(32, 1)

def forward(self, state):
hidden = F.relu(self.hidden(state))
return (
self.prob(hidden).flatten(),
self.critic(hidden).flatten(),
)

def select_action(prob):
m = torch.distributions.Categorical(logits=prob)
action_pt = m.sample()
return action_pt.numpy(), m.log_prob(action_pt)

def ewma(a, alpha=0.99):
ans = []
acc = a[0]
for x in a:
acc = acc * alpha + x * (1 - alpha)
ans.append(acc)
return ans

#%%
policy = ActorCritic()
eps = np.finfo(np.float32).eps.item()
# make cartpole environment
env = gym.make("CartPole-v1")

all_rewards = []

with trange(5000) as pbar:
for i_episode in pbar:
state, _ = env.reset()
rewards = []
log_probs = []
values = []

for t in range(500):  # Don't infinite loop while learning
state = torch.from_numpy(state).float()
prob, value = policy(state)
action, log_prob = select_action(prob)
log_probs.append(log_prob)
state, reward, done, trunc, _ = env.step(action)
if done or trunc:
break
rewards.append(np.array([reward]).astype(np.float32))
values.append(value)
avg_reward = np.sum(rewards, axis=0).mean()
all_rewards.append(avg_reward)
# format float to 2 decimal places and left align with 5 spaces
pbar.set_description(
f"Episode {i_episode + 1} reward: {ewma(all_rewards)[-1]:.2f}"
)
if ewma(all_rewards)[-1] > env.spec.reward_threshold:
print("Solved!")
break
R = 0
policy_loss = []
critic_loss = []
all_R = []
for r, log_prob, V in zip(
reversed(rewards), reversed(log_probs), reversed(values)
):
R = R * gamma + torch.from_numpy(r)
all_R.append(R)
A = R - V

policy_loss.append(torch.mean(-log_prob * A.detach()))
critic_loss.append(F.huber_loss(V, R))

loss = torch.stack(policy_loss).sum() + torch.stack(critic_loss).sum()
loss.backward()
optimizer.step()

# %%
import matplotlib.pyplot as plt

plt.plot(ewma(all_rewards, 0.99))


Here is what the code of actor critic looks like:

#%%
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import gym

from tqdm.auto import trange

gamma = 0.99

class ActorCritic(nn.Module):
def __init__(self):
super().__init__()
self.hidden = nn.Linear(4, 32)
self.prob = nn.Linear(32, 2)
self.critic = nn.Linear(32, 1)

def forward(self, state):
hidden = F.relu(self.hidden(state))
return (
self.prob(hidden).flatten(),
self.critic(hidden).flatten(),
)

def select_action(prob):
m = torch.distributions.Categorical(logits=prob)
action_pt = m.sample()
return action_pt.numpy(), m.log_prob(action_pt), m.entropy()

def ewma(a, alpha=0.99):
ans = []
acc = a[0]
for x in a:
acc = acc * alpha + x * (1 - alpha)
ans.append(acc)
return ans

#%%
policy = ActorCritic()
eps = np.finfo(np.float32).eps.item()
# make cartpole environment
env = gym.make("CartPole-v1")

all_rewards = []

with trange(5000) as pbar:
for i_episode in pbar:
state, _ = env.reset()
rewards = []
log_probs = []
entropies = []
values = []
for t in range(500):  # Don't infinite loop while learning
state = torch.from_numpy(state).float()
prob, value = policy(state)
action, log_prob, entropy = select_action(prob)
log_probs.append(log_prob)
entropies.append(entropy)
state, reward, done, trunc, _ = env.step(action)
if done or trunc:
break
rewards.append(np.array([reward]).astype(np.float32))
values.append(value)
avg_reward = np.sum(rewards, axis=0).mean()
all_rewards.append(avg_reward)
# format float to 2 decimal places and left align with 5 spaces
pbar.set_description(
f"Episode {i_episode + 1} reward: {ewma(all_rewards)[-1]:.2f}"
)
if ewma(all_rewards)[-1] > env.spec.reward_threshold:
print("Solved!")
break
policy_loss = []
critic_loss = []
next_V = torch.zeros_like(values[-1])
for r, log_prob, V in zip(
reversed(rewards), reversed(log_probs), reversed(values)
):
target_V = torch.from_numpy(r) + gamma * next_V
A = target_V - V

# trick to reduce variance
policy_loss.append(torch.mean(-log_prob * A.detach()))
critic_loss.append(F.huber_loss(V, target_V))
next_V = V

loss = (
torch.stack(policy_loss).sum()
+ torch.stack(critic_loss).sum()
- torch.stack(entropies).sum() * 0.01
)
loss.backward()
optimizer.step()

# %%
import matplotlib.pyplot as plt

plt.plot(ewma(all_rewards, 0.99))

# %%


Here is the output of the first code:

And the second code:

The only difference I see in the code is that in the first case, the target is R * gamma + r and in the second case, the target is next_V * gamma + r.

There is another possible issue, that next_V is not detached from the computation graph.

Therefore, I also tried this line: critic_loss.append(F.huber_loss(V, target_V.detach())).

But it also did not work.