TL;DR

I trained a GPT-2 model on Stockfish self-played games in the most naive way, with no search, and it can play decently. The model is trained to output the next move given the FEN string of the board (single state). While I present some gotchas and caveats, the results are quite acceptable for the amount of work and computing invested. I also present a basic attention visualiser parsing the attention of the text tokens into the board.

Table of content

• Context
  • DeepMind Paper
  • Personal Background
• Training Method
  • Technical Details
  • Straight-Forward Improvements
• Inspecting the Model
  • Attention as an Interpretation
  • Parsing Text Attention in the Board
  • Simple Interpretation
• Resources

# Context

DeepMind Paper

Deepming recently released a paper where they achieved grand-master level without search @1. They used a decoder only trained on the board state (FEN string) to predict whether the value of the board, the value of the action ($Q$ values) or the best action. I decided to take a similar approach but make it as simple as possible; I only trained the model to predict the next action (not necessarily the best).

With this approach, there is no need to compute tons of evals using the Stockfish engine, no need to filter out illegal moves and no need to use the fancy RL scheme (yes, I have been traumatised). There are some catches that come with this oversimplification that I present throughout the post, but I think it was a nice reframing.

Personal Background

For the record, this project only took me 1 full day of work to set up and around 3 more days to release everything. The overhead is mostly making this blog post, but it’s a good exercise for me, so it’s definitely worth it. I have already trained LLMs in the past, although I am far from being an expert. I have also worked on chess models, most recently lc0, and I am familiar with XAI / MI.

This project was not really a challenge for me; it was more of a hobby project I wanted to share.

# Training Method

Technical Details

General Framing: Fine-tuning of a pre-trained model.

The training code is available on the associated GitHub repository training-gpt2-on-stockfish-games (mostly quick and dirty). It is really straightforward and self-explanatory, as I relied on the Hugging Face Trainer (all the complexity is under the hood). Maybe the contribution you’ll find more worth is the Apptainer code to launch the training (not that difficult), which is cluster-friendly.

For the dataset I used the simple format of prompt + completion:

{"prompt": "FEN: {fen}\nMOVE:", "completion": " {move}"}

The run details:

• 1 epoch of the dataset
• LR of $10^{-5}$
• 12h on one A100 (40GB)

Tokenisation Gotcha

For simplicity, I didn’t change the tokeniser, so the tokenisation, BPE, is not tailored for chess modelling. For example, _r (_ for whitespace) and r get encoded differently, and nb is encoded as a single token while representing a knight and a bishop. Also note that I used the raw FEN string, so 8 actually represents 8 empty squares; definitely problematic for the model to grasp a spatial representation of the board with positional encodings (it still manages somewhat).

As discussed in the next section, changing the tokenisation is a low-hanging fruit to improve the model.

Straightforward Improvements

Prompt Engineering: Similarly to using your favourite chatbot, prompt engineering could be a simple way of making the model better. The model could be prompted (during training) with the strength, ELO, of the player to move and the final result. On inference time, the model can then be prompted with the desired ELO and outcome. It obviously is limited; don’t expect your model to reach ELO 10,000 with this method.

Conditioned Model

Recent framing of RL as a (self-)supervised learning problem showed that training an agent on sub-optimal trajectories can lead to optimal strategies should the agent be conditioned on a reward signal. For example, this is the setup of Imitation Learning @2 and Decision Transformer @3.

Tokenisation: I think one the best way to encode the board would be to only represent the board tokens, with 7 different tokens (P, N, B, R, Q, K and X for empty squares), then using vertical and horizontal positional encodings (like in ViT), as well as for meta information. A piece’s colour encodings are different for each piece, and common encodings for castling rights are player colour, en passant, half-clock move, and full move number.

Model: A text encoder with a special pooling token, [CLS], and an input sequence of length 64 representing the board using the aforementioned tokenisation. The output of the encoder would be fed to a classification head of size 1858, minimal move encodings but still complete (DeepMind didn’t get this right ^^ (neither in Alpha Zero for the record ^^)).

