RL Series — Ep 8

A3C and A2C. The eighth part of the RL episodes. If you made it here kudos⁵. Short Notes, not intended for complete beginners. You can actually skip this episode if you are in a rush and this just goes into more concreteness about what was learned in EP7.

Photo by Jon Tyson on Unsplash

A3C: Parallel policy updates

We saw Vanilla Policy Gradient in the last episode and VPG is a pretty robust method for simple problems. This Algo is largely unbiased becuase of Monte Carlo returns, which are complete actual returns experienced directly in the environment, but due to the Neural Network, some bias is injected. To further reduce that we propose asynchronous advantage actor-critic (A3C), which does 2 things

1. It uses n-step returns with bootstrapping to learn the policy and value function
2. It uses concurrent actors to generate a broad set of experience samples in parallel

Using Actor-Workers

Problem: One of the main sources of variance in DRL algorithms is how correlated and non-stationary online samples are. We can use Replay Buffer only in off-Policy methods, what will we do here?.

Solution: Have multiple workers generating experience in parallel and asynchronously updating the policy and value function. Having multiple workers generating experience on multiple instances of the environment in parallel decorates the data used for training and reduces the variance of the algorithm.

Monte Carlo (full return unbiased returns) needs the full trajectory and not batches from the Replay buffer (and hence on policy). For independent or TD Lambda we do not need the full thing and hence replay buffer can be used (hence off policy). This is one way to remember if it’s an on or off-policy method. (Take it with a pinch of salt as this statement might be false).

Grokking ch11

Till this point, it is just like having multiple agents do VPG. VPG used Monte Carlo, do these agents also use Monte Carlo/full trajectory updates ?Will see in the next section.

Using n-step estimates

We go out for n-steps collecting rewards, and then bootstrap after that nth state, or before if we land on a terminal state, whichever comes first. A3C takes advantage of the lower variance of n-step returns when compared to Monte Carlo returns

Grokking ch11

Non-blocking model updates. One of the most critical aspects of A3C is that its network updates are asynchronous and lock free, you can read and dig up more.

The final aspect is in how the global model is updated. It is done by accumulating local gradient updates.

At the core it is just parallel programming basics with worker-agent and uses TD updates instead of Monte Carlo.

[embed]

Skipping the Code section in this, might come back later and add a snippet. Code from viewers is also encouraged :’)

A2C: Synchronous policy updates

This was proposed after A3C and showed to perform comparably to A3C. The main features are.

Weight-sharing model

One change to our current algorithm is to use a single neural network for both the policy and the value function.

Good thing — less computationally expensive.

Bad thing — Policy and Value are in different scales, hence could result in some problems learning.

Restoring order in policy updates

We do not have multiple agents, but multiple environemnts and 1 agent. so instead of having multiple actor-learners (as in A3C), we have multiple actors with a single learner.

Grokking ch11

What is the Loss?

Combined Loss = Policy Loss (Policy Gradient Loss)+ Value Loss (return — value)² (return via Monte Carlo)

Lets code this up. A more rigorous version is present in Grokking. Lets do a simpler version.

import torch
import torch.nn as nn
import torch.optim as optim
import gym
import numpy as np

# # Shared Network
class ActorCritic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(ActorCritic, self).__init__()

        # Shared layers (weight-sharing)
        self.shared_layers = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU()
        )

        # Actor head
        self.actor = nn.Sequential(
            nn.Linear(64, action_dim),
            nn.Softmax(dim=-1)
        )

        # Critic head
        self.critic = nn.Sequential(
            nn.Linear(64, 1)
        )

    def forward(self, x):
        shared_out = self.shared_layers(x)
        actor_probs = self.actor(shared_out)
        value_estimate = self.critic(shared_out)
        return actor_probs, value_estimate

# Hyperparameters
learning_rate = 0.001
gamma = 0.99
num_environments = 4
num_steps = 5
num_episodes = 1000

# Create multiple environments
env_name = 'CartPole-v1'
envs = [gym.make(env_name) for _ in range(num_environments)]
state_dim = envs[0].observation_space.shape[0]
action_dim = envs[0].action_space.n

# Initialize the actor-critic network
actor_critic = ActorCritic(state_dim, action_dim)
optimizer = optim.Adam(actor_critic.parameters(), lr=learning_rate)

# Training loop
for episode in range(num_episodes):
    # Initialize episode variables
    episode_states = []
    episode_actions = []
    episode_rewards = []
    episode_advantages = []

    for _ in range(num_steps):
        state_batch = []

        for env in envs:
            state = env.reset()
            state_batch.append(state)

        state_batch = torch.tensor(state_batch, dtype=torch.float32)
        actor_probs, value_estimates = actor_critic(state_batch)

        # Sample actions from the actor's distribution for all environments
        actions = torch.distributions.Categorical(actor_probs).sample().numpy()

        # Take actions in all environments in parallel
        next_states, rewards, dones, _ = zip(*[env.step(action) for env, action in zip(envs, actions)])

        episode_states.extend(state_batch)
        episode_actions.extend(actions)
        episode_rewards.extend(rewards)
        episode_advantages.append(advantage)  # Store advantage for this time step

        if any(dones):
            # When any environment episode ends, update state for the new episode.
            state = [next_state if done else s for s, next_state, done in zip(state, next_states, dones)]

    # Compute returns for this episode
    R = 0
    returns = []

    for r in reversed(episode_rewards):
        R = r + gamma * R
        returns.insert(0, R)

    # Compute actor and critic losses
    action_log_probs = torch.log(actor_probs)
    actor_loss = -torch.mean(action_log_probs * torch.tensor(episode_advantages, dtype=torch.float32))
    critic_loss = nn.MSELoss()(value_estimates, torch.tensor(returns, dtype=torch.float32))

    # Backpropagation and optimization
    optimizer.zero_grad()
    total_loss = actor_loss + critic_loss
    total_loss.backward()
    optimizer.step()

    # Logging
    total_reward = sum(episode_rewards)
    print(f"Episode {episode}: Total Reward: {total_reward}")

# Close environments
for env in envs:
    env.close()

Take this code with a pinch of salt. The environment thing can be skipped too. Want this in comic form — read this.

We will look into DDPG, SAC, PPO etc soon.

What next?

Find the Full Series here!! Do comment clap and give a follow as it cheers me up. Improvements are welcome.

References

1. Grokking reinforcement learning book (mainly ch11)
2. Reinforcement Learning: An Introduction. Richard and Andrew
3. David Silverman RL Series (mainly Lecture 7)
4. Jerry John Thomas Github
5. https://medium.com/hackernoon/intuitive-rl-intro-to-advantage-actor-critic-a2c-4ff545978752