Upload 35 files

delta-iris/src/models/__init__.py

@@ -0,0 +1 @@
+from .tokenizer import Tokenizer

delta-iris/src/models/convnet.py

@@ -7,25 +7,16 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 
-@dataclass
-class FrameCnnConfig:
-    image_channels: int
-    latent_dim: int
-    num_channels: int
-    mult: List[int]
-    down: List[int]
-
-
 class FrameEncoder(nn.Module):
-    def __init__(self, config: FrameCnnConfig) -> None:
+    def __init__(self, config: dict) -> None:
         super().__init__()
 
-        assert len(config.mult) == len(config.down)
-        encoder_layers = [nn.Conv2d(config.image_channels, config.num_channels, kernel_size=3, stride=1, padding=1)]
-        input_channels = config.num_channels
+        assert len(config["mult"]) == len(config["down"])
+        encoder_layers = [nn.Conv2d(config["image_channels"], config["num_channels"], kernel_size=3, stride=1, padding=1)]
+        input_channels = config["num_channels"]
 
-        for m, d in zip(config.mult, config.down):
-            output_channels = m * config.num_channels
+        for m, d in zip(config["mult"], config["down"]):
+            output_channels = m * config["num_channels"]
             encoder_layers.append(ResidualBlock(input_channels, output_channels))
             input_channels = output_channels
             if d:
@@ -33,7 +24,7 @@ class FrameEncoder(nn.Module):
         encoder_layers.extend([
             nn.GroupNorm(num_groups=32, num_channels=input_channels),
             nn.SiLU(inplace=True),
-            nn.Conv2d(input_channels, config.latent_dim, kernel_size=3, stride=1, padding=1)
+            nn.Conv2d(input_channels, config["latent_dim"], kernel_size=3, stride=1, padding=1)
         ])
         self.encoder = nn.Sequential(*encoder_layers)
 
@@ -47,25 +38,25 @@ class FrameEncoder(nn.Module):
 
 
 class FrameDecoder(nn.Module):
-    def __init__(self, config: FrameCnnConfig) -> None:
+    def __init__(self, config: dict) -> None:
         super().__init__()
 
-        assert len(config.mult) == len(config.down)
+        assert len(config["mult"]) == len(config["down"])
         decoder_layers = []
-        output_channels = config.num_channels
+        output_channels = config["num_channels"]
 
-        for m, d in zip(config.mult, config.down):
-            input_channels = m * config.num_channels
+        for m, d in zip(config["mult"], config["down"]):
+            input_channels = m * config["num_channels"]
             decoder_layers.append(ResidualBlock(input_channels, output_channels))
             output_channels = input_channels
             if d:
                 decoder_layers.append(Upsample(input_channels))
         decoder_layers.reverse()
-        decoder_layers.insert(0, nn.Conv2d(config.latent_dim, input_channels, kernel_size=3, stride=1, padding=1))
+        decoder_layers.insert(0, nn.Conv2d(config["latent_dim"], input_channels, kernel_size=3, stride=1, padding=1))
         decoder_layers.extend([
-            nn.GroupNorm(num_groups=32, num_channels=config.num_channels),
+            nn.GroupNorm(num_groups=32, num_channels=config["num_channels"]),
             nn.SiLU(inplace=True),
-            nn.Conv2d(config.num_channels, config.image_channels, kernel_size=3, stride=1, padding=1)
+            nn.Conv2d(config["num_channels"], config["image_channels"], kernel_size=3, stride=1, padding=1)
         ])
         self.decoder = nn.Sequential(*decoder_layers)
 

delta-iris/src/models/tokenizer/__init__.py

@@ -1 +0,0 @@
-

delta-iris/src/models/transformer.py

@@ -11,58 +11,36 @@ import torch.nn as nn
 
 from .kv_caching import KeysValues, KVCache
 
-
-@dataclass
-class TransformerConfig:
-
-    tokens_per_block: int
-    max_blocks: int
-
-    num_layers: int
-    num_heads: int
-    embed_dim: int
-
-    attention: str
-
-    embed_pdrop: float
-    resid_pdrop: float
-    attn_pdrop: float
-
-    @property
-    def max_tokens(self):
-        return self.tokens_per_block * self.max_blocks
-
-
 class TransformerEncoder(nn.Module):
