Skip to main content

Matrix games in JAX

Project description

Matrax logo

Matrix Games in JAX

Python Versions PyPI version Tests Code Style MyPy License

Enter the Matrax! 😎

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+

Matrax has been tested on Python 3.8 and 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.

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.

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

matrax-0.0.4.tar.gz (16.3 kB view hashes)

Uploaded Source

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page