Matrix Games in JAX

Enter the Matrax! 😎

Installation | Quickstart

Matrax is a lightweight suite of 2-player matrix game environments written in JAX. It is a direct re-implementation of the matrix games provided in this repository from the AARG. It follows the Jumanji RL environment suite API developed by InstaDeep.

2-Player Matrix Games 🧑‍🤝‍🧑

Category	Shape (action space)	Registered Version(s)	Source
🔻 Penalty Game	3 x 3	Penalty-{k}-{state}-v0	code
🧗‍♀️ Climbing Game	3 x 3	Climbing-{state}-v0	code
🤝 No-Conflict Games	2 x 2	NoConflict-{id}-{state}-v0	code
💣 Conflict Games	2 x 2	Conflict-{id}-{state}-v0	code

Installation 🎬

You can install the latest release of Matrax from PyPI:

pip install matrax

Alternatively, you can install the latest development version directly from GitHub:

pip install git+https://github.com/instadeepai/matrax.git

Matrax has been tested on Python 3.9. Note that because the installation of JAX differs depending on your hardware accelerator, we advise users to explicitly install the correct JAX version (see the official installation guide).

Quickstart ⚡

import jax
import matrax

# Instantiate a matrix game environment using the registry
env = matrax.make("Penalty-25-stateless-v0")

# Reset your (jit-able) environment
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)

# Interact with the (jit-able) environment
action = env.action_spec.generate_value()          # Action selection (dummy value here)
state, timestep = jax.jit(env.step)(state, action)   # Take a step and observe the next state and time step

Registry and Versioning 📖

Like Jumanji, Matrax keeps a strict versioning of its environments for reproducibility reasons. We maintain a registry of standard environments with their configuration. For each environment, a version suffix is appended, e.g. EnvironmentName-v1. When changes are made to environments that might impact learning results, the version number is incremented by one to prevent potential confusion.

Citing Matrax 📝

If you use any of the environment implementations in Matrax please consider citing this library.

@software{matrax2023github,
  author = {Arnu Pretorius},
  title = {Matrax: Matrix Games in JAX},
  url = {http://github.com/instadeepai/matrax},
  version = {v0.0.4-alpha},
  year = {2023},
}

See Also 🔎

Other works have embraced the approach of writing RL environments in JAX. In particular, we suggest users check out the following sister repositories:

• 🌴 Jumanji is a diverse suite of scalable reinforcement learning environments.
• 🦾 Brax is a differentiable physics engine that simulates environments made up of rigid bodies, joints, and actuators.
• 🏋️‍ Gymnax implements classic environments including classic control, bsuite, MinAtar and a collection of meta RL tasks.
• 🎲 Pgx provides classic board game environments like Backgammon, Shogi, and Go.