-    def __init__(self, config: TransformerConfig) -> None:
+    def __init__(self, config: dict) -> None:
         super().__init__()
         self.config = config
+        self.config["max_tokens"] = config["tokens_per_block"] * config["max_blocks"]
+        self.pos_emb = nn.Embedding(config["max_tokens"], config["embed_dim"])
+        self.emb_drop = nn.Dropout(config["embed_pdrop"])
+        self.ln = nn.LayerNorm(config["embed_dim"])
 
-        self.pos_emb = nn.Embedding(config.max_tokens, config.embed_dim)
-        self.emb_drop = nn.Dropout(config.embed_pdrop)
-        self.ln = nn.LayerNorm(config.embed_dim)
-
-        assert config.attention in ('causal', 'block_causal')
-        k, m = config.tokens_per_block, config.max_blocks
+        assert config["attention"] in ('causal', 'block_causal')
+        k, m = config["tokens_per_block"], config["max_blocks"]
         mask_sa = torch.tril(torch.ones(k * m, k * m))
-        if config.attention == 'block_causal':
+        if config["attention"] == 'block_causal':
             mask_sa = torch.max(mask_sa, torch.block_diag(*[torch.ones(k, k) for _ in range(m)]))
         mask_sa = mask_sa.bool()
 
-        self.blocks = nn.ModuleList([EncoderLayer(config, mask_sa) for _ in range(config.num_layers)])
+        self.blocks = nn.ModuleList([EncoderLayer(config, mask_sa) for _ in range(config["num_layers"])])
         self.keys_values = None
 
     @property
     def num_blocks_left_in_kv_cache(self) -> float:
         assert self.keys_values is not None
-        return (self.config.max_tokens - self.keys_values.size) / self.config.tokens_per_block
+        return (self.config["max_tokens"] - self.keys_values.size) / self.config["tokens_per_block"]
 
     def reset_kv_cache(self, n: int) -> None:
         device = self.ln.weight.device
-        self.keys_values = KeysValues(n, self.config.max_tokens, self.config.embed_dim, self.config.num_layers, device)
+        self.keys_values = KeysValues(n, self.config["max_tokens"], self.config["embed_dim"], self.config["num_layers"], device)
 
     def forward(self, x: torch.FloatTensor, use_kv_cache: bool = False) -> torch.FloatTensor:
-        assert x.ndim == 3 and x.size(2) == self.config.embed_dim   # (B, TK, E)
+        assert x.ndim == 3 and x.size(2) == self.config["embed_dim"]   # (B, TK, E)
 
         prev_steps = self.keys_values.size if use_kv_cache else 0
         inputs = x + self.pos_emb(prev_steps + torch.arange(x.size(1), device=x.device))
@@ -76,7 +54,7 @@ class TransformerEncoder(nn.Module):
 
 
 class EncoderLayer(nn.Module):
-    def __init__(self, config: TransformerConfig, mask_sa: torch.LongTensor) -> None:
+    def __init__(self, config: dict, mask_sa: torch.LongTensor) -> None:
         super().__init__()
         self.sa = SelfAttentionLayer(config, mask=mask_sa)
         self.mlp = MLPLayer(config)
@@ -86,14 +64,14 @@ class EncoderLayer(nn.Module):
 
 
 class MLPLayer(nn.Module):
-    def __init__(self, config: TransformerConfig) -> None:
+    def __init__(self, config: dict) -> None:
         super().__init__()
-        self.ln = nn.LayerNorm(config.embed_dim)
+        self.ln = nn.LayerNorm(config["embed_dim"])
         self.mlp = nn.Sequential(
-            nn.Linear(config.embed_dim, 4 * config.embed_dim),
+            nn.Linear(config["embed_dim"], 4 * config["embed_dim"]),
             nn.GELU(),
-            nn.Linear(4 * config.embed_dim, config.embed_dim),
-            nn.Dropout(config.resid_pdrop),
+            nn.Linear(4 * config["embed_dim"], config["embed_dim"]),
+            nn.Dropout(config["resid_pdrop"]),
         )
 
     def forward(self, inputs: torch.FloatTensor) -> torch.FloatTensor:
