Advantage Actor-Critic (A2C)

[Core Actor-Critic concepts]

• Combine policy-based and value-based methods
  • Actor (policy determination): Parameterize policy $\pi_\theta(a|s)$ and output an action probability distribution from a state.
  • Critic (value estimation): Parameterize value function $V_\psi(s)$ and estimate the current state's value (expected return).
• Advantages
  • Real-time response: updates are possible before an episode ends
    • No episode-level return calculation
    • Immediate response to environmental changes > reduce REINFORCE's high variance problem

[Advantage Actor-Critic (A2C) algorithm]

• Use the advantage function
  • $A(s, a) = Q(s, a) - V(s)$: additional gain from choosing an action in a particular state
  • The critic estimates $V(s)$, while the actor learns to maximize $A(s, a)$
• Update methods
  • Actor: $\theta \leftarrow \theta + \alpha \nabla_\theta \log \pi_\theta(a|s) A(s, a)$
    • Update the policy using the policy gradient.\
    • $\alpha$: Learning rate
  • Critic: $\psi \leftarrow \psi + \beta (r + \gamma V(s') - V(s)) \nabla_\psi V(s)$
    • Update the value function to minimize TD (temporal difference) error.
    • $\beta$: Learning rate
    • $r$: Reward
    • $\gamma$: Discount factor
    • $s'$: Next state
• Benefits of A2C
  • Stable training: The critic estimates values, and the actor updates the policy based on the critic's estimates, allowing more stable training than REINFORCE.
  • Efficient exploration: The advantage function identifies actions that can produce better-than-average outcomes, enabling efficient exploration.

[A2C implementation]

• Design the actor network and critic network
  • Actor network: $s \rightarrow \boxed{\text{FC} + \text{ReLU}} \rightarrow \boxed{\text{FC} + \text{Softmax}} \rightarrow \pi(a|s)$
    • Take a state as input and output a probability distribution over actions.
    • Same architecture as the REINFORCE policy network
  • Critic network: $s \rightarrow \boxed{\text{FC} + \text{ReLU}} \rightarrow \boxed{\text{FC}} \rightarrow V(s)$
    • Take a state as input and output its value.
    • It can use an MLP like DQN, but the critic outputs a scalar $V(s)$ for each state, unlike DQN's vector of Q-values for individual actions.
• Collect data through interaction with the environment
• Use the collected data to update the actor and critic networks.

[A2C limitations and directions for improvement]

• On-policy learning: data are collected only under the current policy, leading to low sample efficiency
• Directions for improvement
  • Proximal Policy Optimization (PPO) (Week 9)
    • Limit the difference between new and old policies during policy updates for stable training.
  • Soft Actor-Critic SAC (Week 10)
    • Introduce entropy to encourage exploration and enable effective learning in continuous action spaces.

[Activity: CartPole with Advantage Actor-Critic (A2C)]

1. Import libraries and set hyperparameters

• Import the required libraries (PyTorch, Gym, NumPy, etc.).
• Set the hyperparameters for A2C.
  • gamma: Discount factor (importance of future rewards)
  • learning_rate: Learning rate
  • hidden_size: Hidden layer sizes (choose a 2-layer/3-layer neural network)
  • num_steps: Number of steps per rollout
    • Rollout: Segment used to compute advantages
  • num_episodes: Total number of training episodes
  • check_interval: Training progress reporting interval (in episodes)
  • early_stop_threshold: Early-stopping condition (based on mean reward)
  • early_stop_patience: Number of times the early-stopping condition must be maintained
  • render_gif_path: File path for the rendered GIF

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical
import gymnasium as gym
import numpy as np
import imageio

# Hyperparameters
gamma = 0.99
learning_rate = 0.003
hidden_size = 128  # Default (2-layer neural network)
num_steps = 10
num_episodes = 2000
progress_interval = 50
early_stop_threshold = 400 # CartPole-v1: Generally classified as successful at 200 or higher
early_stop_patience = 10
render_gif_path = "cartpole_simulation_a2c.gif"

2. Define the A2C neural network model (actor–critic)

• Define the network architecture in the __init__ method.
• The use_three_layers argument selects a 2-layer or 3-layer neural network.
• Define fc1, fc2 (optional fc_h), fc_pi, and fc_v layers
• fc_pi outputs the probability distribution representing the policy
• fc_v outputs the state value.
• The forward method takes input and returns the policy and state value.

class ActorCritic(nn.Module):
    def __init__(self, input_size, num_actions, hidden_size=128, use_three_layers=False):
        super(ActorCritic, self).__init__()
        self.use_three_layers = use_three_layers
        self.fc1 = nn.Linear(input_size, hidden_size)
        if self.use_three_layers:
            self.fc_h = nn.Linear(hidden_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc_pi = nn.Linear(hidden_size, num_actions)
        self.fc_v = nn.Linear(hidden_size, 1)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        if self.use_three_layers:
            x = F.relu(self.fc_h(x))
        x = F.relu(self.fc2(x))
        policy = F.softmax(self.fc_pi(x), dim=0)
        value = self.fc_v(x)
        return policy, value

3. Define the training loop function

Policy/value prediction and action selection

• state_tensor = torch.from_numpy(state).float()
  • Convert the current state (a NumPy array) to a PyTorch tensor.
• policy, value = model(state_tensor)
  • Feed the current state into the model to obtain the policy and state value.
• action_dist = Categorical(policy)
  • Create a Categorical object to sample an action from the policy (probability distribution)
• action = action_dist.sample()
  • Sample an action from the Categorical object

Apply the selected action to interact with the environment

• next_state, reward, done, _ = env.step(action.item())
  • Apply the selected action to obtain the next state, reward, termination status, and additional information.
  • action.item() converts a tensor action to an integer.
• total_reward += reward: calculate cumulative reward during the episode

Store experiences

• memory.append((state_tensor, action, reward, value, done))
  • Store the current state, selected action, received reward, state value, and termination status as a tuple in list memory
  • These data are used later to calculate advantages and losses

Compute advantages and losses from rollout data

• loss = 0: initialize the loss
• _, last_value = model(torch.from_numpy(next_state).float())
  • Calculate the value function at the rollout's final state (next state)
  • Use this value to calculate returns
• returns = last_value.detach()
  • Initialize Returns with the value function estimate of the final state.
  • .detach() prevents this value from participating in gradient computation.
• for state_tensor, action, reward, value, done in reversed(memory)
  • Traverse the experiences stored in the memory list in reverse order.
    • Start at the end of the rollout and work back to the beginning.
  • returns = reward + gamma * returns * (1 - done)
    • Returns is the sum of the current reward and discounted future rewards.
    • If the episode has ended (done=1), Returns contains only the current reward.
  • advantage = returns - value
    • Advantage is the difference between Returns and the current state's value function estimate.
    • Advantage indicates how much better or worse a particular action is than an average action.
  • loss += ~: Compute the loss function.
    • -Categorical(model(state_tensor)[0]).log_prob(action) * advantage.detach()
      • Actor (policy) loss
      • Uses the policy gradient method
      • Higher advantage increases the probability of choosing that action; lower advantage decreases it.
      • .detach() prevents Advantage from participating in gradient computation.
    • F.smooth_l1_loss(value, returns.detach())
      • Critic (value function) loss
      • Use Huber loss to reduce the difference between value function predictions (value) and actual Returns.
      • .detach() prevents Returns from participating in gradient computation.

Optimization:

• optimizer.zero_grad()
  • Reset previous gradients stored by the optimizer to 0.
• loss.backward()
  • Backpropagate the loss function to calculate gradients for each model parameter
• optimizer.step()
  • Update model parameters using the calculated gradients (gradient descent)
• memory = []
  • Clear memory for the next rollout

def train(model, optimizer, env, num_episodes, num_steps, gamma, check_interval, early_stop_threshold, early_stop_patience):
    scores = []
    for episode in range(num_episodes):
        state, _ = env.reset()
        done = False
        total_reward = 0
        memory = []

        while not done:
            for t in range(num_steps):
                state_tensor = torch.from_numpy(state).float()
                policy, value = model(state_tensor)
                action_dist = Categorical(policy)
                action = action_dist.sample()

                next_state, reward, terminated, truncated, _ = env.step(action.item())
                total_reward += reward
                done = terminated or truncated

                memory.append((state_tensor, action, reward, value, done))
                state = next_state

                if done:
                    break

            # Compute advantages and the loss function from rollout data
            loss = 0
            _, last_value = model(torch.from_numpy(next_state).float())
            returns = last_value.detach()

            for state_tensor, action, reward, value, done in reversed(memory):
                returns = reward + gamma * returns * (1 - done)
                advantage = returns - value
                loss += -Categorical(model(state_tensor)[0]).log_prob(action) * advantage.detach() + \
                        F.smooth_l1_loss(value, returns.detach()) # Huber loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            memory = [] # Clear memory

        scores.append(total_reward)

        # Check progress and early stopping
        if (episode+1) % check_interval == 0:
            avg_score = np.mean(scores[-check_interval:])
            print(f"Episode: {episode+1}, Average Reward: {avg_score:.2f}")

        if np.min(scores[-early_stop_patience:]) >= early_stop_threshold:
            print(f"Early stopping triggered! at Episode: {episode+1}")
            break

4. Run training and check results

• Create a Gym environment. (CartPole-v1)
• Create an instance of ActorCritic (choose 2 or 3 layers).
• Create a optim.Adam optimizer
• Call train to start training

# Create the environment
env = gym.make('CartPole-v1', render_mode = 'rgb_array')
input_size = env.observation_space.shape[0]
num_actions = env.action_space.n

# Create the model (example using a 3-layer neural network)
model = ActorCritic(input_size, num_actions, hidden_size, use_three_layers=True)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Run training
train(model, optimizer, env, num_episodes, num_steps, gamma, progress_interval, early_stop_threshold, early_stop_patience)

env.close()

Episode: 50, Average Reward: 16.78
Episode: 100, Average Reward: 24.02
Episode: 150, Average Reward: 18.66
Episode: 200, Average Reward: 37.02
Episode: 250, Average Reward: 39.84
Episode: 300, Average Reward: 31.26
Episode: 350, Average Reward: 55.00
Episode: 400, Average Reward: 104.56
Early stopping triggered! at Episode: 422

5. Test the trained model and render a GIF

• Run one episode using the trained model
• Collect frames at each step
• Generate a GIF using the imageio library
• Print the total reward obtained during the episode

state, _ = env.reset()
done = False
total_reward = 0
frames = []

with torch.no_grad():  # Disable gradient computation (inference mode)
    while not done:
        state_tensor = torch.from_numpy(state).float()
        policy, _ = model(state_tensor)
        action_dist = Categorical(policy)
        action = action_dist.sample()

        next_state, reward, terminated, truncated, _ = env.step(action.item())
        total_reward += reward
        state = next_state
        done = terminated or truncated

        # Append the rendered image to the frame list
        frames.append(env.render())

# Save as a GIF file
imageio.mimsave(render_gif_path, frames, duration=33)  # Adjust fps as needed: duration = 1000/fps
print(f"Test Episode Reward: {total_reward}")
env.close()

Test Episode Reward: 500.0