Update delta-iris/src/world_model.py

delta-iris/src/world_model.py

@@ -7,11 +7,9 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from .models.convnet import FrameEncoder
-from .data import Batch
 from .models.slicer import  Head
-from .tokenizer import Tokenizer
 from .models.transformer import TransformerEncoder
-from .models.utils import init_weights, LossWithIntermediateLosses, symlog, two_hot
+from .models.utils import init_weights
 
 class WorldModel(nn.Module):
     def __init__(self, config: dict) -> None:
@@ -79,34 +77,6 @@ class WorldModel(nn.Module):
             "logits_ends": logits_ends
         }
 
-    def compute_loss(self, batch: Batch, tokenizer: Tokenizer, **kwargs) -> LossWithIntermediateLosses:
-        assert torch.all(batch.ends.sum(dim=1) <= 1)
-
-        with torch.no_grad():
-            latent_tokens = tokenizer(batch.observations[:, :-1], batch.actions[:, :-1], batch.observations[:, 1:]).tokens
-
-        b, _, k = latent_tokens.size()
-
-        frames_emb = self.frame_cnn(batch.observations)
-        act_tokens_emb = self.act_emb(rearrange(batch.actions, 'b t -> b t 1'))
-        latent_tokens_emb = self.latents_emb(torch.cat((latent_tokens, latent_tokens.new_zeros(b, 1, k)), dim=1))
-        sequence = rearrange(torch.cat((frames_emb, act_tokens_emb, latent_tokens_emb), dim=2), 'b t p1k e -> b (t p1k) e')
-  
-        outputs = self(sequence)
-
-        mask = batch.mask_padding
-
-        labels_latents = latent_tokens[mask[:, :-1]].flatten()
-        logits_latents = outputs.logits_latents[:, :-k][repeat(mask[:, :-1], 'b t -> b (t k)', k=k)]
-        latent_acc = (logits_latents.max(dim=-1)[1] == labels_latents).float().mean()
-        labels_rewards = two_hot(symlog(batch.rewards)) if self.config["two_hot_rews"] else (batch.rewards.sign() + 1).long()
-
-        loss_latents = F.cross_entropy(logits_latents, target=labels_latents) * self.config["latents_weight"]
-        loss_rewards = F.cross_entropy(outputs.logits_rewards[mask], target=labels_rewards[mask]) * self.config["rewards_weight"]
-        loss_ends = F.cross_entropy(outputs.logits_ends[mask], target=batch.ends[mask]) * self.config["ends_weight"]
-
-        return LossWithIntermediateLosses(loss_latents=loss_latents, loss_rewards=loss_rewards, loss_ends=loss_ends), {'latent_accuracy': latent_acc}
-
     @torch.no_grad()
     def burn_in(self, obs: torch.FloatTensor, act: torch.LongTensor, latent_tokens: torch.LongTensor, use_kv_cache: bool = False) -> torch.FloatTensor: 
         assert obs.size(1) == act.size(1) + 1 == latent_tokens.size(1) + 1