Upload 35 files

delta-iris/src/data/__init__.py

@@ -0,0 +1,7 @@
+from .batch import Batch
+from .dataset import EpisodeDataset
+from .episode import Episode
+from .episode_count import EpisodeCountManager
+from .sampler import BatchSampler
+from .segment import SegmentId
+from .utils import collate_segments_to_batch, DatasetTraverser, make_segment

delta-iris/src/data/batch.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import List
+
+import torch
+
+from .segment import SegmentId
+
+
+@dataclass
+class Batch:
+    observations: torch.ByteTensor
+    actions: torch.LongTensor
+    rewards: torch.FloatTensor
+    ends: torch.LongTensor
+    mask_padding: torch.BoolTensor
+    segment_ids: List[SegmentId]
+
+    def pin_memory(self) -> Batch:
+        return Batch(**{k: v if k == 'segment_ids' else v.pin_memory() for k, v in self.__dict__.items()})
+
+    def to(self, device: torch.device) -> Batch:
+        return Batch(**{k: v if k == 'segment_ids' else v.to(device) for k, v in self.__dict__.items()})
+

delta-iris/src/data/dataset.py

@@ -0,0 +1,104 @@
+from pathlib import Path
+import shutil
+from typing import Dict, Optional, Union
+
+import numpy as np
+import torch
+
+from .episode import Episode
+from .segment import Segment, SegmentId
+from .utils import make_segment
+
+
+class EpisodeDataset(torch.utils.data.Dataset):
+    def __init__(self, directory: Path, name: str) -> None:
+        super().__init__()
+        self.name = name
+        self.directory = Path(directory)
+        self.num_episodes, self.num_steps, self.start_idx, self.lengths = None, None, None, None
+
+        if not self.directory.is_dir():
+            self._init_empty()
+        else:
+            self._load_info()
+            print(f'({name}) {self.num_episodes} episodes, {self.num_steps} steps.')
+
+    @property
+    def info_path(self) -> Path:
+        return self.directory / 'info.pt'
+
+    @property
+    def info(self) -> Dict[str, Union[int, np.ndarray]]:
+        return {'num_episodes': self.num_episodes, 'num_steps': self.num_steps, 'start_idx': self.start_idx, 'lengths': self.lengths}
+
+    def __len__(self) -> int:
+        return self.num_steps
+
+    def __getitem__(self, segment_id: SegmentId) -> Segment:
+        return self._load_segment(segment_id)
+
+    def _init_empty(self) -> None:
+        self.directory.mkdir(parents=True, exist_ok=False)
+        self.num_episodes = 0
+        self.num_steps = 0
+        self.start_idx = np.array([], dtype=np.int64)            
+        self.lengths = np.array([], dtype=np.int64)
+        self.save_info()
+
+    def _load_info(self) -> None:
+        info = torch.load(self.info_path)
+        self.num_steps = info['num_steps']
+        self.num_episodes = info['num_episodes']
+        self.start_idx = info['start_idx']
+        self.lengths = info['lengths']
+
+    def save_info(self) -> None:
+        torch.save(self.info, self.info_path)
+
+    def clear(self) -> None:
+        shutil.rmtree(self.directory)
+        self._init_empty()
+
+    def _get_episode_path(self, episode_id: int) -> Path:
+        n = 3 # number of hierarchies
+        powers = np.arange(n)
+        subfolders = list(map(int, np.floor((episode_id % 10 ** (1 + powers)) / 10 ** powers) * 10 ** powers))[::-1]
+        return self.directory / '/'.join(list(map(lambda x: f'{x[1]:0{n - x[0]}d}', enumerate(subfolders)))) / f'{episode_id}.pt'
+
+    def _load_segment(self, segment_id: SegmentId, should_pad: bool = True) -> Segment:
+        episode = self.load_episode(segment_id.episode_id)
+        return make_segment(episode, segment_id, should_pad)
+
+    def load_episode(self, episode_id: int) -> Episode:
+        return Episode(**torch.load(self._get_episode_path(episode_id)))
+
+    def add_episode(self, episode: Episode, *, episode_id: Optional[int] = None) -> int:
+        if episode_id is None:
+            episode_id = self.num_episodes
+            self.start_idx = np.concatenate((self.start_idx, np.array([self.num_steps])))
+            self.lengths = np.concatenate((self.lengths, np.array([len(episode)])))
+            self.num_steps += len(episode)
+            self.num_episodes += 1
+
+        else:
+            assert episode_id < self.num_episodes
+            old_episode = self.load_episode(episode_id) 
+            episode = old_episode.merge(episode)
+            incr_num_steps = len(episode) - len(old_episode)
+            self.lengths[episode_id] = len(episode)
+            self.start_idx[episode_id + 1:] += incr_num_steps            
+            self.num_steps += incr_num_steps
+
+        episode_path = self._get_episode_path(episode_id)
+        episode_path.parent.mkdir(parents=True, exist_ok=True)
+        torch.save(episode.__dict__, episode_path.with_suffix('.tmp'))
+        episode_path.with_suffix('.tmp').rename(episode_path)
+
+        return episode_id
+
+    def get_episode_id_from_global_idx(self, global_idx: np.ndarray) -> np.ndarray:
+        return (np.argmax(self.start_idx.reshape(-1, 1) > global_idx, axis=0) - 1) % self.num_episodes
+
+    def get_global_idx_from_segment_id(self, segment_id: SegmentId) -> np.ndarray:
+        start_idx = self.start_idx[segment_id.episode_id]
+        return np.arange(start_idx + segment_id.start, start_idx + segment_id.stop)