@@ -101,13 +79,13 @@ class MLPLayer(nn.Module):
 
 
 class SelfAttentionLayer(nn.Module):
-    def __init__(self, config: TransformerConfig, mask: torch.BoolTensor) -> None:
+    def __init__(self, config: dict, mask: torch.BoolTensor) -> None:
         super().__init__()
         self.register_buffer('mask', mask)
-        self.ln = nn.LayerNorm(config.embed_dim)
-        self.query = nn.Linear(config.embed_dim, config.embed_dim)
-        self.key = nn.Linear(config.embed_dim, config.embed_dim)
-        self.value = nn.Linear(config.embed_dim, config.embed_dim)
+        self.ln = nn.LayerNorm(config["embed_dim"])
+        self.query = nn.Linear(config["embed_dim"], config["embed_dim"])
+        self.key = nn.Linear(config["embed_dim"], config["embed_dim"])
+        self.value = nn.Linear(config["embed_dim"], config["embed_dim"])
         self.attention = Attention(config)
 
     def forward(self, inputs: torch.FloatTensor, kv_cache: Optional[KVCache] = None) -> torch.FloatTensor:
@@ -134,13 +112,13 @@ class SelfAttentionLayer(nn.Module):
 
 
 class Attention(nn.Module):
-    def __init__(self, config: TransformerConfig) -> None:
+    def __init__(self, config: dict) -> None:
         super().__init__()
-        assert config.embed_dim % config.num_heads == 0
-        self.num_heads = config.num_heads
-        self.attn_pdrop = config.attn_pdrop
-        self.resid_drop = nn.Dropout(config.resid_pdrop)
-        self.proj = nn.Linear(config.embed_dim, config.embed_dim)
+        assert config["embed_dim"] % config["num_heads"] == 0
+        self.num_heads = config["num_heads"]
+        self.attn_pdrop = config["attn_pdrop"]
+        self.resid_drop = nn.Dropout(config["resid_pdrop"])
+        self.proj = nn.Linear(config["embed_dim"], config["embed_dim"])
 
     def forward(self, q: torch.FloatTensor, k: torch.FloatTensor, v: torch.FloatTensor, mask: torch.BoolTensor) -> torch.FloatTensor:
         assert mask.size(0) == q.size(1) and mask.size(1) == k.size(1)

delta-iris/src/tokenizer.py

@@ -6,48 +6,32 @@ from einops import rearrange
 import torch
 import torch.nn as nn
 
-from .models.convnet import FrameCnnConfig, FrameEncoder, FrameDecoder
+from .models.convnet import FrameEncoder, FrameDecoder
 from .data import Batch
 from .models.tokenizer.quantizer import Quantizer, QuantizerOutput
 from .models.utils import init_weights, LossWithIntermediateLosses
 
-
-@dataclass
-class TokenizerConfig:
-    image_channels: int
-    image_size: int
-    num_actions: int
-    num_tokens: int
-    decoder_act_channels: int
-    codebook_size: int
-    codebook_dim: int
-    max_codebook_updates_with_revival: int
-    encoder_config: FrameCnnConfig
-    decoder_config: FrameCnnConfig
-    frame_cnn_config: FrameCnnConfig
-
-
 class Tokenizer(nn.Module):
-    def __init__(self, config: TokenizerConfig) -> None:
+    def __init__(self, config: dict) -> None:
         super().__init__()
         self.config = config
 
-        self.latent_res = config.image_size // 2 ** sum(config.encoder_config.down)
-        self.tokens_grid_res = int(math.sqrt(config.num_tokens))
+        self.latent_res = config["image_size"] // 2 ** sum(config["encoder_config"]["down"])
+        self.tokens_grid_res = int(math.sqrt(config["num_tokens"]))
         self.token_res = self.latent_res // self.tokens_grid_res
 
