Update iris/src/world_model.py

iris/src/world_model.py

@@ -10,9 +10,9 @@ from .models.slicer import Embedder, Head
 from .models.transformer import Transformer
 
 class WorldModel(nn.Module):
-    def __init__(self, obs_vocab_size: int, act_vocab_size: int, config: dict) -> None:
+    def __init__(self, config: dict) -> None:
         super().__init__()
-        self.obs_vocab_size, self.act_vocab_size = obs_vocab_size, act_vocab_size
+        self.obs_vocab_size, self.act_vocab_size = config["obs_vocab_size"], config["act_vocab_size"]
         self.config = config
         self.transformer = Transformer(config)