delta-iris/src/data/episode.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class EpisodeMetrics:
+    episode_length: int
+    episode_return: float
+
+
+@dataclass
+class Episode:
+    observations: torch.ByteTensor
+    actions: torch.LongTensor
+    rewards: torch.FloatTensor
+    ends: torch.LongTensor
+
+    def __post_init__(self):
+        assert len(self.observations) == len(self.actions) == len(self.rewards) == len(self.ends)
+        if self.ends.sum() > 0:
+            idx_end = torch.argmax(self.ends) + 1
+            self.observations = self.observations[:idx_end]
+            self.actions = self.actions[:idx_end]
+            self.rewards = self.rewards[:idx_end]
+            self.ends = self.ends[:idx_end]
+
+    def __len__(self) -> int:
+        return self.observations.size(0)
+
+    def merge(self, other: Episode) -> Episode:
+        return Episode(
+            torch.cat((self.observations, other.observations), dim=0),
+            torch.cat((self.actions, other.actions), dim=0),
+            torch.cat((self.rewards, other.rewards), dim=0),
+            torch.cat((self.ends, other.ends), dim=0),
+        )
+
+    def compute_metrics(self) -> EpisodeMetrics:
+        return EpisodeMetrics(len(self), self.rewards.sum())

delta-iris/src/data/episode_count.py

@@ -0,0 +1,41 @@
+from pathlib import Path
+from typing import Tuple
+
+import numpy as np
+import torch
+
+from .dataset import EpisodeDataset
+
+
+class EpisodeCountManager:
+    def __init__(self, dataset: EpisodeDataset) -> None:
+        self.dataset = dataset
+        self.all_counts = dict()
+
+    def load(self, path_to_checkpoint: Path) -> None:
+        self.all_counts = torch.load(path_to_checkpoint)
+        assert all([counts.shape[0] == self.dataset.num_episodes for counts in self.all_counts.values()])
+
+    def save(self, path_to_checkpoint: Path) -> None:
+        torch.save(self.all_counts, path_to_checkpoint)
+
+    def register(self, *keys: Tuple[str]) -> None:
+        assert all([key not in self.all_counts for key in keys])
+        self.all_counts.update({key: np.zeros(self.dataset.num_episodes, dtype=np.int64) for key in keys}) 
+
+    def add_episode(self, episode_id: int) -> None:
+        for key, counts in self.all_counts.items():
+            assert episode_id <= counts.shape[0]
+            if episode_id == counts.shape[0]:
+                self.all_counts[key] = np.concatenate((counts, np.zeros(1, dtype=np.int64)))
+            assert self.all_counts[key].shape[0] == self.dataset.num_episodes
+
+    def increment_episode_count(self, key: str, episode_id: int) -> None:
+        assert key in self.all_counts
+        self.all_counts[key][episode_id] += 1
+
+    def compute_probabilities(self, key: str, alpha: float) -> np.ndarray:
+        assert key in self.all_counts
+        inverse_counts = 1 / (1 + self.all_counts[key])
+        p = inverse_counts ** alpha
+        return p / p.sum()

delta-iris/src/data/sampler.py

@@ -0,0 +1,42 @@
+from typing import Generator, List
+
+import numpy as np
+import torch
+
+from .dataset import EpisodeDataset
+from .segment import SegmentId
+
+
+class BatchSampler(torch.utils.data.Sampler):
+    def __init__(self, dataset: EpisodeDataset, num_steps_per_epoch: int, batch_size: int, sequence_length: int, can_sample_beyond_end: bool) -> None:
+        super().__init__(dataset)
+        self.dataset = dataset
+        self.probabilities = None
+        self.num_steps_per_epoch = num_steps_per_epoch
+        self.batch_size = batch_size
+        self.sequence_length = sequence_length
+        self.can_sample_beyond_end = can_sample_beyond_end
+
+    def __len__(self) -> int:
+        return self.num_steps_per_epoch
+
+    def __iter__(self) -> Generator[List[SegmentId], None, None]:
+        for _ in range(self.num_steps_per_epoch):
+            yield self.sample()
+
+    def sample(self) -> List[SegmentId]:
+        episode_ids = np.random.choice(np.arange(self.dataset.num_episodes), size=self.batch_size, replace=True, p=self.probabilities)
+        timesteps = np.random.randint(low=0, high=self.dataset.lengths[episode_ids])
+
+        # padding allowed, both before start and after end
+        if self.can_sample_beyond_end:  
+            starts = timesteps - np.random.randint(0, self.sequence_length, len(timesteps))
+            stops = starts + self.sequence_length
+
+        # padding allowed only before start
+        else:                           
+            stops = np.minimum(self.dataset.lengths[episode_ids], timesteps + 1 + np.random.randint(0, self.sequence_length, len(timesteps)))
+            starts = stops - self.sequence_length
+
+        return list(map(lambda x: SegmentId(*x), zip(episode_ids, starts, stops)))
+

