Upload 15 files

iris/atari/mspacman/config.json

@@ -0,0 +1,146 @@
+{
+  "name": "ms_pacman_iris_world_model",
+  "env": "MsPacmanNoFrameskip-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": 9,
+          "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/mspacman/tokenizer.pt

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

iris/atari/mspacman/world_model.pt

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

iris/atari/pong/config.json

@@ -0,0 +1,146 @@
+{
+  "name": "pong_iris_world_model",
+  "env": "PongNoFrameskip-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": 6,
+          "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/pong/tokenizer.pt

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

iris/atari/pong/world_model.pt

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

iris/atari/privateeye/config.json

@@ -0,0 +1,146 @@
+{
+  "name": "private_eye_iris_world_model",
+  "env": "PrivateEyeNoFrameskip-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/privateeye/tokenizer.pt

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

iris/atari/privateeye/world_model.pt

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

iris/atari/qbert/config.json

@@ -0,0 +1,146 @@
+{
+  "name": "qbert_iris_world_model",
+  "env": "QbertNoFrameskip-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": 6,
+          "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/qbert/tokenizer.pt

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

iris/atari/qbert/world_model.pt

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

iris/atari/roadrunner/config.json

@@ -0,0 +1,146 @@
+{
+  "name": "road_runner_iris_world_model",
+  "env": "RoadRunnerNoFrameskip-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/roadrunner/tokenizer.pt

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

iris/atari/roadrunner/world_model.pt

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