Update delta-iris/src/tokenizer.py

delta-iris/src/tokenizer.py

@@ -64,6 +64,22 @@ class Tokenizer(nn.Module):
 
         return r
 
+    def embed_tokens(self, tokens):
+        q = self.quantizer.embed_tokens(tokens)
+        b, t, hw, kle = q.shape
+
+        h = self.tokens_grid_res
+        w = self.tokens_grid_res
+        k = self.token_res
+        l = self.token_res
+        e = kle // (k * l)
+
+        q = q.reshape(b, t, h, w, k, l, e)
+        q = q.transpose(0, 1, 6, 2, 4, 3, 5)
+        q = q.reshape(b, t, e, h * k, w * l)
+
+        return q
+
     @torch.no_grad()
     def burn_in(self, obs: torch.FloatTensor, act: torch.LongTensor) -> torch.LongTensor: 
         assert obs.size(1) == act.size(1) + 1