delta-iris/src/data/segment.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class SegmentId:
+    episode_id: int
+    start: int
+    stop: int
+
+
+@dataclass
+class Segment:
+    observations: torch.ByteTensor
+    actions: torch.LongTensor
+    rewards: torch.FloatTensor
+    ends: torch.LongTensor
+    mask_padding: torch.BoolTensor
+    id: SegmentId
+
+    @property
+    def effective_size(self) -> int:
+        return self.mask_padding.sum().item()

delta-iris/src/data/utils.py

@@ -0,0 +1,69 @@
+import math
+from typing import Generator, List
+
+import torch
+
+from .batch import Batch
+from .episode import Episode
+from .segment import Segment, SegmentId
+
+
+def collate_segments_to_batch(segments: List[Segment]) -> Batch:
+    return Batch(
+        torch.stack(list(map(lambda s: s.observations, segments))).div(255),
+        torch.stack(list(map(lambda s: s.actions, segments))),
+        torch.stack(list(map(lambda s: s.rewards, segments))),
+        torch.stack(list(map(lambda s: s.ends, segments))),
+        torch.stack(list(map(lambda s: s.mask_padding, segments))),
+        list(map(lambda segment: segment.id, segments))
+    )
+
+
+def make_segment(episode: Episode, segment_id: SegmentId, should_pad: bool = True) -> Segment:
+    assert segment_id.start < len(episode) and segment_id.stop > 0 and segment_id.start < segment_id.stop
+    padding_length_right = max(0, segment_id.stop - len(episode))
+    padding_length_left = max(0, -segment_id.start)
+    assert padding_length_right == padding_length_left == 0 or should_pad
+
+    def pad(x):
+        pad_right = torch.nn.functional.pad(x, [0 for _ in range(2 * x.ndim - 1)] + [padding_length_right]) if padding_length_right > 0 else x
+        return torch.nn.functional.pad(pad_right, [0 for _ in range(2 * x.ndim - 2)] + [padding_length_left, 0]) if padding_length_left > 0 else pad_right
+
+    start = max(0, segment_id.start)
+    stop = min(len(episode), segment_id.stop)
+
+    return Segment(
+        pad(episode.observations[start:stop]),
+        pad(episode.actions[start:stop]),
+        pad(episode.rewards[start:stop]),
+        pad(episode.ends[start:stop]),
+        mask_padding=torch.cat((torch.zeros(padding_length_left), torch.ones(stop - start), torch.zeros(padding_length_right))).bool(),
+        id=SegmentId(segment_id.episode_id, start, stop)
+    )
+
+
+class DatasetTraverser:
+    def __init__(self, dataset, batch_num_samples: int, chunk_size: int) -> None:
+        self.dataset = dataset
+        self.batch_num_samples = batch_num_samples
+        self.chunk_size = chunk_size
+        self._num_batches = math.ceil(sum([math.ceil(dataset.lengths[episode_id] / chunk_size) - int(dataset.lengths[episode_id] % chunk_size == 1) for episode_id in range(dataset.num_episodes)]) / batch_num_samples)
+
+    def __len__(self) -> int:
+        return self._num_batches 
+
+    def __iter__(self) -> Generator[Batch, None, None]:
+        chunks = []
+
+        for episode_id in range(self.dataset.num_episodes):
+            episode = self.dataset.load_episode(episode_id)
+            chunks.extend(make_segment(episode, SegmentId(episode_id, start=i * self.chunk_size, stop=(i + 1) * self.chunk_size), should_pad=True) for i in range(math.ceil(len(episode) / self.chunk_size)))
+            if chunks[-1].effective_size < 2:
+                chunks.pop()
+
+            while len(chunks) >= self.batch_num_samples:
+                yield collate_segments_to_batch(chunks[:self.batch_num_samples])
+                chunks = chunks[self.batch_num_samples:]
+
+        if len(chunks) > 0:
+            yield collate_segments_to_batch(chunks)

delta-iris/src/models/__init__.py

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

delta-iris/src/models/convnet.py

