al_1.0_g_0.98_x100

101 RL agent checkpoints trained on the JaxGMG maze environment with alpha=1.0 and discount_rate=0.98. Baseline collection for studying phase transitions in RL training dynamics with full value function weighting.

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

Sweep

Seed sweep: 0-100 (101 runs).

Shared Hyperparams

rl_action=train
alpha=1.0
discount_rate=0.98
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=4
log_optimizer_state=True
deterministic=True
eval_schedule=0:1,250:2,500:5,2000:10
ckpt_dir=al_1.0_g_0.98_x100
wandb_project=jaxgmg2_train
use_wandb=True
use_hf=True

Naming Schema

Checkpoints are named al_1.0_g_0.98_seed_{seed}.

Reproduced with

See train.yaml in this repository. Run with:

timaeus run train.yaml

from the timaeus monorepo.