## Organization

**Reproducing figures.** 
The notebooks used to plot Figures 1 and 2 are located in the directory `figures`;
the figures can be replicated by running the cells.

**Figure data.**
The data used to generate the figures is located in the directory `data`.

**Source code.**
The training script is located in `src/train_star.py`
with default parameters in `configs/train.yaml`.

## Setup
To install requirements, run

```bash
conda create -n coverage python=3.10
conda activate coverage
pip install -r requirements.txt
```

## Figure 1

First, we pre-train a model using the following command,
and use the final checkpoint for finetuning.

```bash
python src/train_star.py +training=offline +dataset=pretrain train.max_iters=200000 eval.save_model_interval=10000
```

For finetuning, run

```bash
python src/train_star.py +training=online +dataset=finetune train.max_iters=50000 optimizer.learning_rate=5e-6 eval.logprob_interval=200 eval.eval_interval=200
```

Parameters like the ratios in the dataset mixture can be set in `configs/dataset/finetune.yaml`.


## Figure 2

We train for tasks with horizons $H \in \{8, 16, 24\}$:

```bash
python src/train_star.py +training=offline +dataset=horizon train.max_iters=50000 data.params.0.layers=8 

python src/train_star.py +training=offline +dataset=horizon train.max_iters=50000 data.params.0.layers=16 

python src/train_star.py +training=offline +dataset=horizon train.max_iters=50000 data.params.0.layers=24 
```


