1

src/models/tokenizer/tokenizer.py

@@ -23,12 +23,12 @@ class TokenizerEncoderOutput:
 
 
 class Tokenizer(nn.Module):
-    def __init__(self, vocab_size: int, embed_dim: int, encoder: Encoder, decoder: Decoder, with_lpips: bool = True) -> None:
+    def __init__(self, embedding, vocab_size: int, embed_dim: int, encoder: Encoder, decoder: Decoder, with_lpips: bool = True) -> None:
         super().__init__()
         self.vocab_size = vocab_size
         self.encoder = encoder
         self.pre_quant_conv = torch.nn.Conv2d(encoder.config.z_channels, embed_dim, 1)
-        self.embedding = nn.Embedding(vocab_size, embed_dim) # TODO: use model embed?
+        self.embedding = embedding # nn.Embedding(vocab_size, embed_dim) # TODO: use model embed?
         self.post_quant_conv = torch.nn.Conv2d(embed_dim, decoder.config.z_channels, 1)
         self.decoder = decoder
         self.embedding.weight.data.uniform_(-1.0 / vocab_size, 1.0 / vocab_size)

src/models/world_model.py

@@ -92,6 +92,7 @@ class WorldModel(nn.Module):
         assert num_steps <= self.config.max_tokens
         prev_steps = 0 if past_keys_values is None else past_keys_values.size
 
+        # TODO: replace wth model embedder?
         sequences = self.embedder(tokens, num_steps, prev_steps) + self.pos_emb(prev_steps + torch.arange(num_steps, device=tokens.device))
         # [batch=8, num_steps=170, embed_size=2048]
         sequences = self.post_embed(sequences)

src/play.py

@@ -33,8 +33,9 @@ def main(cfg: DictConfig):
         keymap = 'episode_replay'
 
     else:
-        tokenizer = instantiate(cfg.tokenizer)
-        world_model = WorldModel(obs_vocab_size=tokenizer.vocab_size, act_vocab_size=test_env.num_actions, config=instantiate(cfg.world_model))
+        # tokenizer = instantiate(cfg.tokenizer)
+        world_model = WorldModel(obs_vocab_size=cfg.tokenizer.vocab_size, act_vocab_size=test_env.num_actions, config=instantiate(cfg.world_model))
+        tokenizer = Tokenizer(embedding=embedding, **cfg.tokenizer.embedding)
         actor_critic = ActorCritic(**cfg.actor_critic, act_vocab_size=test_env.num_actions)
         agent = Agent(tokenizer, world_model, actor_critic).to(device)
         agent.load(Path('checkpoints/last.pt'), device)        

src/trainer.py

@@ -22,6 +22,7 @@ from src.make_reconstructions import make_reconstructions_from_batch
 from src.models.actor_critic import ActorCritic
 from src.models.world_model import WorldModel
 from src.utils import configure_optimizer, EpisodeDirManager, set_seed
+from src.models.tokenizer import Tokenizer
 
 
 class Trainer:
@@ -80,8 +81,10 @@ class Trainer:
         assert self.cfg.training.should or self.cfg.evaluation.should
         env = train_env if self.cfg.training.should else test_env
 
-        tokenizer = instantiate(cfg.tokenizer)
-        world_model = WorldModel(obs_vocab_size=tokenizer.vocab_size, act_vocab_size=env.num_actions, config=instantiate(cfg.world_model))
+        # tokenizer = instantiate(cfg.tokenizer)
+        world_model = WorldModel(obs_vocab_size=cfg.tokenizer.vocab_size, act_vocab_size=env.num_actions, config=instantiate(cfg.world_model))
+        embedding = world_model.transformer.embedding
+        tokenizer = Tokenizer(embedding=embedding, **cfg.tokenizer.embedding)
         actor_critic = ActorCritic(**cfg.actor_critic, act_vocab_size=env.num_actions)
         self.agent = Agent(tokenizer, world_model, actor_critic).to(self.device)
         print(f'{sum(p.numel() for p in self.agent.tokenizer.parameters())} parameters in agent.tokenizer')