# Inspecting the Model

A Gradio space is available to interact with the model and its attention; check it out on Hugging Face ->.

Attention as an Interpretation

Attention is the most simple interpretation you can think of when studying Transformers. It is built into the model, and so it does reflect some internal knowledge of the model. In the autoregressive case (our case), you are often interested in “Why did my model predict this token?”. Basically, looking at a row of the attention coefficients $A$, defined by equation $\ref{eq:attention}$, for the predicted query token and the past key tokens gives you an importance distribution on the past tokens.

Interpretation Gotcha

While attention can be used to do simple interpretation, the claim needs to be measured. Indeed, when interpreting the attention coefficients, you say nothing about the output values $O=AV$. This weighted sum needs to be interpreted as a whole for the conclusion to be more rigorous.

While @4 claimed that attention is not an explanation on its own, @5 conducted a more in-depth study pointing out the limitation of attention. You should especially remember the computation graph as its whole. Namely, attention heads will be merged and, in addition, composed via the residual stream, so the computation graph becomes a mess.

Parsing Text Attention on the board

In order to project attention on the board, I simply parsed the attention per token and kept only the attention on pieces. The attention on the last part of the FEN string is summed into Meta while the rest of the tokens are summed in Dump. This is a simple workaround and proxy to circumvent the unadapted tokenisation.

Measure Gotcha

Always think twice about what you’re actually measuring. Here, the heatmaps only carry information about the piece configuration, so my heuristic would be to leave out heatmaps where less than 50% of attention is put on pieces (Config metric).

Always look at the distribution among the three categories of tokens and the raw attention figure. The raw attention can give you insight into when not to trust grouped attention, i.e. when a token refers to more than one square. This is especially important since it creates illusory patterns, so you should refrain from your deepest instinct to interpret this as a meaningful pattern (everyone loves to see patterns ^^).

Simple Interpretation

Disclaimer

The following section contains rough preliminary results, experiments and takes. Thus, everything should be taken with care.

Configuration attention: In terms of configuration attention, the crux is to predict the first letter of the move. I would mostly aim to interpret the attention of this token. While I didn’t conduct an in-depth analysis, it seems that the latter tokens are mostly doing a pooling of the previous representation.

Discovering simple heads: The most simple heads will be located in the early layers. Interpreting later heads is harder due to polysemanticity, convoluted circuits, and a high degree of composition, you name it.

[!success] Advice If you’re not familiar with chess notation or attention, maybe the better is to interact with the associated Hugging Face space. By copying the PGN (sequence of moves) or FEN (state of the board) you can visualise the plot of this post and more. It’s especially important if you disagree with my claims or don’t trust me (you shouldn’t ^^).

L1 H7: Focuses on the knights but only when it is paired in a token; see the figure 1. This is where this tokenisation gets annoying since it’s hard to say what the model might have learned from these paired tokens.

Figure 1: Configuration attention of the head 7 in layer 1.
[FEN] r1b1k1nr/pppp1ppp/2n1pq2/8/1b1PP3/2NQ4/PPP2PPP/R1B1KBNR w KQkq - 5 5,
[Distribution] Config: 0.78 Meta: 0.10 Dump: 0.13.

e2e3 opening: This opening is by far my favourite (not a strong player here ^^), but out of the 262k training games, only 13k opens with e2e3. The main response to this opening is e7e5 with 4k games. I then like to push my queen to f3, but since this combination is not in the training games, it seems that the model blunders doing c7c5. Pushing the bishop to c4 leads the model to respond, moving the knight to c6, letting the queen go for the checkmate (so in 4). This feels expected since, most likely, as a predictor, the LLM learned in a Bayesian way

since there is no prior on this sequence.

Counting Experiment

I think that counting all the games present in the dataset and comparing them to the learned distribution could be a nice low-hanging fruit experiment. The aim would be to examine the learned distribution on boards where the model has a prior against boards where it has no prior. This toy experiment could be framed in the larger picture of prediction generalisation and coherence.

