jaxrl 0.1.1
Add your description here
Requires Python: >=3.13

# Partially Observable Multi-Agent RL with Transformers

This project provides a JAX-based framework for training Transformer-based agents in custom, partially observable multi-agent reinforcement learning (MARL) environments. The custom grid environments are designed to test memory, planning, multi-task learning and cooperation.

### ✨ Key features

* **Transformer-based Agents**: Uses a Transformer-over-time architecture to handle partial observability. Uses updated transformer architecture such as rope and pre-layer norm. 
* **High Performance**: Single layer transformers train a **10 million+** steps per second and the medium sized transformer trains at **2 million** steps per second on a single 5090. Cudnn, grouped query attention and bfloat16 training significantly boost performance.
* **Custom Environments**: Includes several challenging MARL environments designed to test memory and coordination.
* **PPO Implementation**: A clean and efficient Proximal Policy Optimization (PPO) implementation.
* **Multi-task training**: One model can be trained on multiple environments simultaneously with unique task id embeddings.  
* **Snapshot League self play**: Policy snapshots can be rotated into the opponent poll during training.  
* **Trueskill Evaluations**: Policies for zero sum games can be compared using a trueskill based evaluation system and a model checkpoint pool.  
* **Discretized Value Functions**: Option to either minimize MSE or Hl-Gauss discretized values.  

---

## 🚀 Getting started

### Prerequisites

* Python 3.13+
* [uv](https://github.com/astral-sh/uv) (a fast Python package installer and resolver)

### Installation

Clone the repository and install the dependencies:

```bash
git clone git@github.com:gabe00122/jaxrl.git
cd jaxrl
uv sync --extra cuda
```


## 💻 Usage
### Train an agent
To start a training run, use the train command and provide a configuration file.


```bash
uv run pmarl train --config ./config/return_baseline.json
```
A unique run name will be generated for you (e.g., silly-camel-34). You will need this name to view the results.


### Watch a trained agent
To render an environment with a trained agent, use the enjoy command with the run_name from your training session.
Note: it may take up to 2 minutes to start rendering

```bash
uv run pmarl enjoy young-shark-cff1bi --seed 5 --video-path ./videos/out.mp4
```

### Download the checkpoint files

```bash
uv tool install "huggingface_hub[cli]"

huggingface-cli download gabe00122/grid-pomarl --local-dir ./results/
```

### Test out an environment
You can test out a environment without training a model using the `play` command or play with trained models using the `enjoy` command.

```bash
uv run pmarl enjoy koth --pov --human
```

The `pov` option renders from one agents point of pov
The `human` option gives you keyboard controls for one of the agents

Typical control scheme
Keyboard:
w: up
d: right
s: down
a: left
e: dig a wall
space: attack
n: cycle to the next agent

---

### Evaluating Zero Sum Games

By definition because these games are zero sum the mean of the reward during self play remains zero, regardless of the sill level of the policy.
To track progress and compare different training runs we can use trueskill to create a elo like score using a opponent pool.

To build a trueskill graph from a series of runs you can use the following command.

```bash
uv run pmarl eval --run blue-whale --run red-fish --rounds 1000 --out ./analysis/graph.png
```

Here's an example of a output evaluate of agents trained with different hyper parameters on the king of the hill environment.
<img width="900" height="600" alt="image" src="https://github.com/user-attachments/assets/b05a45be-b65d-4ea7-b65b-fec0ab369b8f" />

---

### Multi-task Training

Is the observation and action space is shared by the environments then one policy can be trained on multiple environments simultaneously.
See the `multitask.json` config for an example.

If multitask training was used then for evaluation and playback the `--env name` argument is required to specify which task you are viewing/evaluating from the multi-task models.

## 🏋️‍♂️ Environments
### Grid Return
A multi-agent 2D grid world task requiring spatial memory.

* Description: One or more goal tiles are placed at random locations. When an agent finds a goal, it receives a `+1` reward, and the agent is moved to a new random location. Agents must remember goal locations and navigate around obstacles to return to them efficiently. The number of goals can be configured via `num_flags` and defaults to one.

* Observation: A small rectangular grid centered on the agent. Agents can see each other. No absolute positions are given.

* Actions: [`move-up`, `move-right`, `move-down`, `move-left`]

* Reward: `+1` for finding a goal tile, `0` for anything else.

* Note: Agents can move through each other but not through obstacle tiles.

Agents can "dig" through obstacle tiles. Moving into an obstacle removes the tile but adds a timeout before the agent can move again.

https://github.com/user-attachments/assets/98cc3318-67ca-44c0-ac6e-e4537bd30ed1

### Scouts
A multi-agent coordination task with two specialized agent types. A Harvester must first unlock a resource tile, which a Scout can then gather.

Scout: Fast-moving agents that can gather resources only after they are "unlocked."
Harvester: Slow agents (can only move every 6th turn). When a Harvester reaches a resource tile, it gets a reward and unlocks the resource, allowing Scouts to gather it.

https://github.com/user-attachments/assets/d566840e-1837-4fc1-8c78-439677f358a8

### Traveling salesman

The several way points are randomly scattered and the agents and the agents get a reward for the first time they get each flag. When all flags are taken they all reset.

https://github.com/user-attachments/assets/af009d24-c65e-4195-99af-0a4e703652cd

### King of the hill

Two teams of agents compete to capture random control points in the center.

A multi-agent gridworld where two teams of Knights and Archers battle to capture and hold flags.

Randomly generated maps with destructible walls and central control points

Knights: melee fighters with higher HP

Archers: ranged fighters with arrows and cooldowns

Teams score points every turn for each flag they control

Agents can move, attack, dig through walls, or fire arrows

Rewards are fully team-shared, encouraging coordination

https://github.com/user-attachments/assets/3483745f-7c53-46e9-b838-3cc76b9e3ee4

## Scaling Results

I found that seemingly preformace scales predictablly with network depth. Interestingly there was not a similar improvement with width scaling. The following graph shows depth scaling on the grid return environment.
<img width="956" height="400" alt="image" src="https://github.com/user-attachments/assets/698dc9d9-3a37-4959-aa2c-5e9f06e3ba59" />

## RNNs

For the grid return task grouped query attention also outpreformed RNN layers while using roughly 1/4th the compute at a BPTT length of 512.
<img width="956" height="400" alt="image" src="https://github.com/user-attachments/assets/bf1849ba-82ee-4f25-b8be-8614f030d1da" />

## Bonus ##

Currently training only supports episodes the entirly fit in context, this makes variable length episodes tricky to train on but you can still train on games like craftax if episodes are trunctated to fit within context.
In this case I truncated episodes to fit within 1024 steps of context and still acheived a score of 17.7% with 1b samples. In the future the kv cache at the start of the rollout could be saved and reused in training with sliding window attention to enable learning with variable length or long episodes.

https://github.com/user-attachments/assets/d667e777-c480-4b40-b190-46946d3548d5

---

## 🙏 Acknowledgements

This project was made possible thanks to:

* **Google Research Cloud TPU Program** — for providing access to TPUs that enabled training and experimentation.  
* **[Urizen Onebit Tileset](https://vurmux.itch.io/urizen-onebit-tileset)** by Vurmux — for the excellent pixel art tileset used in the gridworld environments.