Update delta-iris/src/tokenizer.py

delta-iris/src/tokenizer.py

@@ -7,9 +7,8 @@ import torch
 import torch.nn as nn
 
 from .models.convnet import FrameEncoder, FrameDecoder
-from .data import Batch
 from .models.tokenizer.quantizer import Quantizer
-from .models.utils import init_weights, LossWithIntermediateLosses
+from .models.utils import init_weights
 
 class Tokenizer(nn.Module):
     def __init__(self, config: dict) -> None:
@@ -44,26 +43,6 @@ class Tokenizer(nn.Module):
 
         return self.quantizer(z)
 
-    def compute_loss(self, batch: Batch, **kwargs) -> Tuple[LossWithIntermediateLosses, Dict]:
-        x1 = batch.observations[:, :-1]
-        a = batch.actions[:, :-1]
-        x2 = batch.observations[:, 1:]
-
-        quantizer_outputs = self(x1, a, x2)
-
-        r = self.decode(x1, a, rearrange(quantizer_outputs.q, 'b t (h w) (k l e) -> b t e (h k) (w l)', h=self.tokens_grid_res, k=self.token_res, l=self.token_res))
-        delta = (x2 - r)
-        delta = delta[torch.logical_and(batch.mask_padding[:, 1:], batch.mask_padding[:, :-1])]
-
-        losses = {
-            **quantizer_outputs.loss,
-            'reconstruction_loss_l1': 0.1 * torch.abs(delta).mean(),
-            'reconstruction_loss_l2': delta.pow(2).mean(),
-            'reconstruction_loss_l2_worst_pixel': 0.01 * rearrange(delta, 'b c h w -> b (c h w)').pow(2).max(dim=-1)[0].mean(),
-        }
-
-        return LossWithIntermediateLosses(**losses), quantizer_outputs.metrics
-
     def encode(self, x1: torch.FloatTensor, a: torch.LongTensor, x2: torch.FloatTensor) -> torch.FloatTensor:
         a_emb = rearrange(self.encoder_act_emb(a), 'b t (h w) -> b t 1 h w', h=x1.size(3))
         encoder_input = torch.cat((x1, a_emb, x2), dim=2)