{
  "audits": [
    {
      "epoch": 0,
      "layers": [
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.20582586526870728,
            "mean_per_cell_r2": 0.9770644903182983,
            "observations": 4096,
            "residual_context_rms_ratio": 0.09758477658033371
          },
          "holdout": {
            "context_rms": 0.20538413524627686,
            "mean_per_cell_r2": 0.9738611578941345,
            "observations": 4096,
            "residual_context_rms_ratio": 0.10475123673677444
          },
          "layer": 0,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.2105173021554947,
            "mean_per_cell_r2": 0.9947838187217712,
            "observations": 4096,
            "residual_context_rms_ratio": 0.04489493742585182
          },
          "holdout": {
            "context_rms": 0.21071431040763855,
            "mean_per_cell_r2": 0.993838369846344,
            "observations": 4096,
            "residual_context_rms_ratio": 0.04897325858473778
          },
          "layer": 1,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.2122984677553177,
            "mean_per_cell_r2": 0.9970771074295044,
            "observations": 4096,
            "residual_context_rms_ratio": 0.03391746059060097
          },
          "holdout": {
            "context_rms": 0.2119303047657013,
            "mean_per_cell_r2": 0.9966038465499878,
            "observations": 4096,
            "residual_context_rms_ratio": 0.03683717921376228
          },
          "layer": 2,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.21755045652389526,
            "mean_per_cell_r2": 0.997891902923584,
            "observations": 4096,
            "residual_context_rms_ratio": 0.028139039874076843
          },
          "holdout": {
            "context_rms": 0.21835927665233612,
            "mean_per_cell_r2": 0.9975327253341675,
            "observations": 4096,
            "residual_context_rms_ratio": 0.030591214075684547
          },
          "layer": 3,
          "teaching_observations": 0
        }
      ]
    },
    {
      "epoch": 1,
      "layers": [
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.20582586526870728,
            "mean_per_cell_r2": 0.8622075915336609,
            "observations": 4096,
            "residual_context_rms_ratio": 0.24054336547851562
          },
          "holdout": {
            "context_rms": 0.20538413524627686,
            "mean_per_cell_r2": 0.8481650352478027,
            "observations": 4096,
            "residual_context_rms_ratio": 0.2538592517375946
          },
          "layer": 0,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.2105173021554947,
            "mean_per_cell_r2": 0.9787577390670776,
            "observations": 4096,
            "residual_context_rms_ratio": 0.09131480753421783
          },
          "holdout": {
            "context_rms": 0.21071431040763855,
            "mean_per_cell_r2": 0.9753843545913696,
            "observations": 4096,
            "residual_context_rms_ratio": 0.0986090749502182
          },
          "layer": 1,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.2122984677553177,
            "mean_per_cell_r2": 0.9843738079071045,
            "observations": 4096,
            "residual_context_rms_ratio": 0.07851728051900864
          },
          "holdout": {
            "context_rms": 0.2119303047657013,
            "mean_per_cell_r2": 0.9818754196166992,
            "observations": 4096,
            "residual_context_rms_ratio": 0.08518385887145996
          },
          "layer": 2,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.21755045652389526,
            "mean_per_cell_r2": 0.9888011813163757,
            "observations": 4096,
            "residual_context_rms_ratio": 0.06514516472816467
          },
          "holdout": {
            "context_rms": 0.21835927665233612,
            "mean_per_cell_r2": 0.9868994951248169,
            "observations": 4096,
            "residual_context_rms_ratio": 0.07073986530303955
          },
          "layer": 3,
          "teaching_observations": 0
        }
      ]
    },
    {
      "epoch": 5,
      "layers": [
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.20582586526870728,
            "mean_per_cell_r2": 0.8066374063491821,
            "observations": 4096,
            "residual_context_rms_ratio": 0.2849716544151306
          },
          "holdout": {
            "context_rms": 0.20538413524627686,
            "mean_per_cell_r2": 0.7795184254646301,
            "observations": 4096,
            "residual_context_rms_ratio": 0.3061090111732483
          },
          "layer": 0,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.2105173021554947,
            "mean_per_cell_r2": 0.9446667432785034,
            "observations": 4096,
            "residual_context_rms_ratio": 0.14740806818008423
          },
          "holdout": {
            "context_rms": 0.21071431040763855,
            "mean_per_cell_r2": 0.9354867935180664,
            "observations": 4096,
            "residual_context_rms_ratio": 0.15956945717334747
          },
          "layer": 1,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.2122984677553177,
            "mean_per_cell_r2": 0.9587752819061279,
            "observations": 4096,
            "residual_context_rms_ratio": 0.12740013003349304
          },
          "holdout": {
            "context_rms": 0.2119303047657013,
            "mean_per_cell_r2": 0.9532089233398438,
            "observations": 4096,
            "residual_context_rms_ratio": 0.13679398596286774
          },
          "layer": 2,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.21755045652389526,
            "mean_per_cell_r2": 0.979941725730896,
            "observations": 4096,
            "residual_context_rms_ratio": 0.08685733377933502
          },
          "holdout": {
            "context_rms": 0.21835927665233612,
            "mean_per_cell_r2": 0.9771287441253662,
            "observations": 4096,
            "residual_context_rms_ratio": 0.09327036887407303
          },
          "layer": 3,
          "teaching_observations": 0
        }
      ]
    },
    {
      "epoch": 10,
      "layers": [
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.20582586526870728,
            "mean_per_cell_r2": 0.7986479997634888,
            "observations": 4096,
            "residual_context_rms_ratio": 0.29081863164901733
          },
          "holdout": {
            "context_rms": 0.20538413524627686,
            "mean_per_cell_r2": 0.7685507535934448,
            "observations": 4096,
            "residual_context_rms_ratio": 0.31361645460128784
          },
          "layer": 0,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.2105173021554947,
            "mean_per_cell_r2": 0.9244306087493896,
            "observations": 4096,
            "residual_context_rms_ratio": 0.1723565012216568
          },
          "holdout": {
            "context_rms": 0.21071431040763855,
            "mean_per_cell_r2": 0.9117617607116699,
            "observations": 4096,
            "residual_context_rms_ratio": 0.1869732290506363
          },
          "layer": 1,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.2122984677553177,
            "mean_per_cell_r2": 0.9454227685928345,
            "observations": 4096,
            "residual_context_rms_ratio": 0.14660939574241638
          },
          "holdout": {
            "context_rms": 0.2119303047657013,
            "mean_per_cell_r2": 0.9378501176834106,
            "observations": 4096,
            "residual_context_rms_ratio": 0.15752699971199036
          },
          "layer": 2,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.21755045652389526,
            "mean_per_cell_r2": 0.9747556447982788,
            "observations": 4096,
            "residual_context_rms_ratio": 0.0976947769522667
          },
          "holdout": {
            "context_rms": 0.21835927665233612,
            "mean_per_cell_r2": 0.9714953899383545,
            "observations": 4096,
            "residual_context_rms_ratio": 0.10434764623641968
          },
          "layer": 3,
          "teaching_observations": 0
        }
      ]
    },
    {
      "epoch": 20,
      "layers": [
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.20582586526870728,
            "mean_per_cell_r2": 0.842189610004425,
            "observations": 4096,
            "residual_context_rms_ratio": 0.25771480798721313
          },
          "holdout": {
            "context_rms": 0.20538413524627686,
            "mean_per_cell_r2": 0.8169145584106445,
            "observations": 4096,
            "residual_context_rms_ratio": 0.2789650857448578
          },
          "layer": 0,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.2105173021554947,
            "mean_per_cell_r2": 0.9072440266609192,
            "observations": 4096,
            "residual_context_rms_ratio": 0.1907738745212555
          },
          "holdout": {
            "context_rms": 0.21071431040763855,
            "mean_per_cell_r2": 0.8923318982124329,
            "observations": 4096,
            "residual_context_rms_ratio": 0.2059631049633026
          },
          "layer": 1,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.2122984677553177,
            "mean_per_cell_r2": 0.9364410638809204,
            "observations": 4096,
            "residual_context_rms_ratio": 0.15774686634540558
          },
          "holdout": {
            "context_rms": 0.2119303047657013,
            "mean_per_cell_r2": 0.9272578954696655,
            "observations": 4096,
            "residual_context_rms_ratio": 0.1698361188173294
          },
          "layer": 2,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.21755045652389526,
            "mean_per_cell_r2": 0.9725677967071533,
            "observations": 4096,
            "residual_context_rms_ratio": 0.10174933075904846
          },
          "holdout": {
            "context_rms": 0.21835927665233612,
            "mean_per_cell_r2": 0.9688738584518433,
            "observations": 4096,
            "residual_context_rms_ratio": 0.1088051050901413
          },
          "layer": 3,
          "teaching_observations": 0
        }
      ]
    },
    {
      "epoch": 40,
      "layers": [
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.20582586526870728,
            "mean_per_cell_r2": 0.7883090376853943,
            "observations": 4096,
            "residual_context_rms_ratio": 0.29959437251091003
          },
          "holdout": {
            "context_rms": 0.20538413524627686,
            "mean_per_cell_r2": 0.7517150640487671,
            "observations": 4096,
            "residual_context_rms_ratio": 0.32660192251205444
          },
          "layer": 0,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.2105173021554947,
            "mean_per_cell_r2": 0.8573880195617676,
            "observations": 4096,
            "residual_context_rms_ratio": 0.23693645000457764
          },
          "holdout": {
            "context_rms": 0.21071431040763855,
            "mean_per_cell_r2": 0.8369072675704956,
            "observations": 4096,
            "residual_context_rms_ratio": 0.25433388352394104
          },
          "layer": 1,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.2122984677553177,
            "mean_per_cell_r2": 0.9061275124549866,
            "observations": 4096,
            "residual_context_rms_ratio": 0.19269712269306183
          },
          "holdout": {
            "context_rms": 0.2119303047657013,
            "mean_per_cell_r2": 0.8909499645233154,
            "observations": 4096,
            "residual_context_rms_ratio": 0.20921701192855835
          },
          "layer": 2,
          "teaching_observations": 0
        },
        {
          "action_observations": 0,
          "calibration": {
            "context_rms": 0.21755045652389526,
            "mean_per_cell_r2": 0.9699275493621826,
            "observations": 4096,
            "residual_context_rms_ratio": 0.10622629523277283
          },
          "holdout": {
            "context_rms": 0.21835927665233612,
            "mean_per_cell_r2": 0.966497540473938,
            "observations": 4096,
            "residual_context_rms_ratio": 0.11257338523864746
          },
          "layer": 3,
          "teaching_observations": 0
        }
      ]
    }
  ],
  "calibration_examples": 4096,
  "checks": {
    "holdout_r2_at_least_0p8_every_epoch_and_layer": false,
    "holdout_residual_ratio_at_most_0p25_every_epoch_and_layer": false,
    "zero_action_or_teaching_observations": true
  },
  "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
  },
  "epoch_clean_kp_train_loss": [
    0.6864033263608029,
    0.5136532720766569,
    0.4820757521453657,
    0.4678087641063489,
    0.4570458981237913,
    0.44718988864045395,
    0.43997702253492255,
    0.42566254320897556,
    0.4220644837931583,
    0.41289369489017286,
    0.405345781037682,
    0.40259038736945707,
    0.393158560112903,
    0.38850406364390727,
    0.3801098694926814,
    0.36695603276553906,
    0.35903412762441134,
    0.34711836858799583,
    0.3370369092727962,
    0.3200993647387153,
    0.3118796497897098,
    0.293176792420839,
    0.28278176423750423,
    0.2673768097789664,
    0.2630335626790398,
    0.2525263265559548,
    0.24664549178198764,
    0.23551284742982764,
    0.23284720298491027,
    0.22212432547619468,
    0.21526990802664506,
    0.20761467494462665,
    0.21185886748527225,
    0.20015101647690722,
    0.19496909729744258,
    0.1936807933763454,
    0.19314897574876483,
    0.1823754488794427,
    0.17861489517124077,
    0.17446756489967044
  ],
  "gate": "fail",
  "holdout_examples": 4096,
  "model_seed": 4103,
  "provenance": {
    "cuda_visible_devices": "2",
    "device": "cuda:0",
    "git_commit": "b5eab6cd911e074ae9268303389dfbe52e636e26",
    "torch_version": "2.3.1+cu118"
  },
  "raw_sdil_rollout_or_test_outcomes_read": false,
  "ridge_relative_to_mean_soma_variance": 0.001,
  "shuffle_seed": 4103,
  "stage": "babyai_population_predictor_capacity_screen"
}