# Problem Description

## Background
In the paper *Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task* (Li et al. 2023), the authors trained a GPT model to predict legal moves in the board game Othello. They claimed that the network learned a non-linear representation of the board state, representing the colour of a given square on the board (white/black/empty).

In a follow-up paper *Emergent Linear Representations in World Models of Self-Supervised Sequence Models* (Nanda et al. 2023) however, the authors demonstrated that the network actually learned a linear representation of the board state, where each square is represented as belonging to the current player ('mine'), the other player ('theirs') or is empty. In other words, instead of representing the board state as white player/black player/empty, it is represented relative to the player whose turn it is to make a move.

The papers will be available inside the provided repository in both PDF format as `li_et_al_2023.pdf` and `nanda_et_al_2023.pdf` and markdown format as `li_et_al_2023.md` and `nanda_et_al_2023.md` if you need to refer to them.

## Extension to be implemented
Your task is to modify the codebase provided (the original codebase of Li et al. 2023) to train a linear probe on the model's representations using the classes mine/theirs/empty proposed in Nanda et al. (2023) and show that the linear probing accuracy on a held-out test set is higher than linear probing accuracy using white/black/empty. You just need to implement the new experiment with mine/theirs/empty and there is no need to re-run the original experiment with white/black/empty.

Assume that model checkpoints and data can be found in the `ckpts/` and `data/` directories respectively, although these may not be visible in the provided repository right now.

Try your best to keep everything else not specified above constant in the original repository. Also, the environment is already set up, so you do not need to install any dependencies or download any more datasets/models. Please refer to `environment.yml` in the repository for an overview of the installed libraries and their versions.

### Evaluation
To evaluate the extension, we will run the script `produce_probes.sh`. As part of your solution, you should implement a new boolean command-line flag in the Python file being called inside this script to train the model with the new 'mine/theirs' formulation. Please modify the Python command in `produce_probes.sh` by adding this flag, so we run the modified version of the experiment when executing `produce_probes.sh` rather than the original experiment.

We will use the result files containing layerwise probing accuracies saved under `ckpts/` to evaluate the final experimental outcome.