visualise_malt_rollouts.py

Visualise the rollouts of a MALT experiment as a forest of trees.

The MALT [MSD+24] trainer samples a set of trees of responses. These are stored in a flat array. This script provides a command line interface to visualise the trees in a forest.

By default, the visualisation is saved to the ‘analysis’ subdirectory of the checkpoint.

scripts/visualise_malt_rollouts.py

Visualise the rollouts of a MALT experiment as a forest of trees.

usage: scripts/visualise_malt_rollouts.py [-h] [--stdout]
                                          [--format {yaml,jsonl}]
                                          [--wandb_project WANDB_PROJECT]
                                          [--wandb_entity WANDB_ENTITY]
                                          [--overwrite]
                                          run_id iteration

run_id

The id of the wandb run to download the rollouts from.

iteration

The experiment iteration of the rollouts. Either an integer or ‘test’.

-h, --help

show this help message and exit

--stdout, -s

Whether to print the visualisation to stdout. Otherwise, it is saved to the ‘analysis’ subdirectory of the checkpoint.

--format {yaml,jsonl}, -f {yaml,jsonl}

The format to visualise the MALT forest in.

--wandb_project <wandb_project>

The wandb project to use.

--wandb_entity <wandb_entity>

The wandb entity to use.

--overwrite, -o

Whether to overwrite an existing visualisation.

The MALT [MSD+24] trainer samples a set of trees of responses. These are stored in a flat array. This script provides a command line interface to visualise the trees in a forest.

By default, the visualisation is saved to the ‘analysis’ subdirectory of the checkpoint.