@@ -0,0 +1,114 @@
+from dataclasses import dataclass
+from typing import List
+
+from einops import rearrange
+import torch
+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:
+        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
+
+        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:
+                encoder_layers.append(Downsample(output_channels))
+        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)
+        ])
+        self.encoder = nn.Sequential(*encoder_layers)
+
+    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
+        b, t, _, _, _ = x.size()
+        x = rearrange(x, 'b t c h w -> (b t) c h w')
+        x = self.encoder(x)
+        x = rearrange(x, '(b t) c h w -> b t c h w', b=b, t=t)
+
+        return x
+
+
+class FrameDecoder(nn.Module):
+    def __init__(self, config: FrameCnnConfig) -> None:
+        super().__init__()
+
+        assert len(config.mult) == len(config.down)
+        decoder_layers = []
+        output_channels = 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.extend([
+            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)
+        ])
+        self.decoder = nn.Sequential(*decoder_layers)
+
+    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
+        b, t, _, _, _ = x.size()
+        x = rearrange(x, 'b t c h w -> (b t) c h w')
+        x = self.decoder(x)
+        x = rearrange(x, '(b t) c h w -> b t c h w', b=b, t=t)
+        return x
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, num_groups_norm: int = 32) -> None:
+        super().__init__()
+
+        self.f = nn.Sequential(
+            nn.GroupNorm(num_groups_norm, in_channels), 
+            nn.SiLU(inplace=True),
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),
+            nn.GroupNorm(num_groups_norm, out_channels),
+            nn.SiLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
+        )
+        self.skip_projection = nn.Identity() if in_channels == out_channels else torch.nn.Conv2d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.skip_projection(x) + self.f(x) 
+
+
+class Downsample(nn.Module):
+    def __init__(self, num_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(num_channels, num_channels, kernel_size=2, stride=2, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
+class Upsample(nn.Module):
+    def __init__(self, num_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(num_channels, num_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+        return self.conv(x)

delta-iris/src/models/kv_caching.py

@@ -0,0 +1,106 @@
+from typing import Tuple
+
+import numpy as np
+import torch
+
+
+class Cache:
+    def __init__(self, num_samples: int, max_tokens: int, embed_dim: int, device: torch.device) -> None:
+        self._n, self._cache, self._size = num_samples, None, None
+        self._reset = lambda n: torch.empty(n, max_tokens, embed_dim, device=device)  # (B, T, E)
+        self.reset()
+
+    @property
+    def shape(self) -> Tuple[int, int, int]:
+        n, _, embed_dim = self._cache.shape
+
+        return n, self._size, embed_dim
+
+    def reset(self) -> None:
+        self._cache = self._reset(self._n)
+        self._size = 0
+
+    def prune(self, mask: np.ndarray) -> None:
+        assert mask.ndim == 1 and mask.shape[0] == self.shape[0]
+        self._cache = self._cache[mask]
+        self._n = self._cache.shape[0]
+
+    def get(self) -> torch.Tensor:
+        return self._cache[:, :self._size, :]
+
+    def update(self, x: torch.Tensor) -> None:
+        assert (x.ndim == self._cache.ndim) and all([x.size(i) == self._cache.size(i) for i in (0, 2)])
+        assert self._size + x.size(1) <= self._cache.shape[1]
+        self._cache = AssignWithoutInplaceCheck.apply(self._cache, x, 1, self._size, self._size + x.size(1))
+        self._size += x.size(1)
+
+
+class KVCache:
+    def __init__(self, n: int, max_tokens: int, embed_dim: int, device: torch.device) -> None:
+        self._k_cache = Cache(n, max_tokens, embed_dim, device)
+        self._v_cache = Cache(n, max_tokens, embed_dim, device)
+
+    @property
+    def shape(self) -> Tuple[int, int, int]:
+        return self._k_cache.shape
+
+    def reset(self) -> None:
+        self._k_cache.reset()
+        self._v_cache.reset()
+
+    def prune(self, mask: np.ndarray) -> None:
+        self._k_cache.prune(mask)
+        self._v_cache.prune(mask)
+
+    def get(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        return self._k_cache.get(), self._v_cache.get()
+
+    def update(self, k: torch.Tensor, v: torch.Tensor):
+        self._k_cache.update(k)
+        self._v_cache.update(v)
+
+
+class KeysValues:
+    def __init__(self, n: int, max_tokens: int, embed_dim: int, num_layers: int, device: torch.device) -> None:
+        self._keys_values = tuple([KVCache(n, max_tokens, embed_dim, device) for _ in range(num_layers)])
+
+    def __getitem__(self, key: int) -> KVCache:
+        return self._keys_values[key]
+
+    def __len__(self):
+        return len(self._keys_values)
+
+    @property
+    def size(self):
+        return self._keys_values[0].shape[1]
+
+    def reset(self) -> None:
+        for kv_cache in self._keys_values:
+            kv_cache.reset()
+
+    def prune(self, mask: np.ndarray) -> None:
+        for kv_cache in self._keys_values:
+            kv_cache.prune(mask)
+
+
+class AssignWithoutInplaceCheck(torch.autograd.Function):
+    """
+    Inspired from : https://discuss.pytorch.org/t/disable-in-place-correctness-version-check-any-other-workaround/90738/4
+    Warning : do not use it to overwrite a slice twice.
+    """
+
+    @staticmethod
+    def get_slice(dim: int, start: int, stop: int) -> Tuple[slice]:
+        return tuple([slice(None), ] * dim + [slice(start, stop)])
+
+    @staticmethod
+    def forward(ctx, input: torch.Tensor, value: torch.Tensor, dim: int, start: int, stop: int) -> torch.Tensor:
+        ctx.dim = dim
+        ctx.start = start
+        ctx.stop = stop
+        input.data[AssignWithoutInplaceCheck.get_slice(dim, start, stop)] = value
+        return input
+
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor) -> Tuple[torch.Tensor]:
+        return grad_out, grad_out[AssignWithoutInplaceCheck.get_slice(ctx.dim, ctx.start, ctx.stop)], None, None, None

delta-iris/src/models/slicer.py

@@ -0,0 +1,55 @@
+import math
+from typing import List
+
+import torch
+import torch.nn as nn
+
+
+class Slicer(nn.Module):
+    def __init__(self, max_blocks: int, block_mask: torch.Tensor) -> None:
+        super().__init__()
+        self.block_size = block_mask.size(0)
+        self.num_kept_tokens = block_mask.sum().long().item()
+        kept_indices = torch.where(block_mask)[0].repeat(max_blocks)
+        offsets = torch.arange(max_blocks).repeat_interleave(self.num_kept_tokens)
+        self.register_buffer('indices', kept_indices + block_mask.size(0) * offsets)
+
+    def compute_slice(self, num_steps: int, prev_steps: int = 0) -> torch.Tensor:
+        total_steps = num_steps + prev_steps
+        num_blocks = math.ceil(total_steps / self.block_size)
+        indices = self.indices[:num_blocks * self.num_kept_tokens]
+        return indices[torch.logical_and(prev_steps <= indices, indices < total_steps)] - prev_steps
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+class Head(Slicer):
+    def __init__(self, max_blocks: int, block_mask: torch.Tensor, head_module: nn.Module) -> None:
+        super().__init__(max_blocks, block_mask)
+        assert isinstance(head_module, nn.Module)
+        self.head_module = head_module
+
+    def forward(self, x: torch.Tensor, num_steps: int, prev_steps: int) -> torch.Tensor:
+        x_sliced = x[:, self.compute_slice(num_steps, prev_steps)]  # x is (B, T, E)
+        return self.head_module(x_sliced)
+
+
+class Embedder(nn.Module):
+    def __init__(self, max_blocks: int, block_masks: List[torch.Tensor], embedding_tables: List[nn.Embedding]) -> None:
+        super().__init__()
+        assert len(block_masks) == len(embedding_tables)
+        assert (sum(block_masks) == 1).all()  # block mask are a partition of a block
+        self.embedding_dim = embedding_tables[0].embedding_dim
+        assert all([e.embedding_dim == self.embedding_dim for e in embedding_tables])
+        self.embedding_tables = embedding_tables
+        self.slicers = [Slicer(max_blocks, block_mask) for block_mask in block_masks]
+
+    def forward(self, tokens: torch.LongTensor, num_steps: int, prev_steps: int) -> torch.FloatTensor:
+        assert tokens.ndim == 2  # x is (B, T)
+        output = torch.zeros(*tokens.size(), self.embedding_dim, device=tokens.device)
+        for slicer, emb in zip(self.slicers, self.embedding_tables):
+            s = slicer.compute_slice(num_steps, prev_steps)
+            output[:, s] = emb(tokens[:, s])
+
+        return output

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

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

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

@@ -0,0 +1,95 @@
+from dataclasses import dataclass
+import math
+from typing import Dict, Optional
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class QuantizerOutput:
+    q: torch.FloatTensor
+    tokens: torch.LongTensor
+    loss: Dict[str, torch.FloatTensor]
+    metrics: Dict[str, float]
+
+
+class Quantizer(nn.Module):
+    def __init__(self, codebook_size: int, codebook_dim: int, input_dim: int, max_codebook_updates_with_revival: Optional[int] = None) -> None:
+        super().__init__()
+        assert math.log2(codebook_size).is_integer()
+        self.revival_entropy_threshold = int(math.log2(codebook_size)) - 2
+        self.max_codebook_updates_with_revival = max_codebook_updates_with_revival
+        self.pre_quant_proj = nn.Linear(input_dim, codebook_dim)
+        self.post_quant_proj = nn.Linear(codebook_dim, input_dim)
+        codebook = torch.empty(codebook_size, codebook_dim, requires_grad=False).uniform_(-1.0 / codebook_size, 1.0 / codebook_size)
+        self.register_buffer('num_codebook_updates', torch.tensor(0))
+        self.register_buffer('codebook', codebook)
+        self.register_buffer('codewords_freqs', torch.ones(codebook_size).div(codebook_size))
+
+    def forward(self, z: torch.Tensor) -> QuantizerOutput:
+        z = self.pre_quant_proj(z)
+        z = F.normalize(z, dim=-1)
+        b, k = z.size(0), z.size(2)
+        z = rearrange(z, 'b t k e -> (b t k) e')
+
+        cosine_similarity = torch.einsum('n e, c e -> n c', z, self.codebook)
+        tokens = cosine_similarity.argmax(dim=-1)
+        q = self.codebook[tokens]
+
+        losses = {'commitment_loss': 0.02 * (z - q.detach()).pow(2).mean()}
+
+        if self.training:
+            metrics = {**self.update_codebook(z, tokens), 'codebook_entropy': self.compute_codebook_entropy()}
+        else:
+            metrics = {}
+
+        q = z + (q - z).detach()
+        q = self.post_quant_proj(q)
+
+        q = rearrange(q, '(b t k) e -> b t k e', b=b, k=k)
+        tokens = rearrange(tokens, '(b t k) -> b t k', b=b, k=k)
+
+        return QuantizerOutput(q, tokens, losses, metrics)
+
+    @torch.no_grad()
+    def update_codebook(self, z: torch.Tensor, tokens: torch.LongTensor) -> None:
+        tokens_one_hot = F.one_hot(tokens, self.codebook.size(0)).float()  # (N, C)
+
+        # Update codebook
+        counts = tokens_one_hot.sum(dim=0) 
+        codebook_update = torch.einsum('n e, n c -> c e', z, tokens_one_hot) / torch.clamp(counts.unsqueeze(-1), min=1)
+        codebook_update = F.normalize(codebook_update, dim=-1)
+        self.codebook.lerp_(codebook_update, 1 - 0.99)
+
+        # Update counts and revive dead codewords
+        freqs = counts / tokens_one_hot.size(0)
+        self.codewords_freqs.lerp_(freqs, 1 - 0.98)
+
+        can_revive = (self.compute_codebook_entropy() < 1) or (self.max_codebook_updates_with_revival is None) or (self.num_codebook_updates.item() < self.max_codebook_updates_with_revival) 
+        if can_revive and (self.compute_codebook_entropy() < self.revival_entropy_threshold):
+            expired = torch.where(self.codewords_freqs < 1 / (10 * self.codewords_freqs.size(0)))[0]
+            num_expired = expired.size(0)
+            expired = expired[torch.randperm(num_expired)[:z.size(0)]]
+            idx_revived = torch.randperm(z.size(0), device=z.device)[:expired.size(0)]
+            self.codebook[expired] = z[idx_revived]
+            self.codewords_freqs[expired] = 1 / self.codewords_freqs.size(0)
+        else:
+            num_expired = 0
+
+        self.codebook = F.normalize(self.codebook, dim=-1)
+
+        self.num_codebook_updates += 1
+        metrics = {'codewords_revived': num_expired}
+
+        return metrics
+
+    def compute_codebook_entropy(self) -> float:
+        probs = self.codewords_freqs[self.codewords_freqs != 0]
+        return -(torch.log2(probs) * probs).sum().item()
+
+    @torch.no_grad()
+    def embed_tokens(self, tokens: torch.LongTensor) -> torch.FloatTensor:
+        return self.post_quant_proj(self.codebook[tokens])

delta-iris/src/models/transformer.py

@@ -0,0 +1,157 @@
+"""
+Inspired from https://github.com/karpathy/minGPT
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+from einops import rearrange
+import torch
+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:
+        super().__init__()
+        self.config = config
+
+        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
+        mask_sa = torch.tril(torch.ones(k * m, k * m))
+        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.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
+
+    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)
+
+    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)
+
+        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))
+
+        y = self.emb_drop(inputs)
+        for i, block in enumerate(self.blocks):
+            y = block(y, self.keys_values[i] if use_kv_cache else None)
+        y = self.ln(y)
+
+        return y
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, config: TransformerConfig, mask_sa: torch.LongTensor) -> None:
+        super().__init__()
+        self.sa = SelfAttentionLayer(config, mask=mask_sa)
+        self.mlp = MLPLayer(config)
+
+    def forward(self, x: torch.FloatTensor, kv_cache: Optional[KVCache] = None) -> torch.FloatTensor:
+        return self.mlp(self.sa(x, kv_cache))   
+
+
+class MLPLayer(nn.Module):
+    def __init__(self, config: TransformerConfig) -> None:
+        super().__init__()
+        self.ln = nn.LayerNorm(config.embed_dim)
+        self.mlp = nn.Sequential(
+            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),
+        )
+
+    def forward(self, inputs: torch.FloatTensor) -> torch.FloatTensor:
+        return inputs + self.mlp(self.ln(inputs)) 
+
+
+class SelfAttentionLayer(nn.Module):
+    def __init__(self, config: TransformerConfig, 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.attention = Attention(config)
+
+    def forward(self, inputs: torch.FloatTensor, kv_cache: Optional[KVCache] = None) -> torch.FloatTensor:
+        B, T, C = inputs.size()
+        if kv_cache is not None:
+            b, L, c = kv_cache.shape
+            assert b == B and c == C
+        else:
+            L = 0
+
+        x = self.ln(inputs)
+
+        q = self.query(x)
+        k = self.key(x)
+        v = self.value(x)
+
+        if kv_cache is not None:
+            kv_cache.update(k, v)
+            k, v = kv_cache.get()
+
+        y = inputs + self.attention(q, k, v, self.mask[L:L + T, :L + T])
+
+        return y
+
+
+class Attention(nn.Module):
+    def __init__(self, config: TransformerConfig) -> 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)
+
+    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)
+
+        q = rearrange(q, 'b q (h e) -> b h q e', h=self.num_heads)
+        k = rearrange(k, 'b k (h e) -> b h k e', h=self.num_heads)
+        v = rearrange(v, 'b k (h d) -> b h k d', h=self.num_heads)
+
+        y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.attn_pdrop, is_causal=False) if q.size(2) != 0 else q.new_empty(*q.shape[:-1], v.size(-1))
+
+        y = rearrange(y, 'b h q d -> b q (h d)')
+        y = self.resid_drop(self.proj(y))
+
+        return y

delta-iris/src/models/utils.py

@@ -0,0 +1,198 @@
+from collections import OrderedDict
+import cv2
+from pathlib import Path
+import random
+import shutil
+from typing import Callable, Dict
+
+import matplotlib.pyplot as plt
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.optim import AdamW
+
+from data import Episode
+
+
+def configure_optimizer(model: nn.Module, learning_rate: float, weight_decay: float, *blacklist_module_names) -> AdamW:
+    """Credits to https://github.com/karpathy/minGPT"""
+    # separate out all parameters to those that will and won't experience regularizing weight decay
+    decay = set()
+    no_decay = set()
+    whitelist_weight_modules = (torch.nn.Linear, torch.nn.Conv1d)
+    blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding, nn.Conv2d, nn.GroupNorm)
+    for mn, m in model.named_modules():
+        for pn, p in m.named_parameters():
+            fpn = '%s.%s' % (mn, pn) if mn else pn  # full param name
+            if any([fpn.startswith(module_name) for module_name in blacklist_module_names]):
+                no_decay.add(fpn)
+            elif 'bias' in pn:
+                # all biases will not be decayed
+                no_decay.add(fpn)
+            elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
+                # weights of whitelist modules will be weight decayed
+                decay.add(fpn)
+            elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
+                # weights of blacklist modules will NOT be weight decayed
+                no_decay.add(fpn)
+
+    # validate that we considered every parameter
+    param_dict = {pn: p for pn, p in model.named_parameters()}
+    inter_params = decay & no_decay
+    union_params = decay | no_decay
+    assert len(inter_params) == 0, f"parameters {str(inter_params)} made it into both decay/no_decay sets!"
+    assert len(param_dict.keys() - union_params) == 0, f"parameters {str(param_dict.keys() - union_params)} were not separated into either decay/no_decay set!"
+
+    # create the pytorch optimizer object
+    optim_groups = [
+        {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": weight_decay},
+        {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
+    ]
+    optimizer = AdamW(optim_groups, lr=learning_rate)
+
+    return optimizer
+
+
+def init_weights(module: nn.Module) -> None:
+    if isinstance(module, (nn.Linear, nn.Embedding)):
+        module.weight.data.normal_(mean=0.0, std=0.02)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+    elif isinstance(module, nn.LayerNorm):
+        module.bias.data.zero_()
+        module.weight.data.fill_(1.0)
+
+
+def extract_state_dict(state_dict: Dict, module_name: str) -> OrderedDict:
+    return OrderedDict({k.split('.', 1)[1]: v for k, v in state_dict.items() if k.startswith(module_name)})
+
+
+def set_seed(seed: int) -> None:
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    random.seed(seed)
+
+
+@torch.no_grad()
+def compute_discounted_returns(rewards: torch.FloatTensor, gamma: float) -> torch.FloatTensor:
+    assert 0 < gamma <= 1 and rewards.ndim == 2  # (B, T)
+    gammas = gamma ** torch.arange(rewards.size(1))
+    r = rewards * gammas
+
+    return (r + r.sum(dim=1, keepdim=True) - r.cumsum(dim=1)) / gammas
+
+
+class LossWithIntermediateLosses:
+    def __init__(self, **kwargs) -> None:
+        self.loss_total = sum(kwargs.values())
+        self.intermediate_losses = {k: v.item() for k, v in kwargs.items()}
+
+
+class EpisodeDirManager:
+    def __init__(self, episode_dir: Path, max_num_episodes: int) -> None:
+        self.episode_dir = episode_dir
+        self.episode_dir.mkdir(parents=False, exist_ok=True)
+        self.max_num_episodes = max_num_episodes
+        self.best_return = float('-inf')
+
+    def save(self, episode: Episode, episode_id: int, epoch: int) -> None:
+        if self.max_num_episodes is not None and self.max_num_episodes > 0:
+            self._save(episode, episode_id, epoch)
+
+    def _save(self, episode: Episode, episode_id: int, epoch: int) -> None:
+        ep_paths = [p for p in self.episode_dir.iterdir() if p.stem.startswith('episode_')]
+        assert len(ep_paths) <= self.max_num_episodes
+        if len(ep_paths) == self.max_num_episodes:
+            to_remove = min(ep_paths, key=lambda ep_path: int(ep_path.stem.split('_')[1]))
+            to_remove.unlink()
+        torch.save(episode.__dict__, self.episode_dir / f'episode_{episode_id}_epoch_{epoch}.pt')
+
+        ep_return = episode.compute_metrics().episode_return
+        if ep_return > self.best_return:
+            self.best_return = ep_return
+            path_best_ep = [p for p in self.episode_dir.iterdir() if p.stem.startswith('best_')]
+            assert len(path_best_ep) in (0, 1)
+            if len(path_best_ep) == 1:
+                path_best_ep[0].unlink()
+            torch.save(episode.__dict__, self.episode_dir / f'best_episode_{episode_id}_epoch_{epoch}.pt')
+
+
+class RandomHeuristic:
+    def __init__(self, num_actions):
+        self.num_actions = num_actions
+
+    def act(self, obs):
+        assert obs.ndim == 4  # (N, H, W, C)
+        n = obs.size(0)
+
+        return torch.randint(low=0, high=self.num_actions, size=(n,))
+
+
+def make_video(fname, fps, frames):
+    assert frames.ndim == 4  # (T, H, W, C)
+    _, h, w, c = frames.shape
+    assert c == 3
+
+    video = cv2.VideoWriter(str(fname), cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+    for frame in frames:
+        video.write(frame[:, :, ::-1])
+    video.release()
+
+
+def try_until_no_except(fn: Callable):
+    while True:
+        try:
+            fn()
+        except:
+            continue
+        else:
+            break
+
+
+def symlog(x: torch.Tensor) -> torch.Tensor:
+    return torch.sign(x) * torch.log(torch.abs(x) + 1)
+
+
+def symexp(x: torch.Tensor) -> torch.Tensor:
+    return torch.sign(x) * (torch.exp(torch.abs(x)) - 1)
+
+
+def two_hot(x: torch.FloatTensor, x_min: int = -20, x_max: int = 20, num_buckets: int = 255) -> torch.FloatTensor:
+    x.clamp_(x_min, x_max - 1e-5)
+    buckets = torch.linspace(x_min, x_max, num_buckets).to(x.device)
+    k = torch.searchsorted(buckets, x) - 1
+    values = torch.stack((buckets[k + 1] - x, x - buckets[k]), dim=-1) / (buckets[k + 1] - buckets[k]).unsqueeze(-1)  
+    two_hots = torch.scatter(x.new_zeros(*x.size(), num_buckets), dim=-1, index=torch.stack((k, k + 1), dim=-1), src=values)
+
+    return two_hots
+
+
+def compute_softmax_over_buckets(logits: torch.FloatTensor, x_min: int = -20, x_max: int = 20, num_buckets: int = 255) -> torch.FloatTensor:
+    buckets = torch.linspace(x_min, x_max, num_buckets).to(logits.device)
+    probs = F.softmax(logits, dim=-1)
+
+    return probs @ buckets
+
+
+def plot_counts(counts: np.ndarray) -> Image:
+    fig, ax = plt.subplots(figsize=(14, 7))
+    ax.plot(counts)
+    p = Path('priorities.png')
+    fig.savefig(p)
+    plt.close(fig)
+    im = Image.open(p)
+    p.unlink()
+
+    return im
+
+
+def compute_mask_after_first_done(ends: torch.LongTensor) -> torch.BoolTensor:
+    assert ends.ndim == 2
+    first_one_index = torch.argmax(ends, dim=1)
+    mask = torch.arange(ends.size(1), device=ends.device).unsqueeze(0) <= first_one_index.unsqueeze(1)
+    mask = torch.logical_or(mask, ends.sum(dim=1, keepdim=True) == 0)
+
+    return mask

delta-iris/src/tokenizer.py

@@ -0,0 +1,115 @@
+from dataclasses import dataclass
+import math
+from typing import Dict, Tuple
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+
+from .models.convnet import FrameCnnConfig, 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:
+        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.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.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
+        )
+
+        self.encoder = FrameEncoder(config.encoder_config)
+        self.decoder = FrameDecoder(config.decoder_config)
+        self.frame_cnn = FrameEncoder(config.frame_cnn_config)
+
+        self.apply(init_weights)
+
+    def __repr__(self) -> str:
+        return "tokenizer"
+
+    def forward(self, x1: torch.FloatTensor, a: torch.LongTensor, x2: torch.FloatTensor) -> QuantizerOutput:
+        z = self.encode(x1, a, x2)
+        z = rearrange(z, 'b t c (h k) (w l) -> b t (h w) (k l c)', h=self.tokens_grid_res, w=self.tokens_grid_res)
+
+        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)
+        z = self.encoder(encoder_input)
+
+        return z
+
+    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))
+
+        decoder_input = torch.cat((x1_emb, a_emb, q2), dim=2)
+
+        r = self.decoder(decoder_input)
+        r = torch.clamp(r, 0, 1).mul(255).round().div(255) if should_clamp else r
+
+        return r
+
+    @torch.no_grad()
+    def encode_decode(self, x1: torch.FloatTensor, a: torch.LongTensor, x2: torch.FloatTensor) -> torch.Tensor:
+        z = self.encode(x1, a, x2)
+        z = rearrange(z, 'b t c (h k) (w l) -> b t (h w) (k l c)', k=self.token_res, l=self.token_res)
+        q = rearrange(self.quantizer(z).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)
+        r = self.decode(x1, a, q, should_clamp=True)
+
+        return r
+
+    @torch.no_grad()
+    def burn_in(self, obs: torch.FloatTensor, act: torch.LongTensor) -> torch.LongTensor: 
+        assert obs.size(1) == act.size(1) + 1
+        quantizer_output = self(obs[:, :-1], act, obs[:, 1:])
+
+        return quantizer_output.tokens

delta-iris/src/world_model.py

@@ -0,0 +1,139 @@
+from dataclasses import dataclass
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .models.convnet import FrameCnnConfig, FrameEncoder
+from .data import Batch
+from .models.slicer import  Head
+from .tokenizer import Tokenizer
+from .models.transformer import TransformerEncoder, TransformerConfig
+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:
+        super().__init__()
+        self.config = 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))
+
+        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[1] = 1
+        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,
+            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)
+            )
+        )
+
+        self.head_rewards = Head(
+            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)
+            )
+        )
+
+        self.head_ends = Head(
+            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)
+            )
+        )
+
+        self.apply(init_weights)
+
+    def __repr__(self) -> str:
+        return "world_model"
+
+    def forward(self, sequence: torch.FloatTensor, use_kv_cache: bool = False) -> WorldModelOutput:      
+        prev_steps = self.transformer.keys_values.size if use_kv_cache else 0
+        num_steps = sequence.size(1)
+
+        outputs = self.transformer(sequence, use_kv_cache=use_kv_cache)
+
+        logits_latents = self.head_latents(outputs, num_steps, prev_steps)
+        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)
+
+    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
+
+        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