Add Figure 5 Ant flat-ground and maze experiments (part 4)

configs/figure5/medium/reservoir_ro/seed_17/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: reservoir_ro
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 17
+seed_roles:
+  optimization: 17
+  sampling: 17
+  collection: 1700000
+  evaluation: 9300000

configs/figure5/medium/reservoir_ro/seed_18/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: reservoir_ro
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 18
+seed_roles:
+  optimization: 18
+  sampling: 18
+  collection: 1800000
+  evaluation: 9300000

configs/figure5/medium/reservoir_ro/seed_19/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: reservoir_ro
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 19
+seed_roles:
+  optimization: 19
+  sampling: 19
+  collection: 1900000
+  evaluation: 9300000

configs/figure5/medium/reservoir_ro/seed_2/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: reservoir_ro
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 2
+seed_roles:
+  optimization: 2
+  sampling: 2
+  collection: 200000
+  evaluation: 9300000

configs/figure5/medium/reservoir_ro/seed_3/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: reservoir_ro
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 3
+seed_roles:
+  optimization: 3
+  sampling: 3
+  collection: 300000
+  evaluation: 9300000

configs/figure5/medium/reservoir_ro/seed_4/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: reservoir_ro
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 4
+seed_roles:
+  optimization: 4
+  sampling: 4
+  collection: 400000
+  evaluation: 9300000

configs/figure5/medium/reservoir_ro/seed_5/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: reservoir_ro
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 5
+seed_roles:
+  optimization: 5
+  sampling: 5
+  collection: 500000
+  evaluation: 9300000

configs/figure5/medium/reservoir_ro/seed_6/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: reservoir_ro
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 6
+seed_roles:
+  optimization: 6
+  sampling: 6
+  collection: 600000
+  evaluation: 9300000

configs/figure5/medium/reservoir_ro/seed_7/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: reservoir_ro
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 7
+seed_roles:
+  optimization: 7
+  sampling: 7
+  collection: 700000
+  evaluation: 9300000

configs/figure5/medium/reservoir_ro/seed_8/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: reservoir_ro
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 8
+seed_roles:
+  optimization: 8
+  sampling: 8
+  collection: 800000
+  evaluation: 9300000

configs/figure5/medium/reservoir_ro/seed_9/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: reservoir_ro
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 9
+seed_roles:
+  optimization: 9
+  sampling: 9
+  collection: 900000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_0/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 0
+seed_roles:
+  optimization: 0
+  sampling: 0
+  collection: 0
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_1/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 1
+seed_roles:
+  optimization: 1
+  sampling: 1
+  collection: 100000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_10/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 10
+seed_roles:
+  optimization: 10
+  sampling: 10
+  collection: 1000000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_11/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 11
+seed_roles:
+  optimization: 11
+  sampling: 11
+  collection: 1100000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_12/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 12
+seed_roles:
+  optimization: 12
+  sampling: 12
+  collection: 1200000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_13/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 13
+seed_roles:
+  optimization: 13
+  sampling: 13
+  collection: 1300000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_14/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 14
+seed_roles:
+  optimization: 14
+  sampling: 14
+  collection: 1400000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_15/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 15
+seed_roles:
+  optimization: 15
+  sampling: 15
+  collection: 1500000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_16/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 16
+seed_roles:
+  optimization: 16
+  sampling: 16
+  collection: 1600000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_17/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 17
+seed_roles:
+  optimization: 17
+  sampling: 17
+  collection: 1700000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_18/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 18
+seed_roles:
+  optimization: 18
+  sampling: 18
+  collection: 1800000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_19/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 19
+seed_roles:
+  optimization: 19
+  sampling: 19
+  collection: 1900000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_2/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 2
+seed_roles:
+  optimization: 2
+  sampling: 2
+  collection: 200000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_3/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 3
+seed_roles:
+  optimization: 3
+  sampling: 3
+  collection: 300000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_4/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 4
+seed_roles:
+  optimization: 4
+  sampling: 4
+  collection: 400000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_5/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 5
+seed_roles:
+  optimization: 5
+  sampling: 5
+  collection: 500000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_6/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 6
+seed_roles:
+  optimization: 6
+  sampling: 6
+  collection: 600000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_7/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 7
+seed_roles:
+  optimization: 7
+  sampling: 7
+  collection: 700000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_8/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 8
+seed_roles:
+  optimization: 8
+  sampling: 8
+  collection: 800000
+  evaluation: 9300000

configs/figure5/medium/transformer/seed_9/final.yaml

