Using PettingZoo with RLlib for Multi-Agent Deep Reinforcement Learning | by J K Terry | Towards Data Science

Thank you Yuri Plotkin, Rohan Potdar, Ben Black and Kaan Ozdogru, who each created or edited large parts of this article.
This tutorial provides an overview for using the RLlib Python library with PettingZoo environments for multi-agent deep reinforcement learning. More details about using the PettingZoo environment can be found in this recent blog post. In the following article, we will cover how to train and evaluate an RLlib policy using two PettingZoo environments:
Pistonball — No illegal actions exist, all agents step simultaneously
Leduc Hold’em — Illegal action masking, turn based actions
PettingZoo and Pistonball
PettingZoo is a Python library developed for multi-agent reinforcement-learning simulations. The current software provides a standard API to train on environments using other well-known open source reinforcement learning libraries. You can think of the library as analogous to OpenAI’s Gym library, however, tailored for multi-agent reinforcement learning. Similar to others, the basic API usage is as follows:
from pettingzoo.butterfly import pistonball_v5env = pistonball_v5.env() env.reset() for agent in env.agent_iter(): observation, reward, done, info = env.last() action = policy(observation, agent) env.step(action)
Pistonball is a cooperative PettingZoo environment, and is visualized in the rendering below:
The goal of the environment is to train the pistons to cooperatively work together to move the ball to the left as quickly as possible.
Each piston acts as an independent agent controlled by a policy π trained with function approximation techniques such as neural networks (hence deep reinforcement learning). The observation space of each agent is a window above and to the side of each piston (see below).
We assume full observability, and the policy π returns an action that serves to either raise or lower the piston from +4 to -4 pixels (image dimension is 84x84 pixels).
Following an action by each piston, the environment outputs both a global reward of
(ΔX/Xe) * 100 + τt
Where ΔX is the change in the ball’s x-position, Xe is the ball’s starting position and τ is the time penalty (with default value of 0.1) times the length of time t.
For more details about the PettingZoo environment, please check out the following description.
The Code
Let’s step through the code in greater detail.
First, to run the reinforcement learning environment, we import the required libraries:
from ray import tunefrom ray.rllib.models import ModelCatalogfrom ray.rllib.models.torch.torch_modelv2 import TorchModelV2from ray.tune.registry import register_envfrom ray.rllib.env.wrappers.pettingzoo_env import ParallelPettingZooEnvfrom pettingzoo.butterfly import pistonball_v5import supersuit as ssimport torchfrom torch import nn
The multi-agent reinforcement learning environment requires distributed training. To set up the environment, we use the open-source library Ray. Ray is a framework developed to provide a universal API for building distributed applications. Tune is a library built on top of Ray for scalable hyperparameter tuning in distributed reinforcement learning. We therefore only use Tune to execute a single training run in RLlib. Since we will require the use of a custom model to train our policy π, we first register the model in RLlib’s ModelCatalog. To create a custom model, we subclass the TorchModelV2 class from RLlib.
To use the PettingZoo environment with Tune, we first register the environment using the register_env function. ParallelPettingZooEnv is a wrapper used with PettingZoo environments such as Pistonball to interface with RLlib’s multi-agent API. SuperSuit is a library that provides preprocessing functions for both Gym and PettingZoo environments, as we’ll see below.
Initially, we create a convolutional neural network model in Pytorch for training our policy π:
class CNNModelV2(TorchModelV2, nn.Module):def __init__(self, obs_space, act_space, num_outputs, *args, **kwargs):TorchModelV2.__init__(self, obs_space, act_space, num_outputs, *args, **kwargs)nn.Module.__init__(self)self.model = nn.Sequential(nn.Conv2d( 3, 32, [8, 8], stride=(4, 4)),nn.ReLU(),nn.Conv2d( 32, 64, [4, 4], stride=(2, 2)),nn.ReLU(),nn.Conv2d( 64, 64, [3, 3], stride=(1, 1)),nn.ReLU(),nn.Flatten(),(nn.Linear(3136,512)),nn.ReLU(),)self.policy_fn = nn.Linear(512, num_outputs)self.value_fn = nn.Linear(512, 1)def forward(self, input_dict, state, seq_lens):model_out = self.model(input_dict[“obs”].permute(0, 3, 1, 2))self._value_out = self.value_fn(model_out)return self.policy_fn(model_out), statedef value_function(self):return self._value_out.flatten()
You can find more information on how to use custom models with the RLlib library here. We then need to define a function to create and return the environment:
def env_creator(args):env = pistonball_v4.parallel_env(n_pistons=20, local_ratio=0, time_penalty=-0.1, continuous=True, random_drop=True, random_rotate=True, ball_mass=0.75, ball_friction=0.3, ball_elasticity=1.5, max_cycles=125)env = ss.color_reduction_v0(env, mode=’B’)env = ss.dtype_v0(env, ‘float32’)env = ss.resize_v0(env, x_size=84, y_size=84)env = ss.frame_stack_v1(env, 3)env = ss.normalize_obs_v0(env, env_min=0, env_max=1)return env
We use PettingZoo’s parallel API to create the environment. The arguments to the parallel_env function control the environment’s properties and behavior. We additionally use SuperSuit’s wrapper functions for preprocessing operations.
The first function is employed to convert the full-color observations images produced by the environment to grayscale for reducing the computational complexity and cost. Subsequently, we cast the pixel image data type from uint8 to float32, downscale it to a 84x84 grid representation and normalize pixel intensity. Lastly, in order to measure the change in the ball’s direction ΔX (used in calculating the total reward), 3 consecutive frames of observations are stacked together to give the policy an easy way to learn.
A more detailed explanation of these preprocessing operations can be found in the previous tutorial.
After defining our model and environment, we can run the trainer using the tune.run() function using the parameters in the config dictionary. You can read about these hyperparameters in detail here. Of note, we instantiate a multi-agent specific configuration wherein we specify our policies using a dictionary mapping:
policy_id strings → tuples of (policy_cls, obs_space, act_space, config).
The policy_mapping_fn is a function mapping agent_ids to policy_ids. In our particular case, we use parameter sharing for training the pistons as described in the previous tutorial i.e. all the pistons are controlled by the same policy π with id ‘policy_0’.
if __name__ == “__main__”:env_name = “pistonball_v4”register_env(env_name, lambda config: ParallelPettingZooEnv(env_creator(config)))test_env = ParallelPettingZooEnv(env_creator({}))obs_space = test_env.observation_spaceact_space = test_env.action_spaceModelCatalog.register_custom_model(“CNNModelV2”, CNNModelV2)def gen_policy(i):config = {“model”: {“custom_model”: “CNNModelV2”,},“gamma”: 0.99,}return (None, obs_space, act_space, config)policies = {“policy_0”: gen_policy(0)}policy_ids = list(policies.keys())tune.run(“PPO”,name=”PPO”,stop={“timesteps_total”: 5000000},checkpoint_freq=10,local_dir=”~/ray_results/”+env_name,config={# Environment specific“env”: env_name,# General“log_level”: “ERROR”,“framework”: “torch”,“num_gpus”: 1,“num_workers”: 4,“num_envs_per_worker”: 1,“compress_observations”: False,“batch_mode”: ‘truncate_episodes’,# ‘use_critic’: True,‘use_gae’: True,“lambda”: 0.9,“gamma”: .99,# “kl_coeff”: 0.001,# “kl_target”: 1000.,“clip_param”: 0.4,‘grad_clip’: None,“entropy_coeff”: 0.1,‘vf_loss_coeff’: 0.25,“sgd_minibatch_size”: 64,“num_sgd_iter”: 10, # epoc‘rollout_fragment_length’: 512,“train_batch_size”: 512*4,‘lr’: 2e-05,“clip_actions”: True,# Method specific“multiagent”: {“policies”: policies,“policy_mapping_fn”: (lambda agent_id: policy_ids[0]),},},)
Now, we can watch our trained policy execute itself in the environment. During training, the policy is saved at each checkpoint, according to the frequency specified by checkpoint_freq. By default, RLlib stores the checkpoints in ~/ray_results. We first specify the path to our checkpoints:
checkpoint_path = “foo”
We can then load and restore our trained policy π, and use the policy to choose actions while rendering the environment at each timestep. We save the rendered video as a GIF:
with open(params_path, “rb”) as f:config = pickle.load(f)# num_workers not needed since we are not trainingdel config[‘num_workers’]del config[‘num_gpus’]ray.init(num_cpus=8, num_gpus=1)PPOagent = PPOTrainer(env=env_name, config=config)PPOagent.restore(checkpoint_path)reward_sum = 0frame_list = []i = 0env.reset()for agent in env.agent_iter():observation, reward, done, info = env.last()reward_sum += rewardif done:action = Noneelse:action, _, _ = PPOagent.get_policy(“policy_0”).compute_single_action(observation)env.step(action)i += 1if i % (len(env.possible_agents)+1) == 0:frame_list.append(PIL.Image.fromarray(env.render(mode=’rgb_array’)))env.close()print(reward_sum)frame_list[0].save(“out.gif”, save_all=True, append_images=frame_list[1:], duration=3, loop=0)
A gif of the solved environment can be seen above.
The entire training code detailed in this tutorial can be located here. And the code for rendering can be found here. For any improvements to the codebase or for any issues encountered with this article, please create an issue in the PettingZoo repository.
Leduc Hold’em
Leduc Hold’em is a poker variant popular in AI research detailed here and here; we’ll be using the two player variant. This environment is notable in that it is a purely turn based game and some actions are illegal (e.g. action masking is required).
To begin, the initial code we need looks much the same:
from copy import deepcopyimport osimport rayfrom ray import tunefrom ray.rllib.agents.registry import get_agent_classfrom ray.rllib.env import PettingZooEnvfrom pettingzoo.classic import leduc_holdem_v2from ray.rllib.models import ModelCatalogfrom ray.tune.registry import register_envfrom gym.spaces import Boxfrom ray.rllib.agents.dqn.dqn_torch_model import DQNTorchModelfrom ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFCfrom ray.rllib.utils.framework import try_import_torchfrom ray.rllib.utils.torch_ops import FLOAT_MAXtorch, nn = try_import_torch()
Initializing a suitable policy with action masking via parametric action masking then looks like this:
class TorchMaskedActions(DQNTorchModel):“””PyTorch version of above ParametricActionsModel.”””def __init__(self,obs_space,action_space,num_outputs,model_config,name,**kw):DQNTorchModel.__init__(self, obs_space, action_space, num_outputs,model_config, name, **kw)obs_len = obs_space.shape[0]-action_space.norig_obs_space = Box(shape=(obs_len,), low=obs_space.low[:obs_len], high=obs_space.high[:obs_len])self.action_embed_model = TorchFC(orig_obs_space, action_space, action_space.n, model_config, name + “_action_embed”)def forward(self, input_dict, state, seq_lens):# Extract the available actions tensor from the observation.action_mask = input_dict[“obs”][“action_mask”]# Compute the predicted action embeddingaction_logits, _ = self.action_embed_model({“obs”: input_dict[“obs”][‘observation’]})# turns probit action mask into logit action maskinf_mask = torch.clamp(torch.log(action_mask), -1e10, FLOAT_MAX)return action_logits + inf_mask, statedef value_function(self):return self.action_embed_model.value_function()
Rendering in Leduc Hold’em functions similarly to that of Pistonball, using this code snippet:
import rayimport pickle5 as picklefrom ray.tune.registry import register_envfrom ray.rllib.agents.dqn import DQNTrainerfrom pettingzoo.classic import leduc_holdem_v4import supersuit as ssfrom ray.rllib.env.wrappers.pettingzoo_env import PettingZooEnvimport PILfrom ray.rllib.models import ModelCatalogimport numpy as npimport osfrom ray.rllib.agents.registry import get_agent_classfrom copy import deepcopyimport argparsefrom pathlib import Pathfrom rllib_leduc_holdem import TorchMaskedActionsos.environ[“SDL_VIDEODRIVER”] = “dummy”parser = argparse.ArgumentParser(description=’Render pretrained policy loaded from checkpoint’)parser.add_argument(“checkpoint_path”, help=”Path to the checkpoint. This path will likely be something like this: `~/ray_results/pistonball_v4/PPO/PPO_pistonball_v4_660ce_00000_0_2021–06–11_12–30–57/checkpoint_000050/checkpoint-50`”)args = parser.parse_args()checkpoint_path = os.path.expanduser(args.checkpoint_path)params_path = Path(checkpoint_path).parent.parent/”params.pkl”alg_name = “DQN”ModelCatalog.register_custom_model(“pa_model”, TorchMaskedActions)# function that outputs the environment you wish to register.def env_creator():env = leduc_holdem_v4.env()return envnum_cpus = 1config = deepcopy(get_agent_class(alg_name)._default_config)register_env(“leduc_holdem”,lambda config: PettingZooEnv(env_creator()))env = (env_creator())# obs_space = env.observation_space# print(obs_space)# act_space = test_env.action_spacewith open(params_path, “rb”) as f:config = pickle.load(f)# num_workers not needed since we are not trainingdel config[‘num_workers’]del config[‘num_gpus’]ray.init(num_cpus=8, num_gpus=0)DQNAgent = DQNTrainer(env=”leduc_holdem”, config=config)DQNAgent.restore(checkpoint_path)reward_sums = {a:0 for a in env.possible_agents}i = 0env.reset()for agent in env.agent_iter():observation, reward, done, info = env.last()obs = observation[‘observation’]reward_sums[agent] += rewardif done:action = Noneelse:print(DQNAgent.get_policy(agent))policy = DQNAgent.get_policy(agent)batch_obs = {‘obs’:{‘observation’: np.expand_dims(observation[‘observation’], 0),‘action_mask’: np.expand_dims(observation[‘action_mask’],0)}}batched_action, state_out, info = policy.compute_actions_from_input_dict(batch_obs)single_action = batched_action[0]action = single_actionenv.step(action)i += 1env.render()print(“rewards:”)print(reward_sums)
Similar to before, the full training code used in this tutorial can be found here and the rendering code can be found here. If you have any improvements to this code or encounter any issues with this article, please create an issue on the PettingZoo repository and we’d be happy to help.