Add model config

config.json

@@ -7,7 +7,7 @@
   "lowrank_dim": 0,
   "bInclude_glyph_bag": true,
   "bInclude_hero": true,
-  "dropout_rate": 0.2,
+  "dropout_rate": 0.1,
   "enable_dropout_on_latent": true,
   "enable_dropout_on_decoder": true,
   "architecture": "Multi-modal Variational Autoencoder for NetHack game states",
@@ -27,9 +27,11 @@
   ],
   "training_config": {
     "epochs": 15,
-    "batch_size": 32,
-    "learning_rate": 0.0005,
+    "batch_size": 1024,
+    "max_learning_rate": 0.001,
     "sequence_size": 32,
+    "shuffle_batches": true,
+    "shuffle_within_batch": true,
     "adaptive_weighting": {
       "initial_weight_emb": 1.5,
       "final_weight_emb": 0.0,
@@ -37,22 +39,34 @@
       "initial_weight_raw": 0.4,
       "final_weight_raw": 1.0,
       "weight_raw_shape": "linear",
-      "initial_kl_beta": 0.0001,
-      "final_kl_beta": 0.6,
-      "kl_beta_shape": "cosine",
-      "warmup_epoch_ratio": 0.4
+      "initial_mi_beta": 0.0,
+      "final_mi_beta": 0.0,
+      "mi_beta_shape": "constant",
+      "initial_tc_beta": 5.0,
+      "final_tc_beta": 5.0,
+      "tc_beta_shape": "constant",
+      "initial_dw_beta": 0.02,
+      "final_dw_beta": 0.3,
+      "dw_beta_shape": "custom",
+      "warmup_epoch_ratio": 0.2
     },
-    "total_correlation_beta_multiplier": 10.0,
     "free_bits": 0.15,
     "focal_loss_alpha": 0.75,
     "focal_loss_gamma": 2.0,
-    "dropout_rate": 0.2,
+    "dropout_rate": 0.1,
     "enable_dropout_on_latent": true,
-    "enable_dropout_on_decoder": true
+    "enable_dropout_on_decoder": true,
+    "early_stopping": {
+      "enabled": true,
+      "patience": 3,
+      "min_delta": 0.01,
+      "triggered": true,
+      "best_epoch": 3
+    }
   },
-  "final_train_loss": 549.4707604980468,
-  "final_test_loss": 8138.669360351562,
-  "best_train_loss": 522.560433959961,
-  "best_test_loss": 3383.009033203125,
+  "final_train_loss": 2710.9047668457033,
+  "final_test_loss": 2245.631726074219,
+  "best_train_loss": 2710.9047668457033,
+  "best_test_loss": 2245.631726074219,
   "total_epochs": 15
 }