Implementing A3C Algorithm with Tensorflow 2.x
Understanding Actor-Critic Models: A2C and the Significance of A3C
Reinforcement learning (RL) is a powerful paradigm in machine learning where an agent learns to make decisions by interacting with an environment. Among the many RL algorithms, actor-critic models stand out for their balance of stability and efficiency. This blog shares my experience of implementing A3C algorithm to solve the Lunarlander-v3 problem.
What Are Actor-Critic Models?
Actor-critic models combine two key components: the actor, which selects actions based on a policy, and the critic, which evaluates the chosen actions by estimating the value function. The actor follows a policy (a probability distribution over actions) to decide what action to take in a given state, while the critic assesses the action’s quality by estimating the expected future rewards (the value or Q-value). This synergy allows the actor to improve its policy based on the critic’s feedback, making actor-critic methods a hybrid of value-based and policy-based RL approaches.
The actor-critic framework addresses limitations of standalone policy gradient methods (like REINFORCE), which can suffer from high variance, and value-based methods (like Q-learning), which can be computationally expensive. By leveraging the critic’s value estimates, the actor reduces variance in policy updates, leading to more stable learning.
Advantage Actor-Critic (A2C)
The Advantage Actor-Critic (A2C) algorithm is a synchronous version of the actor-critic framework. It builds on the basic actor-critic model by introducing the advantage function, which measures how much better an action is compared to the average action in a given state. The advantage function is defined as:
[ A(s, a) = Q(s, a) — V(s) ]
Here, ( Q(s, a) ) is the action-value function (expected reward for taking action ( a ) in state ( s )), and ( V(s) ) is the state-value function (expected reward from state ( s )). By using the advantage, A2C reduces variance in policy gradients, leading to more efficient learning.
In A2C, multiple environments are run in parallel to collect experiences synchronously. The actor updates the policy using the advantage estimates provided by the critic, and the critic updates its value estimates to minimize the error between predicted and actual rewards. This synchronous approach ensures stable updates, as all environments contribute to the same gradient computation.
The Significance of A3C
The Asynchronous Advantage Actor-Critic (A3C) algorithm, introduced by Google DeepMind, takes A2C a step further by introducing asynchronous updates. Instead of waiting for all environments to synchronize, A3C runs multiple agents in parallel, each with its own copy of the environment and model parameters. These agents independently collect experiences and update a shared global model asynchronously.
A3C’s asynchronous nature offers several key benefits:
- Improved Exploration: Each agent explores the environment independently, increasing the diversity of experiences and reducing the risk of getting stuck in suboptimal policies.
- Efficiency: Asynchronous updates allow A3C to utilize computational resources more effectively, as agents don’t wait for each other, leading to faster training.
- Stability: The uncorrelated experiences from multiple agents act as a form of noise reduction, stabilizing gradient updates without requiring experience replay buffers (used in methods like DQN).
A3C’s ability to scale across multiple CPU cores makes it particularly effective for complex tasks, such as playing Atari games or controlling robots. Its asynchronous updates also reduce the need for extensive hyperparameter tuning, making it more robust across different environments. A pictorial block diagram of A3C model is shown below in Figure 1.
Python Implementation of A3C Model
In the rest of this blog, I will discuss how to implement the A3C algorithm using Python 3.x and Tensorflow 2.x to solve the LunarLander-v3 problem.
I am assuming that you already have a working version of your A2C algorithm. We will use the implementation available at this link for this blog. Details of implementing A2C algorithm is discussed in this blog.
Get Swagat K’s stories in your inbox
Join Medium for free to get updates from this writer.
We will use the Python’s multiprocessing module to run multiple workers in parallel processes. So, we need two functions — one called the “worker” function and the other “run_workers” functions that runs multiple workers. To start with, we load the necessary modules. Also, notice that we are disabling the GPU at the top.
import gymnasium as gym
import numpy as np
import multiprocessing
import time
import wandb
import os
import gc
import queue
from a2c import A2CAgent
from collections import deque
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # Disable GPU for this script
import tensorflow as tf
import tensorflow_probability as tfp
DEBUG = FalseThe worker function
The worker function creates a local A2CAgent referred to as local_network using the class definition provided in the previous section. It also creates its own gym environment for interaction. During each episode, the local_network starts by copying the weights from the global_network and uses these parameters to collects experience tuples (s, a, r, s′ ). It
then computes gradients at the end of the each episode by using this batch of experiences. The computed gradients are pushed to the gradients_queue which can then be used for updating parameters of the global_network. The parameters between multiple parallel processes are shared by using queues. Several steps are taken to ensure that the queues are not empty leading to deadlock situations. This includes pushing more data into the
queue and taking appropriate steps if the queue is empty or full etc. It is also necessary to clear the memory after each thread to prevent running out of memory on the computer. This code uses Python’s multiprocessing library to run the code on a multi-core CPU and does not require GPUs for execution.
def worker(worker_id, global_weights_queue, gradients_queue,
save_request_queue, env_id, create_actor_func=None,
create_critic_func=None, max_episodes=1500,
max_score = 500, min_score = -500,
max_steps = None, wandb_log=False):
# Set random seed for reproducibility in each process
tf.random.set_seed(worker_id + 1)
np.random.seed(worker_id + 1)
# create environment
env = gym.make(env_id)
obs_shape = env.observation_space.shape
action_size = env.action_space.n
# Create local network and environment
if create_actor_func is not None:
actor = create_actor_func(obs_shape, action_size)
else:
actor = None
if create_critic_func is not None:
critic = create_critic_func(obs_shape)
else:
critic = None
local_network = A2CAgent(obs_shape, action_size,
a_model=actor, c_model=critic)
# collect experience and compute gradients
# for the local network
episode = 0
ep_scores = []
best_score = -np.inf
for episode in range(max_episodes):
weights_updated = False
for _ in range(3): # Retry up to 3 times
try:
global_weights = global_weights_queue.get_nowait()
local_network.set_weights(*global_weights)
weights_updated = True
if DEBUG:
print(f"Worker {worker_id}: Retrieved global weights from queue")
break
except queue.Empty:
if DEBUG:
print(f"Worker {worker_id}: Global weights queue empty, retrying...")
time.sleep(0.1) # Brief pause before retry
if not weights_updated:
if DEBUG:
print(f"Worker {worker_id}: Using current weights after failed attempts")
# Continue with current weights if queue is empty
continue
state = env.reset()[0]
episode_reward = 0
done = False
step = 0
states = deque(maxlen=max_steps if max_steps is not None else 1000)
actions = deque(maxlen=max_steps if max_steps is not None else 1000)
rewards = deque(maxlen=max_steps if max_steps is not None else 1000)
# Collect trajectory
while not done:
action = int(local_network.policy(state))
next_state, reward, done, truncated, _ = env.step(action)
done = done or truncated
states.append(state) # type: ignore
actions.append(action)
rewards.append(reward)
state = next_state
episode_reward += reward
step += 1
if max_score is not None and episode_reward >= max_score:
done = True
if min_score is not None and episode_reward <= min_score:
done = True
if max_steps is not None and step >= max_steps:
done = True
# end of episode
ep_scores.append(episode_reward)
# update the local network
a_loss, c_loss, actor_grads, critic_grads = local_network.compute_gradients(
states, actions, rewards
)
# Update global network
try:
gradients_queue.put_nowait((actor_grads, critic_grads))
except queue.Full:
if DEBUG:
print(f"Worker {worker_id}: Gradients queue is full, skipping update.")
continue
if episode_reward > best_score:
best_score = episode_reward
try:
save_request_queue.put_nowait((worker_id, episode, episode_reward))
except queue.Full:
if DEBUG:
print(f"Worker {worker_id}: Save request queue is full, skipping save request.")
continue
# free memory
tf.keras.backend.clear_session() # Clear TensorFlow session
states.clear()
actions.clear()
rewards.clear()
gc.collect() # Collect garbage to free memory
# end of for-loop
env.close()
Running multiple workers in parallel
The code for running multiple workers on parallel threads is provided in the listing below. The function run_workers() uses Python’s multiprocessing module to achieve parallel execution of worker function. It creates a global A2CAgent referred to as global_network whose parameters are shared with the local network of each worker through the multiprocessing queue called global_weights_queue. The gradients
computed by the local networks are received by the global network by using another multiprocessing queue called gradients_queue. The gradients received are averaged and then applied to
update the parameters of the global_network.
def run_workers(env_name, max_num_workers=5, max_episodes=1500,
wandb_log=True, max_score=500, min_score=-200,
max_steps=1000,
create_actor_func=None,
create_critic_func=None):
# Set random seed for reproducibility
tf.random.set_seed(42)
np.random.seed(42)
N_WORKERS = min(multiprocessing.cpu_count(), max_num_workers)
print('N_WORKERS:', N_WORKERS)
# Initialize environment to get state and action sizes
env = gym.make(env_name)
# check if the environment has a discrete action space
if not isinstance(env.action_space, gym.spaces.Discrete):
raise ValueError("A3C Agent currently supports discrete action spaces only")
obs_shape = env.observation_space.shape
action_size = env.action_space.n
env_id = env.spec.id
env.close()
if create_actor_func is not None:
a_model = create_actor_func(obs_shape, action_size)
else:
a_model = None
if create_critic_func is not None:
c_model = create_critic_func(obs_shape)
else:
c_model = None
# Initialize global network
global_network = A2CAgent(obs_shape, action_size,
a_model=a_model, c_model=c_model)
# Create a queue to share weights between processes
manager = multiprocessing.Manager()
global_weights_queue = manager.Queue(maxsize=2*N_WORKERS)
gradients_queue = manager.Queue(maxsize=2*N_WORKERS)
save_request_queue = manager.Queue(maxsize=N_WORKERS)
save_lock = manager.Lock()
# initial weights for the global network
for _ in range(2*N_WORKERS):
# Initialize the global weights queue with the global network's weights
global_weights_queue.put(global_network.get_weights())
# Create and start worker processes
processes = []
for i in range(N_WORKERS):
p = multiprocessing.Process(
target=worker,
args=(i, global_weights_queue, gradients_queue,
save_request_queue, env_id,
create_actor_func, create_critic_func,
max_episodes,
max_score, min_score, max_steps, wandb_log)
)
p.start()
processes.append(p)
print(f'Started worker {p.name}')
try:
last_refill_time = time.time()
best_reward = -np.inf
while(any(p.is_alive() for p in processes)):
gradients = []
for _ in range(N_WORKERS):
if not gradients_queue.empty():
gradients.append(gradients_queue.get())
if gradients: # process gradients if available
# Filter out None gradients
valid_gradients = [
(actor_grads, critic_grads)
for actor_grads, critic_grads in gradients
if actor_grads is not None and critic_grads is not None
]
if valid_gradients: # Only proceed if there are valid gradients
actor_grads_avg = []
critic_grads_avg = []
for actor_grads, critic_grads in valid_gradients:
if not actor_grads_avg:
actor_grads_avg = [tf.convert_to_tensor(g) for g in actor_grads]
critic_grads_avg = [tf.convert_to_tensor(g) for g in critic_grads]
else:
for i, g in enumerate(actor_grads):
actor_grads_avg[i] = tf.add(actor_grads_avg[i], tf.convert_to_tensor(g))
for i, g in enumerate(critic_grads):
critic_grads_avg[i] = tf.add(critic_grads_avg[i], tf.convert_to_tensor(g))
n = len(valid_gradients)
actor_grads_avg = [tf.math.truediv(g, n) for g in actor_grads_avg]
critic_grads_avg = [tf.math.truediv(g, n) for g in critic_grads_avg]
global_network.apply_gradients(actor_grads_avg, critic_grads_avg)
# Synchronize global weights with all workers
updated_weights = global_network.get_weights()
for _ in range(N_WORKERS):
try:
# Synchronize global weights with all workers
global_weights_queue.put_nowait(updated_weights)
except queue.Full:
if DEBUG:
print("Global weights queue is full, skipping synchronization.")
while not global_weights_queue.empty():
try:
global_weights_queue.get_nowait()
except queue.Empty:
break
global_weights_queue.put_nowait(updated_weights)
# periodically refill the global weights queue
if time.time() - last_refill_time > 5:
try:
global_weights_queue.put_nowait(global_network.get_weights())
last_refill_time = time.time()
except queue.Full:
if DEBUG:
print("Global weights queue is full, skipping refill.")
# saving weights
while not save_request_queue.empty():
try:
worker_id, episode, episode_reward = save_request_queue.get_nowait()
print(f"Worker: {worker_id}, episode: {episode}, reward: {episode_reward:.2f}")
if episode_reward > best_reward:
best_reward = episode_reward
with save_lock:
try:
global_network.save_weights(
actor_wt_file='a3c_actor.weights.h5',
critic_wt_file='a3c_critic.weights.h5',
)
print(f"Best Score: {best_reward}. Saved weights!")
except Exception as e:
print(f"Error saving weights: {e}")
except queue.Empty:
break
except Exception as e:
print(f"An error occurred: {e}")
finally:
print("Shutting down workers...")
#Wait for all processes to complete
for p in processes:
p.join()
print(f'Worker {p.name} has finished.')
Separating gradient computation from parameter update
The A2CAgent class should provide additional methods to facilitate working with multiple workers as shown below. This primarily includes separating the gradient computation from the parameter update steps. The method compute_gradients() is used by the local_network
agent of each worker to compute gradients. The method apply_gradients() is used by the global_network to update its parameters. The methods get_weights() and set_weights() allow copying of parameters between the global network and local networks. Gradient clipping is applied to avoid extreme values leading to NaN values.
def compute_gradients(self, states, actions, rewards):
states, actions, discnt_rewards = self.compute_discounted_rewards(states, actions, rewards)
with tf.GradientTape() as tape1, tf.GradientTape() as tape2:
values = self.critic(states)
td_error = tf.math.subtract(discnt_rewards, values)
#actor_loss = self.actor.compute_actor_loss(states, actions, td_error)
actor_loss = self.compute_actor_loss(states, actions, td_error)
critic_loss = tf.reduce_mean(tf.square(td_error))
# compute gradients
actor_grads = tape1.gradient(actor_loss, self.actor.model.trainable_variables)
critic_grads = tape2.gradient(critic_loss, self.critic.model.trainable_variables)
# clip gradients
actor_grads = [tf.clip_by_norm(grad, self.grad_clip_norm) \
if grad is not None else None for grad in actor_grads ]
critic_grads = [tf.clip_by_norm(grad, self.grad_clip_norm) \
if grad is not None else None for grad in critic_grads]
# Check for NaN gradients
actor_has_nan = any(tf.reduce_any(tf.math.is_nan(grad)) for grad in actor_grads if grad is not None)
critic_has_nan = any(tf.reduce_any(tf.math.is_nan(grad)) for grad in critic_grads if grad is not None)
if actor_has_nan or critic_has_nan:
print(f"NaN gradients detected! "
f"Actor NaN: {actor_has_nan}, Critic NaN: {critic_has_nan}, "
f"Actor Loss: {actor_loss:.4f}, Critic Loss: {critic_loss:.4f}")
# Skip sending gradients to avoid corrupting global network
return None, None, None, None
return actor_loss.numpy(), critic_loss.numpy(), actor_grads, critic_grads
def apply_gradients(self, actor_grads, critic_grads):
self.actor.optimizer.apply_gradients(zip(actor_grads, self.actor.model.trainable_variables))
self.critic.optimizer.apply_gradients(zip(critic_grads, self.critic.model.trainable_variables))
def get_weights(self):
return self.actor.model.get_weights(), self.critic.model.get_weights()
def set_weights(self, actor_weights, critic_weights):
self.actor.model.set_weights(actor_weights)
self.critic.model.set_weights(critic_weights)Steps to improve implementation performance
Running multiple workers can take up a lot of resources. Some steps taken to improve the implementation performance are as follows:
- Instead of using open lists for storing states, actions, rewards, you can use
dequeswith fixed sizes. - Release allocated memories after the end of each episode.
- Adequate steps should be taken to ensure that the multiprocessing queues are not empty when accessed by multiple processes in quick succession. This would lead to a deadlock situations where processes wait indefinitely to receive data. This is achieved by having a larger queue size than the number of workers and pushing data into the queues more frequently.
- Reducing the number of workers if required.
Solving LunarLander-V3 Problem using A3C Algorithm
We use a 3 layer dense network for both actor and critic models with 128 neurons in the hidden layers. We run 20 parallel workers in this case. The resulting performance is shown below in Figure 2 and 3 respectively. As we can see that the A3C clearly provides superior performance compared to a single-worker A2C agent. avg100score is the average of last 100 episodes. The A3C_v1 uses the difference between discounted return and the current value function V(s) as the estimate for the advantage function. While the model A3C_v2 uses the standard TD error involving current state s and next state s’ to estimate the advantage function. We will discuss this in more detail in a separate post about A2C algorithm.
from a3c import run_workers
import multiprocessing
import tensorflow as tf
# create actor & Critic models
def create_actor_model(obs_shape, n_actions):
s_input = tf.keras.layers.Input(shape=obs_shape)
x = tf.keras.layers.Dense(128, activation='relu')(s_input)
x = tf.keras.layers.Dense(128, activation='relu')(x)
a = tf.keras.layers.Dense(n_actions, activation='softmax')(x)
model = tf.keras.models.Model(s_input, a, name='actor_network')
model.summary()
return model
def create_critic_model(obs_shape):
s_input = tf.keras.layers.Input(shape=obs_shape)
x = tf.keras.layers.Dense(128, activation='relu')(s_input)
x = tf.keras.layers.Dense(128, activation='relu')(x)
v = tf.keras.layers.Dense(1, activation=None)(x)
model = tf.keras.models.Model(s_input, v, name='critic_network')
model.summary()
return model
if __name__ == '__main__':
multiprocessing.set_start_method('spawn')
run_workers(
env_name='LunarLander-v3',
max_num_workers=20,
max_episodes=1500,
wandb_log=True,
max_score=500,
min_score=-200,
max_steps=1000,
create_actor_func=create_actor_model,
create_critic_func=create_critic_model,
)
Conclusion
Actor-critic models like A2C and A3C represent a powerful approach to reinforcement learning, blending the strengths of policy-based and value-based methods. A2C provides a stable, synchronous framework for learning policies with reduced variance, while A3C’s asynchronous updates enhance efficiency and exploration. The significance of A3C lies in its scalability and robustness, making it a cornerstone of modern RL applications. As RL continues to evolve, understanding these algorithms provides a foundation for tackling complex decision-making problems in AI.
The full source code for A3C algorithm is available here.