c7c5 blunder: I am not giving the reason for the aforementioned blunder here but mostly pointing out the limitation of doing interpretability with attention. Inspecting the attention of the last layer after 1. e3 e5 2. Qf3 most heads are strongly attending to the pawn in e5, see figure 2, but in the end, the pawn c7 is moved. By the way, L12 H2 seems to focus on “player to move best-placed pawn”.

Figure 2: Configuration attention of the head 2 in layer 12.
[FEN] rnbqkbnr/pppp1ppp/8/4p3/8/4PQ2/PPPP1PPP/RNB1KBNR b KQkq - 1 2,
[Distribution] Config: 0.95 Meta: 0.01 Dump: 0.04.

b8c6 second blunder: First, it seems to fail to see any danger or threat and continue to focus on pieces that would develop his position (like the knight move). Then, as a white, before the kill, the model doesn’t find the checkmate with the queen (obviously out-of-distribution here). It is important to note that all training games are long since the self-play means two players of approximately the same level.

Adaptative Strength Experiment

It could be interesting to see if such a predictor model adapts its level to the adversary. Did it blunder because I played such weak moves (w.r.t. to the meta)? Does it play better if you use the book openings? This might not be straightforward since this correlates with out-of-distribution evaluation.

This analysis is only preliminary, so feel free to continue it and send me feedback; the discussion happens in this Discord thread on Eleuther AI.

# Resources

With this blog post, I release everything needed to reproduce, explore, understand and extend this work. In summary:

• The model ->
• The dataset ->
• The interactive space ->
• The code for the training and the space ->

References

1. Ruoss, Anian, et al. “Grandmaster-Level Chess Without Search.” ArXiv, 2024, /abs/2402.04494.
2. Jonathan Ho et al. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
3. Chen, Lili, et al. “Decision Transformer: Reinforcement Learning via Sequence Modeling.” ArXiv, 2021, /abs/2106.01345.
4. Jain, Sarthak, and Byron C. Wallace. “Attention Is Not Explanation.” ArXiv, 2019, /abs/1902.10186.
5. Wiegreffe, Sarah, and Yuval Pinter. “Attention Is Not Not Explanation.” ArXiv, 2019, /abs/1908.04626.

TL;DR

Let’s dive into MARL training with BenchMARL and how to scale up your experiments to a cluster. We’ll cover environment and model customisation, as well as how to run your experiments on a cluster even without internet access.

Table of Contents

• Context
  • BenchMARL
• MARL Training
  • Basic Setup
  • Custom Task
  • Custom Model
• Cluster Training
  • Cluster Setup
  • Running Experiments
  • No Internet
  • Results
• Resources

# Context

As a part of my PhD, I’m working on Multi-Agent Reinforcement Learning (MARL), more precisely on interpretability of MARL, as I outlined in this article. Yet to be able to interpret any agent it’s important to master MARL training. My primary goal is to spend as little time as possible training MARL agents. I also want to use classical MARL algorithms and environments to be able to contextualise my results with the state of the art, and after testing a few tools, I chose to go with BenchMARL.

BenchMARL

BenchMARL is a specialised library designed to ease MARL training. It provides a standardised interface that enables reproducibility and fair benchmarking across various MARL algorithms and environments.

Extra

BenchMARL is really well packaged, easily extensible and already embeds configs’ defaults.

BenchMARL’s backends:

• TorchRL: provides a standardised interface for MARL algorithms and environments
• Hydra: provides a flexible and modular configuration system
• marl-eval: provides a standardised evaluation system

I encourage you to check out the BenchMARL documentation for more details.

# MARL Training

Basic Setup

Along with this blog post, I have prepared a repository with a basic setup of BenchMARL. Please refer to the README for installation instructions and details. You can also check some results on my wandb report.

Feedback

Feel free to open an issue if you have any questions, problems or remarks you want to share.

The first thing you might want to do is train supported algorithms on a supported environment. Let’s say you want to compare MAPPO and IPPO on Multiwalker. Such a benchmark is made of 2 independent experiments, that you can test individually using the script experiments/run_local.py, inspired by the run.py script provided by BenchMARL.