-        self.encoder_act_emb = nn.Embedding(config.num_actions, config.image_size ** 2)
-        self.decoder_act_emb = nn.Embedding(config.num_actions, config.decoder_act_channels * self.latent_res ** 2)
+        self.encoder_act_emb = nn.Embedding(config["num_actions"], config["image_size"] ** 2)
+        self.decoder_act_emb = nn.Embedding(config["num_actions"], config["decoder_act_channels"] * self.latent_res ** 2)
 
         self.quantizer = Quantizer(
-            config.codebook_size, config.codebook_dim,
-            input_dim=config.encoder_config.latent_dim * self.token_res ** 2,
-            max_codebook_updates_with_revival=config.max_codebook_updates_with_revival
+            config["codebook_size"], config["codebook_dim"],
+            input_dim=config["encoder_config"]["latent_dim"] * self.token_res ** 2,
+            max_codebook_updates_with_revival=config["max_codebook_updates_with_revival"]
         )
 
-        self.encoder = FrameEncoder(config.encoder_config)
-        self.decoder = FrameDecoder(config.decoder_config)
-        self.frame_cnn = FrameEncoder(config.frame_cnn_config)
+        self.encoder = FrameEncoder(config["encoder_config"])
+        self.decoder = FrameDecoder(config["decoder_config"])
+        self.frame_cnn = FrameEncoder(config["frame_cnn_config"])
 
         self.apply(init_weights)
 
@@ -89,7 +73,7 @@ class Tokenizer(nn.Module):
 
     def decode(self, x1: torch.FloatTensor, a: torch.LongTensor, q2: torch.FloatTensor, should_clamp: bool = False) -> torch.FloatTensor:
         x1_emb = self.frame_cnn(x1)
-        a_emb = rearrange(self.decoder_act_emb(a), 'b t (c h w) -> b t c h w', c=self.config.decoder_act_channels, h=x1_emb.size(3))
+        a_emb = rearrange(self.decoder_act_emb(a), 'b t (c h w) -> b t c h w', c=self.config["decoder_act_channels"], h=x1_emb.size(3))
 
         decoder_input = torch.cat((x1_emb, a_emb, q2), dim=2)
 

delta-iris/src/world_model.py

@@ -6,77 +6,54 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from .models.convnet import FrameCnnConfig, FrameEncoder
+from .models.convnet import FrameEncoder
 from .data import Batch
 from .models.slicer import  Head
 from .tokenizer import Tokenizer
-from .models.transformer import TransformerEncoder, TransformerConfig
+from .models.transformer import TransformerEncoder
 from .models.utils import init_weights, LossWithIntermediateLosses, symlog, two_hot
 
-
-@dataclass
-class WorldModelOutput:
-    output_sequence: torch.FloatTensor
-    logits_latents: torch.FloatTensor
-    logits_rewards: torch.FloatTensor
-    logits_ends: torch.FloatTensor
-
-
-@dataclass
-class WorldModelConfig:
-    latent_vocab_size: int
-    num_actions: int
-    image_channels: int
-    image_size: int
-    latents_weight: float
-    rewards_weight: float
-    ends_weight: float
-    two_hot_rews: bool
-    transformer_config: TransformerConfig
-    frame_cnn_config: FrameCnnConfig
-
-
 class WorldModel(nn.Module):
-    def __init__(self, config: WorldModelConfig) -> None:
+    def __init__(self, config: dict) -> None:
         super().__init__()
         self.config = config
-        self.transformer = TransformerEncoder(config.transformer_config)
+        self.transformer = TransformerEncoder(config["transformer_config"])
 
