Asynchronous Advantage Actor-Critic (A3C) Implementation in PyTorch

This implementation builds Asynchronous Advantage Actor-Critic (A3C) from scratch. It includes a shared Actor-Critic network with a softmax policy head and a value head, multiple Worker processes that run independent CartPole environments, compute n-step returns and advantages, and sync gradients into a global model that lives in shared memory. A greedy test harness runs the trained policy in render_mode='human' for visual sanity checking.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.multiprocessing as mp
import torch.optim as optim
import numpy as np
import gym
from collections import deque
import time
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning, message=r".*np\.bool8.*")

torch.manual_seed(123)
np.random.seed(123)

# Hyperparameters
GAMMA = 0.99
ENTROPY_BETA = 0.01
LEARNING_RATE = 0.0001
MAX_EPISODES = 5000
MAX_STEPS = 200
NUM_PROCESSES = 4

class ActorCritic(nn.Module):
    def __init__(self, input_dim, n_actions):
        super(ActorCritic, self).__init__()
        self.shared = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU()
        )

        # Actor head
        self.actor = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, n_actions)
        )

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

    def forward(self, x):
        x = self.shared(x)

        policy_logits = self.actor(x)
        policy_probs = F.softmax(policy_logits, dim=-1)
        
        value = self.critic(x)
        
        return policy_probs, value
    
    def evaluate_actions(self, state):
        policy_probs, value = self.forward(state)
        dist = torch.distributions.Categorical(policy_probs)
        return dist, value
    
# Worker class for each parallel agent
class Worker(mp.Process):
    def __init__(self, global_model, optimizer, global_episode_counter, global_reward_queue, process_idx):
        super(Worker, self).__init__()
        self.process_idx = process_idx
        self.env = gym.make('CartPole-v1')
        self.input_dim = self.env.observation_space.shape[0]
        self.n_actions = self.env.action_space.n

        # Global model and optimizer
        self.global_model = global_model
        self.optimizer = optimizer

        # Local model for each worker
        self.local_model = ActorCritic(self.input_dim, self.n_actions)
        self.local_model.load_state_dict(self.global_model.state_dict())
        
        # Shared counters and queue
        self.global_episode_counter = global_episode_counter
        self.global_reward_queue = global_reward_queue

        # Initialize episode reward
        self.episode_reward = 0

    def run(self):
        """worker loop"""
        while self.global_episode_counter.value < MAX_EPISODES:
            # Reset gradients
            self.local_model.load_state_dict(self.global_model.state_dict())
            
            states = []
            actions = []
            rewards = []
            values = []

            # Reset episode
            state, _ = self.env.reset()
            state = torch.FloatTensor(state)
            done = False
            self.episode_reward = 0
            step = 0

            # Collect trajectory
            while not done and step < MAX_STEPS:
                policy_probs, value = self.local_model(state)
                dist = torch.distributions.Categorical(policy_probs)
                action = dist.sample()
                
                next_state, reward, terminated, truncated, _ = self.env.step(action.item())
                done = terminated or truncated
                
                states.append(state)
                actions.append(action)
                rewards.append(reward)
                values.append(value)
                
                state = torch.FloatTensor(next_state)
                self.episode_reward += reward
                step += 1

            # Compute returns and advantages
            R = torch.zeros(1, 1)
            if not done:
                _, value = self.local_model(state)
                R = value.detach()

            returns = []
            advantages = []
            
            for i in reversed(range(len(rewards))):
                R = rewards[i] + GAMMA * R
                advantage = R - values[i]
                returns.insert(0, R)
                advantages.insert(0, advantage)

            states = torch.stack(states)
            actions = torch.stack(actions)
            returns = torch.cat(returns).detach()
            advantages = torch.cat(advantages).detach()

            # update global network
            self.update_global_network(states, actions, returns, advantages)
            
            # Report results
            with self.global_episode_counter.get_lock():
                self.global_episode_counter.value += 1
            
            self.global_reward_queue.put(self.episode_reward)
            
            if self.process_idx == 0 and self.global_episode_counter.value % 10 == 0:
                print(f"Episode: {self.global_episode_counter.value}, Avg Reward: {np.mean(self.get_average_reward()):.2f}")

    def update_global_network(self, states, actions, returns, advantages):
        """Update the global network using gradients from the local network"""

        dists, values = self.local_model.evaluate_actions(states)

        # Calculate policy loss
        log_probs = dists.log_prob(actions)
        policy_loss = -(log_probs * advantages.detach()).mean()
        
        # Calculate value loss
        value_loss = F.mse_loss(values.squeeze(-1), returns.squeeze(-1).detach())
        
        # Calculate entropy 
        entropy = dists.entropy().mean()

        # Total loss
        total_loss = policy_loss + 0.5 * value_loss - ENTROPY_BETA * entropy
        
        self.optimizer.zero_grad()
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.local_model.parameters(), 40.0)
        
        # Sync the gradients to the global model
        for global_param, local_param in zip(self.global_model.parameters(), self.local_model.parameters()):
            if global_param.grad is None:
                global_param._grad = local_param.grad
            else:
                global_param._grad += local_param.grad

        # Update global model
        self.optimizer.step()

    def get_average_reward(self):
        """Get average reward from the global reward queue"""
        rewards = []
        while not self.global_reward_queue.empty():
            rewards.append(self.global_reward_queue.get())
        return rewards if rewards else [0]
    
# Training function
def train():
    # Initialize environment to get dimensions
    env = gym.make('CartPole-v1')
    input_dim = env.observation_space.shape[0]
    n_actions = env.action_space.n
    env.close()

    # Create global model
    global_model = ActorCritic(input_dim, n_actions)
    global_model.share_memory()  # Share model parameters across processes
    
    optimizer = optim.Adam(global_model.parameters(), lr=LEARNING_RATE)
    
    # Create shared counters and queue
    global_episode_counter = mp.Value('i', 0)
    global_reward_queue = mp.Queue()

    workers = [Worker(global_model, optimizer, global_episode_counter, global_reward_queue, i) 
               for i in range(NUM_PROCESSES)]
    
    for worker in workers:
        worker.start()
    
    for worker in workers:
        worker.join()
    
    return global_model

# Test the trained agent
def test(model, num_episodes=5):
    env = gym.make('CartPole-v1', render_mode='human')

    for episode in range(num_episodes):
        state, _ = env.reset()
        state = torch.FloatTensor(state)
        done = False
        total_reward = 0
        
        while not done:
            policy_probs, _ = model(state)
            action = torch.argmax(policy_probs).item()
            next_state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            state = torch.FloatTensor(next_state)
            total_reward += reward
        
        print(f"Test Episode: {episode+1}, Reward: {total_reward}")
    
    env.close()

if __name__ == "__main__":
    # Set method for starting processes
    mp.set_start_method('spawn')
    
    print("Starting A3C training...")
    start_time = time.time()
    trained_model = train()
    end_time = time.time()
    print(f"Training completed in {end_time - start_time:.2f} seconds")
    
    torch.save(trained_model.state_dict(), "a3c_model.pth")
    print("Model saved as a3c_model.pth")
    
    print("Testing the trained agent...")
    test(trained_model)