Delete iris/world_model.py

iris/world_model.py

@@ -1,93 +0,0 @@
-from typing import Any, Optional, Tuple
-
-from einops import rearrange
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from models.kv_caching import KeysValues
-from models.slicer import Embedder, Head
-from models.transformer import Transformer
-
-class WorldModel(nn.Module):
-    def __init__(self, obs_vocab_size: int, act_vocab_size: int, config: dict) -> None:
-        super().__init__()
-        self.obs_vocab_size, self.act_vocab_size = obs_vocab_size, act_vocab_size
-        self.config = config
-        self.transformer = Transformer(config)
-
-        all_but_last_obs_tokens_pattern = torch.ones(config["tokens_per_block"])
-        all_but_last_obs_tokens_pattern[-2] = 0
-        act_tokens_pattern = torch.zeros(self.config["tokens_per_block"])
-        act_tokens_pattern[-1] = 1
-        obs_tokens_pattern = 1 - act_tokens_pattern
-
-        self.pos_emb = nn.Embedding(config["max_tokens"], config["embed_dim"])
-
-        self.embedder = Embedder(
-            max_blocks=config["max_blocks"],
-            block_masks=[act_tokens_pattern, obs_tokens_pattern],
-            embedding_tables=nn.ModuleList([nn.Embedding(act_vocab_size, config["embed_dim"]), nn.Embedding(obs_vocab_size, config["embed_dim"])])
-        )
-
-        self.head_observations = Head(
-            max_blocks=config["max_blocks"],
-            block_mask=all_but_last_obs_tokens_pattern,
-            head_module=nn.Sequential(
-                nn.Linear(config["embed_dim"], config["embed_dim"]),
-                nn.ReLU(),
-                nn.Linear(config["embed_dim"], obs_vocab_size)
-            )
-        )
-
-        self.head_rewards = Head(
-            max_blocks=config["max_blocks"],
-            block_mask=act_tokens_pattern,
-            head_module=nn.Sequential(
-                nn.Linear(config["embed_dim"], config["embed_dim"]),
-                nn.ReLU(),
-                nn.Linear(config["embed_dim"], 3)
-            )
-        )
-
-        self.head_ends = Head(
-            max_blocks=config["max_blocks"],
-            block_mask=act_tokens_pattern,
-            head_module=nn.Sequential(
-                nn.Linear(config["embed_dim"], config["embed_dim"]),
-                nn.ReLU(),
-                nn.Linear(config["embed_dim"], 2)
-            )
-        )
-
-    def __repr__(self) -> str:
-        return "world_model"
-
-    def forward(self, tokens: torch.LongTensor, past_keys_values: Optional[KeysValues] = None) -> dict:
-
-        num_steps = tokens.size(1)  # (B, T)
-        assert num_steps <= self.config["max_tokens"]
-        prev_steps = 0 if past_keys_values is None else past_keys_values.size
-
-        sequences = self.embedder(tokens, num_steps, prev_steps) + self.pos_emb(prev_steps + torch.arange(num_steps, device=tokens.device))
-
-        x = self.transformer(sequences, past_keys_values)
-
-        logits_observations = self.head_observations(x, num_steps=num_steps, prev_steps=prev_steps)
-        logits_rewards = self.head_rewards(x, num_steps=num_steps, prev_steps=prev_steps)
-        logits_ends = self.head_ends(x, num_steps=num_steps, prev_steps=prev_steps)
-        return {
-            "output_sequence": x, 
-            "logits_observations": logits_observations,
-            "logits_rewards": logits_rewards,
-            "logits_ends": logits_ends
-
-        }
-    
-    def compute_labels_world_model(self, obs_tokens: torch.Tensor, rewards: torch.Tensor, ends: torch.Tensor, mask_padding: torch.BoolTensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        assert torch.all(ends.sum(dim=1) <= 1)  # at most 1 done
-        mask_fill = torch.logical_not(mask_padding)
-        labels_observations = rearrange(obs_tokens.masked_fill(mask_fill.unsqueeze(-1).expand_as(obs_tokens), -100), 'b t k -> b (t k)')[:, 1:]
-        labels_rewards = (rewards.sign() + 1).masked_fill(mask_fill, -100).long()  # Rewards clipped to {-1, 0, 1}
-        labels_ends = ends.masked_fill(mask_fill, -100)
-        return labels_observations.reshape(-1), labels_rewards.reshape(-1), labels_ends.reshape(-1)