-        assert ((config.image_size // 2 ** sum(config.frame_cnn_config.down)) ** 2) * config.frame_cnn_config.latent_dim == config.transformer_config.embed_dim
-        self.frame_cnn = nn.Sequential(FrameEncoder(config.frame_cnn_config), Rearrange('b t c h w -> b t 1 (h w c)'), nn.LayerNorm(config.transformer_config.embed_dim))
+        assert ((config["image_size"] // 2 ** sum(config["frame_cnn_config"]["down"])) ** 2) * config["frame_cnn_config"]["latent_dim"] == config["transformer_config"]["embed_dim"]
+        self.frame_cnn = nn.Sequential(FrameEncoder(config["frame_cnn_config"]), Rearrange('b t c h w -> b t 1 (h w c)'), nn.LayerNorm(config["transformer_config"]["embed_dim"]))
 
-        self.act_emb = nn.Embedding(config.num_actions, config.transformer_config.embed_dim)
-        self.latents_emb = nn.Embedding(config.latent_vocab_size, config.transformer_config.embed_dim)
+        self.act_emb = nn.Embedding(config["num_actions"], config["transformer_config"]["embed_dim"])
+        self.latents_emb = nn.Embedding(config["latent_vocab_size"], config["transformer_config"]["embed_dim"])
 
-        act_pattern = torch.zeros(config.transformer_config.tokens_per_block)
+        act_pattern = torch.zeros(config["transformer_config"]["tokens_per_block"])
         act_pattern[1] = 1
-        act_and_latents_but_last_pattern = torch.zeros(config.transformer_config.tokens_per_block) 
+        act_and_latents_but_last_pattern = torch.zeros(config["transformer_config"]["tokens_per_block"]) 
         act_and_latents_but_last_pattern[1:-1] = 1
 
         self.head_latents = Head(
-            max_blocks=config.transformer_config.max_blocks,
+            max_blocks=config["transformer_config"]["max_blocks"],
             block_mask=act_and_latents_but_last_pattern,
             head_module=nn.Sequential(
-                nn.Linear(config.transformer_config.embed_dim, config.transformer_config.embed_dim), nn.ReLU(),
-                nn.Linear(config.transformer_config.embed_dim, config.latent_vocab_size)
+                nn.Linear(config["transformer_config"]["embed_dim"], config["transformer_config"]["embed_dim"]), nn.ReLU(),
+                nn.Linear(config["transformer_config"]["embed_dim"], config["latent_vocab_size"])
             )
         )
 
         self.head_rewards = Head(
-            max_blocks=config.transformer_config.max_blocks,
+            max_blocks=config["transformer_config"]["max_blocks"],
             block_mask=act_pattern,
             head_module=nn.Sequential(
-                nn.Linear(config.transformer_config.embed_dim, config.transformer_config.embed_dim), nn.ReLU(),
-                nn.Linear(config.transformer_config.embed_dim, 255 if config.two_hot_rews else 3)
+                nn.Linear(config["transformer_config"]["embed_dim"], config["transformer_config"]["embed_dim"]), nn.ReLU(),
+                nn.Linear(config["transformer_config"]["embed_dim"], 255 if config["two_hot_rews"] else 3)
             )
         )
 
         self.head_ends = Head(
-            max_blocks=config.transformer_config.max_blocks,
+            max_blocks=config["transformer_config"]["max_blocks"],
             block_mask=act_pattern,
             head_module=nn.Sequential(
-                nn.Linear(config.transformer_config.embed_dim, config.transformer_config.embed_dim), nn.ReLU(),
-                nn.Linear(config.transformer_config.embed_dim, 2)
+                nn.Linear(config["transformer_config"]["embed_dim"], config["transformer_config"]["embed_dim"]), nn.ReLU(),
+                nn.Linear(config["transformer_config"]["embed_dim"], 2)
             )
         )
 
@@ -85,7 +62,7 @@ class WorldModel(nn.Module):
     def __repr__(self) -> str:
         return "world_model"
 
-    def forward(self, sequence: torch.FloatTensor, use_kv_cache: bool = False) -> WorldModelOutput:      
+    def forward(self, sequence: torch.FloatTensor, use_kv_cache: bool = False) -> dict:      
         prev_steps = self.transformer.keys_values.size if use_kv_cache else 0
         num_steps = sequence.size(1)
 
@@ -95,7 +72,12 @@ class WorldModel(nn.Module):
         logits_rewards = self.head_rewards(outputs, num_steps, prev_steps)
         logits_ends = self.head_ends(outputs, num_steps, prev_steps)
 
-        return WorldModelOutput(outputs, logits_latents, logits_rewards, logits_ends)
+        return {
+            "output_sequence": outputs, 
+            "logits_latents": logits_latents,
+            "logits_rewards": logits_rewards,
+            "logits_ends": logits_ends
+        }
 
     def compute_loss(self, batch: Batch, tokenizer: Tokenizer, **kwargs) -> LossWithIntermediateLosses:
         assert torch.all(batch.ends.sum(dim=1) <= 1)
@@ -117,11 +99,11 @@ class WorldModel(nn.Module):
         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()
+        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
+        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}