VMAS: a vectorized, differentiable 2D simulator for multi-agent reinforcement learning

The bottleneck is the environment loop, not the policy

VMAS is described as a vectorized differentiable simulator designed for efficient MARL benchmarking, and it is made up of a fully-differentiable vectorized 2D physics engine written in PyTorch plus a set of challenging multi-robot scenarios. The reason that matters is the standard loop in most multi-agent reinforcement learning code. You write a step function, and you run it once per environment, usually in a Python for loop, and every one of those calls is interpreted Python that the compiler never sees. With a few hundred environments that loop dominates the wall clock, and the GPU sits mostly idle waiting for it. The vectorization approach in PyTorch, as the README puts it, allows VMAS to perform simulations in a batch, scaling to tens of thousands of parallel environments on accelerated hardware. The audience is therefore researchers who have hit that wall and want to compare algorithms on the same scenarios rather than on their own hardware.

A physics engine written in tensors, with a limited motion model

The engine is not a general rigid body simulator, and the README is straightforward about what it does. VMAS simulates agents and landmarks of different shapes, and supports rotations, elastic collisions, joints and custom gravity. Agents use holonomic motion models, which the README says are used to simplify simulation. That is a real simplification and it is the main thing to understand before choosing this over a conventional simulator, because a holonomic agent can move in any direction without the turning constraints that a differential drive robot has, and a lot of interesting multi-agent navigation behaviour comes from exactly those constraints. Custom sensors such as LIDARs are available, and the simulator supports inter-agent communication, which is what makes it usable for tasks where agents have to exchange information rather than just avoid colliding.

One simulator, many training frameworks

What VMAS actually provides is an environment, and it provides it in a shape that several training libraries can consume. The README lists an interface compatible with OpenAI Gym, with Gymnasium, with RLlib, with torchrl and with its MARL training library BenchMARL, which the authors released and which lets you train VMAS tasks using TorchRL. That last one is the recommended path in the notebook links, where a notebook titled Training VMAS in BenchMARL is marked suggested, alongside a notebook that creates a VMAS scenario and trains it with MAPPO and MLP and GNN policies. The practical consequence for an adopting team is that the environment is not tied to one algorithm. You can benchmark the same task under RLlib and under BenchMARL and know the environment was identical, which is the property that makes a benchmark a benchmark.

MPE scenarios, ported and vectorized

The implementation is inspired by OpenAI's MPE, the multi-agent particle environments, and the README states that alongside VMAS's own scenarios the project ports and vectorizes all the scenarios in MPE. There is a mpe_comparison/ directory at the top level of the repository, which is presumably where that comparison lives, and the environment list in the table of contents is split into a VMAS section with main scenarios and debug scenarios, and an MPE section. The scenarios are described as challenging, and the ones referenced in the notebooks make that concrete: a scenario where multiple robots with different embodiments need to navigate to their goals while avoiding each other and obstacles. The multi embodiment part is the interesting bit, since it is a scenario property that most multi-agent benchmarks do not offer, and it forces a policy to handle agents with different capabilities in the same world.

Installing and running a scenario

The package is published on PyPI as vmas, and the install extras in setup.py tell you what you need. The base install requires numpy, torch, pyglet at 1.5.27 or earlier, gym, gym-notices 0.0.8 and six. There are four optional extras. The gymnasium extra adds gymnasium and shimmy, the rllib extra adds ray with the rllib component at 2.2 or earlier, the render extra adds opencv-python, moviepy, matplotlib and opencv-python again for video output, and the test extra adds pytest, pytest-instafail, pyyaml and tqdm. The requirements file repeats the base set. The version number is not in setup.py itself, it is read out of vmas/__init__.py at build time, with a RuntimeError if the version data is missing.

numpy
torch
pyglet<=1.5.27
gym
gym-notices==0.0.8
six

The example scripts in the examples folder are the fastest way in, since the notebook on using a VMAS environment reproduces the use_vmas_env.py script, which creates, steps and renders a scenario. The repository also carries a docs/ directory published through Read the Docs, notebooks/, tests/ and a pre-commit configuration.

Rendering, and the problem of running it on a server

Rendering is a separate concern in VMAS, and the README's table of contents gives it its own section, including a subsection on rendering on server machines. The visual output of a 2D simulator is a genuine engineering problem in a vectorized setting, because a batch of tens of thousands of environments cannot be drawn one window at a time, so the renderer has to pick a subset of the batch and draw those. The extras reflect the same concern, since the render extra pulls in opencv-python for image handling, moviepy for encoding the video and matplotlib for plotting. A plot function under rendering is called out separately in the contents, which suggests the rendering path is a function you can override per scenario. If your experiments only need numbers, the base install avoids the whole rendering stack, which is worth knowing when you are provisioning a training cluster.

Where VMAS is the wrong tool

There are boundaries, and the first one is the physics model. If your task needs 3D, if it needs non-holonomic agents with realistic kinematics, or if you need deformable bodies, this engine is not it, and the holonomic motion model means the dynamics you simulate are not the dynamics of a real platform. The second is that the scenario set, however rich, is a fixed set. A benchmark with a shipped set of hard scenarios is a benchmark, and evaluating a policy on a task nobody has implemented in VMAS means writing that scenario yourself, which is exactly what the modular interface invites but also means you are then on your own for comparability. The third is the dependency pinning, which is old enough to matter. gym rather than gymnasium in the base install, with a gymnasium extra that adds the shimmy compatibility layer, and ray with rllib at 2.2 or earlier, mean you are on a generation-old RLlib if you use the rllib extra, and the pyglet ceiling of 1.5.27 is there for the same reason. The alternative worth naming is a discrete event simulator with a mature 3D or physics engine underneath, which is the right choice when physical fidelity matters more than throughput, and VMAS's advantage evaporates the moment your scenario is not expressible in its simplified 2D model. The paper behind the project is on arXiv at 2207.03530, published at the 16th International Symposium on Distributed Autonomous Robotic Systems, and the README asks that you cite it if you use VMAS in your research.

Editorial conclusion

Adopt VMAS if you are benchmarking multi-agent reinforcement learning policies and your current loop over environments is the bottleneck, since the whole design is about batching those environments into tensor operations. Do not adopt it if your task needs 3D physics, non-holonomic agents, or a simulation that is not roughly 2D, because the engine and its motion models are built for that case. Verify first that your intended scenarios are in the shipped set, that your training library has an adapter for it, and that the rendering extras are installed if you need video output.

Frequently asked questions

Official sources

1. License: GPL-3.0
2. Project website
3. proroklab/VectorizedMultiAgentSimulator on GitHub
4. README
5. Releases