{
  "best_val_loss": 0.0074350039093196395,
  "test_loss": 0.007343360532224178,
  "history": [
    {
      "epoch": 1,
      "train_loss": 0.03299331233888864,
      "val_loss": 0.025188268708586693
    },
    {
      "epoch": 2,
      "train_loss": 0.015773713049292565,
      "val_loss": 0.01214492101430893
    },
    {
      "epoch": 3,
      "train_loss": 0.011268323176503182,
      "val_loss": 0.011055553800463677
    },
    {
      "epoch": 4,
      "train_loss": 0.010149860578477382,
      "val_loss": 0.009853113369047642
    },
    {
      "epoch": 5,
      "train_loss": 0.009246495521962642,
      "val_loss": 0.00933172680824995
    },
    {
      "epoch": 6,
      "train_loss": 0.008860349697470665,
      "val_loss": 0.009004901867508889
    },
    {
      "epoch": 7,
      "train_loss": 0.008579391563683748,
      "val_loss": 0.008856775346696377
    },
    {
      "epoch": 8,
      "train_loss": 0.008333331797942519,
      "val_loss": 0.00840215255945921
    },
    {
      "epoch": 9,
      "train_loss": 0.008022411068752408,
      "val_loss": 0.008290996616780759
    },
    {
      "epoch": 10,
      "train_loss": 0.007867376347407699,
      "val_loss": 0.008192000594735146
    },
    {
      "epoch": 11,
      "train_loss": 0.007745014687851071,
      "val_loss": 0.008125807121396066
    },
    {
      "epoch": 12,
      "train_loss": 0.007642358961328864,
      "val_loss": 0.0080087585195899
    },
    {
      "epoch": 13,
      "train_loss": 0.0075519825259596106,
      "val_loss": 0.007978914509415626
    },
    {
      "epoch": 14,
      "train_loss": 0.0074750817969441415,
      "val_loss": 0.007824904982745648
    },
    {
      "epoch": 15,
      "train_loss": 0.007405085674375296,
      "val_loss": 0.007741572364121676
    },
    {
      "epoch": 16,
      "train_loss": 0.0073448564377427104,
      "val_loss": 0.00778407304406166
    },
    {
      "epoch": 17,
      "train_loss": 0.007286443270593882,
      "val_loss": 0.007653026497960091
    },
    {
      "epoch": 18,
      "train_loss": 0.007236691383421421,
      "val_loss": 0.007664755893498659
    },
    {
      "epoch": 19,
      "train_loss": 0.007191714821234345,
      "val_loss": 0.007630967005193234
    },
    {
      "epoch": 20,
      "train_loss": 0.007150140590667725,
      "val_loss": 0.007592904810905457
    },
    {
      "epoch": 21,
      "train_loss": 0.007111764398515225,
      "val_loss": 0.007567611115574837
    },
    {
      "epoch": 22,
      "train_loss": 0.007073474420830607,
      "val_loss": 0.007488598871529102
    },
    {
      "epoch": 23,
      "train_loss": 0.0070404944337159395,
      "val_loss": 0.00750829313442111
    },
    {
      "epoch": 24,
      "train_loss": 0.007008424286395311,
      "val_loss": 0.007458318319022656
    },
    {
      "epoch": 25,
      "train_loss": 0.006979450356587768,
      "val_loss": 0.0074350039093196395
    }
  ],
  "args": {
    "dataset_path": "le-wm/models/datasets/pusht_expert_train.h5",
    "checkpoint_cache_dir": "le-wm/models",
    "checkpoint": "hf_pusht/weights.pt",
    "output_dir": "models/latent_decoder/pusht_lewm_1_000_000",
    "max_samples": 1000000,
    "sample_block_size": 16,
    "batch_size": 128,
    "num_workers": 0,
    "epochs": 25,
    "patience": 6,
    "lr": 0.001,
    "weight_decay": 0.0001,
    "hidden_dim": 256,
    "depth": 4,
    "heads": 8,
    "mlp_ratio": 4.0,
    "patch_size": 16,
    "image_size": 224,
    "mse_weight": 0.25,
    "foreground_weight": 10.0,
    "foreground_threshold": 0.05,
    "seed": 3072,
    "device": "cuda",
    "cache_latents": true
  },
  "split": {
    "train_fraction": 0.8,
    "val_fraction": 0.1,
    "seed": 3072
  },
  "num_samples": {
    "train": 800000,
    "val": 100000,
    "test": 100000
  }
}