results: reject diagonal SDIL on PickupLoc

13 files changed, 3417 insertions, 0 deletions

diff --git a/results/babyai_shared/pickup_p2/seed4101_bp.json b/results/babyai_shared/pickup_p2/seed4101_bp.json
new file mode 100644
index 0000000..371cc33
--- /dev/null
+++ b/results/babyai_shared/pickup_p2/seed4101_bp.json
@@ -0,0 +1,262 @@
+{
+  "condition": "bp",
+  "config": {
+    "action_dim": 7,
+    "context_gain": 1.0,
+    "hidden_layers": 4,
+    "input_dim": 3964,
+    "learning_rate": 0.03,
+    "mission_dim": 23,
+    "momentum": 0.9,
+    "reciprocal_learning_rate": 0.03,
+    "weight_decay": 0.0001,
+    "width": 256
+  },
+  "data": {
+    "color_cardinality": 6,
+    "elapsed_seconds": 81.82925486564636,
+    "env_id": "BabyAI-PickupLoc-v0",
+    "minigrid_version": "3.1.0",
+    "object_cardinality": 11,
+    "protocol": "babyai_shared_feedback_b0",
+    "rollout_episodes": 500,
+    "rollout_seed_start": 200000,
+    "state_cardinality": 3,
+    "train_episodes": 20000,
+    "train_steps": 121467,
+    "validation_episodes": 2000,
+    "validation_steps": 12347,
+    "vocabulary": [
+      "<unk>",
+      "a",
+      "ball",
+      "behind",
+      "blue",
+      "box",
+      "front",
+      "green",
+      "grey",
+      "in",
+      "key",
+      "left",
+      "of",
+      "on",
+      "pick",
+      "purple",
+      "red",
+      "right",
+      "the",
+      "up",
+      "yellow",
+      "you",
+      "your"
+    ]
+  },
+  "epoch_history": [
+    {
+      "epoch": 1,
+      "train_loss": 0.617732492434351
+    },
+    {
+      "epoch": 2,
+      "train_loss": 0.4960194272116611
+    },
+    {
+      "epoch": 3,
+      "train_loss": 0.46592118231873764
+    },
+    {
+      "epoch": 4,
+      "train_loss": 0.4488015041225835
+    },
+    {
+      "epoch": 5,
+      "train_loss": 0.43406150341033933
+    },
+    {
+      "epoch": 6,
+      "train_loss": 0.40718337021375955
+    },
+    {
+      "epoch": 7,
+      "train_loss": 0.3807649670776568
+    },
+    {
+      "epoch": 8,
+      "train_loss": 0.3537148001319484
+    },
+    {
+      "epoch": 9,
+      "train_loss": 0.32643107909905283
+    },
+    {
+      "epoch": 10,
+      "train_loss": 0.3063848673356207
+    },
+    {
+      "epoch": 11,
+      "train_loss": 0.28984395764375986
+    },
+    {
+      "epoch": 12,
+      "train_loss": 0.2761740297392795
+    },
+    {
+      "epoch": 13,
+      "train_loss": 0.2647564570841036
+    },
+    {
+      "epoch": 14,
+      "train_loss": 0.25636729594908264
+    },
+    {
+      "epoch": 15,
+      "train_loss": 0.24771365335113124
+    },
+    {
+      "epoch": 16,
+      "train_loss": 0.2439696610287616
+    },
+    {
+      "epoch": 17,
+      "train_loss": 0.2316527129474439
+    },
+    {
+      "epoch": 18,
+      "train_loss": 0.22913515367006002
+    },
+    {
+      "epoch": 19,
+      "train_loss": 0.21849248581811
+    },
+    {
+      "epoch": 20,
+      "train_loss": 0.21554886654803626
+    },
+    {
+      "epoch": 21,
+      "train_loss": 0.20946172889910247
+    },
+    {
+      "epoch": 22,
+      "train_loss": 0.20310457860168657
+    },
+    {
+      "epoch": 23,
+      "train_loss": 0.2004946485946053
+    },
+    {
+      "epoch": 24,
+      "train_loss": 0.19238460020015113
+    },
+    {
+      "epoch": 25,
+      "train_loss": 0.19454976933567147
+    },
+    {
+      "epoch": 26,
+      "train_loss": 0.1896667920915704
+    },
+    {
+      "epoch": 27,
+      "train_loss": 0.17798410443883192
+    },
+    {
+      "epoch": 28,
+      "train_loss": 0.18044299861318186
+    },
+    {
+      "epoch": 29,
+      "train_loss": 0.17217574811295458
+    },
+    {
+      "epoch": 30,
+      "train_loss": 0.1684879423128931
+    },
+    {
+      "epoch": 31,
+      "train_loss": 0.16110639711743907
+    },
+    {
+      "epoch": 32,
+      "train_loss": 0.1620375840914877
+    },
+    {
+      "epoch": 33,
+      "train_loss": 0.15719828299785915
+    },
+    {
+      "epoch": 34,
+      "train_loss": 0.1522270987222069
+    },
+    {
+      "epoch": 35,
+      "train_loss": 0.14935507338297996
+    },
+    {
+      "epoch": 36,
+      "train_loss": 0.1444174110889435
+    },
+    {
+      "epoch": 37,
+      "train_loss": 0.14394060233705921
+    },
+    {
+      "epoch": 38,
+      "train_loss": 0.13934059754798286
+    },
+    {
+      "epoch": 39,
+      "train_loss": 0.13100762662918944
+    },
+    {
+      "epoch": 40,
+      "train_loss": 0.1349288523667737
+    }
+  ],
+  "epochs_completed": 40,
+  "finite": true,
+  "first_nonfinite_epoch": null,
+  "mission_lesion_rollout": {
+    "episodes": 500,
+    "mean_length": 41.77,
+    "mean_return": 0.36140937500000003,
+    "success": 0.488
+  },
+  "mission_lesion_validation": {
+    "accuracy": 0.7446343241273183,
+    "loss": 0.9336490147585154
+  },
+  "predictor": [],
+  "provenance": {
+    "cuda_device_name": "NVIDIA GeForce GTX 1080",
+    "cuda_peak_allocated_bytes": 130533376,
+    "cuda_visible_devices": "0",
+    "device": "cuda:0",
+    "git_commit": "7b7d7a08a5e78c3379397f01f0c9af97d5d84fbd",
+    "git_dirty_tracked": false,
+    "minigrid_version": "3.1.0",
+    "torch_version": "2.3.1+cu118"
+  },
+  "rollout": {
+    "episodes": 500,
+    "mean_length": 22.676,
+    "mean_return": 0.66191875,
+    "success": 0.808
+  },
+  "stage": "babyai_pickup_p2",
+  "training": {
+    "batch_size": 256,
+    "history_steps": 4,
+    "model_seed": 4101,
+    "neutral_examples_per_epoch": 0,
+    "parameter_count_including_reciprocal_context_predictor": 1438215,
+    "shuffle_seed": 4101,
+    "total_wall_seconds_including_rollouts": 337.6041464805603,
+    "training_wall_seconds": 322.6370050907135
+  },
+  "validation": {
+    "accuracy": 0.878108042439459,
+    "loss": 0.3454075009463489
+  }
+}
diff --git a/results/babyai_shared/pickup_p2/seed4101_clean_kp.json b/results/babyai_shared/pickup_p2/seed4101_clean_kp.json
new file mode 100644
index 0000000..fc7d8a2
--- /dev/null
+++ b/results/babyai_shared/pickup_p2/seed4101_clean_kp.json
@@ -0,0 +1,262 @@
+{
+  "condition": "clean_kp",
+  "config": {
+    "action_dim": 7,
+    "context_gain": 1.0,
+    "hidden_layers": 4,
+    "input_dim": 3964,
+    "learning_rate": 0.03,
+    "mission_dim": 23,
+    "momentum": 0.9,
+    "reciprocal_learning_rate": 0.03,
+    "weight_decay": 0.0001,
+    "width": 256
+  },
+  "data": {
+    "color_cardinality": 6,
+    "elapsed_seconds": 81.82925486564636,
+    "env_id": "BabyAI-PickupLoc-v0",
+    "minigrid_version": "3.1.0",
+    "object_cardinality": 11,
+    "protocol": "babyai_shared_feedback_b0",
+    "rollout_episodes": 500,
+    "rollout_seed_start": 200000,
+    "state_cardinality": 3,
+    "train_episodes": 20000,
+    "train_steps": 121467,
+    "validation_episodes": 2000,
+    "validation_steps": 12347,
+    "vocabulary": [
+      "<unk>",
+      "a",
+      "ball",
+      "behind",
+      "blue",
+      "box",
+      "front",
+      "green",
+      "grey",
+      "in",
+      "key",
+      "left",
+      "of",
+      "on",
+      "pick",
+      "purple",
+      "red",
+      "right",
+      "the",
+      "up",
+      "yellow",
+      "you",
+      "your"
+    ]
+  },
+  "epoch_history": [
+    {
+      "epoch": 1,
+      "train_loss": 0.6793887242517973
+    },
+    {
+      "epoch": 2,
+      "train_loss": 0.5070055065029546
+    },
+    {
+      "epoch": 3,
+      "train_loss": 0.4780606180743167
+    },
+    {
+      "epoch": 4,
+      "train_loss": 0.46464681211270786
+    },
+    {
+      "epoch": 5,
+      "train_loss": 0.45651326932405173
+    },
+    {
+      "epoch": 6,
+      "train_loss": 0.4414772361203244
+    },
+    {
+      "epoch": 7,
+      "train_loss": 0.4295881901916705
+    },
+    {
+      "epoch": 8,
+      "train_loss": 0.4249380529554267
+    },
+    {
+      "epoch": 9,
+      "train_loss": 0.4178563620542225
+    },
+    {
+      "epoch": 10,
+      "train_loss": 0.40975865175849513
+    },
+    {
+      "epoch": 11,
+      "train_loss": 0.40546370631770084
+    },
+    {
+      "epoch": 12,
+      "train_loss": 0.3982252176811821
+    },
+    {
+      "epoch": 13,
+      "train_loss": 0.3858343799490678
+    },
+    {
+      "epoch": 14,
+      "train_loss": 0.379007538933503
+    },
+    {
+      "epoch": 15,
+      "train_loss": 0.3656208556576779
+    },
+    {
+      "epoch": 16,
+      "train_loss": 0.35379995308424295
+    },
+    {
+      "epoch": 17,
+      "train_loss": 0.33681938353337737
+    },
+    {
+      "epoch": 18,
+      "train_loss": 0.32208908896697197
+    },
+    {
+      "epoch": 19,
+      "train_loss": 0.30617407221543164
+    },
+    {
+      "epoch": 20,
+      "train_loss": 0.2926546273419732
+    },
+    {
+      "epoch": 21,
+      "train_loss": 0.28012341188757045
+    },
+    {
+      "epoch": 22,
+      "train_loss": 0.27170656586948194
+    },
+    {
+      "epoch": 23,
+      "train_loss": 0.2618562983211718
+    },
+    {
+      "epoch": 24,
+      "train_loss": 0.25525567559819473
+    },
+    {
+      "epoch": 25,
+      "train_loss": 0.24340780433855558
+    },
+    {
+      "epoch": 26,
+      "train_loss": 0.241100468886526
+    },
+    {
+      "epoch": 27,
+      "train_loss": 0.22936643666342685
+    },
+    {
+      "epoch": 28,
+      "train_loss": 0.22153453873960596
+    },
+    {
+      "epoch": 29,
+      "train_loss": 0.2213413413574821
+    },
+    {
+      "epoch": 30,
+      "train_loss": 0.21846851916689622
+    },
+    {
+      "epoch": 31,
+      "train_loss": 0.2104246342495868
+    },
+    {
+      "epoch": 32,
+      "train_loss": 0.20888692896617087
+    },
+    {
+      "epoch": 33,
+      "train_loss": 0.19663872825471979
+    },
+    {
+      "epoch": 34,
+      "train_loss": 0.19428577112524134
+    },
+    {
+      "epoch": 35,
+      "train_loss": 0.1947261790539089
+    },
+    {
+      "epoch": 36,
+      "train_loss": 0.18938347912148426
+    },
+    {
+      "epoch": 37,
+      "train_loss": 0.1853264942294673
+    },
+    {
+      "epoch": 38,
+      "train_loss": 0.17872829263147555
+    },
+    {
+      "epoch": 39,
+      "train_loss": 0.17811942146012658
+    },
+    {
+      "epoch": 40,
+      "train_loss": 0.1693246652264344
+    }
+  ],
+  "epochs_completed": 40,
+  "finite": true,
+  "first_nonfinite_epoch": null,
+  "mission_lesion_rollout": {
+    "episodes": 500,
+    "mean_length": 43.54,
+    "mean_return": 0.33751875,
+    "success": 0.498
+  },
+  "mission_lesion_validation": {
+    "accuracy": 0.744796306795173,
+    "loss": 0.739418601047426
+  },
+  "predictor": [],
+  "provenance": {
+    "cuda_device_name": "NVIDIA GeForce GTX 1080",
+    "cuda_peak_allocated_bytes": 130533376,
+    "cuda_visible_devices": "1",
+    "device": "cuda:0",
+    "git_commit": "7b7d7a08a5e78c3379397f01f0c9af97d5d84fbd",
+    "git_dirty_tracked": false,
+    "minigrid_version": "3.1.0",
+    "torch_version": "2.3.1+cu118"
+  },
+  "rollout": {
+    "episodes": 500,
+    "mean_length": 28.318,
+    "mean_return": 0.568978125,
+    "success": 0.672
+  },
+  "stage": "babyai_pickup_p2",
+  "training": {
+    "batch_size": 256,
+    "history_steps": 4,
+    "model_seed": 4101,
+    "neutral_examples_per_epoch": 0,
+    "parameter_count_including_reciprocal_context_predictor": 1438215,
+    "shuffle_seed": 4101,
+    "total_wall_seconds_including_rollouts": 341.9231562614441,
+    "training_wall_seconds": 327.4489586353302
+  },
+  "validation": {
+    "accuracy": 0.8613428363165141,
+    "loss": 0.33315312845982875
+  }
+}
diff --git a/results/babyai_shared/pickup_p2/seed4101_raw_shared.json b/results/babyai_shared/pickup_p2/seed4101_raw_shared.json
new file mode 100644
index 0000000..344bad4
--- /dev/null
+++ b/results/babyai_shared/pickup_p2/seed4101_raw_shared.json
@@ -0,0 +1,262 @@
+{
+  "condition": "raw_shared",
+  "config": {
+    "action_dim": 7,
+    "context_gain": 1.0,
+    "hidden_layers": 4,
+    "input_dim": 3964,
+    "learning_rate": 0.03,
+    "mission_dim": 23,
+    "momentum": 0.9,
+    "reciprocal_learning_rate": 0.03,
+    "weight_decay": 0.0001,
+    "width": 256
+  },
+  "data": {
+    "color_cardinality": 6,
+    "elapsed_seconds": 81.82925486564636,
+    "env_id": "BabyAI-PickupLoc-v0",
+    "minigrid_version": "3.1.0",
+    "object_cardinality": 11,
+    "protocol": "babyai_shared_feedback_b0",
+    "rollout_episodes": 500,
+    "rollout_seed_start": 200000,
+    "state_cardinality": 3,
+    "train_episodes": 20000,
+    "train_steps": 121467,
+    "validation_episodes": 2000,
+    "validation_steps": 12347,
+    "vocabulary": [
+      "<unk>",
+      "a",
+      "ball",
+      "behind",
+      "blue",
+      "box",
+      "front",
+      "green",
+      "grey",
+      "in",
+      "key",
+      "left",
+      "of",
+      "on",
+      "pick",
+      "purple",
+      "red",
+      "right",
+      "the",
+      "up",
+      "yellow",
+      "you",
+      "your"
+    ]
+  },
+  "epoch_history": [
+    {
+      "epoch": 1,
+      "train_loss": 1.3772296137558786
+    },
+    {
+      "epoch": 2,
+      "train_loss": 1.7021908586903622
+    },
+    {
+      "epoch": 3,
+      "train_loss": 1.4057510167674014
+    },
+    {
+      "epoch": 4,
+      "train_loss": 1.3790689074365716
+    },
+    {
+      "epoch": 5,
+      "train_loss": 1.4092068611948114
+    },
+    {
+      "epoch": 6,
+      "train_loss": 1.4024050559495624
+    },
+    {
+      "epoch": 7,
+      "train_loss": 1.448081522489849
+    },
+    {
+      "epoch": 8,
+      "train_loss": 1.3555105929625662
+    },
+    {
+      "epoch": 9,
+      "train_loss": 1.314540177897403
+    },
+    {
+      "epoch": 10,
+      "train_loss": 1.3105172322925769
+    },
+    {
+      "epoch": 11,
+      "train_loss": 1.3240358893494857
+    },
+    {
+      "epoch": 12,
+      "train_loss": 1.108769882478212
+    },
+    {
+      "epoch": 13,
+      "train_loss": 0.9759558450548272
+    },
+    {
+      "epoch": 14,
+      "train_loss": 1.1254934888137014
+    },
+    {
+      "epoch": 15,
+      "train_loss": 1.3993210215317575
+    },
+    {
+      "epoch": 16,
+      "train_loss": 1.473675866628948
+    },
+    {
+      "epoch": 17,
+      "train_loss": 1.4721564543874641
+    },
+    {
+      "epoch": 18,
+      "train_loss": 1.4467827192105744
+    },
+    {
+      "epoch": 19,
+      "train_loss": 1.3536395237320349
+    },
+    {
+      "epoch": 20,
+      "train_loss": 1.494690821170807
+    },
+    {
+      "epoch": 21,
+      "train_loss": 1.6195249751994485
+    },
+    {
+      "epoch": 22,
+      "train_loss": 1.3495978566219933
+    },
+    {
+      "epoch": 23,
+      "train_loss": 1.562763890969126
+    },
+    {
+      "epoch": 24,
+      "train_loss": 1.3056548786163331
+    },
+    {
+      "epoch": 25,
+      "train_loss": 1.510417148188541
+    },
+    {
+      "epoch": 26,
+      "train_loss": 1.2385981110522621
+    },
+    {
+      "epoch": 27,
+      "train_loss": 1.416951956372512
+    },
+    {
+      "epoch": 28,
+      "train_loss": 1.324134414321498
+    },
+    {
+      "epoch": 29,
+      "train_loss": 1.6655939639242072
+    },
+    {
+      "epoch": 30,
+      "train_loss": 1.3269542169570923
+    },
+    {
+      "epoch": 31,
+      "train_loss": 1.198546481634441
+    },
+    {
+      "epoch": 32,
+      "train_loss": 1.6666638913907503
+    },
+    {
+      "epoch": 33,
+      "train_loss": 1.427612681388855
+    },
+    {
+      "epoch": 34,
+      "train_loss": 1.2676852522398296
+    },
+    {
+      "epoch": 35,
+      "train_loss": 1.4735501365912589
+    },
+    {
+      "epoch": 36,
+      "train_loss": 1.1051569257284466
+    },
+    {
+      "epoch": 37,
+      "train_loss": 1.2862973846887287
+    },
+    {
+      "epoch": 38,
+      "train_loss": 1.035027784046374
+    },
+    {
+      "epoch": 39,
+      "train_loss": 1.1972457239502354
+    },
+    {
+      "epoch": 40,
+      "train_loss": 1.2982738676824068
+    }
+  ],
+  "epochs_completed": 40,
+  "finite": true,
+  "first_nonfinite_epoch": null,
+  "mission_lesion_rollout": {
+    "episodes": 500,
+    "mean_length": 64.0,
+    "mean_return": 0.0,
+    "success": 0.0
+  },
+  "mission_lesion_validation": {
+    "accuracy": 0.49890661699198185,
+    "loss": 1.3048100115036958
+  },
+  "predictor": [],
+  "provenance": {
+    "cuda_device_name": "NVIDIA GeForce GTX 1080",
+    "cuda_peak_allocated_bytes": 130029568,
+    "cuda_visible_devices": "2",
+    "device": "cuda:0",
+    "git_commit": "7b7d7a08a5e78c3379397f01f0c9af97d5d84fbd",
+    "git_dirty_tracked": false,
+    "minigrid_version": "3.1.0",
+    "torch_version": "2.3.1+cu118"
+  },
+  "rollout": {
+    "episodes": 500,
+    "mean_length": 64.0,
+    "mean_return": 0.0,
+    "success": 0.0
+  },
+  "stage": "babyai_pickup_p2",
+  "training": {
+    "batch_size": 256,
+    "history_steps": 4,
+    "model_seed": 4101,
+    "neutral_examples_per_epoch": 0,
+    "parameter_count_including_reciprocal_context_predictor": 1438215,
+    "shuffle_seed": 4101,
+    "total_wall_seconds_including_rollouts": 349.1060309410095,
+    "training_wall_seconds": 327.87563133239746
+  },
+  "validation": {
+    "accuracy": 0.49890661699198185,
+    "loss": 1.304808247974742
+  }
+}
diff --git a/results/babyai_shared/pickup_p2/seed4101_sdil.json b/results/babyai_shared/pickup_p2/seed4101_sdil.json
new file mode 100644
index 0000000..6be8f4b
--- /dev/null
+++ b/results/babyai_shared/pickup_p2/seed4101_sdil.json
@@ -0,0 +1,295 @@
+{
+  "condition": "sdil",
+  "config": {
+    "action_dim": 7,
+    "context_gain": 1.0,
+    "hidden_layers": 4,
+    "input_dim": 3964,
+    "learning_rate": 0.03,
+    "mission_dim": 23,
+    "momentum": 0.9,
+    "reciprocal_learning_rate": 0.03,
+    "weight_decay": 0.0001,
+    "width": 256
+  },
+  "data": {
+    "color_cardinality": 6,
+    "elapsed_seconds": 81.82925486564636,
+    "env_id": "BabyAI-PickupLoc-v0",
+    "minigrid_version": "3.1.0",
+    "object_cardinality": 11,
+    "protocol": "babyai_shared_feedback_b0",
+    "rollout_episodes": 500,
+    "rollout_seed_start": 200000,
+    "state_cardinality": 3,
+    "train_episodes": 20000,
+    "train_steps": 121467,
+    "validation_episodes": 2000,
+    "validation_steps": 12347,
+    "vocabulary": [
+      "<unk>",
+      "a",
+      "ball",
+      "behind",
+      "blue",
+      "box",
+      "front",
+      "green",
+      "grey",
+      "in",
+      "key",
+      "left",
+      "of",
+      "on",
+      "pick",
+      "purple",
+      "red",
+      "right",
+      "the",
+      "up",
+      "yellow",
+      "you",
+      "your"
+    ]
+  },
+  "epoch_history": [
+    {
+      "epoch": 1,
+      "train_loss": 1.4496507589440597
+    },
+    {
+      "epoch": 2,
+      "train_loss": 1.4869623249455501
+    },
+    {
+      "epoch": 3,
+      "train_loss": 1.4099679946899415
+    },
+    {
+      "epoch": 4,
+      "train_loss": 1.4097933264782554
+    },
+    {
+      "epoch": 5,
+      "train_loss": 1.5673320049988595
+    },
+    {
+      "epoch": 6,
+      "train_loss": 1.3529665086143896
+    },
+    {
+      "epoch": 7,
+      "train_loss": 1.6511782859501085
+    },
+    {
+      "epoch": 8,
+      "train_loss": 1.4052464414897718
+    },
+    {
+      "epoch": 9,
+      "train_loss": 1.5705512192374782
+    },
+    {
+      "epoch": 10,
+      "train_loss": 1.3675857054559808
+    },
+    {
+      "epoch": 11,
+      "train_loss": 1.384743598887795
+    },
+    {
+      "epoch": 12,
+      "train_loss": 1.3992685817417345
+    },
+    {
+      "epoch": 13,
+      "train_loss": 1.3667162430913824
+    },
+    {
+      "epoch": 14,
+      "train_loss": 1.4492891484812687
+    },
+    {
+      "epoch": 15,
+      "train_loss": 1.4339373264814679
+    },
+    {
+      "epoch": 16,
+      "train_loss": 1.4766567792390521
+    },
+    {
+      "epoch": 17,
+      "train_loss": 1.415901366534986
+    },
+    {
+      "epoch": 18,
+      "train_loss": 1.4241125407971833
+    },
+    {
+      "epoch": 19,
+      "train_loss": 1.4059961216073287
+    },
+    {
+      "epoch": 20,
+      "train_loss": 1.6259911338906539
+    },
+    {
+      "epoch": 21,
+      "train_loss": 1.5026495936042383
+    },
+    {
+      "epoch": 22,
+      "train_loss": 1.492478413832815
+    },
+    {
+      "epoch": 23,
+      "train_loss": 1.3337757732993678
+    },
+    {
+      "epoch": 24,
+      "train_loss": 1.4029922046159442
+    },
+    {
+      "epoch": 25,
+      "train_loss": 1.4238720023004632
+    },
+    {
+      "epoch": 26,
+      "train_loss": 1.6601720375763742
+    },
+    {
+      "epoch": 27,
+      "train_loss": 1.3956078170475208
+    },
+    {
+      "epoch": 28,
+      "train_loss": 1.3906170870128431
+    },
+    {
+      "epoch": 29,
+      "train_loss": 1.5026050532491584
+    },
+    {
+      "epoch": 30,
+      "train_loss": 1.356335889414737
+    },
+    {
+      "epoch": 31,
+      "train_loss": 1.4696443364494725
+    },
+    {
+      "epoch": 32,
+      "train_loss": 1.4274836944278917
+    },
+    {
+      "epoch": 33,
+      "train_loss": 1.3983684896167956
+    },
+    {
+      "epoch": 34,
+      "train_loss": 1.416925176821257
+    },
+    {
+      "epoch": 35,
+      "train_loss": 1.6563161368119088
+    },
+    {
+      "epoch": 36,
+      "train_loss": 1.4001352036626715
+    },
+    {
+      "epoch": 37,
+      "train_loss": 1.425227477926957
+    },
+    {
+      "epoch": 38,
+      "train_loss": 1.3887689311880815
+    },
+    {
+      "epoch": 39,
+      "train_loss": 1.364363015074479
+    },
+    {
+      "epoch": 40,
+      "train_loss": 1.377758667845475
+    }
+  ],
+  "epochs_completed": 40,
+  "finite": true,
+  "first_nonfinite_epoch": null,
+  "mission_lesion_rollout": {
+    "episodes": 500,
+    "mean_length": 64.0,
+    "mean_return": 0.0,
+    "success": 0.0
+  },
+  "mission_lesion_validation": {
+    "accuracy": 0.49890661699198185,
+    "loss": 1.3227376229581325
+  },
+  "predictor": [
+    {
+      "action_observations": 0,
+      "context_rms": 0.20984283089637756,
+      "mean_per_cell_r2": 0.07823440432548523,
+      "neutral_observations": 4096,
+      "residual_context_rms_ratio": 0.6220114827156067,
+      "teaching_observations": 0
+    },
+    {
+      "action_observations": 0,
+      "context_rms": 0.2008315920829773,
+      "mean_per_cell_r2": 0.3734843134880066,
+      "neutral_observations": 4096,
+      "residual_context_rms_ratio": 0.5552217364311218,
+      "teaching_observations": 0
+    },
+    {
+      "action_observations": 0,
+      "context_rms": 0.19726455211639404,
+      "mean_per_cell_r2": 0.33756017684936523,
+      "neutral_observations": 4096,
+      "residual_context_rms_ratio": 0.5612375736236572,
+      "teaching_observations": 0
+    },
+    {
+      "action_observations": 0,
+      "context_rms": 0.19451536238193512,
+      "mean_per_cell_r2": 0.4691427946090698,
+      "neutral_observations": 4096,
+      "residual_context_rms_ratio": 0.5088029503822327,
+      "teaching_observations": 0
+    }
+  ],
+  "provenance": {
+    "cuda_device_name": "NVIDIA GeForce GTX 1080",
+    "cuda_peak_allocated_bytes": 435783168,
+    "cuda_visible_devices": "3",
+    "device": "cuda:0",
+    "git_commit": "7b7d7a08a5e78c3379397f01f0c9af97d5d84fbd",
+    "git_dirty_tracked": false,
+    "minigrid_version": "3.1.0",
+    "torch_version": "2.3.1+cu118"
+  },
+  "rollout": {
+    "episodes": 500,
+    "mean_length": 64.0,
+    "mean_return": 0.0,
+    "success": 0.0
+  },
+  "stage": "babyai_pickup_p2",
+  "training": {
+    "batch_size": 256,
+    "history_steps": 4,
+    "model_seed": 4101,
+    "neutral_examples_per_epoch": 4096,
+    "parameter_count_including_reciprocal_context_predictor": 1438215,
+    "shuffle_seed": 4101,
+    "total_wall_seconds_including_rollouts": 346.5359559059143,
+    "training_wall_seconds": 322.71134519577026
+  },
+  "validation": {
+    "accuracy": 0.49890661699198185,
+    "loss": 1.3229210194647723
+  }
+}
diff --git a/results/babyai_shared/pickup_p2/seed4102_bp.json b/results/babyai_shared/pickup_p2/seed4102_bp.json
new file mode 100644
index 0000000..294fb8c
--- /dev/null
+++ b/results/babyai_shared/pickup_p2/seed4102_bp.json
@@ -0,0 +1,262 @@
+{
+  "condition": "bp",
+  "config": {
+    "action_dim": 7,
+    "context_gain": 1.0,
+    "hidden_layers": 4,
+    "input_dim": 3964,
+    "learning_rate": 0.03,
+    "mission_dim": 23,
+    "momentum": 0.9,
+    "reciprocal_learning_rate": 0.03,
+    "weight_decay": 0.0001,
+    "width": 256
+  },
+  "data": {
+    "color_cardinality": 6,
+    "elapsed_seconds": 81.82925486564636,
+    "env_id": "BabyAI-PickupLoc-v0",
+    "minigrid_version": "3.1.0",
+    "object_cardinality": 11,
+    "protocol": "babyai_shared_feedback_b0",
+    "rollout_episodes": 500,
+    "rollout_seed_start": 200000,
+    "state_cardinality": 3,
+    "train_episodes": 20000,
+    "train_steps": 121467,
+    "validation_episodes": 2000,
+    "validation_steps": 12347,
+    "vocabulary": [
+      "<unk>",
+      "a",
+      "ball",
+      "behind",
+      "blue",
+      "box",
+      "front",
+      "green",
+      "grey",
+      "in",
+      "key",
+      "left",
+      "of",
+      "on",
+      "pick",
+      "purple",
+      "red",
+      "right",
+      "the",
+      "up",
+      "yellow",
+      "you",
+      "your"
+    ]
+  },
+  "epoch_history": [
+    {
+      "epoch": 1,
+      "train_loss": 0.6358707290574124
+    },
+    {
+      "epoch": 2,
+      "train_loss": 0.4983976507814307
+    },
+    {
+      "epoch": 3,
+      "train_loss": 0.4675803022008193
+    },
+    {
+      "epoch": 4,
+      "train_loss": 0.4492099952697754
+    },
+    {
+      "epoch": 5,
+      "train_loss": 0.42707399907865023
+    },
+    {
+      "epoch": 6,
+      "train_loss": 0.39973109395880446
+    },
+    {
+      "epoch": 7,
+      "train_loss": 0.3692683135835748
+    },
+    {
+      "epoch": 8,
+      "train_loss": 0.34918095908666913
+    },
+    {
+      "epoch": 9,
+      "train_loss": 0.32295236201662764
+    },
+    {
+      "epoch": 10,
+      "train_loss": 0.3064782916558416
+    },
+    {
+      "epoch": 11,
+      "train_loss": 0.2867251298615807
+    },
+    {
+      "epoch": 12,
+      "train_loss": 0.2729152385498348
+    },
+    {
+      "epoch": 13,
+      "train_loss": 0.2690554935994901
+    },
+    {
+      "epoch": 14,
+      "train_loss": 0.2533571675576662
+    },
+    {
+      "epoch": 15,
+      "train_loss": 0.2478961335985284
+    },
+    {
+      "epoch": 16,
+      "train_loss": 0.24331268574062145
+    },
+    {
+      "epoch": 17,
+      "train_loss": 0.23489129813093887
+    },
+    {
+      "epoch": 18,
+      "train_loss": 0.22351061075925827
+    },
+    {
+      "epoch": 19,
+      "train_loss": 0.2218321235242643
+    },
+    {
+      "epoch": 20,
+      "train_loss": 0.21205236274945108
+    },
+    {
+      "epoch": 21,
+      "train_loss": 0.21070943317915264
+    },
+    {
+      "epoch": 22,
+      "train_loss": 0.2010072124474927
+    },
+    {
+      "epoch": 23,
+      "train_loss": 0.1968216976366545
+    },
+    {
+      "epoch": 24,
+      "train_loss": 0.19213551078972063
+    },
+    {
+      "epoch": 25,
+      "train_loss": 0.1923028072401097
+    },
+    {
+      "epoch": 26,
+      "train_loss": 0.18429007967835978
+    },
+    {
+      "epoch": 27,
+      "train_loss": 0.17929337776020954
+    },
+    {
+      "epoch": 28,
+      "train_loss": 0.1737828024908116
+    },
+    {
+      "epoch": 29,
+      "train_loss": 0.17116386283385127
+    },
+    {
+      "epoch": 30,
+      "train_loss": 0.16565450434622012
+    },
+    {
+      "epoch": 31,
+      "train_loss": 0.16759385447753103
+    },
+    {
+      "epoch": 32,
+      "train_loss": 0.1587610994514666
+    },
+    {
+      "epoch": 33,
+      "train_loss": 0.15462220474293356
+    },
+    {
+      "epoch": 34,
+      "train_loss": 0.15275993905569377
+    },
+    {
+      "epoch": 35,
+      "train_loss": 0.1518227168447093
+    },
+    {
+      "epoch": 36,
+      "train_loss": 0.14997457358397936
+    },
+    {
+      "epoch": 37,
+      "train_loss": 0.13869191097585779
+    },
+    {
+      "epoch": 38,
+      "train_loss": 0.14246609537225022
+    },
+    {
+      "epoch": 39,
+      "train_loss": 0.13483418814445797
+    },
+    {
+      "epoch": 40,
+      "train_loss": 0.1340645107156352
+    }
+  ],
+  "epochs_completed": 40,
+  "finite": true,
+  "first_nonfinite_epoch": null,
+  "mission_lesion_rollout": {
+    "episodes": 500,
+    "mean_length": 42.626,
+    "mean_return": 0.351171875,
+    "success": 0.506
+  },
+  "mission_lesion_validation": {
+    "accuracy": 0.743905402121973,
+    "loss": 1.0319232535939609
+  },
+  "predictor": [],
+  "provenance": {
+    "cuda_device_name": "NVIDIA GeForce GTX 1080",
+    "cuda_peak_allocated_bytes": 130533376,
+    "cuda_visible_devices": "4",
+    "device": "cuda:0",
+    "git_commit": "7b7d7a08a5e78c3379397f01f0c9af97d5d84fbd",
+    "git_dirty_tracked": false,
+    "minigrid_version": "3.1.0",
+    "torch_version": "2.3.1+cu118"
+  },
+  "rollout": {
+    "episodes": 500,
+    "mean_length": 21.354,
+    "mean_return": 0.683309375,
+    "success": 0.836
+  },
+  "stage": "babyai_pickup_p2",
+  "training": {
+    "batch_size": 256,
+    "history_steps": 4,
+    "model_seed": 4102,
+    "neutral_examples_per_epoch": 0,
+    "parameter_count_including_reciprocal_context_predictor": 1438215,
+    "shuffle_seed": 4102,
+    "total_wall_seconds_including_rollouts": 331.3091230392456,
+    "training_wall_seconds": 311.8552315235138
+  },
+  "validation": {
+    "accuracy": 0.8787559731108772,
+    "loss": 0.35747272120397217
+  }
+}
diff --git a/results/babyai_shared/pickup_p2/seed4102_clean_kp.json b/results/babyai_shared/pickup_p2/seed4102_clean_kp.json
new file mode 100644
index 0000000..9cd9956
--- /dev/null
+++ b/results/babyai_shared/pickup_p2/seed4102_clean_kp.json
@@ -0,0 +1,262 @@
+{
+  "condition": "clean_kp",
+  "config": {
+    "action_dim": 7,
+    "context_gain": 1.0,
+    "hidden_layers": 4,
+    "input_dim": 3964,
+    "learning_rate": 0.03,
+    "mission_dim": 23,
+    "momentum": 0.9,
+    "reciprocal_learning_rate": 0.03,
+    "weight_decay": 0.0001,
+    "width": 256
+  },
+  "data": {
+    "color_cardinality": 6,
+    "elapsed_seconds": 81.82925486564636,
+    "env_id": "BabyAI-PickupLoc-v0",
+    "minigrid_version": "3.1.0",
+    "object_cardinality": 11,
+    "protocol": "babyai_shared_feedback_b0",
+    "rollout_episodes": 500,
+    "rollout_seed_start": 200000,
+    "state_cardinality": 3,
+    "train_episodes": 20000,
+    "train_steps": 121467,
+    "validation_episodes": 2000,
+    "validation_steps": 12347,
+    "vocabulary": [
+      "<unk>",
+      "a",
+      "ball",
+      "behind",
+      "blue",
+      "box",
+      "front",
+      "green",
+      "grey",
+      "in",
+      "key",
+      "left",
+      "of",
+      "on",
+      "pick",
+      "purple",
+      "red",
+      "right",
+      "the",
+      "up",
+      "yellow",
+      "you",
+      "your"
+    ]
+  },
+  "epoch_history": [
+    {
+      "epoch": 1,
+      "train_loss": 0.681839566669966
+    },
+    {
+      "epoch": 2,
+      "train_loss": 0.5106817564838811
+    },
+    {
+      "epoch": 3,
+      "train_loss": 0.4866123088410026
+    },
+    {
+      "epoch": 4,
+      "train_loss": 0.46724821322842647
+    },
+    {
+      "epoch": 5,
+      "train_loss": 0.45275609010144285
+    },
+    {
+      "epoch": 6,
+      "train_loss": 0.444956036241431
+    },
+    {
+      "epoch": 7,
+      "train_loss": 0.4342330169050317
+    },
+    {
+      "epoch": 8,
+      "train_loss": 0.42846722878907856
+    },
+    {
+      "epoch": 9,
+      "train_loss": 0.41873464232996893
+    },
+    {
+      "epoch": 10,
+      "train_loss": 0.4157650209100623
+    },
+    {
+      "epoch": 11,
+      "train_loss": 0.40508943651851853
+    },
+    {
+      "epoch": 12,
+      "train_loss": 0.3999078395492152
+    },
+    {
+      "epoch": 13,
+      "train_loss": 0.3982298595654337
+    },
+    {
+      "epoch": 14,
+      "train_loss": 0.3850154632016232
+    },
+    {
+      "epoch": 15,
+      "train_loss": 0.3766156781347174
+    },
+    {
+      "epoch": 16,
+      "train_loss": 0.368439452836388
+    },
+    {
+      "epoch": 17,
+      "train_loss": 0.3616289800719211
+    },
+    {
+      "epoch": 18,
+      "train_loss": 0.3508304416505914
+    },
+    {
+      "epoch": 19,
+      "train_loss": 0.3401417228422667
+    },
+    {
+      "epoch": 20,
+      "train_loss": 0.3305343171797301
+    },
+    {
+      "epoch": 21,
+      "train_loss": 0.3165550116802517
+    },
+    {
+      "epoch": 22,
+      "train_loss": 0.3019398627469414
+    },
+    {
+      "epoch": 23,
+      "train_loss": 0.29349281496123264
+    },
+    {
+      "epoch": 24,
+      "train_loss": 0.2876167286697187
+    },
+    {
+      "epoch": 25,
+      "train_loss": 0.26772761492352737
+    },
+    {
+      "epoch": 26,
+      "train_loss": 0.26045800096110294
+    },
+    {
+      "epoch": 27,
+      "train_loss": 0.24848401973122045
+    },
+    {
+      "epoch": 28,
+      "train_loss": 0.2440636997787576
+    },
+    {
+      "epoch": 29,
+      "train_loss": 0.2351069336188467
+    },
+    {
+      "epoch": 30,
+      "train_loss": 0.23047932257777767
+    },
+    {
+      "epoch": 31,
+      "train_loss": 0.22535064524725865
+    },
+    {
+      "epoch": 32,
+      "train_loss": 0.22159839981480647
+    },
+    {
+      "epoch": 33,
+      "train_loss": 0.2125453116078126
+    },
+    {
+      "epoch": 34,
+      "train_loss": 0.20711057270828045
+    },
+    {
+      "epoch": 35,
+      "train_loss": 0.20062021220985213
+    },
+    {
+      "epoch": 36,
+      "train_loss": 0.19858844374355517
+    },
+    {
+      "epoch": 37,
+      "train_loss": 0.1931002575629636
+    },
+    {
+      "epoch": 38,
+      "train_loss": 0.1893660755063358
+    },
+    {
+      "epoch": 39,
+      "train_loss": 0.18631849710878573
+    },
+    {
+      "epoch": 40,
+      "train_loss": 0.18104866024694943
+    }
+  ],
+  "epochs_completed": 40,
+  "finite": true,
+  "first_nonfinite_epoch": null,
+  "mission_lesion_rollout": {
+    "episodes": 500,
+    "mean_length": 41.762,
+    "mean_return": 0.36832187499999997,
+    "success": 0.556
+  },
+  "mission_lesion_validation": {
+    "accuracy": 0.7485219081558273,
+    "loss": 0.7462387038617324
+  },
+  "predictor": [],
+  "provenance": {
+    "cuda_device_name": "NVIDIA GeForce GTX 1080",
+    "cuda_peak_allocated_bytes": 130533376,
+    "cuda_visible_devices": "5",
+    "device": "cuda:0",
+    "git_commit": "7b7d7a08a5e78c3379397f01f0c9af97d5d84fbd",
+    "git_dirty_tracked": false,
+    "minigrid_version": "3.1.0",
+    "torch_version": "2.3.1+cu118"
+  },
+  "rollout": {
+    "episodes": 500,
+    "mean_length": 25.062,
+    "mean_return": 0.6245656249999999,
+    "success": 0.77
+  },
+  "stage": "babyai_pickup_p2",
+  "training": {
+    "batch_size": 256,
+    "history_steps": 4,
+    "model_seed": 4102,
+    "neutral_examples_per_epoch": 0,
+    "parameter_count_including_reciprocal_context_predictor": 1438215,
+    "shuffle_seed": 4102,
+    "total_wall_seconds_including_rollouts": 340.4907646179199,
+    "training_wall_seconds": 326.9207751750946
+  },
+  "validation": {
+    "accuracy": 0.8682270997003321,
+    "loss": 0.3341673115462806
+  }
+}
diff --git a/results/babyai_shared/pickup_p2/seed4102_raw_shared.json b/results/babyai_shared/pickup_p2/seed4102_raw_shared.json
new file mode 100644
index 0000000..18eddd3
--- /dev/null
+++ b/results/babyai_shared/pickup_p2/seed4102_raw_shared.json
@@ -0,0 +1,262 @@
+{
+  "condition": "raw_shared",
+  "config": {
+    "action_dim": 7,
+    "context_gain": 1.0,
+    "hidden_layers": 4,
+    "input_dim": 3964,
+    "learning_rate": 0.03,
+    "mission_dim": 23,
+    "momentum": 0.9,
+    "reciprocal_learning_rate": 0.03,
+    "weight_decay": 0.0001,
+    "width": 256
+  },
+  "data": {
+    "color_cardinality": 6,
+    "elapsed_seconds": 81.82925486564636,
+    "env_id": "BabyAI-PickupLoc-v0",
+    "minigrid_version": "3.1.0",
+    "object_cardinality": 11,
+    "protocol": "babyai_shared_feedback_b0",
+    "rollout_episodes": 500,
+    "rollout_seed_start": 200000,
+    "state_cardinality": 3,
+    "train_episodes": 20000,
+    "train_steps": 121467,
+    "validation_episodes": 2000,
+    "validation_steps": 12347,
+    "vocabulary": [
+      "<unk>",
+      "a",
+      "ball",
+      "behind",
+      "blue",
+      "box",
+      "front",
+      "green",
+      "grey",
+      "in",
+      "key",
+      "left",
+      "of",
+      "on",
+      "pick",
+      "purple",
+      "red",
+      "right",
+      "the",
+      "up",
+      "yellow",
+      "you",
+      "your"
+    ]
+  },
+  "epoch_history": [
+    {
+      "epoch": 1,
+      "train_loss": 1.426944848863702
+    },
+    {
+      "epoch": 2,
+      "train_loss": 1.4811008263889112
+    },
+    {
+      "epoch": 3,
+      "train_loss": 1.537802341109828
+    },
+    {
+      "epoch": 4,
+      "train_loss": 1.6175212340605887
+    },
+    {
+      "epoch": 5,
+      "train_loss": 1.4761648283506694
+    },
+    {
+      "epoch": 6,
+      "train_loss": 1.5078701266489531
+    },
+    {
+      "epoch": 7,
+      "train_loss": 1.228762877489391
+    },
+    {
+      "epoch": 8,
+      "train_loss": 1.3726487501044022
+    },
+    {
+      "epoch": 9,
+      "train_loss": 1.2172372989905509
+    },
+    {
+      "epoch": 10,
+      "train_loss": 1.1436676884952344
+    },
+    {
+      "epoch": 11,
+      "train_loss": 1.3721513145848325
+    },
+    {
+      "epoch": 12,
+      "train_loss": 1.3090164377814846
+    },
+    {
+      "epoch": 13,
+      "train_loss": 1.3085567035173116
+    },
+    {
+      "epoch": 14,
+      "train_loss": 1.24689010431892
+    },
+    {
+      "epoch": 15,
+      "train_loss": 1.438764830388521
+    },
+    {
+      "epoch": 16,
+      "train_loss": 1.4798997766093205
+    },
+    {
+      "epoch": 17,
+      "train_loss": 1.5367788784127487
+    },
+    {
+      "epoch": 18,
+      "train_loss": 1.4143389420760306
+    },
+    {
+      "epoch": 19,
+      "train_loss": 1.4158848059804816
+    },
+    {
+      "epoch": 20,
+      "train_loss": 1.2560336269830403
+    },
+    {
+      "epoch": 21,
+      "train_loss": 1.4611049637041593
+    },
+    {
+      "epoch": 22,
+      "train_loss": 1.3461816581926849
+    },
+    {
+      "epoch": 23,
+      "train_loss": 1.3200064413171066
+    },
+    {
+      "epoch": 24,
+      "train_loss": 1.209783155039737
+    },
+    {
+      "epoch": 25,
+      "train_loss": 1.2179505748497812
+    },
+    {
+      "epoch": 26,
+      "train_loss": 1.3550261970570212
+    },
+    {
+      "epoch": 27,
+      "train_loss": 1.0599337084669815
+    },
+    {
+      "epoch": 28,
+      "train_loss": 0.9845607308337563
+    },
+    {
+      "epoch": 29,
+      "train_loss": 0.9575978085869237
+    },
+    {
+      "epoch": 30,
+      "train_loss": 0.9305678352556731
+    },
+    {
+      "epoch": 31,
+      "train_loss": 0.9234064896483171
+    },
+    {
+      "epoch": 32,
+      "train_loss": 0.9524868351534793
+    },
+    {
+      "epoch": 33,
+      "train_loss": 0.9005411171913147
+    },
+    {
+      "epoch": 34,
+      "train_loss": 0.8903469843613474
+    },
+    {
+      "epoch": 35,
+      "train_loss": 0.9097682004225881
+    },
+    {
+      "epoch": 36,
+      "train_loss": 0.9102116966247559
+    },
+    {
+      "epoch": 37,
+      "train_loss": 0.8369047345613179
+    },
+    {
+      "epoch": 38,
+      "train_loss": 0.9056826622862565
+    },
+    {
+      "epoch": 39,
+      "train_loss": 1.0265569414590534
+    },
+    {
+      "epoch": 40,
+      "train_loss": 0.8987322279026634
+    }
+  ],
+  "epochs_completed": 40,
+  "finite": true,
+  "first_nonfinite_epoch": null,
+  "mission_lesion_rollout": {
+    "episodes": 500,
+    "mean_length": 52.798,
+    "mean_return": 0.176928125,
+    "success": 0.194
+  },
+  "mission_lesion_validation": {
+    "accuracy": 0.692556896412084,
+    "loss": 0.8453130396868214
+  },
+  "predictor": [],
+  "provenance": {
+    "cuda_device_name": "NVIDIA GeForce GTX 1080",
+    "cuda_peak_allocated_bytes": 130029568,
+    "cuda_visible_devices": "6",
+    "device": "cuda:0",
+    "git_commit": "7b7d7a08a5e78c3379397f01f0c9af97d5d84fbd",
+    "git_dirty_tracked": false,
+    "minigrid_version": "3.1.0",
+    "torch_version": "2.3.1+cu118"
+  },
+  "rollout": {
+    "episodes": 500,
+    "mean_length": 50.642,
+    "mean_return": 0.211046875,
+    "success": 0.232
+  },
+  "stage": "babyai_pickup_p2",
+  "training": {
+    "batch_size": 256,
+    "history_steps": 4,
+    "model_seed": 4102,
+    "neutral_examples_per_epoch": 0,
+    "parameter_count_including_reciprocal_context_predictor": 1438215,
+    "shuffle_seed": 4102,
+    "total_wall_seconds_including_rollouts": 344.8074402809143,
+    "training_wall_seconds": 326.14055609703064
+  },
+  "validation": {
+    "accuracy": 0.6942577144245565,
+    "loss": 0.8413812175901195
+  }
+}
diff --git a/results/babyai_shared/pickup_p2/seed4102_sdil.json b/results/babyai_shared/pickup_p2/seed4102_sdil.json
new file mode 100644
index 0000000..d06ebe4
--- /dev/null
+++ b/results/babyai_shared/pickup_p2/seed4102_sdil.json
@@ -0,0 +1,295 @@
+{
+  "condition": "sdil",
+  "config": {
+    "action_dim": 7,
+    "context_gain": 1.0,
+    "hidden_layers": 4,
+    "input_dim": 3964,
+    "learning_rate": 0.03,
+    "mission_dim": 23,
+    "momentum": 0.9,
+    "reciprocal_learning_rate": 0.03,
+    "weight_decay": 0.0001,
+    "width": 256
+  },
+  "data": {
+    "color_cardinality": 6,
+    "elapsed_seconds": 81.82925486564636,
+    "env_id": "BabyAI-PickupLoc-v0",
+    "minigrid_version": "3.1.0",
+    "object_cardinality": 11,
+    "protocol": "babyai_shared_feedback_b0",
+    "rollout_episodes": 500,
+    "rollout_seed_start": 200000,
+    "state_cardinality": 3,
+    "train_episodes": 20000,
+    "train_steps": 121467,
+    "validation_episodes": 2000,
+    "validation_steps": 12347,
+    "vocabulary": [
+      "<unk>",
+      "a",
+      "ball",
+      "behind",
+      "blue",
+      "box",
+      "front",
+      "green",
+      "grey",
+      "in",
+      "key",
+      "left",
+      "of",
+      "on",
+      "pick",
+      "purple",
+      "red",
+      "right",
+      "the",
+      "up",
+      "yellow",
+      "you",
+      "your"
+    ]
+  },
+  "epoch_history": [
+    {
+      "epoch": 1,
+      "train_loss": 1.4037305317426982
+    },
+    {
+      "epoch": 2,
+      "train_loss": 1.5199073093815854
+    },
+    {
+      "epoch": 3,
+      "train_loss": 1.490153865814209
+    },
+    {
+      "epoch": 4,
+      "train_loss": 1.3832739104722676
+    },
+    {
+      "epoch": 5,
+      "train_loss": 1.343873203428168
+    },
+    {
+      "epoch": 6,
+      "train_loss": 1.5590096739718788
+    },
+    {
+      "epoch": 7,
+      "train_loss": 1.6177833466780813
+    },
+    {
+      "epoch": 8,
+      "train_loss": 1.3586675656469245
+    },
+    {
+      "epoch": 9,
+      "train_loss": 1.4770930380570262
+    },
+    {
+      "epoch": 10,
+      "train_loss": 1.4020695854488172
+    },
+    {
+      "epoch": 11,
+      "train_loss": 1.660883920067235
+    },
+    {
+      "epoch": 12,
+      "train_loss": 1.4256328241448653
+    },
+    {
+      "epoch": 13,
+      "train_loss": 1.3678666669444033
+    },
+    {
+      "epoch": 14,
+      "train_loss": 1.5727304014406707
+    },
+    {
+      "epoch": 15,
+      "train_loss": 1.4565586441441587
+    },
+    {
+      "epoch": 16,
+      "train_loss": 1.461268330373262
+    },
+    {
+      "epoch": 17,
+      "train_loss": 1.39766593155108
+    },
+    {
+      "epoch": 18,
+      "train_loss": 1.4142666897020841
+    },
+    {
+      "epoch": 19,
+      "train_loss": 1.476966627020585
+    },
+    {
+      "epoch": 20,
+      "train_loss": 1.4425141741100111
+    },
+    {
+      "epoch": 21,
+      "train_loss": 1.457439170134695
+    },
+    {
+      "epoch": 22,
+      "train_loss": 1.3950570189325433
+    },
+    {
+      "epoch": 23,
+      "train_loss": 1.387892014101932
+    },
+    {
+      "epoch": 24,
+      "train_loss": 1.5007062272021645
+    },
+    {
+      "epoch": 25,
+      "train_loss": 1.4848456322519403
+    },
+    {
+      "epoch": 26,
+      "train_loss": 1.3709611882661519
+    },
+    {
+      "epoch": 27,
+      "train_loss": 1.4066063529566715
+    },
+    {
+      "epoch": 28,
+      "train_loss": 1.5965287012802927
+    },
+    {
+      "epoch": 29,
+      "train_loss": 1.4587098232068514
+    },
+    {
+      "epoch": 30,
+      "train_loss": 1.5121034484160574
+    },
+    {
+      "epoch": 31,
+      "train_loss": 1.551018009436758
+    },
+    {
+      "epoch": 32,
+      "train_loss": 1.4567223797346416
+    },
+    {
+      "epoch": 33,
+      "train_loss": 1.3996574670390078
+    },
+    {
+      "epoch": 34,
+      "train_loss": 1.3792390657726088
+    },
+    {
+      "epoch": 35,
+      "train_loss": 1.5840236104162115
+    },
+    {
+      "epoch": 36,
+      "train_loss": 1.4191108206698768
+    },
+    {
+      "epoch": 37,
+      "train_loss": 1.5500210498508653
+    },
+    {
+      "epoch": 38,
+      "train_loss": 1.8080971825750252
+    },
+    {
+      "epoch": 39,
+      "train_loss": 1.5819644702108282
+    },
+    {
+      "epoch": 40,
+      "train_loss": 1.5324563553458765
+    }
+  ],
+  "epochs_completed": 40,
+  "finite": true,
+  "first_nonfinite_epoch": null,
+  "mission_lesion_rollout": {
+    "episodes": 500,
+    "mean_length": 64.0,
+    "mean_return": 0.0,
+    "success": 0.0
+  },
+  "mission_lesion_validation": {
+    "accuracy": 0.49890661699198185,
+    "loss": 1.2736983379198104
+  },
+  "predictor": [
+    {
+      "action_observations": 0,
+      "context_rms": 0.21022486686706543,
+      "mean_per_cell_r2": 0.14354562759399414,
+      "neutral_observations": 4096,
+      "residual_context_rms_ratio": 0.6005861759185791,
+      "teaching_observations": 0
+    },
+    {
+      "action_observations": 0,
+      "context_rms": 0.20877140760421753,
+      "mean_per_cell_r2": 0.3952368497848511,
+      "neutral_observations": 4096,
+      "residual_context_rms_ratio": 0.5114675164222717,
+      "teaching_observations": 0
+    },
+    {
+      "action_observations": 0,
+      "context_rms": 0.20698049664497375,
+      "mean_per_cell_r2": 0.3069768249988556,
+      "neutral_observations": 4096,
+      "residual_context_rms_ratio": 0.5448850393295288,
+      "teaching_observations": 0
+    },
+    {
+      "action_observations": 0,
+      "context_rms": 0.20679785311222076,
+      "mean_per_cell_r2": 0.3696891665458679,
+      "neutral_observations": 4096,
+      "residual_context_rms_ratio": 0.5242287516593933,
+      "teaching_observations": 0
+    }
+  ],
+  "provenance": {
+    "cuda_device_name": "NVIDIA GeForce GTX 1080",
+    "cuda_peak_allocated_bytes": 435783168,
+    "cuda_visible_devices": "7",
+    "device": "cuda:0",
+    "git_commit": "7b7d7a08a5e78c3379397f01f0c9af97d5d84fbd",
+    "git_dirty_tracked": false,
+    "minigrid_version": "3.1.0",
+    "torch_version": "2.3.1+cu118"
+  },
+  "rollout": {
+    "episodes": 500,
+    "mean_length": 64.0,
+    "mean_return": 0.0,
+    "success": 0.0
+  },
+  "stage": "babyai_pickup_p2",
+  "training": {
+    "batch_size": 256,
+    "history_steps": 4,
+    "model_seed": 4102,
+    "neutral_examples_per_epoch": 4096,
+    "parameter_count_including_reciprocal_context_predictor": 1438215,
+    "shuffle_seed": 4102,
+    "total_wall_seconds_including_rollouts": 344.80195808410645,
+    "training_wall_seconds": 318.03023290634155
+  },
+  "validation": {
+    "accuracy": 0.49890661699198185,
+    "loss": 1.2736801557756392
+  }
+}
diff --git a/results/babyai_shared/pickup_p2/seed4103_bp.json b/results/babyai_shared/pickup_p2/seed4103_bp.json
new file mode 100644
index 0000000..4080f66
--- /dev/null
+++ b/results/babyai_shared/pickup_p2/seed4103_bp.json
@@ -0,0 +1,262 @@
+{
+  "condition": "bp",
+  "config": {
+    "action_dim": 7,
+    "context_gain": 1.0,
+    "hidden_layers": 4,
+    "input_dim": 3964,
+    "learning_rate": 0.03,
+    "mission_dim": 23,
+    "momentum": 0.9,
+    "reciprocal_learning_rate": 0.03,
+    "weight_decay": 0.0001,
+    "width": 256
+  },
+  "data": {
+    "color_cardinality": 6,
+    "elapsed_seconds": 81.82925486564636,
+    "env_id": "BabyAI-PickupLoc-v0",
+    "minigrid_version": "3.1.0",
+    "object_cardinality": 11,
+    "protocol": "babyai_shared_feedback_b0",
+    "rollout_episodes": 500,
+    "rollout_seed_start": 200000,
+    "state_cardinality": 3,
+    "train_episodes": 20000,
+    "train_steps": 121467,
+    "validation_episodes": 2000,
+    "validation_steps": 12347,
+    "vocabulary": [
+      "<unk>",
+      "a",
+      "ball",
+      "behind",
+      "blue",
+      "box",
+      "front",
+      "green",
+      "grey",
+      "in",
+      "key",
+      "left",
+      "of",
+      "on",
+      "pick",
+      "purple",
+      "red",
+      "right",
+      "the",
+      "up",
+      "yellow",
+      "you",
+      "your"
+    ]
+  },
+  "epoch_history": [
+    {
+      "epoch": 1,
+      "train_loss": 0.6392265326098392
+    },
+    {
+      "epoch": 2,
+      "train_loss": 0.49516982342067517
+    },
+    {
+      "epoch": 3,
+      "train_loss": 0.4656036344327425
+    },
+    {
+      "epoch": 4,
+      "train_loss": 0.4490726175433711
+    },
+    {
+      "epoch": 5,
+      "train_loss": 0.42668137694659986
+    },
+    {
+      "epoch": 6,
+      "train_loss": 0.3981025429148423
+    },
+    {
+      "epoch": 7,
+      "train_loss": 0.3687292110292535
+    },
+    {
+      "epoch": 8,
+      "train_loss": 0.3368211528815721
+    },
+    {
+      "epoch": 9,
+      "train_loss": 0.3183619886950443
+    },
+    {
+      "epoch": 10,
+      "train_loss": 0.29898343581902354
+    },
+    {
+      "epoch": 11,
+      "train_loss": 0.28272959040968043
+    },
+    {
+      "epoch": 12,
+      "train_loss": 0.26972324841900874
+    },
+    {
+      "epoch": 13,
+      "train_loss": 0.26071007860334294
+    },
+    {
+      "epoch": 14,
+      "train_loss": 0.25385804000653717
+    },
+    {
+      "epoch": 15,
+      "train_loss": 0.2435340386315396
+    },
+    {
+      "epoch": 16,
+      "train_loss": 0.23799892532198053
+    },
+    {
+      "epoch": 17,
+      "train_loss": 0.2289504592355929
+    },
+    {
+      "epoch": 18,
+      "train_loss": 0.22272453988853252
+    },
+    {
+      "epoch": 19,
+      "train_loss": 0.2201314499503688
+    },
+    {
+      "epoch": 20,
+      "train_loss": 0.21376912342874627
+    },
+    {
+      "epoch": 21,
+      "train_loss": 0.20853730493470243
+    },
+    {
+      "epoch": 22,
+      "train_loss": 0.20172592633648923
+    },
+    {
+      "epoch": 23,
+      "train_loss": 0.19995384768435828
+    },
+    {
+      "epoch": 24,
+      "train_loss": 0.19138428857451992
+    },
+    {
+      "epoch": 25,
+      "train_loss": 0.185930360743874
+    },
+    {
+      "epoch": 26,
+      "train_loss": 0.17905265743795193
+    },
+    {
+      "epoch": 27,
+      "train_loss": 0.17856170729586954
+    },
+    {
+      "epoch": 28,
+      "train_loss": 0.17182566167492616
+    },
+    {
+      "epoch": 29,
+      "train_loss": 0.16652865444359027
+    },
+    {
+      "epoch": 30,
+      "train_loss": 0.16592447939671967
+    },
+    {
+      "epoch": 31,
+      "train_loss": 0.16058228897420984
+    },
+    {
+      "epoch": 32,
+      "train_loss": 0.15412096956842825
+    },
+    {
+      "epoch": 33,
+      "train_loss": 0.1573458581692294
+    },
+    {
+      "epoch": 34,
+      "train_loss": 0.14750758492632915
+    },
+    {
+      "epoch": 35,
+      "train_loss": 0.14922092113055682
+    },
+    {
+      "epoch": 36,
+      "train_loss": 0.14258274528541062
+    },
+    {
+      "epoch": 37,
+      "train_loss": 0.1406223738664075
+    },
+    {
+      "epoch": 38,
+      "train_loss": 0.13483527417245664
+    },
+    {
+      "epoch": 39,
+      "train_loss": 0.13803405683291586
+    },
+    {
+      "epoch": 40,
+      "train_loss": 0.1281460135861447
+    }
+  ],
+  "epochs_completed": 40,
+  "finite": true,
+  "first_nonfinite_epoch": null,
+  "mission_lesion_rollout": {
+    "episodes": 500,
+    "mean_length": 43.936,
+    "mean_return": 0.33075,
+    "success": 0.486
+  },
+  "mission_lesion_validation": {
+    "accuracy": 0.7431764801166275,
+    "loss": 0.9334571064749526
+  },
+  "predictor": [],
+  "provenance": {
+    "cuda_device_name": "NVIDIA GeForce GTX 1080",
+    "cuda_peak_allocated_bytes": 130533376,
+    "cuda_visible_devices": "0",
+    "device": "cuda:0",
+    "git_commit": "7b7d7a08a5e78c3379397f01f0c9af97d5d84fbd",
+    "git_dirty_tracked": false,
+    "minigrid_version": "3.1.0",
+    "torch_version": "2.3.1+cu118"
+  },
+  "rollout": {
+    "episodes": 500,
+    "mean_length": 26.046,
+    "mean_return": 0.607328125,
+    "success": 0.736
+  },
+  "stage": "babyai_pickup_p2",
+  "training": {
+    "batch_size": 256,
+    "history_steps": 4,
+    "model_seed": 4103,
+    "neutral_examples_per_epoch": 0,
+    "parameter_count_including_reciprocal_context_predictor": 1438215,
+    "shuffle_seed": 4103,
+    "total_wall_seconds_including_rollouts": 165.793874502182,
+    "training_wall_seconds": 151.43734502792358
+  },
+  "validation": {
+    "accuracy": 0.8766501984287681,
+    "loss": 0.34083314674570847
+  }
+}
diff --git a/results/babyai_shared/pickup_p2/seed4103_clean_kp.json b/results/babyai_shared/pickup_p2/seed4103_clean_kp.json
new file mode 100644
index 0000000..5d862e9
--- /dev/null
+++ b/results/babyai_shared/pickup_p2/seed4103_clean_kp.json
@@ -0,0 +1,262 @@
+{
+  "condition": "clean_kp",
+  "config": {
+    "action_dim": 7,
+    "context_gain": 1.0,
+    "hidden_layers": 4,
+    "input_dim": 3964,
+    "learning_rate": 0.03,
+    "mission_dim": 23,
+    "momentum": 0.9,
+    "reciprocal_learning_rate": 0.03,
+    "weight_decay": 0.0001,
+    "width": 256
+  },
+  "data": {
+    "color_cardinality": 6,
+    "elapsed_seconds": 81.82925486564636,
+    "env_id": "BabyAI-PickupLoc-v0",
+    "minigrid_version": "3.1.0",
+    "object_cardinality": 11,
+    "protocol": "babyai_shared_feedback_b0",
+    "rollout_episodes": 500,
+    "rollout_seed_start": 200000,
+    "state_cardinality": 3,
+    "train_episodes": 20000,
+    "train_steps": 121467,
+    "validation_episodes": 2000,
+    "validation_steps": 12347,
+    "vocabulary": [
+      "<unk>",
+      "a",
+      "ball",
+      "behind",
+      "blue",
+      "box",
+      "front",
+      "green",
+      "grey",
+      "in",
+      "key",
+      "left",
+      "of",
+      "on",
+      "pick",
+      "purple",
+      "red",
+      "right",
+      "the",
+      "up",
+      "yellow",
+      "you",
+      "your"
+    ]
+  },
+  "epoch_history": [
+    {
+      "epoch": 1,
+      "train_loss": 0.6864033263608029
+    },
+    {
+      "epoch": 2,
+      "train_loss": 0.5136532720766569
+    },
+    {
+      "epoch": 3,
+      "train_loss": 0.4820757521453657
+    },
+    {
+      "epoch": 4,
+      "train_loss": 0.4678087641063489
+    },
+    {
+      "epoch": 5,
+      "train_loss": 0.4570458981237913
+    },
+    {
+      "epoch": 6,
+      "train_loss": 0.44718988864045395
+    },
+    {
+      "epoch": 7,
+      "train_loss": 0.43997702253492255
+    },
+    {
+      "epoch": 8,
+      "train_loss": 0.42566254320897556
+    },
+    {
+      "epoch": 9,
+      "train_loss": 0.4220644837931583
+    },
+    {
+      "epoch": 10,
+      "train_loss": 0.41289369489017286
+    },
+    {
+      "epoch": 11,
+      "train_loss": 0.405345781037682
+    },
+    {
+      "epoch": 12,
+      "train_loss": 0.40259038736945707
+    },
+    {
+      "epoch": 13,
+      "train_loss": 0.393158560112903
+    },
+    {
+      "epoch": 14,
+      "train_loss": 0.38850406364390727
+    },
+    {
+      "epoch": 15,
+      "train_loss": 0.3801098694926814
+    },
+    {
+      "epoch": 16,
+      "train_loss": 0.36695603276553906
+    },
+    {
+      "epoch": 17,
+      "train_loss": 0.35903412762441134
+    },
+    {
+      "epoch": 18,
+      "train_loss": 0.34711836858799583
+    },
+    {
+      "epoch": 19,
+      "train_loss": 0.3370369092727962
+    },
+    {
+      "epoch": 20,
+      "train_loss": 0.3200993647387153
+    },
+    {
+      "epoch": 21,
+      "train_loss": 0.3118796497897098
+    },
+    {
+      "epoch": 22,
+      "train_loss": 0.293176792420839
+    },
+    {
+      "epoch": 23,
+      "train_loss": 0.28278176423750423
+    },
+    {
+      "epoch": 24,
+      "train_loss": 0.2673768097789664
+    },
+    {
+      "epoch": 25,
+      "train_loss": 0.2630335626790398
+    },
+    {
+      "epoch": 26,
+      "train_loss": 0.2525263265559548
+    },
+    {
+      "epoch": 27,
+      "train_loss": 0.24664549178198764
+    },
+    {
+      "epoch": 28,
+      "train_loss": 0.23551284742982764
+    },
+    {
+      "epoch": 29,
+      "train_loss": 0.23284720298491027
+    },
+    {
+      "epoch": 30,
+      "train_loss": 0.22212432547619468
+    },
+    {
+      "epoch": 31,
+      "train_loss": 0.21526990802664506
+    },
+    {
+      "epoch": 32,
+      "train_loss": 0.20761467494462665
+    },
+    {
+      "epoch": 33,
+      "train_loss": 0.21185886748527225
+    },
+    {
+      "epoch": 34,
+      "train_loss": 0.20015101647690722
+    },
+    {
+      "epoch": 35,
+      "train_loss": 0.19496909729744258
+    },
+    {
+      "epoch": 36,
+      "train_loss": 0.1936807933763454
+    },
+    {
+      "epoch": 37,
+      "train_loss": 0.19314897574876483
+    },
+    {
+      "epoch": 38,
+      "train_loss": 0.1823754488794427
+    },
+    {
+      "epoch": 39,
+      "train_loss": 0.17861489517124077
+    },
+    {
+      "epoch": 40,
+      "train_loss": 0.17446756489967044
+    }
+  ],
+  "epochs_completed": 40,
+  "finite": true,
+  "first_nonfinite_epoch": null,
+  "mission_lesion_rollout": {
+    "episodes": 500,
+    "mean_length": 45.232,
+    "mean_return": 0.305125,
+    "success": 0.412
+  },
+  "mission_lesion_validation": {
+    "accuracy": 0.7619664695877542,
+    "loss": 0.7881599735844866
+  },
+  "predictor": [],
+  "provenance": {
+    "cuda_device_name": "NVIDIA GeForce GTX 1080",
+    "cuda_peak_allocated_bytes": 130533376,
+    "cuda_visible_devices": "1",
+    "device": "cuda:0",
+    "git_commit": "7b7d7a08a5e78c3379397f01f0c9af97d5d84fbd",
+    "git_dirty_tracked": false,
+    "minigrid_version": "3.1.0",
+    "torch_version": "2.3.1+cu118"
+  },
+  "rollout": {
+    "episodes": 500,
+    "mean_length": 28.356,
+    "mean_return": 0.56964375,
+    "success": 0.684
+  },
+  "stage": "babyai_pickup_p2",
+  "training": {
+    "batch_size": 256,
+    "history_steps": 4,
+    "model_seed": 4103,
+    "neutral_examples_per_epoch": 0,
+    "parameter_count_including_reciprocal_context_predictor": 1438215,
+    "shuffle_seed": 4103,
+    "total_wall_seconds_including_rollouts": 166.57137203216553,
+    "training_wall_seconds": 153.0613570213318
+  },
+  "validation": {
+    "accuracy": 0.8635296023325504,
+    "loss": 0.35984307350374345
+  }
+}
diff --git a/results/babyai_shared/pickup_p2/seed4103_raw_shared.json b/results/babyai_shared/pickup_p2/seed4103_raw_shared.json
new file mode 100644
index 0000000..2b57c87
--- /dev/null
+++ b/results/babyai_shared/pickup_p2/seed4103_raw_shared.json
@@ -0,0 +1,262 @@
+{
+  "condition": "raw_shared",
+  "config": {
+    "action_dim": 7,
+    "context_gain": 1.0,
+    "hidden_layers": 4,
+    "input_dim": 3964,
+    "learning_rate": 0.03,
+    "mission_dim": 23,
+    "momentum": 0.9,
+    "reciprocal_learning_rate": 0.03,
+    "weight_decay": 0.0001,
+    "width": 256
+  },
+  "data": {
+    "color_cardinality": 6,
+    "elapsed_seconds": 81.82925486564636,
+    "env_id": "BabyAI-PickupLoc-v0",
+    "minigrid_version": "3.1.0",
+    "object_cardinality": 11,
+    "protocol": "babyai_shared_feedback_b0",
+    "rollout_episodes": 500,
+    "rollout_seed_start": 200000,
+    "state_cardinality": 3,
+    "train_episodes": 20000,
+    "train_steps": 121467,
+    "validation_episodes": 2000,
+    "validation_steps": 12347,
+    "vocabulary": [
+      "<unk>",
+      "a",
+      "ball",
+      "behind",
+      "blue",
+      "box",
+      "front",
+      "green",
+      "grey",
+      "in",
+      "key",
+      "left",
+      "of",
+      "on",
+      "pick",
+      "purple",
+      "red",
+      "right",
+      "the",
+      "up",
+      "yellow",
+      "you",
+      "your"
+    ]
+  },
+  "epoch_history": [
+    {
+      "epoch": 1,
+      "train_loss": 1.4279862976074218
+    },
+    {
+      "epoch": 2,
+      "train_loss": 1.6384362813046105
+    },
+    {
+      "epoch": 3,
+      "train_loss": 1.4120025303489283
+    },
+    {
+      "epoch": 4,
+      "train_loss": 1.5674286854894537
+    },
+    {
+      "epoch": 5,
+      "train_loss": 1.5401874258643702
+    },
+    {
+      "epoch": 6,
+      "train_loss": 1.4283890882291292
+    },
+    {
+      "epoch": 7,
+      "train_loss": 1.5478344112948368
+    },
+    {
+      "epoch": 8,
+      "train_loss": 1.6749039321196706
+    },
+    {
+      "epoch": 9,
+      "train_loss": 1.3940781457800615
+    },
+    {
+      "epoch": 10,
+      "train_loss": 1.3927083012932224
+    },
+    {
+      "epoch": 11,
+      "train_loss": 1.5125436822991623
+    },
+    {
+      "epoch": 12,
+      "train_loss": 1.6823845085344817
+    },
+    {
+      "epoch": 13,
+      "train_loss": 1.4184414868605764
+    },
+    {
+      "epoch": 14,
+      "train_loss": 1.4393992233276367
+    },
+    {
+      "epoch": 15,
+      "train_loss": 1.5265272190696315
+    },
+    {
+      "epoch": 16,
+      "train_loss": 1.6243808332242464
+    },
+    {
+      "epoch": 17,
+      "train_loss": 1.3540453097694798
+    },
+    {
+      "epoch": 18,
+      "train_loss": 1.499637666501497
+    },
+    {
+      "epoch": 19,
+      "train_loss": 1.3106657308026364
+    },
+    {
+      "epoch": 20,
+      "train_loss": 1.5018967904542622
+    },
+    {
+      "epoch": 21,
+      "train_loss": 1.3504434702270909
+    },
+    {
+      "epoch": 22,
+      "train_loss": 1.6013833044704637
+    },
+    {
+      "epoch": 23,
+      "train_loss": 1.6672487831115723
+    },
+    {
+      "epoch": 24,
+      "train_loss": 1.4507962395015517
+    },
+    {
+      "epoch": 25,
+      "train_loss": 1.3700003109480205
+    },
+    {
+      "epoch": 26,
+      "train_loss": 1.3767117627043473
+    },
+    {
+      "epoch": 27,
+      "train_loss": 1.4459412696487026
+    },
+    {
+      "epoch": 28,
+      "train_loss": 1.3630849577251234
+    },
+    {
+      "epoch": 29,
+      "train_loss": 1.6651583372919183
+    },
+    {
+      "epoch": 30,
+      "train_loss": 1.3656946495959634
+    },
+    {
+      "epoch": 31,
+      "train_loss": 1.410263954965692
+    },
+    {
+      "epoch": 32,
+      "train_loss": 1.3972627205597727
+    },
+    {
+      "epoch": 33,
+      "train_loss": 1.5788899810690629
+    },
+    {
+      "epoch": 34,
+      "train_loss": 1.5010775119379947
+    },
+    {
+      "epoch": 35,
+      "train_loss": 1.4368829202651978
+    },
+    {
+      "epoch": 36,
+      "train_loss": 1.3498843838039198
+    },
+    {
+      "epoch": 37,
+      "train_loss": 1.3490342864237335
+    },
+    {
+      "epoch": 38,
+      "train_loss": 1.3266521142658434
+    },
+    {
+      "epoch": 39,
+      "train_loss": 1.3577745432602732
+    },
+    {
+      "epoch": 40,
+      "train_loss": 1.2795040179553785
+    }
+  ],
+  "epochs_completed": 40,
+  "finite": true,
+  "first_nonfinite_epoch": null,
+  "mission_lesion_rollout": {
+    "episodes": 500,
+    "mean_length": 59.866,
+    "mean_return": 0.06493437499999999,
+    "success": 0.068
+  },
+  "mission_lesion_validation": {
+    "accuracy": 0.6101077184741233,
+    "loss": 1.1676633622739598
+  },
+  "predictor": [],
+  "provenance": {
+    "cuda_device_name": "NVIDIA GeForce GTX 1080",
+    "cuda_peak_allocated_bytes": 130029568,
+    "cuda_visible_devices": "2",
+    "device": "cuda:0",
+    "git_commit": "7b7d7a08a5e78c3379397f01f0c9af97d5d84fbd",
+    "git_dirty_tracked": false,
+    "minigrid_version": "3.1.0",
+    "torch_version": "2.3.1+cu118"
+  },
+  "rollout": {
+    "episodes": 500,
+    "mean_length": 59.886,
+    "mean_return": 0.064653125,
+    "success": 0.068
+  },
+  "stage": "babyai_pickup_p2",
+  "training": {
+    "batch_size": 256,
+    "history_steps": 4,
+    "model_seed": 4103,
+    "neutral_examples_per_epoch": 0,
+    "parameter_count_including_reciprocal_context_predictor": 1438215,
+    "shuffle_seed": 4103,
+    "total_wall_seconds_including_rollouts": 171.7450816631317,
+    "training_wall_seconds": 154.85608053207397
+  },
+  "validation": {
+    "accuracy": 0.5960962177047056,
+    "loss": 1.1890508834481712
+  }
+}
diff --git a/results/babyai_shared/pickup_p2/seed4103_sdil.json b/results/babyai_shared/pickup_p2/seed4103_sdil.json
new file mode 100644
index 0000000..acbef8b
--- /dev/null
+++ b/results/babyai_shared/pickup_p2/seed4103_sdil.json
@@ -0,0 +1,295 @@
+{
+  "condition": "sdil",
+  "config": {
+    "action_dim": 7,
+    "context_gain": 1.0,
+    "hidden_layers": 4,
+    "input_dim": 3964,
+    "learning_rate": 0.03,
+    "mission_dim": 23,
+    "momentum": 0.9,
+    "reciprocal_learning_rate": 0.03,
+    "weight_decay": 0.0001,
+    "width": 256
+  },
+  "data": {
+    "color_cardinality": 6,
+    "elapsed_seconds": 81.82925486564636,
+    "env_id": "BabyAI-PickupLoc-v0",
+    "minigrid_version": "3.1.0",
+    "object_cardinality": 11,
+    "protocol": "babyai_shared_feedback_b0",
+    "rollout_episodes": 500,
+    "rollout_seed_start": 200000,
+    "state_cardinality": 3,
+    "train_episodes": 20000,
+    "train_steps": 121467,
+    "validation_episodes": 2000,
+    "validation_steps": 12347,
+    "vocabulary": [
+      "<unk>",
+      "a",
+      "ball",
+      "behind",
+      "blue",
+      "box",
+      "front",
+      "green",
+      "grey",
+      "in",
+      "key",
+      "left",
+      "of",
+      "on",
+      "pick",
+      "purple",
+      "red",
+      "right",
+      "the",
+      "up",
+      "yellow",
+      "you",
+      "your"
+    ]
+  },
+  "epoch_history": [
+    {
+      "epoch": 1,
+      "train_loss": 1.4704136446902627
+    },
+    {
+      "epoch": 2,
+      "train_loss": 1.4733195139232436
+    },
+    {
+      "epoch": 3,
+      "train_loss": 1.6606671679647345
+    },
+    {
+      "epoch": 4,
+      "train_loss": 1.6223741378282246
+    },
+    {
+      "epoch": 5,
+      "train_loss": 1.380738465158563
+    },
+    {
+      "epoch": 6,
+      "train_loss": 1.4576800007569162
+    },
+    {
+      "epoch": 7,
+      "train_loss": 1.3935401956658615
+    },
+    {
+      "epoch": 8,
+      "train_loss": 1.5138149407035426
+    },
+    {
+      "epoch": 9,
+      "train_loss": 1.5381833568372225
+    },
+    {
+      "epoch": 10,
+      "train_loss": 1.602822558503402
+    },
+    {
+      "epoch": 11,
+      "train_loss": 1.4934680660147417
+    },
+    {
+      "epoch": 12,
+      "train_loss": 1.462449248213517
+    },
+    {
+      "epoch": 13,
+      "train_loss": 1.5229607145409836
+    },
+    {
+      "epoch": 14,
+      "train_loss": 1.4338193808103863
+    },
+    {
+      "epoch": 15,
+      "train_loss": 1.7763569934744583
+    },
+    {
+      "epoch": 16,
+      "train_loss": 1.4266947914424695
+    },
+    {
+      "epoch": 17,
+      "train_loss": 1.4615587412683588
+    },
+    {
+      "epoch": 18,
+      "train_loss": 1.3746792825899625
+    },
+    {
+      "epoch": 19,
+      "train_loss": 1.4826162985751503
+    },
+    {
+      "epoch": 20,
+      "train_loss": 1.3551955336018613
+    },
+    {
+      "epoch": 21,
+      "train_loss": 1.6226069836867483
+    },
+    {
+      "epoch": 22,
+      "train_loss": 1.381537736340573
+    },
+    {
+      "epoch": 23,
+      "train_loss": 1.5173151382647063
+    },
+    {
+      "epoch": 24,
+      "train_loss": 1.54268351931321
+    },
+    {
+      "epoch": 25,
+      "train_loss": 1.5190619021967837
+    },
+    {
+      "epoch": 26,
+      "train_loss": 1.66076710274345
+    },
+    {
+      "epoch": 27,
+      "train_loss": 1.6404679448981034
+    },
+    {
+      "epoch": 28,
+      "train_loss": 1.3772655429338154
+    },
+    {
+      "epoch": 29,
+      "train_loss": 1.4842414411745573
+    },
+    {
+      "epoch": 30,
+      "train_loss": 1.382061290239033
+    },
+    {
+      "epoch": 31,
+      "train_loss": 1.3829637015493292
+    },
+    {
+      "epoch": 32,
+      "train_loss": 1.428046850907175
+    },
+    {
+      "epoch": 33,
+      "train_loss": 1.4464430650911833
+    },
+    {
+      "epoch": 34,
+      "train_loss": 1.4072801037838585
+    },
+    {
+      "epoch": 35,
+      "train_loss": 1.3329616769991424
+    },
+    {
+      "epoch": 36,
+      "train_loss": 1.436322965872915
+    },
+    {
+      "epoch": 37,
+      "train_loss": 1.3869416563134445
+    },
+    {
+      "epoch": 38,
+      "train_loss": 1.469730708975541
+    },
+    {
+      "epoch": 39,
+      "train_loss": 1.48009900143272
+    },
+    {
+      "epoch": 40,
+      "train_loss": 1.498611428110223
+    }
+  ],
+  "epochs_completed": 40,
+  "finite": true,
+  "first_nonfinite_epoch": null,
+  "mission_lesion_rollout": {
+    "episodes": 500,
+    "mean_length": 64.0,
+    "mean_return": 0.0,
+    "success": 0.0
+  },
+  "mission_lesion_validation": {
+    "accuracy": 0.49890661699198185,
+    "loss": 1.3523273018484532
+  },
+  "predictor": [
+    {
+      "action_observations": 0,
+      "context_rms": 0.20582586526870728,
+      "mean_per_cell_r2": 0.1324767768383026,
+      "neutral_observations": 4096,
+      "residual_context_rms_ratio": 0.6195260882377625,
+      "teaching_observations": 0
+    },
+    {
+      "action_observations": 0,
+      "context_rms": 0.2105173021554947,
+      "mean_per_cell_r2": 0.40293416380882263,
+      "neutral_observations": 4096,
+      "residual_context_rms_ratio": 0.5007978081703186,
+      "teaching_observations": 0
+    },
+    {
+      "action_observations": 0,
+      "context_rms": 0.2122984677553177,
+      "mean_per_cell_r2": 0.2902011275291443,
+      "neutral_observations": 4096,
+      "residual_context_rms_ratio": 0.5484216213226318,
+      "teaching_observations": 0
+    },
+    {
+      "action_observations": 0,
+      "context_rms": 0.21755045652389526,
+      "mean_per_cell_r2": 0.4205133020877838,
+      "neutral_observations": 4096,
+      "residual_context_rms_ratio": 0.4854767322540283,
+      "teaching_observations": 0
+    }
+  ],
+  "provenance": {
+    "cuda_device_name": "NVIDIA GeForce GTX 1080",
+    "cuda_peak_allocated_bytes": 435783168,
+    "cuda_visible_devices": "3",
+    "device": "cuda:0",
+    "git_commit": "7b7d7a08a5e78c3379397f01f0c9af97d5d84fbd",
+    "git_dirty_tracked": false,
+    "minigrid_version": "3.1.0",
+    "torch_version": "2.3.1+cu118"
+  },
+  "rollout": {
+    "episodes": 500,
+    "mean_length": 64.0,
+    "mean_return": 0.0,
+    "success": 0.0
+  },
+  "stage": "babyai_pickup_p2",
+  "training": {
+    "batch_size": 256,
+    "history_steps": 4,
+    "model_seed": 4103,
+    "neutral_examples_per_epoch": 4096,
+    "parameter_count_including_reciprocal_context_predictor": 1438215,
+    "shuffle_seed": 4103,
+    "total_wall_seconds_including_rollouts": 171.52505588531494,
+    "training_wall_seconds": 152.92210721969604
+  },
+  "validation": {
+    "accuracy": 0.49890661699198185,
+    "loss": 1.351929626687293
+  }
+}
diff --git a/results/babyai_shared/pickup_p2_analysis.json b/results/babyai_shared/pickup_p2_analysis.json
new file mode 100644
index 0000000..845e70c
--- /dev/null
+++ b/results/babyai_shared/pickup_p2_analysis.json
@@ -0,0 +1,174 @@
+{
+  "checks": {
+    "all_runs_finite": true,
+    "clean_kp_success_at_least_60_every_seed": true,
+    "raw_below_clean_kp_every_seed": true,
+    "sdil_above_raw_every_seed": false,
+    "sdil_nonzero_success_every_seed": false
+  },
+  "conditions": {
+    "bp": {
+      "all_finite": true,
+      "expert_action_accuracy_percent": {
+        "mean": 87.78380713263681,
+        "sample_std": 0.1078533752082175,
+        "values": [
+          87.8108042439459,
+          87.87559731108772,
+          87.66501984287682
+        ]
+      },
+      "mission_lesion_rollout_success_percent": {
+        "mean": 49.333333333333336,
+        "sample_std": 1.1015141094572214,
+        "values": [
+          48.8,
+          50.6,
+          48.6
+        ]
+      },
+      "rollout_success_percent": {
+        "mean": 79.33333333333333,
+        "sample_std": 5.1588112325741635,
+        "values": [
+          80.80000000000001,
+          83.6,
+          73.6
+        ]
+      }
+    },
+    "clean_kp": {
+      "all_finite": true,
+      "expert_action_accuracy_percent": {
+        "mean": 86.43665127831322,
+        "sample_std": 0.3517610682148285,
+        "values": [
+          86.13428363165141,
+          86.8227099700332,
+          86.35296023325503
+        ]
+      },
+      "mission_lesion_rollout_success_percent": {
+        "mean": 48.86666666666667,
+        "sample_std": 7.245228314783008,
+        "values": [
+          49.8,
+          55.60000000000001,
+          41.199999999999996
+        ]
+      },
+      "rollout_success_percent": {
+        "mean": 70.86666666666667,
+        "sample_std": 5.34540300944029,
+        "values": [
+          67.2,
+          77.0,
+          68.4
+        ]
+      }
+    },
+    "raw_shared": {
+      "all_finite": true,
+      "expert_action_accuracy_percent": {
+        "mean": 59.642018304041464,
+        "sample_std": 9.767595165741607,
+        "values": [
+          49.89066169919818,
+          69.42577144245566,
+          59.609621770470554
+        ]
+      },
+      "mission_lesion_rollout_success_percent": {
+        "mean": 8.733333333333334,
+        "sample_std": 9.843441132720475,
+        "values": [
+          0.0,
+          19.400000000000002,
+          6.800000000000001
+        ]
+      },
+      "rollout_success_percent": {
+        "mean": 10.000000000000002,
+        "sample_std": 11.92644121270046,
+        "values": [
+          0.0,
+          23.200000000000003,
+          6.800000000000001
+        ]
+      }
+    },
+    "sdil": {
+      "all_finite": true,
+      "expert_action_accuracy_percent": {
+        "mean": 49.89066169919818,
+        "sample_std": 0.0,
+        "values": [
+          49.89066169919818,
+          49.89066169919818,
+          49.89066169919818
+        ]
+      },
+      "mission_lesion_rollout_success_percent": {
+        "mean": 0.0,
+        "sample_std": 0.0,
+        "values": [
+          0.0,
+          0.0,
+          0.0
+        ]
+      },
+      "rollout_success_percent": {
+        "mean": 0.0,
+        "sample_std": 0.0,
+        "values": [
+          0.0,
+          0.0,
+          0.0
+        ]
+      }
+    }
+  },
+  "decision": "Close the diagonal per-cell predictor on PickupLoc. Raw shared feedback fails as hypothesized, but this SDIL implementation does not recover it. Screen the originally specified population predictor using neutral prediction only before another endpoint.",
+  "gate": "fail",
+  "paired_clean_kp_minus_raw_rollout_points": {
+    "mean": 60.866666666666674,
+    "sample_std": 6.730032194078519,
+    "values": [
+      67.2,
+      53.800000000000004,
+      61.60000000000001
+    ]
+  },
+  "paired_sdil_minus_raw_rollout_points": {
+    "mean": -10.000000000000002,
+    "sample_std": 11.92644121270046,
+    "values": [
+      0.0,
+      -23.200000000000003,
+      -6.800000000000001
+    ]
+  },
+  "sdil_predictor": {
+    "action_or_teaching_observations": 0,
+    "maximum_layer_residual_context_rms_ratio": {
+      "mean": 0.614041248957316,
+      "sample_std": 0.011718512673796736,
+      "values": [
+        0.6220114827156067,
+        0.6005861759185791,
+        0.6195260882377625
+      ]
+    },
+    "mean_per_cell_r2": {
+      "mean": 0.30999962737162906,
+      "sample_std": 0.005533015870761776,
+      "values": [
+        0.3146054223179817,
+        0.3038621172308922,
+        0.31153134256601334
+      ]
+    }
+  },
+  "stage": "babyai_pickup_p2_analysis",
+  "test_split_generated_or_read": false
+}