AgileRL：实现 MATD3¶
本教程展示了如何在 简单说话者听众 多智能体粒子环境中训练一个 MATD3 智能体。
什么是 MATD3？¶
MATD3（多智能体孪生延迟深度确定性策略梯度）通过使用第二组评论家网络和延迟更新策略网络，扩展了 MADDPG（多智能体深度确定性策略梯度）算法，以减少多智能体领域中的过高估计偏差。与 MADDPG 相比，这使得性能更优越。有关 MATD3 的更多信息，请查看 AgileRL 的文档。
我可以使用它吗？¶
动作空间 |
观测空间 |
|
|---|---|---|
离散 |
✔️ |
✔️ |
连续 |
✔️ |
✔️ |
环境设置¶
要遵循本教程，您需要安装下面显示的依赖项。建议使用新创建的虚拟环境以避免依赖冲突。
agilerl==2.2.1; python_version >= '3.10' and python_version < '3.12'
pettingzoo[classic,atari,mpe]>=1.23.1
AutoROM>=0.6.1
SuperSuit>=3.9.0
torch>=2.0.1
numpy>=1.24.2
tqdm>=4.65.0
fastrand==1.3.0
gymnasium>=0.28.1
imageio>=2.31.1
Pillow>=9.5.0
PyYAML>=5.4.1
代码¶
使用 MADDPG 训练多个智能体¶
以下代码应该可以顺利运行。注释旨在帮助您了解如何将 PettingZoo 与 AgileRL 一起使用。如果您有任何疑问，请随时在 Discord 服务器中提问。
"""This tutorial shows how to train an MATD3 agent on the simple speaker listener multi-particle environment.
Authors: Michael (https://github.com/mikepratt1), Nickua (https://github.com/nicku-a), Jaime (https://github.com/jaimesabalbermudez)
"""
import os
import numpy as np
import torch
from agilerl.algorithms.core.registry import HyperparameterConfig, RLParameter
from agilerl.components.multi_agent_replay_buffer import MultiAgentReplayBuffer
from agilerl.hpo.mutation import Mutations
from agilerl.hpo.tournament import TournamentSelection
from agilerl.utils.algo_utils import obs_channels_to_first
from agilerl.utils.utils import create_population, observation_space_channels_to_first
from agilerl.vector.pz_async_vec_env import AsyncPettingZooVecEnv
from tqdm import trange
from pettingzoo.mpe import simple_speaker_listener_v4
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("===== AgileRL Online Multi-Agent Demo =====")
# Define the network configuration
NET_CONFIG = {
"encoder_config": {
"hidden_size": [32, 32], # Actor hidden size
}
}
# Define the initial hyperparameters
INIT_HP = {
"POPULATION_SIZE": 4,
"ALGO": "MATD3", # Algorithm
# Swap image channels dimension from last to first [H, W, C] -> [C, H, W]
"CHANNELS_LAST": False,
"BATCH_SIZE": 32, # Batch size
"O_U_NOISE": True, # Ornstein Uhlenbeck action noise
"EXPL_NOISE": 0.1, # Action noise scale
"MEAN_NOISE": 0.0, # Mean action noise
"THETA": 0.15, # Rate of mean reversion in OU noise
"DT": 0.01, # Timestep for OU noise
"LR_ACTOR": 0.001, # Actor learning rate
"LR_CRITIC": 0.001, # Critic learning rate
"GAMMA": 0.95, # Discount factor
"MEMORY_SIZE": 100000, # Max memory buffer size
"LEARN_STEP": 100, # Learning frequency
"TAU": 0.01, # For soft update of target parameters
"POLICY_FREQ": 2, # Policy frequnecy
}
num_envs = 8
# Define the simple speaker listener environment as a parallel environment
env = simple_speaker_listener_v4.parallel_env(continuous_actions=True)
env = AsyncPettingZooVecEnv([lambda: env for _ in range(num_envs)])
env.reset()
# Configure the multi-agent algo input arguments
observation_spaces = [env.single_observation_space(agent) for agent in env.agents]
action_spaces = [env.single_action_space(agent) for agent in env.agents]
if INIT_HP["CHANNELS_LAST"]:
observation_spaces = [
observation_space_channels_to_first(obs) for obs in observation_spaces
]
# Append number of agents and agent IDs to the initial hyperparameter dictionary
INIT_HP["N_AGENTS"] = env.num_agents
INIT_HP["AGENT_IDS"] = env.agents
# Mutation config for RL hyperparameters
hp_config = HyperparameterConfig(
lr_actor=RLParameter(min=1e-4, max=1e-2),
lr_critic=RLParameter(min=1e-4, max=1e-2),
batch_size=RLParameter(min=8, max=512, dtype=int),
learn_step=RLParameter(
min=20, max=200, dtype=int, grow_factor=1.5, shrink_factor=0.75
),
)
# Create a population ready for evolutionary hyper-parameter optimisation
pop = create_population(
INIT_HP["ALGO"],
observation_spaces,
action_spaces,
NET_CONFIG,
INIT_HP,
hp_config=hp_config,
population_size=INIT_HP["POPULATION_SIZE"],
num_envs=num_envs,
device=device,
)
# Configure the multi-agent replay buffer
field_names = ["state", "action", "reward", "next_state", "done"]
memory = MultiAgentReplayBuffer(
INIT_HP["MEMORY_SIZE"],
field_names=field_names,
agent_ids=INIT_HP["AGENT_IDS"],
device=device,
)
# Instantiate a tournament selection object (used for HPO)
tournament = TournamentSelection(
tournament_size=2, # Tournament selection size
elitism=True, # Elitism in tournament selection
population_size=INIT_HP["POPULATION_SIZE"], # Population size
eval_loop=1, # Evaluate using last N fitness scores
)
# Instantiate a mutations object (used for HPO)
mutations = Mutations(
no_mutation=0.2, # Probability of no mutation
architecture=0.2, # Probability of architecture mutation
new_layer_prob=0.2, # Probability of new layer mutation
parameters=0.2, # Probability of parameter mutation
activation=0, # Probability of activation function mutation
rl_hp=0.2, # Probability of RL hyperparameter mutation
mutation_sd=0.1, # Mutation strength
rand_seed=1,
device=device,
)
# Define training loop parameters
max_steps = 13000 # Max steps (default: 2000000)
learning_delay = 0 # Steps before starting learning
evo_steps = 1000 # Evolution frequency
eval_steps = None # Evaluation steps per episode - go until done
eval_loop = 1 # Number of evaluation episodes
elite = pop[0] # Assign a placeholder "elite" agent
total_steps = 0
# TRAINING LOOP
print("Training...")
pbar = trange(max_steps, unit="step")
while np.less([agent.steps[-1] for agent in pop], max_steps).all():
pop_episode_scores = []
for agent in pop: # Loop through population
state, info = env.reset() # Reset environment at start of episode
scores = np.zeros(num_envs)
completed_episode_scores = []
steps = 0
if INIT_HP["CHANNELS_LAST"]:
state = {
agent_id: obs_channels_to_first(s) for agent_id, s in state.items()
}
for idx_step in range(evo_steps // num_envs):
# Get next action from agent
cont_actions, discrete_action = agent.get_action(
state, training=True, infos=info
)
if agent.discrete_actions:
action = discrete_action
else:
action = cont_actions
# Act in environment
next_state, reward, termination, truncation, info = env.step(action)
scores += np.sum(np.array(list(reward.values())).transpose(), axis=-1)
total_steps += num_envs
steps += num_envs
# Image processing if necessary for the environment
if INIT_HP["CHANNELS_LAST"]:
next_state = {
agent_id: obs_channels_to_first(ns)
for agent_id, ns in next_state.items()
}
# Save experiences to replay buffer
memory.save_to_memory(
state,
cont_actions,
reward,
next_state,
termination,
is_vectorised=True,
)
# Learn according to learning frequency
# Handle learn steps > num_envs
if agent.learn_step > num_envs:
learn_step = agent.learn_step // num_envs
if (
idx_step % learn_step == 0
and len(memory) >= agent.batch_size
and memory.counter > learning_delay
):
# Sample replay buffer
experiences = memory.sample(agent.batch_size)
# Learn according to agent's RL algorithm
agent.learn(experiences)
# Handle num_envs > learn step; learn multiple times per step in env
elif (
len(memory) >= agent.batch_size and memory.counter > learning_delay
):
for _ in range(num_envs // agent.learn_step):
# Sample replay buffer
experiences = memory.sample(agent.batch_size)
# Learn according to agent's RL algorithm
agent.learn(experiences)
state = next_state
# Calculate scores and reset noise for finished episodes
reset_noise_indices = []
term_array = np.array(list(termination.values())).transpose()
trunc_array = np.array(list(truncation.values())).transpose()
for idx, (d, t) in enumerate(zip(term_array, trunc_array)):
if np.any(d) or np.any(t):
completed_episode_scores.append(scores[idx])
agent.scores.append(scores[idx])
scores[idx] = 0
reset_noise_indices.append(idx)
agent.reset_action_noise(reset_noise_indices)
pbar.update(evo_steps // len(pop))
agent.steps[-1] += steps
pop_episode_scores.append(completed_episode_scores)
# Evaluate population
fitnesses = [
agent.test(
env,
swap_channels=INIT_HP["CHANNELS_LAST"],
max_steps=eval_steps,
loop=eval_loop,
)
for agent in pop
]
mean_scores = [
(
np.mean(episode_scores)
if len(episode_scores) > 0
else "0 completed episodes"
)
for episode_scores in pop_episode_scores
]
print(f"--- Global steps {total_steps} ---")
print(f"Steps {[agent.steps[-1] for agent in pop]}")
print(f"Scores: {mean_scores}")
print(f'Fitnesses: {["%.2f"%fitness for fitness in fitnesses]}')
print(
f'5 fitness avgs: {["%.2f"%np.mean(agent.fitness[-5:]) for agent in pop]}'
)
# Tournament selection and population mutation
elite, pop = tournament.select(pop)
pop = mutations.mutation(pop)
# Update step counter
for agent in pop:
agent.steps.append(agent.steps[-1])
# Save the trained algorithm
path = "./models/MATD3"
filename = "MATD3_trained_agent.pt"
os.makedirs(path, exist_ok=True)
save_path = os.path.join(path, filename)
elite.save_checkpoint(save_path)
pbar.close()
env.close()
观看训练好的智能体玩游戏¶
以下代码允许您从之前的训练块加载已保存的 MATD3 算法，测试算法性能，然后将一些回合可视化为 GIF。
import os
import imageio
import numpy as np
import torch
from agilerl.algorithms.matd3 import MATD3
from PIL import Image, ImageDraw
from pettingzoo.mpe import simple_speaker_listener_v4
# Define function to return image
def _label_with_episode_number(frame, episode_num):
im = Image.fromarray(frame)
drawer = ImageDraw.Draw(im)
if np.mean(frame) < 128:
text_color = (255, 255, 255)
else:
text_color = (0, 0, 0)
drawer.text(
(im.size[0] / 20, im.size[1] / 18), f"Episode: {episode_num+1}", fill=text_color
)
return im
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Configure the environment
env = simple_speaker_listener_v4.parallel_env(
continuous_actions=True, render_mode="rgb_array"
)
env.reset()
observation_spaces = [env.observation_space(agent) for agent in env.agents]
action_spaces = [env.action_space(agent) for agent in env.agents]
# Append number of agents and agent IDs to the initial hyperparameter dictionary
agent_ids = env.agents
# Instantiate an MADDPG object
matd3 = MATD3(
observation_spaces=observation_spaces,
action_spaces=action_spaces,
agent_ids=agent_ids,
device=device,
)
# Load the saved algorithm into the MADDPG object
path = "./models/MATD3/MATD3_trained_agent.pt"
matd3.load_checkpoint(path)
# Define test loop parameters
episodes = 10 # Number of episodes to test agent on
max_steps = 25 # Max number of steps to take in the environment in each episode
rewards = [] # List to collect total episodic reward
frames = [] # List to collect frames
indi_agent_rewards = {
agent_id: [] for agent_id in agent_ids
} # Dictionary to collect inidivdual agent rewards
rewards = [] # List to collect total episodic reward
frames = [] # List to collect frames
indi_agent_rewards = {
agent_id: [] for agent_id in agent_ids
} # Dictionary to collect inidivdual agent rewards
# Test loop for inference
for ep in range(episodes):
state, info = env.reset()
agent_reward = {agent_id: 0 for agent_id in agent_ids}
score = 0
for _ in range(max_steps):
agent_mask = info["agent_mask"] if "agent_mask" in info.keys() else None
env_defined_actions = (
info["env_defined_actions"]
if "env_defined_actions" in info.keys()
else None
)
# Get next action from agent
cont_actions, discrete_action = matd3.get_action(state, training=False)
if matd3.discrete_actions:
action = discrete_action
else:
action = cont_actions
# Save the frame for this step and append to frames list
frame = env.render()
frames.append(_label_with_episode_number(frame, episode_num=ep))
# Take action in environment
action = {agent_id: action[agent_id].reshape(-1) for agent_id in agent_ids}
state, reward, termination, truncation, info = env.step(action)
# Save agent's reward for this step in this episode
for agent_id, r in reward.items():
agent_reward[agent_id] += r
# Determine total score for the episode and then append to rewards list
score = sum(agent_reward.values())
# Stop episode if any agents have terminated
if any(truncation.values()) or any(termination.values()):
break
rewards.append(score)
# Record agent specific episodic reward
for agent_id in agent_ids:
indi_agent_rewards[agent_id].append(agent_reward[agent_id])
print("-" * 15, f"Episode: {ep}", "-" * 15)
print("Episodic Reward: ", rewards[-1])
for agent_id, reward_list in indi_agent_rewards.items():
print(f"{agent_id} reward: {reward_list[-1]}")
env.close()
# Save the gif to specified path
gif_path = "./videos/"
os.makedirs(gif_path, exist_ok=True)
imageio.mimwrite(
os.path.join("./videos/", "speaker_listener.gif"), frames, duration=10
)