Update delta-iris/src/world_model.py

delta-iris/src/world_model.py

@@ -70,16 +70,3 @@ class WorldModel(nn.Module):
             "logits_rewards": logits_rewards,
             "logits_ends": logits_ends
         }
-
-    @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
-
-        x_emb = self.frame_cnn(obs)
-        act_emb = rearrange(self.act_emb(act), 'b t e -> b t 1 e')
-        q_emb = self.latents_emb(latent_tokens)
-        x_a_q = rearrange(torch.cat((x_emb[:, :-1], act_emb, q_emb), dim=2), 'b t k2 e -> b (t k2) e')
-        wm_input_sequence = torch.cat((x_a_q, x_emb[:, -1]), dim=1)
-        wm_output_sequence = self(wm_input_sequence, use_kv_cache=use_kv_cache).output_sequence
-
-        return wm_output_sequence