Upload 12 files

iris/atari/bankheist/config.json

@@ -0,0 +1,146 @@
+{
+  "name": "bankheist_iris_world_model",
+  "env": "BankHeistNoFrameskip-v4",
+  "model_type": "iris",
+  "metadata": {
+    "latent_dim": [1, 16],
+    "num_tokens": 340
+  },
+  "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": [
+        {
+          "vocab_size": 512,
+          "act_vocab_size": 18,
+          "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": "generate_empty_keys_values",
+          "method_name": "generate_empty_keys_values",
+          "signature": 
+          {
+            "inputs": ["n"]
+          }
+        }
+      ]
+    },
+    {
+      "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": "decode",
+          "method_name": "decode",
+          "signature": 
+          {
+            "inputs": ["z", "should_postprocess"]
+          }
+        },
+        {
+          "name": "decode_obs_tokens",
+          "method_name": "decode_obs_tokens",
+          "signature": 
+          {
+            "inputs": ["obs_tokens", "num_observations_tokens"]
+          }
+        },
+        {
+          "name": "encode",
+          "method_name": "encode",
+          "signature": 
+          {
+            "inputs": ["observations", "should_preprocess"]
+          }
+        }
+      ]
+    }
+  ]
+}

iris/atari/bankheist/tokenizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ae0e59bd9e87f0f39faeb3395a20033f75a9d5632f352f5c63f72d1e0db22204
+size 75104314

iris/atari/bankheist/world_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e81d04acf4f11a003340eb158180a2b6b2f45fee80b8e6876f24502c94b2680f
+size 38490814

iris/atari/battlezone/config.json

@@ -0,0 +1,146 @@
+{
+  "name": "battlezone_iris_world_model",
+  "env": "BattleZoneNoFrameskip-v4",
+  "model_type": "iris",
+  "metadata": {
+    "latent_dim": [1, 16],
+    "num_tokens": 340
+  },
+  "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": [
+        {
+          "vocab_size": 512,
+          "act_vocab_size": 18,
+          "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": "generate_empty_keys_values",
+          "method_name": "generate_empty_keys_values",
+          "signature": 
+          {
+            "inputs": ["n"]
+          }
+        }
+      ]
+    },
+    {
+      "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": "decode",
+          "method_name": "decode",
+          "signature": 
+          {
+            "inputs": ["z", "should_postprocess"]
+          }
+        },
+        {
+          "name": "decode_obs_tokens",
+          "method_name": "decode_obs_tokens",
+          "signature": 
+          {
+            "inputs": ["obs_tokens", "num_observations_tokens"]
+          }
+        },
+        {
+          "name": "encode",
+          "method_name": "encode",
+          "signature": 
+          {
+            "inputs": ["observations", "should_preprocess"]
+          }
+        }
+      ]
+    }
+  ]
+}

iris/atari/battlezone/tokenizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:903ecacfc9a8b2328f38ee3a6afc74c0724fb16949ecddc53648ca925be8e12c
+size 75104314

iris/atari/battlezone/world_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2abb50dfd0e5adef304544da9f3951c7f4a3b976e9eb76096fea8969842b3db1
+size 38490814

iris/atari/boxing/config.json

@@ -0,0 +1,146 @@
+{
+  "name": "boxing_iris_world_model",
+  "env": "BoxingNoFrameskip-v4",
+  "model_type": "iris",
+  "metadata": {
+    "latent_dim": [1, 16],
+    "num_tokens": 340
+  },
+  "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": [
+        {
+          "vocab_size": 512,
+          "act_vocab_size": 18,
+          "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": "generate_empty_keys_values",
+          "method_name": "generate_empty_keys_values",
+          "signature": 
+          {
+            "inputs": ["n"]
+          }
+        }
+      ]
+    },
+    {
+      "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": "decode",
+          "method_name": "decode",
+          "signature": 
+          {
+            "inputs": ["z", "should_postprocess"]
+          }
+        },
+        {
+          "name": "decode_obs_tokens",
+          "method_name": "decode_obs_tokens",
+          "signature": 
+          {
+            "inputs": ["obs_tokens", "num_observations_tokens"]
+          }
+        },
+        {
+          "name": "encode",
+          "method_name": "encode",
+          "signature": 
+          {
+            "inputs": ["observations", "should_preprocess"]
+          }
+        }
+      ]
+    }
+  ]
+}

iris/atari/boxing/tokenizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c603d6ba26f15ff677945d6065a109e48c62487bedef1b089013e6ec7c0e7e4
+size 75104314

iris/atari/boxing/world_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d88fe6f63f20745a16f2d9a38e4ea1e2584ead0e10c87f41a01aa8e3785f9485
+size 38490814

iris/atari/breakout/config.json

@@ -0,0 +1,146 @@
+{
+  "name": "breakout_iris_world_model",
+  "env": "BreakoutNoFrameskip-v4",
+  "model_type": "iris",
+  "metadata": {
+    "latent_dim": [1, 16],
+    "num_tokens": 340
+  },
+  "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": [
+        {
+          "vocab_size": 512,
+          "act_vocab_size": 4,
+          "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": "generate_empty_keys_values",
+          "method_name": "generate_empty_keys_values",
+          "signature": 
+          {
+            "inputs": ["n"]
+          }
+        }
+      ]
+    },
+    {
+      "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": "decode",
+          "method_name": "decode",
+          "signature": 
+          {
+            "inputs": ["z", "should_postprocess"]
+          }
+        },
+        {
+          "name": "decode_obs_tokens",
+          "method_name": "decode_obs_tokens",
+          "signature": 
+          {
+            "inputs": ["obs_tokens", "num_observations_tokens"]
+          }
+        },
+        {
+          "name": "encode",
+          "method_name": "encode",
+          "signature": 
+          {
+            "inputs": ["observations", "should_preprocess"]
+          }
+        }
+      ]
+    }
+  ]
+}

iris/atari/breakout/tokenizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a702fa66a36fabac44c2f83aca92c8dd4130b7774fa82434b2ed48062720005e
+size 75104314

iris/atari/breakout/world_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b1714bcc2022fb6367629f7e39f746e4708888103b30b8c0a28b5202e729d419
+size 38476478