jaxgmg2_3phase_fast_tight

RL agent checkpoints trained on the JaxGMG maze environment. Contains two sweeps: a broad initial grid sweep (seeds 100-104) and a follow-up extended sweep focused on the two most interesting (alpha, discount_rate) pairs (seeds 105-150).

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

Sweep

Initial sweep (train.yaml)

Grid over alpha x discount_rate x seeds 100-104 (10B steps each):

alpha:         {0.0, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 1.0}
discount_rate: {0.97, 0.98, 0.99}
seeds:         100-104

~165 runs (some combinations missing due to failed jobs).

Extended sweep (train_extended.yaml)

Focused on the two most interesting (alpha, discount_rate) pairs from the initial sweep (5B steps each):

alpha=0.75, discount_rate=0.97  (seeds 105-150)
alpha=0.5,  discount_rate=0.99  (seeds 105-150)

92 runs.

Shared Hyperparams

rl_action=train
lr=5e-05
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=4
eval_schedule=0:1,250:2,500:5,2000:10
f_str_ckpt=al_{alpha}_g_{discount_rate}_seed_{seed}_pa_1
ckpt_dir=jaxgmg2_3phase_fast_tight
wandb_project=jaxgmg2_3phase_fast_tight

Naming Schema

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

Missing runs

5 runs from the initial sweep failed and are not present:

• al_0.0_g_0.97_seed_100_pa_1
• al_0.0_g_0.99_seed_100_pa_1
• al_0.45_g_0.97_seed_103_pa_1
• al_0.45_g_0.98_seed_100_pa_1
• al_0.4_g_0.99_seed_104_pa_1

Reproduced with

See train.yaml (initial sweep), train_extended.yaml (extended sweep), and train_missing.yaml (5 failed runs from initial sweep) in this repository. Run with:

timaeus run train.yaml
timaeus run train_extended.yaml
timaeus run train_missing.yaml

from the timaeus monorepo.