Update delta-iris/atari/boxing/config.json

delta-iris/atari/boxing/config.json

@@ -0,0 +1,155 @@
+{
+  "name": "boxing_delta_iris_world_model",
+  "env": "BoxingNoFrameSkip-v4",
+  "model_type": "delta_iris",
+  "metadata": {
+    "latent_dim": [1, 4, 1024],
+    "two_hot_rews": false,
+    "tokens_per_block": 6,
+    "num_tokens": 4,
+    "tokens_grid_res": 2,
+    "token_res": 16384
+  },
+  "util_folders":{
+    "data": "../src/data",
+    "models": "../src/models"
+  },
+  "requirements":{
+    "-r": "requirements.txt"
+  },
+  "models": [
+    {
+      "name": "world_model",
+      "framework": null,
+      "format": "state_dict",
+      "source": {
+        "weights_path": "world_model.pt",
+        "class_path": "../src/world_model.py",
+        "class_name": "WorldModel",
+        "class_args": [
+        {
+          "latent_vocab_size": 1024,
+          "num_actions": 18,
+          "image_channels": 3,
+          "image_size": 64,
+          "latents_weight": 1.0,
+          "rewards_weight": 1.0,
+          "ends_weight": 1.0,
+          "two_hot_rews": false,
+          "transformer_config": {
+            "tokens_per_block": 6,
+            "max_blocks": 26,
+            "num_layers": 3,
+            "num_heads": 4,
+            "embed_dim": 256,
+            "attention": "causal",
+            "embed_pdrop": 0.0,
+            "resid_pdrop": 0.0,
+            "attn_pdrop": 0.0
+          },
+          "frame_cnn_config": {
+            "image_channels": 3,
+            "latent_dim": 4,
+            "num_channels": 32,
+            "mult": [1, 1, 2, 2, 4],
+            "down": [1, 0, 1, 1, 0]
+          }
+        }]
+      },
+      "signature": {
+        "inputs": ["wm_input_sequence", "use_kv_cache"],
+        "call_mode": "positional"
+      },
+      "sub_models":
+      [
+        {
+          "name": "act_emb",
+          "signature": 
+          {
+            "inputs": ["act"],
+            "call_mode": "positional"
+          }
+        },
+        {
+          "name": "latents_emb",
+          "signature": 
+          {
+            "inputs": ["latent_tokens"],
+            "call_mode": "positional"
+          }
+        },
+        {
+          "name": "transformer",
+          "signature": 
+          {
+            "call_mode": "auto"
+          }
+        },
+        {
+          "name": "frame_cnn",
+          "signature": 
+          {
+            "inputs": ["obs"],
+            "call_mode": "auto"
+          }
+        }
+      ]
+    },
+    {
+      "name": "tokenizer",
+      "framework": null,
+      "format": "state_dict",
+      "source": {
+        "weights_path": "tokenizer.pt",
+        "class_path": "../src/tokenizer.py",
+        "class_name": "Tokenizer",
+        "class_args": [{
+          "image_channels": 3,
+          "image_size": 64,
+          "num_actions": 17,
+          "num_tokens": 4,
+          "decoder_act_channels": 4,
+          "codebook_size": 1024,
+          "codebook_dim": 64,
+          "max_codebook_updates_with_revival": 0,
+          "encoder_config": {
+            "image_channels": 7,
+            "latent_dim": 64,
+            "num_channels": 64,
+            "mult": [1, 1, 2, 2, 4],
+            "down": [1, 0, 1, 1, 0]
+          },
+          "decoder_config": { 
+            "image_channels": 3,
+            "latent_dim": 84,
+            "num_channels": 64,
+            "mult": [1, 1, 2, 2, 4],
+            "down": [1, 0, 1, 1, 0]
+          },
+          "frame_cnn_config": {
+            "image_channels": 3,
+            "latent_dim": 16,
+            "num_channels": 32,
+            "mult": [1, 1, 2, 2, 4],
+            "down": [1, 0, 1, 1, 0]
+          }
+        }]
+      },
+      "signature": {
+        "inputs": ["o1", "a", "o2"],
+        "call_mode": "positional"
+      },
+      "sub_models":
+      [
+        {
+          "name": "decode",
+          "signature": 
+          {
+            "inputs": ["obs", "act", "q", "should_clamp"],
+            "call_mode": "positional"
+          }
+        }
+      ]
+    }
+  ]
+}