The Yokai Learning Environment

This repository implements the Yokai Learning Environment (YLE) based on the cooperative card game Yokai. It accompanies our paper "The Yokai Learning Environment: Tracking Beliefs Over Space and Time".

Organisation

We add our implementation to the JaxMARL project but we do not follow their repository organisation. While the environment is colocated with all other JaxMARL environments our training and evaluation code can be found in src and not in baselines. We give an overview over the most important aspects of the code here:

.
jaxmarl/
|-- ...
|-- environments/
    |-- ...
    |-- yokai/
        |-- __init__.py             # Exports are handeld here
        |-- constants.py            # Helper functions
        |-- yokai_grid.py           # The yokai environment
src/                                # Training and evaluation functions

Environment Versions

We provide two Yokai versions. The one explored in the paper is available under the configuration name yokai_grid_9_game. All versions are yokai_grid_9_game and yokai_grid_16_game. First, yokai_grid_9_game implements the 9 card version in the 9x9 grid. Second, yokai_grid_16_game implements the 16 card version in a 10x10 grid. The two versions can be found at the bottom of the jaxmarl/environments/yokai/yokai_grid.py file.

Training, Evaluation and Reproducability

Running a training itself is easy and boils down to setting the run parameters correctly (refer to the config class in train_*.py). The config already sets reasonable defaults. A example command would be:

CUDA_VISIBLE_DEVICES=${device} python3 -m src/train_*.py

Credit

The implementation is based on several projects:

1. Most importantly, this work is based on the JaxMARL project. The YLE follows its provided interface.
2. Our TransformerXL implementation is based on the transformerXL_PPO_JAX project.
3. The code for running sweeps is mostly from the Sampling for Learnability project.