0
$\begingroup$

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): enter image description here enter image description here

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()
optimizer = optim.Adam(policy.parameters(), lr=1e-2)
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))

        optimizer.zero_grad()
        loss = torch.stack(policy_loss).sum() + torch.stack(critic_loss).sum()
        nn.utils.clip_grad_norm_(policy.parameters(), 100)
        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()
optimizer = optim.Adam(policy.parameters(), lr=1e-2)
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

        optimizer.zero_grad()
        loss = (
            torch.stack(policy_loss).sum()
            + torch.stack(critic_loss).sum()
            - torch.stack(entropies).sum() * 0.01
        )
        nn.utils.clip_grad_norm_(policy.parameters(), 100)
        loss.backward()
        optimizer.step()

# %%
import matplotlib.pyplot as plt

plt.plot(ewma(all_rewards, 0.99))

# %%

Here is the output of the first code:

enter image description here

And the second code:

enter image description here

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.

$\endgroup$

1 Answer 1

1
$\begingroup$

My conclusion so far is that the Actor-Critic algorithm described by Sutton and Barto does not work.

First, all implementations I have seen are actually REINFORE with baseline, like https://github.com/pytorch/examples/blob/main/reinforcement_learning/actor_critic.py and https://github.com/nikhilbarhate99/Actor-Critic-PyTorch/blob/01c833e83006be5762151a29f0719cc9c03c204d/model.py#L33

Second, GAE-Lambda is a method generalizing the two: enter image description here

I tried SpinningUp's Vanilla Policy Gradients with the lambda parameter set to 0, thus reproducing Sutton and Barto's Actor Critic formula. It did not learn.

$\endgroup$

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .