Update iris/atari/alien/config.json

iris/atari/alien/config.json

@@ -0,0 +1,125 @@
+{
+  "name": "alien_iris_world_model",
+  "env": "AlienNoFrameskip-v4",
+  "model_type": "iris",
+  "metadata": {
+    "latent_dim": [1, 4, 1024],
+    "two_hot_rews": true,
+    "tokens_per_block": 6,
+    "num_tokens": 4,
+    "tokens_grid_res": 2,
+    "token_res": 16384
+  },
+  "util_folders":{
+    "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": [
+        {
+          "tokens_per_block": 17,
+          "max_blocks": 20,
+          "attention": "causal",
+          "num_layers": 10,
+          "num_heads": 4,
+          "embed_dim": 256,
+          "embed_pdrop": 0.1,
+          "resid_pdrop": 0.1,
+          "attn_pdrop": 0.1,
+
+        }]
+      },
+      "signature": {
+        "inputs": ["tokens", "past_keys_values"],
+        "call_mode": "positional"
+      },
+      "sub_models":
+      [
+        {
+          "name": "transformer",
+          "sub_model_name": "transformer",
+          "signature": 
+          {
+            "inputs": ["sequences", "past_keys_values"],
+            "call_mode": "positional"
+          }
+        }
+      ],
+      "methods":
+      [
+      ]
+    },
+    {
+      "name": "tokenizer",
+      "framework": null,
+      "format": "state_dict",
+      "source": {
+        "weights_path": "tokenizer.pt",
+        "class_path": "../src/tokenizer.py",
+        "class_name": "Tokenizer",
+        "class_args": [{
+          "vocab_size": 512,
+          "embed_dim": 512,
+          "encoder": {
+            "resolution": 64,
+            "in_channels": 3,
+            "z_channels": 512,
+            "ch": 64,
+            "ch_mult": [1, 1, 1, 1, 1],
+            "num_res_blocks": 2,
+            "attn_resolutions": [8, 16],
+            "out_ch": 3,
+            "dropout": 0.0
+          },
+          "decoder": { 
+            "resolution": 64,
+            "in_channels": 3,
+            "z_channels": 512,
+            "ch": 64,
+            "ch_mult": [1, 1, 1, 1, 1],
+            "num_res_blocks": 2,
+            "attn_resolutions": [8, 16],
+            "out_ch": 3,
+            "dropout": 0.0
+          }
+        }]
+      },
+      "signature": {
+        "inputs": ["x", "should_preprocess", "should_postprocess"],
+        "call_mode": "positional"
+      },
+      "sub_models":
+      [
+        {
+          "name": "embedding",
+          "sub_model_name": "embedding",
+          "signature": 
+          {
+            "call_mode": "auto"
+          }
+        }
+      ],
+      "methods":
+      [
+        {
+          "name": "encode",
+          "method_name": "encode"
+        },
+        {
+          "name": "decode",
+          "method_name": "decode"
+        }
+      ]
+    }
+  ]
+}