# How to Implement Multi-Agent RL Algorithms Using JaxMARL Environments in EvoRL

> Learn to implement multi-agent RL algorithms in EvoRL using JaxMARL environments. Master batching and centralized rollouts for efficient training of shared policy networks.

- Repository: [EMI-Group/evorl](https://github.com/emi-group/evorl)
- Tags: tutorial
- Published: 2026-03-01

---

**You implement multi-agent RL algorithms in EvoRL by wrapping JaxMARL environments with the `JaxMARLAdapter`, transforming agent-specific dictionaries into batched arrays using `batchify` and `unbatchify`, and executing centralized rollouts with `decentralized_rollout_with_shared_model` to train shared policy networks.**

EvoRL provides a unified framework for evolutionary and reinforcement learning research in JAX. When you need to train cooperative or competitive multi-agent systems using environments from the JaxMARL suite, the framework offers specific adapters and utilities to standardize multi-agent interactions. This guide demonstrates how to **implement multi-agent RL algorithms using JaxMARL environments in EvoRL**, covering environment wrapping, batch processing, and loss computation for algorithms like MAPPO.

## Core Architecture and Components

EvoRL's multi-agent support relies on a set of specialized abstractions that handle the complexity of multi-agent interactions. Understanding these components is essential before implementing your algorithm.

The primary components include:

- **`evorl.envs.MultiAgentEnv`** – The abstract base class defined in [`evorl/envs/multi_agent_env.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/multi_agent_env.py) that specifies the API all MARL environments must implement
- **`evorl.envs.jaxmarl.JaxMARLAdapter`** – Located in [`evorl/envs/jaxmarl.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/jaxmarl.py), this adapter converts native JaxMARL environments into EvoRL's standardized format
- **`evorl.utils.ma_utils`** – Provides `batchify` and `unbatchify` functions for converting between dictionary-based agent data and batched JAX arrays
- **`evorl.rollout_ma`** – Contains `decentralized_rollout_with_shared_model` for centralized execution with shared parameters across agents
- **`evorl.utils.rl_toolkits`** – Implements GAE and advantage estimation functions in [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py)

## Wrapping JaxMARL Environments

The first step is adapting the JaxMARL environment to EvoRL's interface using the `JaxMARLAdapter`. This handles automatic conversion of reset, step, and space formats.

```python
from evorl.envs import create_env, AutoresetMode
from evorl.envs.jaxmarl import create_mabrax_env

# Example: a 4-agent Ant environment from JaxMARL

env = create_mabrax_env("ant_4x2", homogenisation_method="max")

```

The `create_mabrax_env` function, implemented in [`evorl/envs/jaxmarl_envs/mabrax.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/jaxmarl_envs/mabrax.py) (lines 48-61), returns a `JaxMARLAdapter` instance that conforms to the `MultiAgentEnv` API. You can also instantiate this via `create_env` with a Hydra configuration specifying `"jaxmarl"` as the backend.

## Handling Batched Observations and Actions

When using a **shared policy network** for all agents, you must convert between environment dictionaries (`{agent_id: obs}`) and batched arrays suitable for neural network processing.

```python
from evorl.utils.ma_utils import batchify, unbatchify
from functools import partial

obs_batchify_fn = partial(batchify, agent_list=env.agents)
action_unbatchify_fn = partial(unbatchify, agent_list=env.agents)

```

These utilities, defined in [`evorl/utils/ma_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/ma_utils.py), ensure observations from multiple agents stack into a single array with shape `[num_agents, ...]`, while network outputs split back into per-agent action dictionaries for environment stepping.

## Running Centralized-Execution Rollouts

For algorithms like MAPPO or MATD3, use centralized execution with decentralized training (CTDE). EvoRL provides `decentralized_rollout_with_shared_model` for this paradigm.

```python
from evorl.rollout_ma import decentralized_rollout_with_shared_model

env_state, trajectory = decentralized_rollout_with_shared_model(
    env=env,
    agent=shared_agent,
    env_state=init_state,
    agent_state=shared_state,
    key=key,
    rollout_length=128,
    obs_batchify_fn=obs_batchify_fn,
    action_unbatchify_fn=action_unbatchify_fn,
    env_extra_fields=("autoreset", "episode_return", "termination"),
)

```

This function, located in [`evorl/rollout_ma.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout_ma.py) (starting at line 54), executes a single shared policy for all agents. The returned `trajectory` is a `SampleBatch` with shape `[T, B, ...]` where `T` is the rollout length and `B` represents parallel environment batch dimensions.

## Computing Advantages and Value Targets

After collecting trajectories, compute advantages using Generalized Advantage Estimation (GAE). The `compute_gae_with_horizon` function in [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py) (lines 18-59) handles bootstrapping from value estimates.

```python
from evorl.utils.rl_toolkits import compute_gae_with_horizon

# Concatenate last observation for bootstrapping

obs_concat = jax.tree_util.tree_map(
    lambda o, no: jnp.concatenate([o, no[-1:]], axis=0),
    trajectory.obs,
    trajectory.next_obs,
)

# Compute V-values using the value network

values = shared_agent.compute_values(
    shared_state, SampleBatch(obs=obs_concat)
)

v_targets, advantages = compute_gae_with_horizon(
    rewards=trajectory.rewards,
    values=values,
    dones=trajectory.dones,
    terminations=trajectory.extras.env_extras.termination,
    gae_horizon=0,
    gae_lambda=0.95,
    discount=0.99,
)

```

Assign the computed values back to the trajectory for use in the loss function:

```python
trajectory = trajectory.replace(
    extras=trajectory.extras.replace(
        v_targets=jax.lax.stop_gradient(v_targets),
        advantages=jax.lax.stop_gradient(advantages),
    )
)

```

## Implementing the MAPPO Loss Function

Multi-Agent PPO reuses the single-agent PPO loss structure but operates on batched agent data. The reference implementation in [`evorl/algorithms/ppo.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/ppo.py) (lines 60-102) serves as your template.

```python
def loss(self, agent_state, sample_batch, key):
    obs = sample_batch.obs
    # Batched policy forward pass

    raw_actions = self.policy_network.apply(
        agent_state.params.policy_params, obs
    )
    actions_dist = get_tanh_norm_dist(*jnp.split(raw_actions, 2, axis=-1))
    
    # Compute PPO surrogate loss

    logp = actions_dist.log_prob(sample_batch.actions)
    logp_old = sample_batch.extras.policy_extras.logp
    rho = jnp.exp(logp - logp_old)
    
    advantages = sample_batch.extras.advantages
    if self.normalize_gae:
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
    
    clipped_rho = jnp.clip(rho, 1 - self.clip_epsilon, 1 + self.clip_epsilon)
    actor_loss = -jnp.mean(
        jnp.minimum(rho * advantages, clipped_rho * advantages)
    )
    
    # Value loss on batched observations

    values = self.value_network.apply(
        agent_state.params.value_params, obs
    )
    critic_loss = optax.l2_loss(values, sample_batch.extras.v_targets).mean()
    
    return PyTreeDict(
        actor_loss=actor_loss,
        critic_loss=critic_loss,
        actor_entropy=actions_dist.entropy(seed=key).mean(),
    )

```

## Complete Working Example: Multi-Agent PPO

Here is a complete implementation combining environment setup, rollout collection, and optimization:

```python
import jax
import optax
from evorl.envs.jaxmarl import create_mabrax_env
from evorl.utils.ma_utils import batchify, unbatchify
from evorl.utils.rl_toolkits import compute_gae_with_horizon
from evorl.algorithms.ppo import make_mlp_ppo_agent
from evorl.rollout_ma import decentralized_rollout_with_shared_model
from functools import partial

# 1. Create environment

env = create_mabrax_env("ant_4x2", homogenisation_method="max")
key = jax.random.PRNGKey(0)
env_state = env.reset(key)

# 2. Initialize shared agent

dummy_obs = env.obs_space[env.agents[0]].sample(key)
agent = make_mlp_ppo_agent(
    action_space=env.action_space[env.agents[0]],
    clip_epsilon=0.2,
    normalize_gae=True,
)
agent_state = agent.init(dummy_obs, dummy_obs, key)

# 3. Setup batching functions

obs_batchify_fn = partial(batchify, agent_list=env.agents)
action_unbatchify_fn = partial(unbatchify, agent_list=env.agents)

# 4. Collect rollout

key, rollout_key = jax.random.split(key)
env_state, traj = decentralized_rollout_with_shared_model(
    env=env,
    agent=agent,
    env_state=env_state,
    agent_state=agent_state,
    key=rollout_key,
    rollout_length=128,
    obs_batchify_fn=obs_batchify_fn,
    action_unbatchify_fn=action_unbatchify_fn,
    env_extra_fields=("autoreset", "episode_return", "termination"),
)

# 5. Compute GAE advantages

obs_concat = jax.tree_util.tree_map(
    lambda o, no: jnp.concatenate([o, no[-1:]], axis=0),
    traj.obs,
    traj.next_obs,
)
values = agent.compute_values(agent_state, SampleBatch(obs=obs_concat))
v_t, adv = compute_gae_with_horizon(
    rewards=traj.rewards,
    values=values,
    dones=traj.dones,
    terminations=traj.extras.env_extras.termination,
    gae_lambda=0.95,
    discount=0.99,
)

traj = traj.replace(
    extras=traj.extras.replace(
        v_targets=jax.lax.stop_gradient(v_t),
        advantages=jax.lax.stop_gradient(adv),
    )
)

# 6. Optimization step

optimizer = optax.adam(3e-4)
opt_state = optimizer.init(agent_state.params)

def loss_fn(params, batch, key):
    tmp_state = agent_state.replace(params=params)
    loss_dict = agent.loss(tmp_state, batch, key)
    return loss_dict.actor_loss + loss_dict.critic_loss, loss_dict

grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
(loss, loss_dict), grads = grad_fn(agent_state.params, traj, key)
updates, opt_state = optimizer.update(grads, opt_state)
new_params = optax.apply_updates(agent_state.params, updates)
agent_state = agent_state.replace(params=new_params)

```

## Summary

To successfully **implement multi-agent RL algorithms using JaxMARL environments in EvoRL**:

- Wrap JaxMARL environments using `create_mabrax_env` or `JaxMARLAdapter` from [`evorl/envs/jaxmarl.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/jaxmarl.py)
- Use `batchify` and `unbatchify` from [`evorl/utils/ma_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/ma_utils.py) to transform agent dictionaries into batched arrays
- Execute training rollouts with `decentralized_rollout_with_shared_model` from [`evorl/rollout_ma.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout_ma.py) for centralized execution
- Compute advantages using `compute_gae_with_horizon` from [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py)
- Adapt single-agent loss functions like PPO to process batched multi-agent data

## Frequently Asked Questions

### What is the difference between decentralized_rollout and decentralized_rollout_with_shared_model?

The `decentralized_rollout` function executes policies with separate parameters for each agent, while `decentralized_rollout_with_shared_model` uses a single shared network for all agents. According to [`evorl/rollout_ma.py`](https://github.com/emi-group/evorl/blob/main/evorl/rollout_ma.py), the shared model variant automatically handles observation batching and action unbatching using the provided helper functions, making it suitable for CTDE (Centralized Training with Decentralized Execution) algorithms like MAPPO that rely on parameter sharing across agents.

### How does JaxMARLAdapter handle environment resets and steps?

The `JaxMARLAdapter` class in [`evorl/envs/jaxmarl.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/jaxmarl.py) automatically converts JaxMARL environment outputs into EvoRL's standardized `MultiAgentEnv` format. It standardizes the observation and action spaces, manages dictionary-based agent IDs, and handles autoreset logic. When you call `env.reset(key)`, it returns a standardized `EnvState` compatible with EvoRL's rollout functions and utility modules.

### Why do I need to use batchify and unbatchify functions?

Multi-agent environments return observations as dictionaries mapping agent IDs to individual arrays, but neural networks expect contiguous batched arrays. The `batchify` function in [`evorl/utils/ma_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/ma_utils.py) stacks these dictionaries into arrays of shape `[num_agents, ...]`, while `unbatchify` splits network outputs back into per-agent dictionaries required by the environment's step function. This transformation is essential for shared-parameter policies that process all agents simultaneously in a single forward pass.

### Can I adapt existing single-agent algorithms for multi-agent training?

Yes, EvoRL's single-agent algorithms in [`evorl/algorithms/ppo.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/ppo.py) can be adapted for multi-agent use by ensuring the loss function handles batched inputs containing multiple agents. The primary implementation changes involve using `decentralized_rollout_with_shared_model` for data collection and ensuring your network architecture processes the additional agent dimension. The core loss computation remains structurally identical to the single-agent case, as the framework handles the multi-agent batching transparently.