• IPPO & Multiwalker:
  uv run -m scripts.experiments.run_local algorithm=ippo task=pettingzoo/multiwalker
• MAPPO & Multiwalker:
  uv run -m scripts.experiments.run_local algorithm=mappo task=pettingzoo/multiwalker

These scripts are based on Hydra’s configuration system, which allows you to easily modify the configuration of your experiment in a YAML file or through the command line. This is especially important when you want to run multiple experiments with different configurations, e.g. for hyperparameter search. Additionally, the defaults can be loaded directly from BenchMARL since the script’s config (exp:run_local.yaml) adds it to the Hydra search path.

Tweaking

You can easily tweak an algorithm by creating a new file in the algorithm config group and iterating on it manually or with a hyperparameter search. You can similarly tweak experiments and models.

But the strength of BenchMARL is to run benchmarks, i.e. a group of reproducible experiments with a similar config. You can start by running the benchmarks/multiwalker.py script, which is equivalent to running the previous experiments (with multiple seeds) but fully baked with the powerful plots from marl-eval to compare the experiments.

• Run the benchmark: uv run -m scripts.benchmarks.multiwalker

As this becomes tedious to run on a personal computer, the next step is to run it on a cluster. Jump to the Cluster Training section for more details.

Benchmark Config

I proposed a config for the multiwalker benchmark based on object packing similar to how layers are composed. You can use it as a template to create your own benchmark configs.

Custom Tas

One of BenchMARL’s strengths is its ability to integrate custom tasks. First, let’s say I want to create a supported task variation, e.g. Multiwalker with a shared reward.

Composition bug

Unfortunately, unlike algorithms, tasks are nested in the framework folder (e.g. pettingzoo). Because of this and a bug in Hydra, it is not possible to compose defaults from nested tasks, see this issue.

But it is still possible to derive a custom task in a few steps:

• Create a config file: multiwalker/shared.yaml
• Register the task (before creating the experiment object):

• Run the experiment:
  uv run -m scripts.experiments.run_local algorithm=mappo task=multiwalker/shared

Hack

Contrary to other groups (algorithms, experiments or models), tasks need to be registered as they are not spawned directly. Indeed, they are spawned through their factored environment class.

Another kind of custom tasks are unsupported tasks from a supported environment class, for example, KAZ (Knights Archers Zombies) from PettingZoo. First, you need to create a custom task class:

Which is used in the task:

Then you can register the task:

And additionally, you can validate the task using a dataclass to serve as a schema, which you would add in a ConfigStore:

For truly custom environments, see the examples. You might also want to dig into the torchrl documentation first, to understand how to create your own environments.

Question

How to handle the KAZ vector state (B, N, F)?

Answer

As N is not an agent dimension, you cannot directly use the Mlp ((B, F) inputs) model nor the Cnn ((B, F) or (B, H, W, C) inputs) model. You’ll need either to modify the environment to output the correct shape or use a custom model based on a Cnn, Mlp or a flattening layer.

Custom Model

As noted in the previous section, KAZ vector state requires a custom model. The easiest will be to modify the MlP model, flattening any extra dimension. The basic idea is to introduce a new num_extra_dims parameter in the model config, which will be used to flatten the input tensor.

This parameter will first be used in the _perform_checks method to check that the input tensor has the correct shape, and finally in the _forward method to simplify the input tensor:

It then needs to come with a config file extra_mlp.yaml and should be registered in the model_config_registry:

And that’s it! You can now use your custom model in your experiments.

Extra Dims

See the refined PR here.

# Cluster Training

Now that we can run the experiments we want locally, let’s scale up our experiments to a cluster.

Tweaks and tricks

The more optimisations you’ll want on the training process, the more you’ll need to dig into the torchrl backend for more control over the environments and objects, and in the torch backend for more control over the models and tensors.

Cluster Setup

After having tested a few different setups, I ended up settling on a full uv config. Here is my opinionated setup:

• [Optional] Setup git in your project
• [Optional] Link a GitHub repository
• Sync your project to the cluster (I prefer using a git remote for easy bidirectional edits but rsync or scp can be simpler)
• Use uv sync to install the dependencies on the cluster (this step can be delegated to the job, should you have an internet connection, see No Internet)
• Run your jobs or notebooks on the cluster

You should do this on a Work or Draft partition, avoid installing in your home directory which might have a limited space (in terms of disk space and inodes). Always be aware of your cluster configuration and check with your admin in case of doubts.

Pros

• It’s super easy to use once you’re used to uv
• The configuration is the same as the local one
• Fully compatible with slurm jobs and jupyter notebooks/hubs
• Works without internet on the nodes (see No Internet)

Cons

• You might miss package optimisations tailored for your cluster
• It can consume a lot of inodes (you might need to remove old virtual environments)
• Not super compatible with classical cuda echosystem

Running Experiments

The easy part with the setup I presented is that the same script can be used to run your experiments locally or on a cluster. So, except for the slurm arguments, which might differ depending on your cluster, you can use the same script.

A typical slurm script, that you can launch using sbatch launch/bench:multiwalker-jz.sh, would look like this:

And to make use of the GPU you can just switch the experiment config by adding the argument experiment=gpu. It will simply load the default experiment config (base_experiment) and override the gpu config:

Now you’re ready to launch a bunch of jobs doing wild hyperparameter search, with bigger models, bigger batch sizes, etc.

JupyterHub

You can find an example of how to run a notebook on a cluster JupyterHub in my project template, this can vary depending on your cluster (example made for JeanZay).

No Internet

Some clusters, for security reasons, don’t allow internet access on the nodes (e.g. JeanZay in France). This can be a problem if you want to set up your environment (e.g. with uv) directly on the nodes (which can be easier). So let’s see how we can easily transpose what we’ve seen so far to a cluster without an internet connection.

As noted in the Cluster Config Setup section, you should install your dependencies before starting your jobs. If you don’t have an internet connection on the head node (never seen this though), you might try to transfer your local setup, or if allowed, use a Docker image.

Then, the only thing you need to do is to remove or disable the tools that require internet access. In our experiments, you just need to use wandb in “offline” mode (e.g. using experiment=gpu_offline in the script):

And when running a script with uv, you should use the --no-sync flag to avoid syncing your dependencies again. Depending on your use case, you might need to download your datasets to a special partition beforehand.

Wandb offline

Unfortunately, wandb has a bug related to the config logging for offline runs, see this issue. In order to circumvent this issue, you’ll need to slightly modify the logger of the Experiment class in BenchMARL to pass the config to wandb.init, see this PR.

Results

You’ll find more detailed script launchers in the launch folder and in the commit history (I did quite a few tries). I was able to match results from the original paper @1, and for some runs, get higher returns. I’ll try to make a clean report on wandb when I find some time, which will be here.

Some considerations if you want to go deeper. In this post, I only scratched the surface of the possibilities of BenchMARL. When using it to run pure experiments from YAML, you’ll find some limitations about customisation, which are totally overcome by manipulating the Python classes directly. One of the most powerful ways to customise Experiment is through callbacks that can enable parameter scheduling (e.g. for lr or batch_size) or add custom pre/post-processing.

Further Customisation

As you saw in the previous sections, BenchMARL is a really powerful tool to train MARL agents. However, there are still some things you might want to customise to fit your experiment’s needs. And the beauty of BenchMARL is that you can keep all the features you want and simply add your secret sauce, whether it’s an algorithm, an environment, or anything else.

# Resources

To learn more about BenchMARL and MARL in general, here are some valuable resources:

• BenchMARL Documentation
• BenchMARL GitHub Repository
• TorchRL Documentation
• Hydra Configuration Framework

BenchMARL’s Discord community is also a great place to ask questions and share experiences with other users. Also, feel free to open an issue or a PR if you want to add or suggest an edit.

References

1. Gupta, Jayesh K. et al. “Cooperative Multi-agent Control Using Deep Reinforcement Learning.” AAMAS Workshops (2017).