update configs

alex-petrenko · Jan 25, 2025 · 683cf40 · 683cf40
1 parent 305eddd
commit 683cf40
Show file tree

Hide file tree

Showing 9 changed files with 31 additions and 12 deletions.
diff --git a/train_configs/progress.py b/train_configs/progress.py
@@ -47,6 +47,7 @@
     "actor_mlp_dim": 64,
     "actor_depth": 1,
     "actor_heads": 8,
+    "model": "vit",
 }
 
 # params different between exps

diff --git a/train_configs/scaled.py b/train_configs/scaled.py
@@ -38,6 +38,7 @@
     "gamma": 0.999,
     "gae_lambda": 0.95,
     "value_loss_coeff": 0.5,
+    "model": "cnn",
 }
 
 # params different between exps

diff --git a/train_configs/vit.py b/train_configs/vit.py
@@ -47,21 +47,14 @@
     "actor_mlp_dim": 64,
     "actor_depth": 1,
     "actor_heads": 8,
+    "model": "vit",
 }
 
 # params different between exps
 params_grid = [
     {
         "seed": list(range(1)),
-        "actor_hidden_dim": [actor_scale],
-        "actor_mlp_dim": [actor_scale * 2],
-        "actor_depth": [1, 2, 3],
-        "critic_hidden_dim": [critic_scale],
-        "critic_mlp_dim": [critic_scale * 2],
-        "critic_depth": [1, 2, 3],
-    }
-    for actor_scale in [32, 64, 128]
-    for critic_scale in [64, 128, 256]
+    },
 ]
 
 experiments_list = create_experiments_helper(

diff --git a/train_configs/vit_accumulation.py b/train_configs/vit_accumulation.py
@@ -47,6 +47,7 @@
     "actor_mlp_dim": 64,
     "actor_depth": 1,
     "actor_heads": 8,
+    "model": "vit",
 }
 
 # params different between exps

diff --git a/train_configs/local.py → train_configs/vit_grid.py b/train_configs/local.py → train_configs/vit_grid.py
@@ -12,9 +12,10 @@
 
 # params for all exps
 config = {
+    "exp_tag": name,
     "env": "nethack_score",
     "run_script": "sf_examples.nethack.train_nethack",
-    "train_for_env_steps": 10_000_000,
+    "train_for_env_steps": 100_000_000,
     "num_workers": num_workers,
     "num_envs_per_worker": num_envs // num_workers,
     "worker_num_splits": 2,
@@ -32,18 +33,36 @@
     "with_wandb": True,
     "decorrelate_envs_on_one_worker": True,
     "max_grad_norm": 40.0,
-    "learning_rate": 2e-4,
+    "learning_rate": 1e-4,
     "exploration_loss_coeff": 0.001,
     "gamma": 0.999,
     "gae_lambda": 0.95,
     "value_loss_coeff": 0.5,
+    "actor_critic_share_weights": False,
+    "critic_hidden_dim": 64,
+    "critic_mlp_dim": 128,
+    "critic_depth": 2,
+    "critic_heads": 8,
+    "actor_hidden_dim": 32,
+    "actor_mlp_dim": 64,
+    "actor_depth": 1,
+    "actor_heads": 8,
+    "model": "vit",
 }
 
 # params different between exps
 params_grid = [
     {
         "seed": list(range(1)),
-    },
+        "actor_hidden_dim": [actor_scale],
+        "actor_mlp_dim": [actor_scale * 2],
+        "actor_depth": [1, 2, 3],
+        "critic_hidden_dim": [critic_scale],
+        "critic_mlp_dim": [critic_scale * 2],
+        "critic_depth": [1, 2, 3],
+    }
+    for actor_scale in [32, 64, 128]
+    for critic_scale in [64, 128, 256]
 ]
 
 experiments_list = create_experiments_helper(

diff --git a/train_configs/vit_rollout.py b/train_configs/vit_rollout.py
@@ -47,6 +47,7 @@
     "actor_mlp_dim": 64,
     "actor_depth": 1,
     "actor_heads": 8,
+    "model": "vit",
 }
 
 # params different between exps

diff --git a/train_configs/vit_rollout_big.py b/train_configs/vit_rollout_big.py
@@ -47,6 +47,7 @@
     "actor_mlp_dim": 128,
     "actor_depth": 1,
     "actor_heads": 8,
+    "model": "vit",
 }
 
 # params different between exps

diff --git a/train_configs/vit_rollout_scale_critic.py b/train_configs/vit_rollout_scale_critic.py
@@ -47,6 +47,7 @@
     "actor_mlp_dim": 128,
     "actor_depth": 1,
     "actor_heads": 8,
+    "model": "vit",
 }
 
 # params different between exps

diff --git a/train_configs/vit_rollout_scale_critic_long.py b/train_configs/vit_rollout_scale_critic_long.py
@@ -47,6 +47,7 @@
     "actor_mlp_dim": 128,
     "actor_depth": 1,
     "actor_heads": 8,
+    "model": "vit",
 }
 
 # params different between exps