Update iris/src/tokenizer.py

iris/src/tokenizer.py

@@ -68,6 +68,12 @@ class Tokenizer(nn.Module):
             rec = self.postprocess_output(rec)
         return rec
 
+    def decode_obs_tokens(self, obs_tokens, num_observations_tokens):
+        embedded_tokens = self.embedding(self.obs_tokens)     # (B, K, E)
+        z = rearrange(embedded_tokens, 'b (h w) e -> b e h w', h=int(np.sqrt(num_observations_tokens)))
+        rec = self.decode(z, should_postprocess=True)         # (B, C, H, W)
+        return torch.clamp(rec, 0, 1)
+
     @torch.no_grad()
     def encode_decode(self, x: torch.Tensor, should_preprocess: bool = False, should_postprocess: bool = False) -> torch.Tensor:
         z_q = self.encode(x, should_preprocess).z_quantized