jaxgmg2_antenna

50 RL agent checkpoints from the antenna-ban experiment on the JaxGMG maze environment. The experiment tests whether banning states identified as important by susceptibility analysis affects phase transition dynamics.

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

Groups

Treatment group (30 runs)

antban2_al_0.6_g_0.98_id_XX_seed_9806XX with run_id 50-79.

These ban the states identified as useful by susceptibility analysis (top-left corner):

env_rule = "ban:c[2,1], m[1,:]; c[1,2], m[:,1]"

See antenna_train.yaml.

Control group (20 of 30 runs)

control_al_0.6_g_0.98_id_XX_seed_9806XXX with run_id 80-109 (10 missing due to failed jobs).

These ban states around the opposite corner (bottom-right), not identified as useful:

env_rule = "ban:c[10,11], m[11,:]; c[11,10], m[:,11]"

See control_train.yaml.

Missing run_ids: 84, 87, 90, 92, 93, 95, 101, 102, 106, 109. To fill these, see control_train_missing.yaml.

Shared Hyperparams

rl_action=train
alpha=0.6
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=5
log_optimizer_state=True
eval_schedule=0:1,250:2,500:5,2000:10
wandb_project=jaxgmg2_antennaban
use_wandb=True
use_hf=True

Reproduced with

timaeus run antenna_train.yaml
timaeus run control_train.yaml
timaeus run control_train_missing.yaml

from the timaeus monorepo.