@@ -0,0 +1,93 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 30.0
+policy:
+  name: transformer
+  actor_kwargs:
+    goal_router_hidden:
+    - 256
+    - 256
+    goal_router_codebook:
+    - - -6.0
+      - 10.0
+    - - -10.0
+      - 6.0
+    - - -6.0
+      - 6.0
+    - - 2.0
+      - 6.0
+    - - 6.0
+      - 6.0
+    - - -6.0
+      - 2.0
+    - - -2.0
+      - 2.0
+    - - 2.0
+      - 2.0
+    - - -10.0
+      - -2.0
+    - - -6.0
+      - -2.0
+    - - 2.0
+      - -2.0
+    - - 6.0
+      - -2.0
+    - - 10.0
+      - -2.0
+    - - 10.0
+      - -6.0
+    - - 10.0
+      - -10.0
+    goal_router_codebook_latch_radius: 0.45
+    goal_router_codebook_active_pairs:
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - -6.0
+    - - - 6.0
+        - -10.0
+      - - -10.0
+        - 10.0
+    - - - -10.0
+        - -6.0
+      - - 10.0
+        - 6.0
+    goal_router_codebook_pair_match_radius: 2.0
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_Medium_Diverse_GR-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 60000
+  dagger_steps: 30000
+  dagger_iterations: 0
+  updates_per_stage: 5000
+  batch_size: 64
+  burn_in: 16
+  learning_rate: 0.0001
+  gradient_clip: 1.0
+  actor_train_mode: base_all
+  goal_input_features: 4
+  loss_space: pre_tanh
+  sampling_mode: episode_uniform
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+  student_goal: final
+  evaluation_waypoint_planner: none
+evaluation:
+  action_smoothing: 0.95
+seed: 9
+seed_roles:
+  optimization: 9
+  sampling: 9
+  collection: 900000
+  evaluation: 9300000

configs/figure5/umaze/adult/seed_0/final.yaml

@@ -0,0 +1,69 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+seed: 0

configs/figure5/umaze/adult/seed_1/final.yaml

@@ -0,0 +1,71 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 1

configs/figure5/umaze/adult/seed_10/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 10
+seed_roles:
+  optimization: 10
+  sampling: 10
+  collection: 1000000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_11/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 11
+seed_roles:
+  optimization: 11
+  sampling: 11
+  collection: 1100000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_12/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 12
+seed_roles:
+  optimization: 12
+  sampling: 12
+  collection: 1200000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_13/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 13
+seed_roles:
+  optimization: 13
+  sampling: 13
+  collection: 1300000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_14/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 14
+seed_roles:
+  optimization: 14
+  sampling: 14
+  collection: 1400000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_15/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 15
+seed_roles:
+  optimization: 15
+  sampling: 15
+  collection: 1500000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_16/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 16
+seed_roles:
+  optimization: 16
+  sampling: 16
+  collection: 1600000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_17/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 17
+seed_roles:
+  optimization: 17
+  sampling: 17
+  collection: 1700000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_18/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 18
+seed_roles:
+  optimization: 18
+  sampling: 18
+  collection: 1800000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_19/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 19
+seed_roles:
+  optimization: 19
+  sampling: 19
+  collection: 1900000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_2/final.yaml

@@ -0,0 +1,71 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 2

configs/figure5/umaze/adult/seed_3/final.yaml

@@ -0,0 +1,71 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 3

configs/figure5/umaze/adult/seed_4/final.yaml

@@ -0,0 +1,71 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 4

configs/figure5/umaze/adult/seed_5/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 5
+seed_roles:
+  optimization: 5
+  sampling: 5
+  collection: 500000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_6/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 6
+seed_roles:
+  optimization: 6
+  sampling: 6
+  collection: 600000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_7/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 7
+seed_roles:
+  optimization: 7
+  sampling: 7
+  collection: 700000
+  evaluation: 9200000

configs/figure5/umaze/adult/seed_8/final.yaml

@@ -0,0 +1,76 @@
+algorithm: antmaze_teacher_distillation
+observation_clip: 10.0
+policy:
+  name: connectome_rate
+  actor_kwargs:
+    dt: 0.01
+    tau_init: 0.02
+    use_lora: false
+    input_target: sensory
+    goal_adapter_hidden:
+    - 256
+    - 128
+    goal_adapter_features: 109
+  value_hidden:
+  - 256
+  - 256
+  value_activation: relu
+connectome:
+  species: adult
+  kind: original
+  adult_dir: larva/data/adult
+  chunksize: 250000
+  min_syn_count: 1
+  normalization: row_l1
+  global_gain: 1.0
+  sensory_selector:
+    source: classification
+    field: flow
+    include:
+    - afferent
+  output_selector:
+    source: classification
+    field: super_class
+    include:
+    - motor
+  polarity:
+    unknown_action: zero
+    nt_signs:
+      ACH: 1.0
+      GLUT: 1.0
+      GABA: -1.0
+      DA: 1.0
+      SER: 1.0
+      OCT: 1.0
+env:
+  class_path: larva.ant_transfer.env.AntMazeTransferEnv
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+resolved_env:
+  env_id: AntMaze_UMaze-v5
+  goal_feature_mode: world
+  continuing_task: false
+imitation:
+  demonstration_steps: 40000
+  dagger_steps: 30000
+  dagger_iterations: 1
+  updates_per_stage: 10000
+  batch_size: 32
+  burn_in: 8
+  learning_rate: 0.0003
+  gradient_clip: 1.0
+  actor_train_mode: goal_adapter_only
+  goal_input_features: 4
+  loss_space: pre_tanh
+  checkpoint_interval_updates: 500
+navigation:
+  waypoint_planner: official_grid_waypoints
+evaluation:
+  action_smoothing: 0.8
+seed: 8
+seed_roles:
+  optimization: 8
+  sampling: 8
+  collection: 800000
+  evaluation: 9200000