jaxgmg2_3phase_unique

224 RL agent checkpoints trained on the JaxGMG maze environment across a grid of alpha and discount_rate values. Without optimizer state saved (see jaxgmg2_3phase_optim_state for the version with optimizer state).

WandB: https://wandb.ai/devinterp/jaxgmg2_3phase_unique

Sweep

Grid over alpha x discount_rate x run_id (0-14):

alpha:         {0.4, 0.5, 0.6, 0.7, 1.0}
discount_rate: {0.97, 0.98, 0.99}
run_id:        0-14

5 x 3 x 15 = 225 combinations. 1 run missing: al_0.7_g_0.98_id_14_seed_980714.

Seed is derived from run_id via: seed = int(discount_rate*100)*10000 + int(alpha*10)*100 + run_id

Shared Hyperparams

rl_action=train
lr=5e-05
num_total_env_steps=10000000000
num_rollout_steps=64
num_levels=9600
cheese_loc=any
env_layout=open
env_size=13
mask_type=first_episode
use_prev_action=False
grad_acc_per_chunk=5
log_optimizer_state=False
eval_schedule=0:1,250:2,500:5,2000:10
seed_formula={int(discount_rate*100):02d}{int(alpha*10):02d}{run_id:02d}
f_str_ckpt=al_{alpha}_g_{discount_rate}_id_{run_id}_seed_{seed}
ckpt_dir=jaxgmg2_3phase_unique
wandb_project=jaxgmg2_3phase_unique
use_wandb=True
use_hf=True

Naming Schema

Checkpoints are named al_{alpha}_g_{discount_rate}_id_{run_id}_seed_{seed}.

Reproduced with

See train.yaml in this repository. Run with:

timaeus run train.yaml

To fill the 1 missing run, see train_missing.yaml:

timaeus run train_missing.yaml

from the timaeus monorepo.