High Performance Simulation for Scalable Multi-Agent Reinforcement Learning
Abstract
Multi-agent reinforcement learning experiments and open-source training environments are typically limited in scale, supporting tens or sometimes up to hundreds of interacting agents.
In this paper we demonstrate the use of Vogue , a high performance agent based model ( ABM ) framework.
Vogue serves as a multi-agent training environment, supporting thousands to tens of thousands of interacting agents while maintaining high training throughput by running both the environment and reinforcement learning ( RL ) agents on the GPU.
High performance multi-agent environments at this scale have the potential to enable the learning of robust and ﬂexible policies for use in
ABM s and simulations of complex systems.
We demonstrate training performance with two newly developed, large scale multi-agent training environments.
Moreover, we show that these environments can train shared RL policies on time-scales of minutes and hours.
1. Introduction
Increases in computational power and available data have led to an increase in the potential scale and ﬁdelity of computer simulations such as agent based models ( ABM s).
In tandem, a desire has developed for complex agent behaviour which cannot be addressed solely by the often error-prone and challenging process of manually designing and implementing agent behaviours that are both ﬂexible and robust to extreme states of the simulation.
In the context of agent-based models multi-agent reinforcement learning ( MARL ) is a means of learning agent behaviours where the reward function is well known but the simulation has a complex state space.
This motivates the
UK being close to other agents, but are penalised for collisions.
This example contains 5,000 agents learning a shared policy.
500 training steps (totalling 625,000 mini batches) took approximately 16 minutes to train on a laptop Nvidia GTX 1650 GPU.
The inset illustrates how agents sense neighbouring agents via a simple vision model, shown in more detail in Figure 3(a).
development of fast and large-scale training environments, enabling the use of MARL to learn behaviours reﬂecting emergent systems such as crowds, ﬂocks, social networks or markets.
2. High Performance MARL with Vogue 2.
Vogue: An Agent Based Modelling Framework Vogue is a commercial Python
framework for high performance ABM s. It is
motivated by some key design principles
: • Usability: Vogue allows users to write ABM logic in native, idiomatic Python, as shown in code snippet 1. a r X i v : . 03945v [ c s .
M A ] J u l This is then compiled by Vogue into high performance bytecode, and can be run in a standard Python process.
• High Performance Execution: Vogue maps model logic over agents using high performance, vectorised update patterns.
This is efﬁcient and obviates the need for modellers to consider the computational complexity and low-level implementation details of ABM s. • Composition: Vogue model logic is described in small pieces of functionality following standard patterns.
This allows modellers to easily combine and reuse model logic, and standardises how models are written.
• Interoperability: Vogue interacts seamlessly with the broader Python data science and machine learning ( ML ) ecosystem.
This allows Vogue models to be easily employed in ML and data science pipelines.
These design choices allow users to quickly implement high performance ABM s without the need to write low-level code, as shown in code snippet 1. In comparison to other numerical Python compilation libraries such as JAX (Bradbury et al., 2018) or Numba (Lam et al., 2015), Vogue speciﬁcally implements common patterns encountered in
ABM s, in particular interactions between spatially embedded agents (based on spatial proximity), and agents interacting via edges on a graph or network.
These base interactions can them be combined into more complex model logic, including interactions between heterogeneous agent types.
Vogue is currently limited to models that update in discrete time-steps. It is designed such that individual interactions update agents simultaneously, in particular the GPU engine applies interactions to entities in parallel.
2.2. End-to-end GPU-based Reinforcement Learning
As well as allowing users to efﬁciently implement models, a major advantage of Vogue is that models can be executed completely on the GPU with minimal changes to the underlying code, leveraging the high degree of parallelism afforded by the GPU to efﬁciently update the state of the model.
Vogue’s implementation allows data to be shared between Vogue and other ML and deep-learning libraries without copying (or transfer via the host) via the CUDA array interface.
This means Vogue can interact directly with any ML framework that implements the CUDA array interface, including popular ML tools like PyTorch (Paszke et al., 2019) and JAX (Bradbury et al., 2018), or GPU based data tools like RAPIDS (Team, 2018).
This allows the execution of both the Vogue training environment/simulation and the deep reinforcement learning ( RL ) agent on the GPU, removing the overhead of moving data to and from the GPU during training. @vogue.interaction.self() def update_opinion( params: Params, person: Person ) -> None: person.opinion = person.new_opinion
@vogue.interaction.graph() def social_influence( params: Params, me: Person, you: Person, edge: Friendship, ) ->
None: d = abs(me.opinion - you.opinion) if d < params.threshold: w = params.strength * edge.weight me.new_opinion = ( (1.0 - w) * me.new_opinion + w * you.opinion ) Code Snippet 1.
A Vogue implementation of a simple opinion dynamics model.
The model is implemented as two interactions .
A self interaction that updates the current opinion of entities, and a graph interaction that updates an entities opinion based on its neighbours opinions.
Model logic in Vogue is expressed as interactions between pairs of agents, or updates of individual agents.
The compilation of these interactions and execution of the model is handled by the Vogue engine.
M ARL has found a wide range of applications including distributed autonomous systems, socioeconomic and game theoretic problems (Busoniu et al., 2008).
The number of learning agents in these systems typically ranges from single-digit to hundreds of interacting agents.
A major obstacle to applying MARL to systems with larger numbers of interacting agents is the lack of high performance simulation environments, as well as the technical challenges and cost associated with their development.
Training environments are often implemented for execution on the CPU with the deep RL agent running on the GPU.
However, the communication of data from the simulation on the host to the agent on the GPU can have a signiﬁcant performance impact.
OpenAI Gym (Brockman et al., 2016) has become a de facto standard API for RL research, allowing training environments and RL agent implementations to be easily combined, though it does not have a standard API for multi-agent training.
There are a number of open-source MARL envi- ronment projects implementing this API, including Petting
Zoo (Terry et al., 2020) and MAgent (Zheng et al., 2018).
MAgent potentially supports environments containing millions of agents, but in contrast to Vogue (proposed here) the environments restrict agents to a discrete grid and the engine is implemented in C++.
Petting Zoo wraps a number of environments (including MAgent), though many of the environments are restricted to tens or hundreds of agents, or have observation spaces that scale with the number of agents, restricting their usable scale.
RLLib (Liang et al., 2017) provides a set of tools for the execution of RL pipelines, including automated hyperparameter tuning and distributed training.
It has the ability to run pipelines with multiple multi-agent policies, and multiple environments executed in parallel.
It does however rely on users to implement the training environment (or providing a pre-existing environment).
Warpdrive (Lan et al., 2021) offers a framework for end-toend MARL on GPU.
It demonstrates high training throughput, though is only tested (in Lan et al., 2021) on a maximum of 1,000 agents (executed across 2,000 environments simultaneously).
Its end-to-end solution contrasts Vogue’s focus on ease of writing high performance models, and interoperability with existing Python tools. NVidia’s Issac
Gym (Makoviychuk et al., 2021) takes a similar approach, sharing data between simulation and RL agent on the GPU directly.
Their focus is on physics simulations for robotics, where Vogue is focused on ABM use cases, and ease of implementing and combining ABM s.
There are a number of ABM frameworks, sharing the broad aim of allowing users to more easily implement ABM s, or provide performance gains.
Mesa (Kazil et al., 2020) is a popular Python ABM framework, though its performance is restricted by Python’s native performance.
Netlogo (Wilensky, 1999) is instead implemented in Java, though this limits its interoperability with Python and requires users to implement models in Netlogo’s scripting language. A comparison of the performance of these libraries for a basic implementation of boids (with hard-coded behaviours) is shown in Figure 2. Both the CPU and GPU version of Vogue are several orders of magnitude faster than Mesa and NetLogo.
The
GPU engine demonstrates improved scaling over the CPU engine with increasing number of agents.
Few frameworks support GPU execution; Flame (Richmond et al., 2021) is a notable example, implemented in C++ with an optional Python interface.
4. Training Environments
At the time of writing, we could not locate any documented training environments at the scale of 1,000+ interacting agents.
Instead, we provide two simple examples by imple- [Figure] menting extended versions of two popular models: ﬂocking and tag .
We note that extending the range and complexity of example environments is a priority for future work.
In both examples, to address the increased overall size of the simulation state space, the agents only observe their local neighbourhood of the simulation via a simple ray casting algorithm that divides an agent’s ﬁeld of view into a ﬁxed number, v , of sectors.
Internally, their view is represented by an array of v values representing the distance to the nearest object at that angle, as shown in Figure 3(a).
An example time series showing an agent’s view during a training step is shown in Figure 3(b).
This local view forms part of the training environment observation space, local to each agent in the environment.
4.1. Flocking
This training environment is based on Reynolds’ Boids model (Reynolds, 1987), designed to mimic natural ﬂocks and shoals.
The training environment consists of a ﬂock of agents navigating a 2d space, with no hard-coded navigation rules.
Previous studies have looked at emergence of ﬂocking from reinforcement learning (for example Durve et al., 2020; Shimada & Bentley, 2018) but in this work we aim to examine the feasibility of efﬁciently training ﬂocking behaviours for 1,000-10,000+ interacting agents.
A snapshot of a simulation state is shown in Figure 1.
The observation space of the environment (for an individual agent) is a 129 length vector containing the agent’s view, with v = 128 sectors, and the agent’s current speed.
The action space is two dimensional.
One dimension accelerates/decelerates the agent, in the range [ − a max , a max ] , updating the speed of the agent as s t +1 = min
( s min , max ( s t + a t , s max )) .
The second dimension in the range [ − θ max , θ max ] [Figure] rotates the heading of the agents up to a maximum rotation rate.
The reward signal for each agent is the sum of contributions due to proximity to other agents, i.e. the reward r i of agent i is
[Equation]
where d ij is the Euclidean distance between agents i, j , d v is the visual range of the agent, and f is the reward contribution, as in Figure 4.
[Figure]
In each environment update, the agents are accelerated and rotated (according to actions sampled from the current policy), after which their local view model and rewards are updated for all pairs of agents in spatial proximity.
4.2. Tag
This environment is based on Petting Zoo’s (in Terry et al., 2020) simple tag environment.
It contains two agent types: runners and chasers .
Chasers are rewarded for touching runners, and runners are rewarded for staying close to other runners, and penalised for touching a chaser.
The view model of the agent has two simulated colour channels allowing agents to distinguish between neighbouring runners or chasers.
The corresponding observation space is alength vector, concatenating the two v = 64 colour channels.
The two dimensional action space rotates the agent’s heading, in the range [ − θ max , θ max ] and moves them along the heading, in the range [0 , s max ] where s max is the maximum speed of the agents.
4.3. Training
For both environments we used proximal policy optimisation ( PPO ) (Schulman et al., 2017) for a continuous action space, with the simulated agents sharing a single policy.
In the tag environment, the runner and chaser types share independent policies, trained simultaneously.
Instead of multiple parallel environments, we treat each individual agent as an independent learner, and gather trajectories across the full set of of agents.
If we have n agents and run t updates of the simulation each training step, we sample a total of n × t trajectories per training step.
Vogue updates the state of the ABM in response to the actions sampled from the PPO policy, then copies the updated observations and rewards to the JAX experience buffer, as shown in Figure 5. [Figure]
5. Performance Results
Benchmarking was performed on an Nvidia A100 GPU.
Each full training run includes T = 200 training steps.
Inside each training step the environment was updated for t = 128 steps, generating m = n × t individual trajectories.
The policy was then updated for p = 2 epochs, across b minibatches sampled from the trajectories.
The number of trajectories scales with the number of simulated agents, so the number of mini-batches is capped as b = min ( m/ | b | , b max ) where b max = 512 is the maximum number of batches, and | b | = 512 is the chosen mini-batch size.
This means for the highest number of entities each training run performs
T × p × b = 204 , 800 mini-batch updates of the PPO policy, and T × t = 25 , 600 updates of the environment.
The agents have a view distance equivalent to 1/10th of the width of the environment and a convex ﬁeld of vision of approximately 250 ◦ .
The RL agent update, and trajectory processing algorithms were implemented in JAX.
The actor and critic PPO networks had two hidden layers with 64 nodes each, with a total of 1035 trainable parameters.
It should also be noted that performance of the simulation can be highly dependent on the parameters of the simulation.
For example the view distance of the agents and their vision angle impact the number of interactions between agents that must be processed.
[Figure]
The total training time for 200 steps using the ﬂock and tag environments are shown in Figure 6.
Values shown are averaged over 10 independent training runs.
With 16, agents, the ﬂock environment takes around 3.5 minutes to complete training, and demonstrates slow scaling with number of agents.
For a tag environment with a total of 17,600 agents (16,000 runners and 1,600 chasers) training takes around 12 minutes.
The tag environment demonstrates roughly linear growth in training time with number of agents, we speculate because of the added complexity of multiple interactions between different agent types.
6. Conclusion
We have demonstrated that joining forward simulation and training of MARL systems into a single end-to-end RL process can lead to drastic improvements in performance.
Using the commercial software Vogue , high performance ABM s can be executed directly on the GPU, while the high level of interoperability enables direct communication between
Vogue and many Python based RL or deep learning frameworks.
Vogue is currently available for academic research – please contact the authors.