Code release.

src/agent.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+
+import torch
+from torch.distributions.categorical import Categorical
+import torch.nn as nn
+
+from models.actor_critic import ActorCritic
+from models.tokenizer import Tokenizer
+from models.world_model import WorldModel
+from utils import extract_state_dict
+
+
+class Agent(nn.Module):
+    def __init__(self, tokenizer: Tokenizer, world_model: WorldModel, actor_critic: ActorCritic):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.world_model = world_model
+        self.actor_critic = actor_critic
+
+    @property
+    def device(self):
+        return self.actor_critic.conv1.weight.device
+
+    def load(self, path_to_checkpoint: Path, device: torch.device, load_tokenizer: bool = True, load_world_model: bool = True, load_actor_critic: bool = True) -> None:
+        agent_state_dict = torch.load(path_to_checkpoint, map_location=device)
+        if load_tokenizer:
+            self.tokenizer.load_state_dict(extract_state_dict(agent_state_dict, 'tokenizer'))
+        if load_world_model:
+            self.world_model.load_state_dict(extract_state_dict(agent_state_dict, 'world_model'))
+        if load_actor_critic:
+            self.actor_critic.load_state_dict(extract_state_dict(agent_state_dict, 'actor_critic'))
+
+    def act(self, obs: torch.FloatTensor, should_sample: bool = True, temperature: float = 1.0) -> torch.LongTensor:
+        input_ac = obs if self.actor_critic.use_original_obs else torch.clamp(self.tokenizer.encode_decode(obs, should_preprocess=True, should_postprocess=True), 0, 1)
+        logits_actions = self.actor_critic(input_ac).logits_actions[:, -1] / temperature
+        act_token = Categorical(logits=logits_actions).sample() if should_sample else logits_actions.argmax(dim=-1)
+        return act_token

src/collector.py

@@ -0,0 +1,125 @@
+import random
+import sys
+from typing import List, Optional, Union
+
+from einops import rearrange
+import numpy as np
+import torch
+from tqdm import tqdm
+import wandb
+
+from agent import Agent
+from dataset import EpisodesDataset
+from envs import SingleProcessEnv, MultiProcessEnv
+from episode import Episode
+from utils import EpisodeDirManager, RandomHeuristic
+
+
+class Collector:
+    def __init__(self, env: Union[SingleProcessEnv, MultiProcessEnv], dataset: EpisodesDataset, episode_dir_manager: EpisodeDirManager) -> None:
+        self.env = env
+        self.dataset = dataset
+        self.episode_dir_manager = episode_dir_manager
+        self.obs = self.env.reset()
+        self.episode_ids = [None] * self.env.num_envs
+        self.heuristic = RandomHeuristic(self.env.num_actions)
+
+    @torch.no_grad()
+    def collect(self, agent: Agent, epoch: int, epsilon: float, should_sample: bool, temperature: float, burn_in: int, *, num_steps: Optional[int] = None, num_episodes: Optional[int] = None):
+        assert self.env.num_actions == agent.world_model.act_vocab_size
+        assert 0 <= epsilon <= 1
+
+        assert (num_steps is None) != (num_episodes is None)
+        should_stop = lambda steps, episodes: steps >= num_steps if num_steps is not None else episodes >= num_episodes
+
+        to_log = []
+        steps, episodes = 0, 0
+        returns = []
+        observations, actions, rewards, dones = [], [], [], []
+
+        burnin_obs_rec, mask_padding = None, None
+        if set(self.episode_ids) != {None} and burn_in > 0:
+            current_episodes = [self.dataset.get_episode(episode_id) for episode_id in self.episode_ids]
+            segmented_episodes = [episode.segment(start=len(episode) - burn_in, stop=len(episode), should_pad=True) for episode in current_episodes]
+            mask_padding = torch.stack([episode.mask_padding for episode in segmented_episodes], dim=0).to(agent.device)
+            burnin_obs = torch.stack([episode.observations for episode in segmented_episodes], dim=0).float().div(255).to(agent.device)
+            burnin_obs_rec = torch.clamp(agent.tokenizer.encode_decode(burnin_obs, should_preprocess=True, should_postprocess=True), 0, 1)
+
+        agent.actor_critic.reset(n=self.env.num_envs, burnin_observations=burnin_obs_rec, mask_padding=mask_padding)
+        pbar = tqdm(total=num_steps if num_steps is not None else num_episodes, desc=f'Experience collection ({self.dataset.name})', file=sys.stdout)
+
+        while not should_stop(steps, episodes):
+
+            observations.append(self.obs)
+            obs = rearrange(torch.FloatTensor(self.obs).div(255), 'n h w c -> n c h w').to(agent.device)
+            act = agent.act(obs, should_sample=should_sample, temperature=temperature).cpu().numpy()
+
+            if random.random() < epsilon:
+                act = self.heuristic.act(obs).cpu().numpy()
+
+            self.obs, reward, done, _ = self.env.step(act)
+
+            actions.append(act)
+            rewards.append(reward)
+            dones.append(done)
+
+            new_steps = len(self.env.mask_new_dones)
+            steps += new_steps
+            pbar.update(new_steps if num_steps is not None else 0)
+
+            # Warning: with EpisodicLifeEnv + MultiProcessEnv, reset is ignored if not a real done.
+            # Thus, segments of experience following a life loss and preceding a general done are discarded.
+            # Not a problem with a SingleProcessEnv.
+
+            if self.env.should_reset():
+                self.add_experience_to_dataset(observations, actions, rewards, dones)
+
+                new_episodes = self.env.num_envs
+                episodes += new_episodes
+                pbar.update(new_episodes if num_episodes is not None else 0)
+
+                for episode_id in self.episode_ids:
+                    episode = self.dataset.get_episode(episode_id)
+                    self.episode_dir_manager.save(episode, episode_id, epoch)
+                    metrics_episode = {k: v for k, v in episode.compute_metrics().__dict__.items()}
+                    metrics_episode['episode_num'] = episode_id
+                    metrics_episode['action_histogram'] = wandb.Histogram(np_histogram=np.histogram(episode.actions.numpy(), bins=np.arange(0, self.env.num_actions + 1) - 0.5, density=True))
+                    to_log.append({f'{self.dataset.name}/{k}': v for k, v in metrics_episode.items()})
+                    returns.append(metrics_episode['episode_return'])
+
+                self.obs = self.env.reset()
+                self.episode_ids = [None] * self.env.num_envs
+                agent.actor_critic.reset(n=self.env.num_envs)
+                observations, actions, rewards, dones = [], [], [], []
+
+        # Add incomplete episodes to dataset, and complete them later.
+        if len(observations) > 0:
+            self.add_experience_to_dataset(observations, actions, rewards, dones)
+
+        agent.actor_critic.clear()
+
+        metrics_collect = {
+            '#episodes': len(self.dataset),
+            '#steps': sum(map(len, self.dataset.episodes)),
+        }
+        if len(returns) > 0:
+            metrics_collect['return'] = np.mean(returns)
+        metrics_collect = {f'{self.dataset.name}/{k}': v for k, v in metrics_collect.items()}
+        to_log.append(metrics_collect)
+
+        return to_log
+
+    def add_experience_to_dataset(self, observations: List[np.ndarray], actions: List[np.ndarray], rewards: List[np.ndarray], dones: List[np.ndarray]) -> None:
+        assert len(observations) == len(actions) == len(rewards) == len(dones)
+        for i, (o, a, r, d) in enumerate(zip(*map(lambda arr: np.swapaxes(arr, 0, 1), [observations, actions, rewards, dones]))):  # Make everything (N, T, ...) instead of (T, N, ...)
+            episode = Episode(
+                observations=torch.ByteTensor(o).permute(0, 3, 1, 2).contiguous(),  # channel-first
+                actions=torch.LongTensor(a),
+                rewards=torch.FloatTensor(r),
+                ends=torch.LongTensor(d),
+                mask_padding=torch.ones(d.shape[0], dtype=torch.bool),
+            )
+            if self.episode_ids[i] is None:
+                self.episode_ids[i] = self.dataset.add_episode(episode)
+            else:
+                self.dataset.update_episode(self.episode_ids[i], episode)

src/dataset.py

@@ -0,0 +1,163 @@
+from collections import deque
+import math
+from pathlib import Path
+import random
+from typing import Dict, List, Optional, Tuple
+
+import psutil
+import torch
+
+from episode import Episode
+
+Batch = Dict[str, torch.Tensor]
+
+
+class EpisodesDataset:
+    def __init__(self, max_num_episodes: Optional[int] = None, name: Optional[str] = None) -> None:
+        self.max_num_episodes = max_num_episodes
+        self.name = name if name is not None else 'dataset'
+        self.num_seen_episodes = 0
+        self.episodes = deque()
+        self.episode_id_to_queue_idx = dict()
+        self.newly_modified_episodes, self.newly_deleted_episodes = set(), set()
+
+    def __len__(self) -> int:
+        return len(self.episodes)
+
+    def clear(self) -> None:
+        self.episodes = deque()
+        self.episode_id_to_queue_idx = dict()
+
+    def add_episode(self, episode: Episode) -> int:
+        if self.max_num_episodes is not None and len(self.episodes) == self.max_num_episodes:
+            self._popleft()
+        episode_id = self._append_new_episode(episode)
+        return episode_id
+
+    def get_episode(self, episode_id: int) -> Episode:
+        assert episode_id in self.episode_id_to_queue_idx
+        queue_idx = self.episode_id_to_queue_idx[episode_id]
+        return self.episodes[queue_idx]
+
+    def update_episode(self, episode_id: int, new_episode: Episode) -> None:
+        assert episode_id in self.episode_id_to_queue_idx
+        queue_idx = self.episode_id_to_queue_idx[episode_id]
+        merged_episode = self.episodes[queue_idx].merge(new_episode)
+        self.episodes[queue_idx] = merged_episode
+        self.newly_modified_episodes.add(episode_id)
+
+    def _popleft(self) -> Episode:
+        id_to_delete = [k for k, v in self.episode_id_to_queue_idx.items() if v == 0]
+        assert len(id_to_delete) == 1
+        self.newly_deleted_episodes.add(id_to_delete[0])
+        self.episode_id_to_queue_idx = {k: v - 1 for k, v in self.episode_id_to_queue_idx.items() if v > 0}
+        return self.episodes.popleft()
+
+    def _append_new_episode(self, episode):
+        episode_id = self.num_seen_episodes
+        self.episode_id_to_queue_idx[episode_id] = len(self.episodes)
+        self.episodes.append(episode)
+        self.num_seen_episodes += 1
+        self.newly_modified_episodes.add(episode_id)
+        return episode_id
+
+    def sample_batch(self, batch_num_samples: int, sequence_length: int, weights: Optional[Tuple[float]] = None, sample_from_start: bool = True) -> Batch:
+        return self._collate_episodes_segments(self._sample_episodes_segments(batch_num_samples, sequence_length, weights, sample_from_start))
+
+    def _sample_episodes_segments(self, batch_num_samples: int, sequence_length: int, weights: Optional[Tuple[float]], sample_from_start: bool) -> List[Episode]:
+        num_episodes = len(self.episodes)
+        num_weights = len(weights) if weights is not None else 0
+
+        if num_weights < num_episodes:
+            weights = [1] * num_episodes
+        else:
+            assert all([0 <= x <= 1 for x in weights]) and sum(weights) == 1
+            sizes = [num_episodes // num_weights + (num_episodes % num_weights) * (i == num_weights - 1) for i in range(num_weights)]
+            weights = [w / s for (w, s) in zip(weights, sizes) for _ in range(s)]
+
+        sampled_episodes = random.choices(self.episodes, k=batch_num_samples, weights=weights)
+
+        sampled_episodes_segments = []
+        for sampled_episode in sampled_episodes:
+            if sample_from_start:
+                start = random.randint(0, len(sampled_episode) - 1)
+                stop = start + sequence_length
+            else:
+                stop = random.randint(1, len(sampled_episode))
+                start = stop - sequence_length
+            sampled_episodes_segments.append(sampled_episode.segment(start, stop, should_pad=True))
+            assert len(sampled_episodes_segments[-1]) == sequence_length
+        return sampled_episodes_segments
+
+    def _collate_episodes_segments(self, episodes_segments: List[Episode]) -> Batch:
+        episodes_segments = [e_s.__dict__ for e_s in episodes_segments]
+        batch = {}
+        for k in episodes_segments[0]:
+            batch[k] = torch.stack([e_s[k] for e_s in episodes_segments])
+        batch['observations'] = batch['observations'].float() / 255.0  # int8 to float and scale
+        return batch
+
+    def traverse(self, batch_num_samples: int, chunk_size: int):
+        for episode in self.episodes:
+            chunks = [episode.segment(start=i * chunk_size, stop=(i + 1) * chunk_size, should_pad=True) for i in range(math.ceil(len(episode) / chunk_size))]
+            batches = [chunks[i * batch_num_samples: (i + 1) * batch_num_samples] for i in range(math.ceil(len(chunks) / batch_num_samples))]
+            for b in batches:
+                yield self._collate_episodes_segments(b)
+
+    def update_disk_checkpoint(self, directory: Path) -> None:
+        assert directory.is_dir()
+        for episode_id in self.newly_modified_episodes:
+            episode = self.get_episode(episode_id)
+            episode.save(directory / f'{episode_id}.pt')
+        for episode_id in self.newly_deleted_episodes:
+            (directory / f'{episode_id}.pt').unlink()
+        self.newly_modified_episodes, self.newly_deleted_episodes = set(), set()
+
+    def load_disk_checkpoint(self, directory: Path) -> None:
+        assert directory.is_dir() and len(self.episodes) == 0
+        episode_ids = sorted([int(p.stem) for p in directory.iterdir()])
+        self.num_seen_episodes = episode_ids[-1] + 1
+        for episode_id in episode_ids:
+            episode = Episode(**torch.load(directory / f'{episode_id}.pt'))
+            self.episode_id_to_queue_idx[episode_id] = len(self.episodes)
+            self.episodes.append(episode)
+
+
+class EpisodesDatasetRamMonitoring(EpisodesDataset):
+    """
+    Prevent episode dataset from going out of RAM.
+    Warning: % looks at system wide RAM usage while G looks only at process RAM usage.
+    """
+    def __init__(self, max_ram_usage: str, name: Optional[str] = None) -> None:
+        super().__init__(max_num_episodes=None, name=name)
+        self.max_ram_usage = max_ram_usage
+        self.num_steps = 0
+        self.max_num_steps = None
+
+        max_ram_usage = str(max_ram_usage)
+        if max_ram_usage.endswith('%'):
+            m = int(max_ram_usage.split('%')[0])
+            assert 0 < m < 100
+            self.check_ram_usage = lambda: psutil.virtual_memory().percent > m
+        else:
+            assert max_ram_usage.endswith('G')
+            m = float(max_ram_usage.split('G')[0])
+            self.check_ram_usage = lambda: psutil.Process().memory_info()[0] / 2 ** 30 > m
+
+    def clear(self) -> None:
+        super().clear()
+        self.num_steps = 0
+
+    def add_episode(self, episode: Episode) -> int:
+        if self.max_num_steps is None and self.check_ram_usage():
+            self.max_num_steps = self.num_steps
+        self.num_steps += len(episode)
+        while (self.max_num_steps is not None) and (self.num_steps > self.max_num_steps):
+            self._popleft()
+        episode_id = self._append_new_episode(episode)
+        return episode_id
+
+    def _popleft(self) -> Episode:
+        episode = super()._popleft()
+        self.num_steps -= len(episode)
+        return episode

src/envs/__init__.py

@@ -0,0 +1,4 @@
+from .multi_process_env import MultiProcessEnv
+from .wrappers import make_atari, ResizeObsWrapper
+from .single_process_env import SingleProcessEnv
+from .world_model_env import WorldModelEnv

src/envs/done_tracker.py

@@ -0,0 +1,27 @@
+import numpy as np
+
+
+class DoneTrackerEnv:
+    def __init__(self, num_envs: int) -> None:
+        """Monitor env dones: 0 when not done, 1 when done, 2 when already done."""
+        self.num_envs = num_envs
+        self.done_tracker = None
+        self.reset_done_tracker()
+
+    def reset_done_tracker(self) -> None:
+        self.done_tracker = np.zeros(self.num_envs, dtype=np.uint8)
+
+    def update_done_tracker(self, done: np.ndarray) -> None:
+        self.done_tracker = np.clip(2 * self.done_tracker + done, 0, 2)
+
+    @property
+    def num_envs_done(self) -> int:
+        return (self.done_tracker > 0).sum()
+
+    @property
+    def mask_dones(self) -> np.ndarray:
+        return np.logical_not(self.done_tracker)
+
+    @property
+    def mask_new_dones(self) -> np.ndarray:
+        return np.logical_not(self.done_tracker[self.done_tracker <= 1])

src/envs/multi_process_env.py

@@ -0,0 +1,94 @@
+from dataclasses import astuple, dataclass
+from enum import Enum
+from multiprocessing import Pipe, Process
+from multiprocessing.connection import Connection
+from typing import Any, Callable, Iterator, List, Optional, Tuple
+
+import numpy as np
+
+from .done_tracker import DoneTrackerEnv
+
+
+class MessageType(Enum):
+    RESET = 0
+    RESET_RETURN = 1
+    STEP = 2
+    STEP_RETURN = 3
+    CLOSE = 4
+
+
+@dataclass
+class Message:
+    type: MessageType
+    content: Optional[Any] = None
+
+    def __iter__(self) -> Iterator:
+        return iter(astuple(self))
+
+
+def child_env(child_id: int, env_fn: Callable, child_conn: Connection) -> None:
+    np.random.seed(child_id + np.random.randint(0, 2 ** 31 - 1))
+    env = env_fn()
+    while True:
+        message_type, content = child_conn.recv()
+        if message_type == MessageType.RESET:
+            obs = env.reset()
+            child_conn.send(Message(MessageType.RESET_RETURN, obs))
+        elif message_type == MessageType.STEP:
+            obs, rew, done, _ = env.step(content)
+            if done:
+                obs = env.reset()
+            child_conn.send(Message(MessageType.STEP_RETURN, (obs, rew, done, None)))
+        elif message_type == MessageType.CLOSE:
+            child_conn.close()
+            return
+        else:
+            raise NotImplementedError
+
+
+class MultiProcessEnv(DoneTrackerEnv):
+    def __init__(self, env_fn: Callable, num_envs: int, should_wait_num_envs_ratio: float) -> None:
+        super().__init__(num_envs)
+        self.num_actions = env_fn().env.action_space.n
+        self.should_wait_num_envs_ratio = should_wait_num_envs_ratio
+        self.processes, self.parent_conns = [], []
+        for child_id in range(num_envs):
+            parent_conn, child_conn = Pipe()
+            self.parent_conns.append(parent_conn)
+            p = Process(target=child_env, args=(child_id, env_fn, child_conn), daemon=True)
+            self.processes.append(p)
+        for p in self.processes:
+            p.start()
+
+    def should_reset(self) -> bool:
+        return (self.num_envs_done / self.num_envs) >= self.should_wait_num_envs_ratio
+
+    def _receive(self, check_type: Optional[MessageType] = None) -> List[Any]:
+        messages = [parent_conn.recv() for parent_conn in self.parent_conns]
+        if check_type is not None:
+            assert all([m.type == check_type for m in messages])
+        return [m.content for m in messages]
+
+    def reset(self) -> np.ndarray:
+        self.reset_done_tracker()
+        for parent_conn in self.parent_conns:
+            parent_conn.send(Message(MessageType.RESET))
+        content = self._receive(check_type=MessageType.RESET_RETURN)
+        return np.stack(content)
+
+    def step(self, actions: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Any]:
+        for parent_conn, action in zip(self.parent_conns, actions):
+            parent_conn.send(Message(MessageType.STEP, action))
+        content = self._receive(check_type=MessageType.STEP_RETURN)
+        obs, rew, done, _ = zip(*content)
+        done = np.stack(done)
+        self.update_done_tracker(done)
+        return np.stack(obs), np.stack(rew), done, None
+
+    def close(self) -> None:
+        for parent_conn in self.parent_conns:
+            parent_conn.send(Message(MessageType.CLOSE))
+        for p in self.processes:
+            p.join()
+        for parent_conn in self.parent_conns:
+            parent_conn.close()

src/envs/single_process_env.py

@@ -0,0 +1,32 @@
+from typing import Any, Tuple
+
+import numpy as np
+
+from .done_tracker import DoneTrackerEnv
+
+
+class SingleProcessEnv(DoneTrackerEnv):
+    def __init__(self, env_fn):
+        super().__init__(num_envs=1)
+        self.env = env_fn()
+        self.num_actions = self.env.action_space.n
+
+    def should_reset(self) -> bool:
+        return self.num_envs_done == 1
+
+    def reset(self) -> np.ndarray:
+        self.reset_done_tracker()
+        obs = self.env.reset()
+        return obs[None, ...]
+
+    def step(self, action) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Any]:
+        obs, reward, done, _ = self.env.step(action[0])  # action is supposed to be ndarray (1,)
+        done = np.array([done])
+        self.update_done_tracker(done)
+        return obs[None, ...], np.array([reward]), done, None
+
+    def render(self) -> None:
+        self.env.render()
+
+    def close(self) -> None:
+        self.env.close()

src/envs/world_model_env.py

@@ -0,0 +1,113 @@
+import random
+from typing import List, Optional, Union
+
+import gym
+from einops import rearrange
+import numpy as np
+from PIL import Image
+import torch
+from torch.distributions.categorical import Categorical
+import torchvision
+
+from utils import extract_state_dict
+
+
+class WorldModelEnv:
+
+    def __init__(self, tokenizer: torch.nn.Module, world_model: torch.nn.Module, device: Union[str, torch.device], pretrained_agent_path: Optional[str] = None, env: Optional[gym.Env] = None) -> None:
+
+        self.device = torch.device(device)
+        self.world_model = world_model.to(self.device).eval()
+        self.tokenizer = tokenizer.to(self.device).eval()
+
+        if pretrained_agent_path is not None:
+            agent_state_dict = torch.load(pretrained_agent_path)
+            self.world_model.load_state_dict(extract_state_dict(agent_state_dict, 'world_model'))
+            incompatible_keys = self.tokenizer.load_state_dict(extract_state_dict(agent_state_dict, 'tokenizer'), strict=False)
+            assert not incompatible_keys.missing_keys
+            assert (not incompatible_keys.unexpected_keys) or all([k.startswith('lpips.') for k in incompatible_keys.unexpected_keys])
+
+        self.keys_values_wm, self.obs_tokens, self._num_observations_tokens = None, None, None
+
+        self.env = env
+
+    @property
+    def num_observations_tokens(self) -> int:
+        return self._num_observations_tokens
+
+    @torch.no_grad()
+    def reset(self) -> torch.FloatTensor:
+        assert self.env is not None
+        obs = torchvision.transforms.functional.to_tensor(self.env.reset()).to(self.device).unsqueeze(0)  # (1, C, H, W) in [0., 1.]
+        return self.reset_from_initial_observations(obs)
+
+    @torch.no_grad()
+    def reset_from_initial_observations(self, observations: torch.FloatTensor) -> torch.FloatTensor:
+        obs_tokens = self.tokenizer.encode(observations, should_preprocess=True).tokens    # (B, C, H, W) -> (B, K)
+        _, num_observations_tokens = obs_tokens.shape
+        if self.num_observations_tokens is None:
+            self._num_observations_tokens = num_observations_tokens
+
+        _ = self.refresh_keys_values_with_initial_obs_tokens(obs_tokens)
+        self.obs_tokens = obs_tokens
+
+        return self.decode_obs_tokens()
+
+    @torch.no_grad()
+    def refresh_keys_values_with_initial_obs_tokens(self, obs_tokens: torch.LongTensor) -> torch.FloatTensor:
+        n, num_observations_tokens = obs_tokens.shape
+        assert num_observations_tokens == self.num_observations_tokens
+        self.keys_values_wm = self.world_model.transformer.generate_empty_keys_values(n=n, max_tokens=self.world_model.config.max_tokens)
+        outputs_wm = self.world_model(obs_tokens, past_keys_values=self.keys_values_wm)
+        return outputs_wm.output_sequence  # (B, K, E)
+
+    @torch.no_grad()
+    def step(self, action: Union[int, np.ndarray, torch.LongTensor], should_predict_next_obs: bool = True) -> None:
+        assert self.keys_values_wm is not None and self.num_observations_tokens is not None
+
+        num_passes = 1 + self.num_observations_tokens if should_predict_next_obs else 1
+
+        output_sequence, obs_tokens = [], []
+
+        if self.keys_values_wm.size + num_passes > self.world_model.config.max_tokens:
+            _ = self.refresh_keys_values_with_initial_obs_tokens(self.obs_tokens)
+
+        token = action.clone().detach() if isinstance(action, torch.Tensor) else torch.tensor(action, dtype=torch.long)
+        token = token.reshape(-1, 1).to(self.device)  # (B, 1)
+
+        for k in range(num_passes):  # assumption that there is only one action token.
+
+            outputs_wm = self.world_model(token, past_keys_values=self.keys_values_wm)
+            output_sequence.append(outputs_wm.output_sequence)
+
+            if k == 0:
+                reward = Categorical(logits=outputs_wm.logits_rewards).sample().float().cpu().numpy().reshape(-1) - 1   # (B,)
+                done = Categorical(logits=outputs_wm.logits_ends).sample().cpu().numpy().astype(bool).reshape(-1)       # (B,)
+
+            if k < self.num_observations_tokens:
+                token = Categorical(logits=outputs_wm.logits_observations).sample()
+                obs_tokens.append(token)
+
+        output_sequence = torch.cat(output_sequence, dim=1)   # (B, 1 + K, E)
+        self.obs_tokens = torch.cat(obs_tokens, dim=1)        # (B, K)
+
+        obs = self.decode_obs_tokens() if should_predict_next_obs else None
+        return obs, reward, done, None
+
+    @torch.no_grad()
+    def render_batch(self) -> List[Image.Image]:
+        frames = self.decode_obs_tokens().detach().cpu()
+        frames = rearrange(frames, 'b c h w -> b h w c').mul(255).numpy().astype(np.uint8)
+        return [Image.fromarray(frame) for frame in frames]
+
+    @torch.no_grad()
+    def decode_obs_tokens(self) -> List[Image.Image]:
+        embedded_tokens = self.tokenizer.embedding(self.obs_tokens)     # (B, K, E)
+        z = rearrange(embedded_tokens, 'b (h w) e -> b e h w', h=int(np.sqrt(self.num_observations_tokens)))
+        rec = self.tokenizer.decode(z, should_postprocess=True)         # (B, C, H, W)
+        return torch.clamp(rec, 0, 1)
+
+    @torch.no_grad()
+    def render(self):
+        assert self.obs_tokens.shape == (1, self.num_observations_tokens)
+        return self.render_batch()[0]

src/envs/wrappers.py

@@ -0,0 +1,147 @@
+"""
+Credits to https://github.com/openai/baselines/blob/master/baselines/common/atari_wrappers.py
+"""
+
+from typing import Tuple
+
+import gym
+import numpy as np
+from PIL import Image
+
+
+def make_atari(id, size=64, max_episode_steps=None, noop_max=30, frame_skip=4, done_on_life_loss=False, clip_reward=False):
+    env = gym.make(id)
+    assert 'NoFrameskip' in env.spec.id or 'Frameskip' not in env.spec
+    env = ResizeObsWrapper(env, (size, size))
+    if clip_reward:
+        env = RewardClippingWrapper(env)
+    if max_episode_steps is not None:
+        env = gym.wrappers.TimeLimit(env, max_episode_steps=max_episode_steps)
+    if noop_max is not None:
+        env = NoopResetEnv(env, noop_max=noop_max)
+    env = MaxAndSkipEnv(env, skip=frame_skip)
+    if done_on_life_loss:
+        env = EpisodicLifeEnv(env)
+    return env
+
+
+class ResizeObsWrapper(gym.ObservationWrapper):
+    def __init__(self, env: gym.Env, size: Tuple[int, int]) -> None:
+        gym.ObservationWrapper.__init__(self, env)
+        self.size = tuple(size)
+        self.observation_space = gym.spaces.Box(low=0, high=255, shape=(size[0], size[1], 3), dtype=np.uint8)
+        self.original_obs = None
+
+    def resize(self, obs: np.ndarray):
+        img = Image.fromarray(obs)
+        img = img.resize(self.size, Image.BILINEAR)
+        return np.array(img)
+
+    def observation(self, observation: np.ndarray) -> np.ndarray:
+        self.original_obs = observation
+        return self.resize(observation)
+
+
+class RewardClippingWrapper(gym.RewardWrapper):
+    def reward(self, reward):
+        return np.sign(reward)
+
+
+class NoopResetEnv(gym.Wrapper):
+    def __init__(self, env, noop_max=30):
+        """Sample initial states by taking random number of no-ops on reset.
+        No-op is assumed to be action 0.
+        """
+        gym.Wrapper.__init__(self, env)
+        self.noop_max = noop_max
+        self.override_num_noops = None
+        self.noop_action = 0
+        assert env.unwrapped.get_action_meanings()[0] == 'NOOP'
+
+    def reset(self, **kwargs):
+        """ Do no-op action for a number of steps in [1, noop_max]."""
+        self.env.reset(**kwargs)
+        if self.override_num_noops is not None:
+            noops = self.override_num_noops
+        else:
+            noops = self.unwrapped.np_random.randint(1, self.noop_max + 1)
+        assert noops > 0
+        obs = None
+        for _ in range(noops):
+            obs, _, done, _ = self.env.step(self.noop_action)
+            if done:
+                obs = self.env.reset(**kwargs)
+        return obs
+
+    def step(self, action):
+        return self.env.step(action)
+
+
+class EpisodicLifeEnv(gym.Wrapper):
+    def __init__(self, env):
+        """Make end-of-life == end-of-episode, but only reset on true game over.
+        Done by DeepMind for the DQN and co. since it helps value estimation.
+        """
+        gym.Wrapper.__init__(self, env)
+        self.lives = 0
+        self.was_real_done = True
+
+    def step(self, action):
+        obs, reward, done, info = self.env.step(action)
+        self.was_real_done = done
+        # check current lives, make loss of life terminal,
+        # then update lives to handle bonus lives
+        lives = self.env.unwrapped.ale.lives()
+        if lives < self.lives and lives > 0:
+            # for Qbert sometimes we stay in lives == 0 condition for a few frames
+            # so it's important to keep lives > 0, so that we only reset once
+            # the environment advertises done.
+            done = True
+        self.lives = lives
+        return obs, reward, done, info
+
+    def reset(self, **kwargs):
+        """Reset only when lives are exhausted.
+        This way all states are still reachable even though lives are episodic,
+        and the learner need not know about any of this behind-the-scenes.
+        """
+        if self.was_real_done:
+            obs = self.env.reset(**kwargs)
+        else:
+            # no-op step to advance from terminal/lost life state
+            obs, _, _, _ = self.env.step(0)
+        self.lives = self.env.unwrapped.ale.lives()
+        return obs
+
+
+class MaxAndSkipEnv(gym.Wrapper):
+    def __init__(self, env, skip=4):
+        """Return only every `skip`-th frame"""
+        gym.Wrapper.__init__(self, env)
+        assert skip > 0
+        # most recent raw observations (for max pooling across time steps)
+        self._obs_buffer = np.zeros((2,) + env.observation_space.shape, dtype=np.uint8)
+        self._skip = skip
+        self.max_frame = np.zeros(env.observation_space.shape, dtype=np.uint8)
+
+    def step(self, action):
+        """Repeat action, sum reward, and max over last observations."""
+        total_reward = 0.0
+        done = None
+        for i in range(self._skip):
+            obs, reward, done, info = self.env.step(action)
+            if i == self._skip - 2:
+                self._obs_buffer[0] = obs
+            if i == self._skip - 1:
+                self._obs_buffer[1] = obs
+            total_reward += reward
+            if done:
+                break
+        # Note that the observation on the done=True frame
+        # doesn't matter
+        self.max_frame = self._obs_buffer.max(axis=0)
+
+        return self.max_frame, total_reward, done, info
+
+    def reset(self, **kwargs):
+        return self.env.reset(**kwargs)

src/episode.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from pathlib import Path
+
+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
+    mask_padding: torch.BoolTensor
+
+    def __post_init__(self):
+        assert len(self.observations) == len(self.actions) == len(self.rewards) == len(self.ends) == len(self.mask_padding)
+        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]
+            self.mask_padding = self.mask_padding[: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),
+            torch.cat((self.mask_padding, other.mask_padding), dim=0),
+        )
+
+    def segment(self, start: int, stop: int, should_pad: bool = False) -> Episode:
+        assert start < len(self) and stop > 0 and start < stop
+        padding_length_right = max(0, stop - len(self))
+        padding_length_left = max(0, -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, start)
+        stop = min(len(self), stop)
+        segment = Episode(
+            self.observations[start:stop],
+            self.actions[start:stop],
+            self.rewards[start:stop],
+            self.ends[start:stop],
+            self.mask_padding[start:stop],
+        )
+
+        segment.observations = pad(segment.observations)
+        segment.actions = pad(segment.actions)
+        segment.rewards = pad(segment.rewards)
+        segment.ends = pad(segment.ends)
+        segment.mask_padding = torch.cat((torch.zeros(padding_length_left, dtype=torch.bool), segment.mask_padding, torch.zeros(padding_length_right, dtype=torch.bool)), dim=0)
+
+        return segment
+
+    def compute_metrics(self) -> EpisodeMetrics:
+        return EpisodeMetrics(len(self), self.rewards.sum())
+
+    def save(self, path: Path) -> None:
+        torch.save(self.__dict__, path)

src/game/__init__.py

@@ -0,0 +1,3 @@
+from .agent_env import AgentEnv
+from .episode_replay_env import EpisodeReplayEnv
+from .game import Game

src/game/agent_env.py

@@ -0,0 +1,54 @@
+from einops import rearrange
+import numpy as np
+from PIL import Image
+import torch
+
+from agent import Agent
+from envs import SingleProcessEnv
+from game.keymap import get_keymap_and_action_names
+
+
+class AgentEnv:
+    def __init__(self, agent: Agent, env: SingleProcessEnv, keymap_name: str) -> None:
+        assert isinstance(env, SingleProcessEnv)
+        self.agent = agent
+        self.env = env
+        _, self.action_names = get_keymap_and_action_names(keymap_name)
+        self.obs = None
+        self._t = None
+        self._return = None
+
+    def _to_tensor(self, obs: np.ndarray):
+        return rearrange(torch.FloatTensor(obs).div(255), 'n h w c -> n c h w').to(self.agent.device)
+
+    def _to_array(self, obs: torch.FloatTensor):
+        assert obs.ndim == 4 and obs.size(0) == 1
+        return obs[0].mul(255).permute(1, 2, 0).cpu().numpy().astype(np.uint8)
+
+    def reset(self):
+        obs = self.env.reset()
+        self.obs = self._to_tensor(obs)
+        self.agent.actor_critic.reset(1)
+        self._t = 0
+        self._return = 0
+        return obs
+
+    def step(self, *args, **kwargs) -> torch.FloatTensor:
+        with torch.no_grad():
+            act = self.agent.act(self.obs, should_sample=True).cpu().numpy()
+        obs, reward, done, _ = self.env.step(act)
+        self.obs = self._to_tensor(obs)
+        self._t += 1
+        self._return += reward[0]
+        info = {
+            'timestep': self._t,
+            'action': self.action_names[act[0]],
+            'return': self._return,
+        }
+        return obs, reward, done, info
+
+    def render(self) -> Image.Image:
+        assert self.obs.size() == (1, 3, 64, 64)
+        rec = torch.clamp(self.agent.tokenizer.encode_decode(self.obs, should_preprocess=True, should_postprocess=True), 0, 1)
+        arr = self._to_array(torch.cat((self.obs, rec), dim=-1))
+        return Image.fromarray(arr)

src/game/episode_replay_env.py

@@ -0,0 +1,111 @@
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+import torch
+
+from episode import Episode
+from game.keymap import get_keymap_and_action_names
+
+
+class EpisodeReplayEnv:
+    def __init__(self, replay_keymap_name: str, episode_dir: Path) -> None:
+        _, self.action_names = get_keymap_and_action_names(replay_keymap_name)
+        assert episode_dir.is_dir()
+        self._paths = {}
+        for mode in ['train', 'test', 'imagination']:
+            directory = episode_dir / mode
+            if directory.is_dir():
+                self._paths[mode] = sorted([p for p in directory.iterdir() if 'episode_' in p.stem and p.suffix == '.pt'])
+                print(f'Found {len(self._paths[mode])} {mode} episodes.')
+            else:
+                print(f'No {mode} episodes.')
+
+        self._t, self._episode = None, None
+        self._ep_idx = 0
+        self._mode = 'train'
+        self.load()
+
+    def load(self):
+        self._episode = Episode(**torch.load(self.paths[self._ep_idx]))
+        self._t = 0
+
+    def load_next(self):
+        self._ep_idx = (self._ep_idx + 1) % len(self.paths)
+        self.load()
+
+    def load_previous(self):
+        self._ep_idx = (self._ep_idx - 1) % len(self.paths)
+        self.load()
+
+    def set_mode(self, mode):
+        assert mode in ['train', 'test', 'imagination']
+        if mode in self._paths:
+            self._mode = mode
+            self._ep_idx = 0
+            self.load()
+        else:
+            print(f'No {mode} episodes.')
+
+    def __len__(self):
+        return len(self.ends)
+
+    @property
+    def paths(self):
+        return self._paths[self._mode]
+
+    @property
+    def observations(self):
+        return self._episode.observations
+
+    @property
+    def actions(self):
+        return self._episode.actions
+
+    @property
+    def rewards(self):
+        return self._episode.rewards
+
+    @property
+    def ends(self):
+        return self._episode.ends
+
+    def reset(self):
+        return self.observations[self._t]
+
+    def step(self, action) -> torch.FloatTensor:
+        if action == 1:
+            self._t = (self._t - 1) % len(self)
+        elif action == 2:
+            self._t = (self._t + 1) % len(self)
+        if action == 3:
+            self._t = (self._t - 10) % len(self)
+        elif action == 4:
+            self._t = (self._t + 10) % len(self)
+        elif action == 5:
+            self._t = 0
+        elif action == 6:
+            self.load_previous()
+        elif action == 7:
+            self.load_next()
+        elif action == 8:
+            self.set_mode('train')
+        elif action == 9:
+            self.set_mode('test')
+        elif action == 10:
+            self.set_mode('imagination')
+        act = self.actions[self._t]
+        reward = self.rewards[self._t].item()
+        done = self.ends[self._t].item()
+        info = {
+            'ep_name': f'[{self._mode}] {self.paths[self._ep_idx].stem}',
+            'timestep': self._t,
+            'action': self.action_names[act],
+            'cum_reward': f'{sum(self.rewards[:self._t + 1]):.3f}'
+        }
+        return self.observations[self._t], reward, done, info
+
+    def render(self) -> Image.Image:
+        obs = self.observations[self._t]  # (C, H, W) in [0., 1.]
+        arr = obs.permute(1, 2, 0).numpy().astype(np.uint8)
+        return Image.fromarray(arr)

src/game/game.py

@@ -0,0 +1,111 @@
+from typing import Tuple, Union
+
+import gym
+import numpy as np
+import pygame
+from PIL import Image
+
+from envs import WorldModelEnv
+from game.keymap import get_keymap_and_action_names
+
+
+class Game:
+    def __init__(self, env: Union[gym.Env, WorldModelEnv], keymap_name: str, size: Tuple[int, int], fps: int, verbose: bool) -> None:
+        self.env = env
+        self.height, self.width = size
+        self.fps = fps
+        self.verbose = verbose
+
+        self.keymap, self.action_names = get_keymap_and_action_names(keymap_name)
+
+        print('Actions:')
+        for key, idx in self.keymap.items():
+            print(f'{pygame.key.name(key)}: {self.action_names[idx]}')
+
+    def run(self) -> None:
+        pygame.init()
+
+        header_height = 100 if self.verbose else 0
+        font_size = 24
+        screen = pygame.display.set_mode((self.width, self.height + header_height))
+        clock = pygame.time.Clock()
+        font = pygame.font.SysFont(None, font_size)
+        header_rect = pygame.Rect(0, 0, self.width, header_height)
+
+        def clear_header():
+            pygame.draw.rect(screen, pygame.Color('black'), header_rect)
+            pygame.draw.rect(screen, pygame.Color('white'), header_rect, 1)
+
+        def draw_text(text, idx_line, idx_column=0):
+            pos = (5 + idx_column * int(self.width // 4), 5 + idx_line * font_size)
+            assert (0 <= pos[0] <= self.width) and (0 <= pos[1] <= header_height)
+            screen.blit(font.render(text, True, pygame.Color('white')), pos)
+
+        def draw_game(image):
+            if isinstance(image, np.ndarray):
+                image = Image.fromarray(image)
+            else:
+                assert isinstance(image, Image.Image)
+            pygame_image = np.array(image.resize((self.width, self.height), resample=Image.NEAREST)).transpose((1, 0, 2))
+            surface = pygame.surfarray.make_surface(pygame_image)
+            screen.blit(surface, (0, header_height))
+
+        if isinstance(self.env, gym.Env):
+            _, info = self.env.reset(return_info=True)
+            draw_game(info['rgb'])
+        else:
+            self.env.reset()
+            draw_game(self.env.render())
+
+        clear_header()
+        pygame.display.flip()
+
+        do_reset, do_wait = False, False
+        should_stop = False
+        while not should_stop:
+
+            action = 0  # noop
+            pygame.event.pump()
+            for event in pygame.event.get():
+                if event.type == pygame.QUIT:
+                    should_stop = True
+                if event.type == pygame.KEYDOWN and event.key in self.keymap.keys():
+                    action = self.keymap[event.key]
+                if event.type == pygame.KEYDOWN and event.key == pygame.K_RETURN:
+                    do_reset = True
+                if event.type == pygame.KEYDOWN and event.key == pygame.K_PERIOD:
+                    do_wait = not do_wait
+            if action == 0:
+                pressed = pygame.key.get_pressed()
+                for key, action in self.keymap.items():
+                    if pressed[key]:
+                        break
+                else:
+                    action = 0
+
+            if do_wait:
+                continue
+
+            _, reward, done, info = self.env.step(action)
+
+            img = info['rgb'] if isinstance(self.env, gym.Env) else self.env.render()
+            draw_game(img)
+
+            if self.verbose:
+                clear_header()
+                draw_text(f'Action: {self.action_names[action]}', idx_line=0)
+                draw_text(f'Reward: {reward if isinstance(reward, float) else reward.item(): .2f}', idx_line=1)
+                draw_text(f'Done: {done}', idx_line=2)
+                if info is not None:
+                    assert isinstance(info, dict)
+                    for i, (k, v) in enumerate(info.items()):
+                        draw_text(f'{k}: {v}', idx_line=i, idx_column=1)
+
+            pygame.display.flip()   # update screen
+            clock.tick(self.fps)    # ensures game maintains the given frame rate
+
+            if do_reset or done:
+                self.env.reset()
+                do_reset = False
+
+        pygame.quit()

src/game/keymap.py

@@ -0,0 +1,103 @@
+import gym
+import pygame
+
+
+def get_keymap_and_action_names(name):
+
+    if name == 'empty':
+        return EMPTY_KEYMAP, EMPTY_ACTION_NAMES
+
+    if name == 'episode_replay':
+        return EPISODE_REPLAY_KEYMAP, EPISODE_REPLAY_ACTION_NAMES
+
+    if name == 'atari':
+        return ATARI_KEYMAP, ATARI_ACTION_NAMES
+
+    assert name.startswith('atari/')
+    env_id = name.split('atari/')[1]
+    action_names = [x.lower() for x in gym.make(env_id).get_action_meanings()]
+    keymap = {}
+    for key, value in ATARI_KEYMAP.items():
+        if ATARI_ACTION_NAMES[value] in action_names:
+            keymap[key] = action_names.index(ATARI_ACTION_NAMES[value])
+    return keymap, action_names
+
+
+ATARI_ACTION_NAMES = [
+    'noop',
+    'fire',
+    'up',
+    'right',
+    'left',
+    'down',
+    'upright',
+    'upleft',
+    'downright',
+    'downleft',
+    'upfire',
+    'rightfire',
+    'leftfire',
+    'downfire',
+    'uprightfire',
+    'upleftfire',
+    'downrightfire',
+    'downleftfire',
+]
+
+ATARI_KEYMAP = {
+    pygame.K_SPACE: 1,
+
+    pygame.K_w: 2,
+    pygame.K_d: 3,
+    pygame.K_a: 4,
+    pygame.K_s: 5,
+
+    pygame.K_t: 6,
+    pygame.K_r: 7,
+    pygame.K_g: 8,
+    pygame.K_f: 9,
+
+    pygame.K_UP: 10,
+    pygame.K_RIGHT: 11,
+    pygame.K_LEFT: 12,
+    pygame.K_DOWN: 13,
+
+    pygame.K_u: 14,
+    pygame.K_y: 15,
+    pygame.K_j: 16,
+    pygame.K_h: 17,
+}
+
+EPISODE_REPLAY_ACTION_NAMES = [
+    'noop',
+    'previous',
+    'next',
+    'previous_10',
+    'next_10',
+    'go_to_start',
+    'load_previous',
+    'load_next',
+    'go_to_train_episodes',
+    'go_to_test_episodes',
+    'go_to_imagination_episodes',
+]
+
+EPISODE_REPLAY_KEYMAP = {
+    pygame.K_LEFT: 1,
+    pygame.K_RIGHT: 2,
+    pygame.K_PAGEDOWN: 3,
+    pygame.K_PAGEUP: 4,
+    pygame.K_SPACE: 5,
+    pygame.K_DOWN: 6,
+    pygame.K_UP: 7,
+    pygame.K_t: 8,
+    pygame.K_y: 9,
+    pygame.K_i: 10,
+}
+
+EMPTY_ACTION_NAMES = [
+    'noop',
+]
+
+EMPTY_KEYMAP = {
+}

src/main.py

@@ -0,0 +1,14 @@
+import hydra
+from omegaconf import DictConfig
+
+from trainer import Trainer
+
+
+@hydra.main(config_path="../config", config_name="trainer")
+def main(cfg: DictConfig):
+    trainer = Trainer(cfg)
+    trainer.run()
+
+
+if __name__ == "__main__":
+    main()

src/make_reconstructions.py

@@ -0,0 +1,53 @@
+from einops import rearrange
+import numpy as np
+from PIL import Image
+import torch
+
+
+@torch.no_grad()
+def make_reconstructions_from_batch(batch, save_dir, epoch, tokenizer):
+    check_batch(batch)
+
+    original_frames = tensor_to_np_frames(rearrange(batch['observations'], 'b t c h w -> b t h w c'))
+    all = [original_frames]
+
+    rec_frames = generate_reconstructions_with_tokenizer(batch, tokenizer)
+    all.append(rec_frames)
+
+    for i, image in enumerate(map(Image.fromarray, np.concatenate(list(np.concatenate((original_frames, rec_frames), axis=-2)), axis=-3))):
+        image.save(save_dir / f'epoch_{epoch:03d}_t_{i:03d}.png')
+
+    return
+
+
+def check_batch(batch):
+    assert sorted(batch.keys()) == ['actions', 'ends', 'mask_padding', 'observations', 'rewards']
+    b, t, _, _, _ = batch['observations'].shape  # (B, T, C, H, W)
+    assert batch['actions'].shape == batch['rewards'].shape == batch['ends'].shape == batch['mask_padding'].shape == (b, t)
+
+
+def tensor_to_np_frames(inputs):
+    check_float_btw_0_1(inputs)
+    return inputs.mul(255).cpu().numpy().astype(np.uint8)
+
+
+def check_float_btw_0_1(inputs):
+    assert inputs.is_floating_point() and (inputs >= 0).all() and (inputs <= 1).all()
+
+
+@torch.no_grad()
+def generate_reconstructions_with_tokenizer(batch, tokenizer):
+    check_batch(batch)
+    inputs = rearrange(batch['observations'], 'b t c h w -> (b t) c h w')
+    outputs = reconstruct_through_tokenizer(inputs, tokenizer)
+    b, t, _, _, _ = batch['observations'].size()
+    outputs = rearrange(outputs, '(b t) c h w -> b t h w c', b=b, t=t)
+    rec_frames = tensor_to_np_frames(outputs)
+    return rec_frames
+
+
+@torch.no_grad()
+def reconstruct_through_tokenizer(inputs, tokenizer):
+    check_float_btw_0_1(inputs)
+    reconstructions = tokenizer.encode_decode(inputs, should_preprocess=True, should_postprocess=True)
+    return torch.clamp(reconstructions, 0, 1)

src/models/__init__.py

@@ -0,0 +1 @@
+from .transformer import Transformer, TransformerConfig

src/models/actor_critic.py

@@ -0,0 +1,166 @@
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+import sys
+
+from einops import rearrange
+import numpy as np
+import torch
+from torch.distributions.categorical import Categorical
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+
+from dataset import Batch
+from envs.world_model_env import WorldModelEnv
+from models.tokenizer import Tokenizer
+from models.world_model import WorldModel
+from utils import compute_lambda_returns, LossWithIntermediateLosses
+
+
+@dataclass
+class ActorCriticOutput:
+    logits_actions: torch.FloatTensor
+    means_values: torch.FloatTensor
+
+
+@dataclass
+class ImagineOutput:
+    observations: torch.ByteTensor
+    actions: torch.LongTensor
+    logits_actions: torch.FloatTensor
+    values: torch.FloatTensor
+    rewards: torch.FloatTensor
+    ends: torch.BoolTensor
+
+
+class ActorCritic(nn.Module):
+    def __init__(self, act_vocab_size, use_original_obs: bool = False) -> None:
+        super().__init__()
+        self.use_original_obs = use_original_obs
+        self.conv1 = nn.Conv2d(3, 32, 3, stride=1, padding=1)
+        self.maxp1 = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(32, 32, 3, stride=1, padding=1)
+        self.maxp2 = nn.MaxPool2d(2, 2)
+        self.conv3 = nn.Conv2d(32, 64, 3, stride=1, padding=1)
+        self.maxp3 = nn.MaxPool2d(2, 2)
+        self.conv4 = nn.Conv2d(64, 64, 3, stride=1, padding=1)
+        self.maxp4 = nn.MaxPool2d(2, 2)
+
+        self.lstm_dim = 512
+        self.lstm = nn.LSTMCell(1024, self.lstm_dim)
+        self.hx, self.cx = None, None
+
+        self.critic_linear = nn.Linear(512, 1)
+        self.actor_linear = nn.Linear(512, act_vocab_size)
+
+    def __repr__(self) -> str:
+        return "actor_critic"
+
+    def clear(self) -> None:
+        self.hx, self.cx = None, None
+
+    def reset(self, n: int, burnin_observations: Optional[torch.Tensor] = None, mask_padding: Optional[torch.Tensor] = None) -> None:
+        device = self.conv1.weight.device
+        self.hx = torch.zeros(n, self.lstm_dim, device=device)
+        self.cx = torch.zeros(n, self.lstm_dim, device=device)
+        if burnin_observations is not None:
+            assert burnin_observations.ndim == 5 and burnin_observations.size(0) == n and mask_padding is not None and burnin_observations.shape[:2] == mask_padding.shape
+            for i in range(burnin_observations.size(1)):
+                if mask_padding[:, i].any():
+                    with torch.no_grad():
+                        self(burnin_observations[:, i], mask_padding[:, i])
+
+    def prune(self, mask: np.ndarray) -> None:
+        self.hx = self.hx[mask]
+        self.cx = self.cx[mask]
+
+    def forward(self, inputs: torch.FloatTensor, mask_padding: Optional[torch.BoolTensor] = None) -> ActorCriticOutput:
+        assert inputs.ndim == 4 and inputs.shape[1:] == (3, 64, 64)
+        assert 0 <= inputs.min() <= 1 and 0 <= inputs.max() <= 1
+        assert mask_padding is None or (mask_padding.ndim == 1 and mask_padding.size(0) == inputs.size(0) and mask_padding.any())
+        x = inputs[mask_padding] if mask_padding is not None else inputs
+
+        x = x.mul(2).sub(1)
+        x = F.relu(self.maxp1(self.conv1(x)))
+        x = F.relu(self.maxp2(self.conv2(x)))
+        x = F.relu(self.maxp3(self.conv3(x)))
+        x = F.relu(self.maxp4(self.conv4(x)))
+        x = torch.flatten(x, start_dim=1)
+
+        if mask_padding is None:
+            self.hx, self.cx = self.lstm(x, (self.hx, self.cx))
+        else:
+            self.hx[mask_padding], self.cx[mask_padding] = self.lstm(x, (self.hx[mask_padding], self.cx[mask_padding]))
+
+        logits_actions = rearrange(self.actor_linear(self.hx), 'b a -> b 1 a')
+        means_values = rearrange(self.critic_linear(self.hx), 'b 1 -> b 1 1')
+
+        return ActorCriticOutput(logits_actions, means_values)
+
+    def compute_loss(self, batch: Batch, tokenizer: Tokenizer, world_model: WorldModel, imagine_horizon: int, gamma: float, lambda_: float, entropy_weight: float, **kwargs: Any) -> LossWithIntermediateLosses:
+        assert not self.use_original_obs
+        outputs = self.imagine(batch, tokenizer, world_model, horizon=imagine_horizon)
+
+        with torch.no_grad():
+            lambda_returns = compute_lambda_returns(
+                rewards=outputs.rewards,
+                values=outputs.values,
+                ends=outputs.ends,
+                gamma=gamma,
+                lambda_=lambda_,
+            )[:, :-1]
+
+        values = outputs.values[:, :-1]
+
+        d = Categorical(logits=outputs.logits_actions[:, :-1])
+        log_probs = d.log_prob(outputs.actions[:, :-1])
+        loss_actions = -1 * (log_probs * (lambda_returns - values.detach())).mean()
+        loss_entropy = - entropy_weight * d.entropy().mean()
+        loss_values = F.mse_loss(values, lambda_returns)
+
+        return LossWithIntermediateLosses(loss_actions=loss_actions, loss_values=loss_values, loss_entropy=loss_entropy)
+
+    def imagine(self, batch: Batch, tokenizer: Tokenizer, world_model: WorldModel, horizon: int, show_pbar: bool = False) -> ImagineOutput:
+        assert not self.use_original_obs
+        initial_observations = batch['observations']
+        mask_padding = batch['mask_padding']
+        assert initial_observations.ndim == 5 and initial_observations.shape[2:] == (3, 64, 64)
+        assert mask_padding[:, -1].all()
+        device = initial_observations.device
+        wm_env = WorldModelEnv(tokenizer, world_model, device)
+
+        all_actions = []
+        all_logits_actions = []
+        all_values = []
+        all_rewards = []
+        all_ends = []
+        all_observations = []
+
+        burnin_observations = torch.clamp(tokenizer.encode_decode(initial_observations[:, :-1], should_preprocess=True, should_postprocess=True), 0, 1) if initial_observations.size(1) > 1 else None
+        self.reset(n=initial_observations.size(0), burnin_observations=burnin_observations, mask_padding=mask_padding[:, :-1])
+
+        obs = wm_env.reset_from_initial_observations(initial_observations[:, -1])
+        for k in tqdm(range(horizon), disable=not show_pbar, desc='Imagination', file=sys.stdout):
+
+            all_observations.append(obs)
+
+            outputs_ac = self(obs)
+            action_token = Categorical(logits=outputs_ac.logits_actions).sample()
+            obs, reward, done, _ = wm_env.step(action_token, should_predict_next_obs=(k < horizon - 1))
+
+            all_actions.append(action_token)
+            all_logits_actions.append(outputs_ac.logits_actions)
+            all_values.append(outputs_ac.means_values)
+            all_rewards.append(torch.tensor(reward).reshape(-1, 1))
+            all_ends.append(torch.tensor(done).reshape(-1, 1))
+
+        self.clear()
+
+        return ImagineOutput(
+            observations=torch.stack(all_observations, dim=1).mul(255).byte(),      # (B, T, C, H, W) in [0, 255]
+            actions=torch.cat(all_actions, dim=1),                                  # (B, T)
+            logits_actions=torch.cat(all_logits_actions, dim=1),                    # (B, T, #actions)
+            values=rearrange(torch.cat(all_values, dim=1), 'b t 1 -> b t'),         # (B, T)
+            rewards=torch.cat(all_rewards, dim=1).to(device),                       # (B, T)
+            ends=torch.cat(all_ends, dim=1).to(device),                             # (B, T)
+        )

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, num_heads: int, max_tokens: int, embed_dim: int, device: torch.device) -> None:
+        assert embed_dim % num_heads == 0
+        self._n, self._cache, self._size = num_samples, None, None
+        self._reset = lambda n: torch.empty(n, num_heads, max_tokens, embed_dim // num_heads, device=device)  # (B, nh, T, hs)
+        self.reset()
+
+    @property
+    def shape(self) -> Tuple[int, int, int, int]:
+        n, num_heads, _, head_dim = self._cache.shape
+        return n, num_heads, self._size, head_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, 1, 3)])
+        assert self._size + x.size(2) <= self._cache.shape[2]
+        self._cache = AssignWithoutInplaceCheck.apply(self._cache, x, 2, self._size, self._size + x.size(2))
+        self._size += x.size(2)
+
+
+class KVCache:
+    def __init__(self, n: int, num_heads: int, max_tokens: int, embed_dim: int, device: torch.device) -> None:
+        self._k_cache = Cache(n, num_heads, max_tokens, embed_dim, device)
+        self._v_cache = Cache(n, num_heads, max_tokens, embed_dim, device)
+
+    @property
+    def shape(self) -> Tuple[int, 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, num_heads: int, max_tokens: int, embed_dim: int, num_layers: int, device: torch.device) -> None:
+        self._keys_values = tuple([KVCache(n, num_heads, 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[2]
+
+    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

src/models/slicer.py

@@ -0,0 +1,54 @@
+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.Tensor, num_steps: int, prev_steps: int) -> torch.Tensor:
+        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

src/models/tokenizer/__init__.py

@@ -0,0 +1,2 @@
+from .nets import Encoder, Decoder, EncoderDecoderConfig
+from .tokenizer import Tokenizer, TokenizerEncoderOutput

src/models/tokenizer/lpips.py

@@ -0,0 +1,167 @@
+"""
+Credits to https://github.com/CompVis/taming-transformers
+"""
+
+from collections import namedtuple
+import hashlib
+import os
+from pathlib import Path
+import requests
+
+import torch
+import torch.nn as nn
+from torchvision import models
+from tqdm import tqdm
+
+
+class LPIPS(nn.Module):
+    # Learned perceptual metric
+    def __init__(self, use_dropout: bool = True):
+        super().__init__()
+        self.scaling_layer = ScalingLayer()
+        self.chns = [64, 128, 256, 512, 512]  # vg16 features
+        self.net = vgg16(pretrained=True, requires_grad=False)
+        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+        self.load_from_pretrained()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def load_from_pretrained(self) -> None:
+        ckpt = get_ckpt_path(name="vgg_lpips", root=Path.home() / ".cache/iris/tokenizer_pretrained_vgg")  # Download VGG if necessary
+        self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
+        outs0, outs1 = self.net(in0_input), self.net(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+        for kk in range(len(self.chns)):
+            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
+
+        res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))]
+        val = res[0]
+        for i in range(1, len(self.chns)):
+            val += res[i]
+        return val
+
+
+class ScalingLayer(nn.Module):
+    def __init__(self) -> None:
+        super(ScalingLayer, self).__init__()
+        self.register_buffer('shift', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
+        self.register_buffer('scale', torch.Tensor([.458, .448, .450])[None, :, None, None])
+
+    def forward(self, inp: torch.Tensor) -> torch.Tensor:
+        return (inp - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+    """ A single linear layer which does a 1x1 conv """
+    def __init__(self, chn_in: int, chn_out: int = 1, use_dropout: bool = False) -> None:
+        super(NetLinLayer, self).__init__()
+        layers = [nn.Dropout(), ] if (use_dropout) else []
+        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False), ]
+        self.model = nn.Sequential(*layers)
+
+
+class vgg16(torch.nn.Module):
+    def __init__(self, requires_grad: bool = False, pretrained: bool = True) -> None:
+        super(vgg16, self).__init__()
+        vgg_pretrained_features = models.vgg16(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
+        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
+        return out
+
+
+def normalize_tensor(x: torch.Tensor, eps: float = 1e-10) -> torch.Tensor:
+    norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True))
+    return x / (norm_factor + eps)
+
+
+def spatial_average(x: torch.Tensor, keepdim: bool = True) -> torch.Tensor:
+    return x.mean([2, 3], keepdim=keepdim)
+
+
+# ********************************************************************
+# *************** Utilities to download pretrained vgg ***************
+# ********************************************************************
+
+
+URL_MAP = {
+    "vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1"
+}
+
+
+CKPT_MAP = {
+    "vgg_lpips": "vgg.pth"
+}
+
+
+MD5_MAP = {
+    "vgg_lpips": "d507d7349b931f0638a25a48a722f98a"
+}
+
+
+def download(url: str, local_path: str, chunk_size: int = 1024) -> None:
+    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+    with requests.get(url, stream=True) as r:
+        total_size = int(r.headers.get("content-length", 0))
+        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+            with open(local_path, "wb") as f:
+                for data in r.iter_content(chunk_size=chunk_size):
+                    if data:
+                        f.write(data)
+                        pbar.update(chunk_size)
+
+
+def md5_hash(path: str) -> str:
+    with open(path, "rb") as f:
+        content = f.read()
+    return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name: str, root: str, check: bool = False) -> str:
+    assert name in URL_MAP
+    path = os.path.join(root, CKPT_MAP[name])
+    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
+        download(URL_MAP[name], path)
+        md5 = md5_hash(path)
+        assert md5 == MD5_MAP[name], md5
+    return path

src/models/tokenizer/nets.py

@@ -0,0 +1,362 @@
+"""
+Credits to https://github.com/CompVis/taming-transformers
+"""
+
+from dataclasses import dataclass
+from typing import List
+
+import torch
+import torch.nn as nn
+
+
+@dataclass
+class EncoderDecoderConfig:
+    resolution: int
+    in_channels: int
+    z_channels: int
+    ch: int
+    ch_mult: List[int]
+    num_res_blocks: int
+    attn_resolutions: List[int]
+    out_ch: int
+    dropout: float
+
+
+class Encoder(nn.Module):
+    def __init__(self, config: EncoderDecoderConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.num_resolutions = len(config.ch_mult)
+        temb_ch = 0  # timestep embedding #channels
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(config.in_channels,
+                                       config.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = config.resolution
+        in_ch_mult = (1,) + tuple(config.ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = config.ch * in_ch_mult[i_level]
+            block_out = config.ch * config.ch_mult[i_level]
+            for i_block in range(self.config.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=temb_ch,
+                                         dropout=config.dropout))
+                block_in = block_out
+                if curr_res in config.attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in, with_conv=True)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=temb_ch,
+                                       dropout=config.dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=temb_ch,
+                                       dropout=config.dropout)
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        config.z_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        temb = None  # timestep embedding
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.config.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(self, config: EncoderDecoderConfig) -> None:
+        super().__init__()
+        self.config = config
+        temb_ch = 0
+        self.num_resolutions = len(config.ch_mult)
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        in_ch_mult = (1,) + tuple(config.ch_mult)
+        block_in = config.ch * config.ch_mult[self.num_resolutions - 1]
+        curr_res = config.resolution // 2 ** (self.num_resolutions - 1)
+        print(f"Tokenizer : shape of latent is {config.z_channels, curr_res, curr_res}.")
+
+        # z to block_in
+        self.conv_in = torch.nn.Conv2d(config.z_channels,
+                                       block_in,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=temb_ch,
+                                       dropout=config.dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=temb_ch,
+                                       dropout=config.dropout)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = config.ch * config.ch_mult[i_level]
+            for i_block in range(config.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=temb_ch,
+                                         dropout=config.dropout))
+                block_in = block_out
+                if curr_res in config.attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, with_conv=True)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        config.out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        temb = None  # timestep embedding
+
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.config.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h, temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+def nonlinearity(x: torch.Tensor) -> torch.Tensor:
+    # swish
+    return x * torch.sigmoid(x)
+
+
+def Normalize(in_channels: int) -> nn.Module:
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int, with_conv: bool) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int, with_conv: bool) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.with_conv:
+            pad = (0, 1, 0, 1)
+            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, *, in_channels: int, out_channels: int = None, conv_shortcut: bool = False,
+                 dropout: float, temb_channels: int = 512) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = torch.nn.Conv2d(in_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels,
+                                             out_channels)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(out_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(in_channels,
+                                                     out_channels,
+                                                     kernel_size=3,
+                                                     stride=1,
+                                                     padding=1)
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(in_channels,
+                                                    out_channels,
+                                                    kernel_size=1,
+                                                    stride=1,
+                                                    padding=0)
+
+    def forward(self, x: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.k = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.v = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=1,
+                                        stride=1,
+                                        padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h * w)
+        q = q.permute(0, 2, 1)      # b,hw,c
+        k = k.reshape(b, c, h * w)  # b,c,hw
+        w_ = torch.bmm(q, k)        # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b, c, h * w)
+        w_ = w_.permute(0, 2, 1)   # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v, w_)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b, c, h, w)
+
+        h_ = self.proj_out(h_)
+
+        return x + h_

src/models/tokenizer/tokenizer.py

@@ -0,0 +1,104 @@
+"""
+Credits to https://github.com/CompVis/taming-transformers
+"""
+
+from dataclasses import dataclass
+from typing import Any, Tuple
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+
+from dataset import Batch
+from .lpips import LPIPS
+from .nets import Encoder, Decoder
+from utils import LossWithIntermediateLosses
+
+
+@dataclass
+class TokenizerEncoderOutput:
+    z: torch.FloatTensor
+    z_quantized: torch.FloatTensor
+    tokens: torch.LongTensor
+
+
+class Tokenizer(nn.Module):
+    def __init__(self, vocab_size: int, embed_dim: int, encoder: Encoder, decoder: Decoder, with_lpips: bool = True) -> None:
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.encoder = encoder
+        self.pre_quant_conv = torch.nn.Conv2d(encoder.config.z_channels, embed_dim, 1)
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.post_quant_conv = torch.nn.Conv2d(embed_dim, decoder.config.z_channels, 1)
+        self.decoder = decoder
+        self.embedding.weight.data.uniform_(-1.0 / vocab_size, 1.0 / vocab_size)
+        self.lpips = LPIPS().eval() if with_lpips else None
+
+    def __repr__(self) -> str:
+        return "tokenizer"
+
+    def forward(self, x: torch.Tensor, should_preprocess: bool = False, should_postprocess: bool = False) -> Tuple[torch.Tensor]:
+        outputs = self.encode(x, should_preprocess)
+        decoder_input = outputs.z + (outputs.z_quantized - outputs.z).detach()
+        reconstructions = self.decode(decoder_input, should_postprocess)
+        return outputs.z, outputs.z_quantized, reconstructions
+
+    def compute_loss(self, batch: Batch, **kwargs: Any) -> LossWithIntermediateLosses:
+        assert self.lpips is not None
+        observations = self.preprocess_input(rearrange(batch['observations'], 'b t c h w -> (b t) c h w'))
+        z, z_quantized, reconstructions = self(observations, should_preprocess=False, should_postprocess=False)
+
+        # Codebook loss. Notes:
+        # - beta position is different from taming and identical to original VQVAE paper
+        # - VQVAE uses 0.25 by default
+        beta = 1.0
+        commitment_loss = (z.detach() - z_quantized).pow(2).mean() + beta * (z - z_quantized.detach()).pow(2).mean()
+
+        reconstruction_loss = torch.abs(observations - reconstructions).mean()
+        perceptual_loss = torch.mean(self.lpips(observations, reconstructions))
+
+        return LossWithIntermediateLosses(commitment_loss=commitment_loss, reconstruction_loss=reconstruction_loss, perceptual_loss=perceptual_loss)
+
+    def encode(self, x: torch.Tensor, should_preprocess: bool = False) -> TokenizerEncoderOutput:
+        if should_preprocess:
+            x = self.preprocess_input(x)
+        shape = x.shape  # (..., C, H, W)
+        x = x.view(-1, *shape[-3:])
+        z = self.encoder(x)
+        z = self.pre_quant_conv(z)
+        b, e, h, w = z.shape
+        z_flattened = rearrange(z, 'b e h w -> (b h w) e')
+        dist_to_embeddings = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + torch.sum(self.embedding.weight**2, dim=1) - 2 * torch.matmul(z_flattened, self.embedding.weight.t())
+
+        tokens = dist_to_embeddings.argmin(dim=-1)
+        z_q = rearrange(self.embedding(tokens), '(b h w) e -> b e h w', b=b, e=e, h=h, w=w).contiguous()
+
+        # Reshape to original
+        z = z.reshape(*shape[:-3], *z.shape[1:])
+        z_q = z_q.reshape(*shape[:-3], *z_q.shape[1:])
+        tokens = tokens.reshape(*shape[:-3], -1)
+
+        return TokenizerEncoderOutput(z, z_q, tokens)
+
+    def decode(self, z_q: torch.Tensor, should_postprocess: bool = False) -> torch.Tensor:
+        shape = z_q.shape  # (..., E, h, w)
+        z_q = z_q.view(-1, *shape[-3:])
+        z_q = self.post_quant_conv(z_q)
+        rec = self.decoder(z_q)
+        rec = rec.reshape(*shape[:-3], *rec.shape[1:])
+        if should_postprocess:
+            rec = self.postprocess_output(rec)
+        return rec
+
+    @torch.no_grad()
+    def encode_decode(self, x: torch.Tensor, should_preprocess: bool = False, should_postprocess: bool = False) -> torch.Tensor:
+        z_q = self.encode(x, should_preprocess).z_quantized
+        return self.decode(z_q, should_postprocess)
+
+    def preprocess_input(self, x: torch.Tensor) -> torch.Tensor:
+        """x is supposed to be channels first and in [0, 1]"""
+        return x.mul(2).sub(1)
+
+    def postprocess_output(self, y: torch.Tensor) -> torch.Tensor:
+        """y is supposed to be channels first and in [-1, 1]"""
+        return y.add(1).div(2)

src/models/transformer.py

@@ -0,0 +1,120 @@
+"""
+Credits to https://github.com/karpathy/minGPT
+"""
+
+from dataclasses import dataclass
+import math
+from typing import Optional
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from .kv_caching import KeysValues, KVCache
+
+
+@dataclass
+class TransformerConfig:
+    tokens_per_block: int
+    max_blocks: int
+    attention: str
+
+    num_layers: int
+    num_heads: int
+    embed_dim: int
+
+    embed_pdrop: float
+    resid_pdrop: float
+    attn_pdrop: float
+
+    @property
+    def max_tokens(self):
+        return self.tokens_per_block * self.max_blocks
+
+
+class Transformer(nn.Module):
+    def __init__(self, config: TransformerConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.drop = nn.Dropout(config.embed_pdrop)
+        self.blocks = nn.ModuleList([Block(config) for _ in range(config.num_layers)])
+        self.ln_f = nn.LayerNorm(config.embed_dim)
+
+    def generate_empty_keys_values(self, n: int, max_tokens: int) -> KeysValues:
+        device = self.ln_f.weight.device  # Assumption that all submodules are on the same device
+        return KeysValues(n, self.config.num_heads, max_tokens, self.config.embed_dim, self.config.num_layers, device)
+
+    def forward(self, sequences: torch.Tensor, past_keys_values: Optional[KeysValues] = None) -> torch.Tensor:
+        assert past_keys_values is None or len(past_keys_values) == len(self.blocks)
+        x = self.drop(sequences)
+        for i, block in enumerate(self.blocks):
+            x = block(x, None if past_keys_values is None else past_keys_values[i])
+
+        x = self.ln_f(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, config: TransformerConfig) -> None:
+        super().__init__()
+        self.ln1 = nn.LayerNorm(config.embed_dim)
+        self.ln2 = nn.LayerNorm(config.embed_dim)
+        self.attn = SelfAttention(config)
+        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, x: torch.Tensor, past_keys_values: Optional[KeysValues] = None) -> torch.Tensor:
+        x_attn = self.attn(self.ln1(x), past_keys_values)
+        x = x + x_attn
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, config: TransformerConfig) -> None:
+        super().__init__()
+        assert config.embed_dim % config.num_heads == 0
+        assert config.attention in ('causal', 'block_causal')
+        self.num_heads = config.num_heads
+        self.key = nn.Linear(config.embed_dim, config.embed_dim)
+        self.query = nn.Linear(config.embed_dim, config.embed_dim)
+        self.value = nn.Linear(config.embed_dim, config.embed_dim)
+        self.attn_drop = nn.Dropout(config.attn_pdrop)
+        self.resid_drop = nn.Dropout(config.resid_pdrop)
+        self.proj = nn.Linear(config.embed_dim, config.embed_dim)
+
+        causal_mask = torch.tril(torch.ones(config.max_tokens, config.max_tokens))
+        block_causal_mask = torch.max(causal_mask, torch.block_diag(*[torch.ones(config.tokens_per_block, config.tokens_per_block) for _ in range(config.max_blocks)]))
+        self.register_buffer('mask', causal_mask if config.attention == 'causal' else block_causal_mask)
+
+    def forward(self, x: torch.Tensor, kv_cache: Optional[KVCache] = None) -> torch.Tensor:
+        B, T, C = x.size()
+        if kv_cache is not None:
+            b, nh, L, c = kv_cache.shape
+            assert nh == self.num_heads and b == B and c * nh == C
+        else:
+            L = 0
+
+        q = self.query(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)   # (B, nh, T, hs)
+        k = self.key(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)     # (B, nh, T, hs)
+        v = self.value(x).view(B, T, self.num_heads, C // self.num_heads).transpose(1, 2)   # (B, nh, T, hs)
+
+        if kv_cache is not None:
+            kv_cache.update(k, v)
+            k, v = kv_cache.get()
+
+        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+        att = att.masked_fill(self.mask[L:L + T, :L + T] == 0, float('-inf'))
+        att = F.softmax(att, dim=-1)
+        att = self.attn_drop(att)
+        y = att @ v
+        y = rearrange(y, 'b h t e -> b t (h e)')
+
+        y = self.resid_drop(self.proj(y))
+
+        return y

src/models/world_model.py

@@ -0,0 +1,122 @@
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from dataset import Batch
+from .kv_caching import KeysValues
+from .slicer import Embedder, Head
+from .tokenizer import Tokenizer
+from .transformer import Transformer, TransformerConfig
+from utils import init_weights, LossWithIntermediateLosses
+
+
+@dataclass
+class WorldModelOutput:
+    output_sequence: torch.FloatTensor
+    logits_observations: torch.FloatTensor
+    logits_rewards: torch.FloatTensor
+    logits_ends: torch.FloatTensor
+
+
+class WorldModel(nn.Module):
+    def __init__(self, obs_vocab_size: int, act_vocab_size: int, config: TransformerConfig) -> 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)
+            )
+        )
+
+        self.apply(init_weights)
+
+    def __repr__(self) -> str:
+        return "world_model"
+
+    def forward(self, tokens: torch.LongTensor, past_keys_values: Optional[KeysValues] = None) -> WorldModelOutput:
+
+        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 WorldModelOutput(x, logits_observations, logits_rewards, logits_ends)
+
+    def compute_loss(self, batch: Batch, tokenizer: Tokenizer, **kwargs: Any) -> LossWithIntermediateLosses:
+
+        with torch.no_grad():
+            obs_tokens = tokenizer.encode(batch['observations'], should_preprocess=True).tokens  # (BL, K)
+
+        act_tokens = rearrange(batch['actions'], 'b l -> b l 1')
+        tokens = rearrange(torch.cat((obs_tokens, act_tokens), dim=2), 'b l k1 -> b (l k1)')  # (B, L(K+1))
+
+        outputs = self(tokens)
+
+        labels_observations, labels_rewards, labels_ends = self.compute_labels_world_model(obs_tokens, batch['rewards'], batch['ends'], batch['mask_padding'])
+
+        logits_observations = rearrange(outputs.logits_observations[:, :-1], 'b t o -> (b t) o')
+        loss_obs = F.cross_entropy(logits_observations, labels_observations)
+        loss_rewards = F.cross_entropy(rearrange(outputs.logits_rewards, 'b t e -> (b t) e'), labels_rewards)
+        loss_ends = F.cross_entropy(rearrange(outputs.logits_ends, 'b t e -> (b t) e'), labels_ends)
+
+        return LossWithIntermediateLosses(loss_obs=loss_obs, loss_rewards=loss_rewards, loss_ends=loss_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)

src/play.py

@@ -0,0 +1,44 @@
+from pathlib import Path
+
+import hydra
+from hydra.utils import instantiate
+from omegaconf import DictConfig
+import torch
+
+from agent import Agent
+from envs import SingleProcessEnv, WorldModelEnv
+from game import AgentEnv, EpisodeReplayEnv, Game
+from models.actor_critic import ActorCritic
+from models.world_model import WorldModel
+
+
+@hydra.main(config_path="../config", config_name="trainer")
+def main(cfg: DictConfig):
+    device = torch.device(cfg.common.device)
+    assert cfg.mode in ('world_model', 'episode_replay', 'agent')
+
+    if cfg.mode in ['world_model', 'agent']:
+        env_fn = lambda: instantiate(cfg.env.test)
+        test_env = SingleProcessEnv(env_fn)
+        tokenizer = instantiate(cfg.tokenizer)
+        world_model = WorldModel(obs_vocab_size=tokenizer.vocab_size, act_vocab_size=test_env.num_actions, config=instantiate(cfg.world_model))
+        if cfg.mode == 'world_model':
+            env = WorldModelEnv(tokenizer=tokenizer, world_model=world_model, pretrained_agent_path=Path('checkpoints/last.pt'), device=device, env=env_fn())
+            keymap = cfg.env.keymap
+        else:
+            actor_critic = ActorCritic(**cfg.actor_critic, act_vocab_size=test_env.num_actions)
+            agent = Agent(tokenizer, world_model, actor_critic).to(device)
+            agent.load(Path('checkpoints/last.pt'), device)
+            env = AgentEnv(agent, test_env, cfg.env.keymap)
+            keymap = 'empty'
+
+    else:
+        env = EpisodeReplayEnv(replay_keymap_name=cfg.env.keymap, episode_dir=Path('media/episodes'))
+        keymap = 'episode_replay'
+
+    game = Game(env, keymap_name=keymap, size=(600, 1200 if cfg.mode == 'agent' else 600), fps=cfg.fps, verbose=bool(cfg.header))
+    game.run()
+
+
+if __name__ == "__main__":
+    main()

src/trainer.py

@@ -0,0 +1,283 @@
+from collections import defaultdict
+from pathlib import Path
+import shutil
+import sys
+import time
+from typing import Any, Dict, Optional, Tuple
+
+import hydra
+from hydra.utils import instantiate
+from omegaconf import DictConfig, OmegaConf
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+import wandb
+
+from agent import Agent
+from collector import Collector
+from envs import SingleProcessEnv, MultiProcessEnv
+from episode import Episode
+from make_reconstructions import make_reconstructions_from_batch
+from models.actor_critic import ActorCritic
+from models.world_model import WorldModel
+from utils import configure_optimizer, EpisodeDirManager, set_seed
+
+
+class Trainer:
+    def __init__(self, cfg: DictConfig) -> None:
+        wandb.init(
+            config=OmegaConf.to_container(cfg, resolve=True),
+            reinit=True,
+            resume=True,
+            **cfg.wandb
+        )
+
+        if cfg.common.seed is not None:
+            set_seed(cfg.common.seed)
+
+        self.cfg = cfg
+        self.start_epoch = 1
+        self.device = torch.device(cfg.common.device)
+
+        self.ckpt_dir = Path('checkpoints')
+        self.media_dir = Path('media')
+        self.episode_dir = self.media_dir / 'episodes'
+        self.reconstructions_dir = self.media_dir / 'reconstructions'
+
+        if not cfg.common.resume:
+            config_dir = Path('config')
+            config_path = config_dir / 'trainer.yaml'
+            config_dir.mkdir(exist_ok=False, parents=False)
+            shutil.copy('.hydra/config.yaml', config_path)
+            wandb.save(str(config_path))
+            shutil.copytree(src=(Path(hydra.utils.get_original_cwd()) / "src"), dst="./src")
+            shutil.copytree(src=(Path(hydra.utils.get_original_cwd()) / "scripts"), dst="./scripts")
+            self.ckpt_dir.mkdir(exist_ok=False, parents=False)
+            self.media_dir.mkdir(exist_ok=False, parents=False)
+            self.episode_dir.mkdir(exist_ok=False, parents=False)
+            self.reconstructions_dir.mkdir(exist_ok=False, parents=False)
+
+        episode_manager_train = EpisodeDirManager(self.episode_dir / 'train', max_num_episodes=cfg.collection.train.num_episodes_to_save)
+        episode_manager_test = EpisodeDirManager(self.episode_dir / 'test', max_num_episodes=cfg.collection.test.num_episodes_to_save)
+        self.episode_manager_imagination = EpisodeDirManager(self.episode_dir / 'imagination', max_num_episodes=cfg.evaluation.actor_critic.num_episodes_to_save)
+
+        def create_env(cfg_env, num_envs):
+            env_fn = lambda: instantiate(cfg_env)
+            return MultiProcessEnv(env_fn, num_envs, should_wait_num_envs_ratio=1.0) if num_envs > 1 else SingleProcessEnv(env_fn)
+
+        if self.cfg.training.should:
+            train_env = create_env(cfg.env.train, cfg.collection.train.num_envs)
+            self.train_dataset = instantiate(cfg.datasets.train)
+            self.train_collector = Collector(train_env, self.train_dataset, episode_manager_train)
+
+        if self.cfg.evaluation.should:
+            test_env = create_env(cfg.env.test, cfg.collection.test.num_envs)
+            self.test_dataset = instantiate(cfg.datasets.test)
+            self.test_collector = Collector(test_env, self.test_dataset, episode_manager_test)
+
+        assert self.cfg.training.should or self.cfg.evaluation.should
+        env = train_env if self.cfg.training.should else test_env
+
+        tokenizer = instantiate(cfg.tokenizer)
+        world_model = WorldModel(obs_vocab_size=tokenizer.vocab_size, act_vocab_size=env.num_actions, config=instantiate(cfg.world_model))
+        actor_critic = ActorCritic(**cfg.actor_critic, act_vocab_size=env.num_actions)
+        self.agent = Agent(tokenizer, world_model, actor_critic).to(self.device)
+        print(f'{sum(p.numel() for p in self.agent.tokenizer.parameters())} parameters in agent.tokenizer')
+        print(f'{sum(p.numel() for p in self.agent.world_model.parameters())} parameters in agent.world_model')
+        print(f'{sum(p.numel() for p in self.agent.actor_critic.parameters())} parameters in agent.actor_critic')
+
+        self.optimizer_tokenizer = torch.optim.Adam(self.agent.tokenizer.parameters(), lr=cfg.training.learning_rate)
+        self.optimizer_world_model = configure_optimizer(self.agent.world_model, cfg.training.learning_rate, cfg.training.world_model.weight_decay)
+        self.optimizer_actor_critic = torch.optim.Adam(self.agent.actor_critic.parameters(), lr=cfg.training.learning_rate)
+
+        if cfg.initialization.path_to_checkpoint is not None:
+            self.agent.load(**cfg.initialization, device=self.device)
+
+        if cfg.common.resume:
+            self.load_checkpoint()
+
+    def run(self) -> None:
+
+        for epoch in range(self.start_epoch, 1 + self.cfg.common.epochs):
+
+            print(f"\nEpoch {epoch} / {self.cfg.common.epochs}\n")
+            start_time = time.time()
+            to_log = []
+
+            if self.cfg.training.should:
+                if epoch <= self.cfg.collection.train.stop_after_epochs:
+                    to_log += self.train_collector.collect(self.agent, epoch, **self.cfg.collection.train.config)
+                to_log += self.train_agent(epoch)
+
+            if self.cfg.evaluation.should and (epoch % self.cfg.evaluation.every == 0):
+                self.test_dataset.clear()
+                to_log += self.test_collector.collect(self.agent, epoch, **self.cfg.collection.test.config)
+                to_log += self.eval_agent(epoch)
+
+            if self.cfg.training.should:
+                self.save_checkpoint(epoch, save_agent_only=not self.cfg.common.do_checkpoint)
+
+            to_log.append({'duration': (time.time() - start_time) / 3600})
+            for metrics in to_log:
+                wandb.log({'epoch': epoch, **metrics})
+
+        self.finish()
+
+    def train_agent(self, epoch: int) -> None:
+        self.agent.train()
+        self.agent.zero_grad()
+
+        metrics_tokenizer, metrics_world_model, metrics_actor_critic = {}, {}, {}
+
+        cfg_tokenizer = self.cfg.training.tokenizer
+        cfg_world_model = self.cfg.training.world_model
+        cfg_actor_critic = self.cfg.training.actor_critic
+
+        w = self.cfg.training.sampling_weights
+
+        if epoch > cfg_tokenizer.start_after_epochs:
+            metrics_tokenizer = self.train_component(self.agent.tokenizer, self.optimizer_tokenizer, sequence_length=1, sample_from_start=True, sampling_weights=w, **cfg_tokenizer)
+        self.agent.tokenizer.eval()
+
+        if epoch > cfg_world_model.start_after_epochs:
+            metrics_world_model = self.train_component(self.agent.world_model, self.optimizer_world_model, sequence_length=self.cfg.common.sequence_length, sample_from_start=True, sampling_weights=w, tokenizer=self.agent.tokenizer, **cfg_world_model)
+        self.agent.world_model.eval()
+
+        if epoch > cfg_actor_critic.start_after_epochs:
+            metrics_actor_critic = self.train_component(self.agent.actor_critic, self.optimizer_actor_critic, sequence_length=1 + self.cfg.training.actor_critic.burn_in, sample_from_start=False, sampling_weights=w, tokenizer=self.agent.tokenizer, world_model=self.agent.world_model, **cfg_actor_critic)
+        self.agent.actor_critic.eval()
+
+        return [{'epoch': epoch, **metrics_tokenizer, **metrics_world_model, **metrics_actor_critic}]
+
+    def train_component(self, component: nn.Module, optimizer: torch.optim.Optimizer, steps_per_epoch: int, batch_num_samples: int, grad_acc_steps: int, max_grad_norm: Optional[float], sequence_length: int, sampling_weights: Optional[Tuple[float]], sample_from_start: bool, **kwargs_loss: Any) -> Dict[str, float]:
+        loss_total_epoch = 0.0
+        intermediate_losses = defaultdict(float)
+
+        for _ in tqdm(range(steps_per_epoch), desc=f"Training {str(component)}", file=sys.stdout):
+            optimizer.zero_grad()
+            for _ in range(grad_acc_steps):
+                batch = self.train_dataset.sample_batch(batch_num_samples, sequence_length, sampling_weights, sample_from_start)
+                batch = self._to_device(batch)
+
+                losses = component.compute_loss(batch, **kwargs_loss) / grad_acc_steps
+                loss_total_step = losses.loss_total
+                loss_total_step.backward()
+                loss_total_epoch += loss_total_step.item() / steps_per_epoch
+
+                for loss_name, loss_value in losses.intermediate_losses.items():
+                    intermediate_losses[f"{str(component)}/train/{loss_name}"] += loss_value / steps_per_epoch
+
+            if max_grad_norm is not None:
+                torch.nn.utils.clip_grad_norm_(component.parameters(), max_grad_norm)
+
+            optimizer.step()
+
+        metrics = {f'{str(component)}/train/total_loss': loss_total_epoch, **intermediate_losses}
+        return metrics
+
+    @torch.no_grad()
+    def eval_agent(self, epoch: int) -> None:
+        self.agent.eval()
+
+        metrics_tokenizer, metrics_world_model = {}, {}
+
+        cfg_tokenizer = self.cfg.evaluation.tokenizer
+        cfg_world_model = self.cfg.evaluation.world_model
+        cfg_actor_critic = self.cfg.evaluation.actor_critic
+
+        if epoch > cfg_tokenizer.start_after_epochs:
+            metrics_tokenizer = self.eval_component(self.agent.tokenizer, cfg_tokenizer.batch_num_samples, sequence_length=1)
+
+        if epoch > cfg_world_model.start_after_epochs:
+            metrics_world_model = self.eval_component(self.agent.world_model, cfg_world_model.batch_num_samples, sequence_length=self.cfg.common.sequence_length, tokenizer=self.agent.tokenizer)
+
+        if epoch > cfg_actor_critic.start_after_epochs:
+            self.inspect_imagination(epoch)
+
+        if cfg_tokenizer.save_reconstructions:
+            batch = self._to_device(self.test_dataset.sample_batch(batch_num_samples=3, sequence_length=self.cfg.common.sequence_length))
+            make_reconstructions_from_batch(batch, save_dir=self.reconstructions_dir, epoch=epoch, tokenizer=self.agent.tokenizer)
+
+        return [metrics_tokenizer, metrics_world_model]
+
+    @torch.no_grad()
+    def eval_component(self, component: nn.Module, batch_num_samples: int, sequence_length: int, **kwargs_loss: Any) -> Dict[str, float]:
+        loss_total_epoch = 0.0
+        intermediate_losses = defaultdict(float)
+
+        steps = 0
+        pbar = tqdm(desc=f"Evaluating {str(component)}", file=sys.stdout)
+        for batch in self.test_dataset.traverse(batch_num_samples, sequence_length):
+            batch = self._to_device(batch)
+
+            losses = component.compute_loss(batch, **kwargs_loss)
+            loss_total_epoch += losses.loss_total.item()
+
+            for loss_name, loss_value in losses.intermediate_losses.items():
+                intermediate_losses[f"{str(component)}/eval/{loss_name}"] += loss_value
+
+            steps += 1
+            pbar.update(1)
+
+        intermediate_losses = {k: v / steps for k, v in intermediate_losses.items()}
+        metrics = {f'{str(component)}/eval/total_loss': loss_total_epoch / steps, **intermediate_losses}
+        return metrics
+
+    @torch.no_grad()
+    def inspect_imagination(self, epoch: int) -> None:
+        mode_str = 'imagination'
+        batch = self.test_dataset.sample_batch(batch_num_samples=self.episode_manager_imagination.max_num_episodes, sequence_length=1 + self.cfg.training.actor_critic.burn_in, sample_from_start=False)
+        outputs = self.agent.actor_critic.imagine(self._to_device(batch), self.agent.tokenizer, self.agent.world_model, horizon=self.cfg.evaluation.actor_critic.horizon, show_pbar=True)
+
+        to_log = []
+        for i, (o, a, r, d) in enumerate(zip(outputs.observations.cpu(), outputs.actions.cpu(), outputs.rewards.cpu(), outputs.ends.long().cpu())):  # Make everything (N, T, ...) instead of (T, N, ...)
+            episode = Episode(o, a, r, d, torch.ones_like(d))
+            episode_id = (epoch - 1 - self.cfg.training.actor_critic.start_after_epochs) * outputs.observations.size(0) + i
+            self.episode_manager_imagination.save(episode, episode_id, epoch)
+
+            metrics_episode = {k: v for k, v in episode.compute_metrics().__dict__.items()}
+            metrics_episode['episode_num'] = episode_id
+            metrics_episode['action_histogram'] = wandb.Histogram(episode.actions.numpy(), num_bins=self.agent.world_model.act_vocab_size)
+            to_log.append({f'{mode_str}/{k}': v for k, v in metrics_episode.items()})
+
+        return to_log
+
+    def _save_checkpoint(self, epoch: int, save_agent_only: bool) -> None:
+        torch.save(self.agent.state_dict(), self.ckpt_dir / 'last.pt')
+        if not save_agent_only:
+            torch.save(epoch, self.ckpt_dir / 'epoch.pt')
+            torch.save({
+                "optimizer_tokenizer": self.optimizer_tokenizer.state_dict(),
+                "optimizer_world_model": self.optimizer_world_model.state_dict(),
+                "optimizer_actor_critic": self.optimizer_actor_critic.state_dict(),
+            }, self.ckpt_dir / 'optimizer.pt')
+            ckpt_dataset_dir = self.ckpt_dir / 'dataset'
+            ckpt_dataset_dir.mkdir(exist_ok=True, parents=False)
+            self.train_dataset.update_disk_checkpoint(ckpt_dataset_dir)
+            if self.cfg.evaluation.should:
+                torch.save(self.test_dataset.num_seen_episodes, self.ckpt_dir / 'num_seen_episodes_test_dataset.pt')
+
+    def save_checkpoint(self, epoch: int, save_agent_only: bool) -> None:
+        tmp_checkpoint_dir = Path('checkpoints_tmp')
+        shutil.copytree(src=self.ckpt_dir, dst=tmp_checkpoint_dir, ignore=shutil.ignore_patterns('dataset'))
+        self._save_checkpoint(epoch, save_agent_only)
+        shutil.rmtree(tmp_checkpoint_dir)
+
+    def load_checkpoint(self) -> None:
+        assert self.ckpt_dir.is_dir()
+        self.start_epoch = torch.load(self.ckpt_dir / 'epoch.pt') + 1
+        self.agent.load(self.ckpt_dir / 'last.pt', device=self.device)
+        ckpt_opt = torch.load(self.ckpt_dir / 'optimizer.pt', map_location=self.device)
+        self.optimizer_tokenizer.load_state_dict(ckpt_opt['optimizer_tokenizer'])
+        self.optimizer_world_model.load_state_dict(ckpt_opt['optimizer_world_model'])
+        self.optimizer_actor_critic.load_state_dict(ckpt_opt['optimizer_actor_critic'])
+        self.train_dataset.load_disk_checkpoint(self.ckpt_dir / 'dataset')
+        if self.cfg.evaluation.should:
+            self.test_dataset.num_seen_episodes = torch.load(self.ckpt_dir / 'num_seen_episodes_test_dataset.pt')
+        print(f'Successfully loaded model, optimizer and {len(self.train_dataset)} episodes from {self.ckpt_dir.absolute()}.')
+
+    def _to_device(self, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        return {k: batch[k].to(self.device) for k in batch}
+
+    def finish(self) -> None:
+        wandb.finish()

src/utils.py

@@ -0,0 +1,142 @@
+from collections import OrderedDict
+from pathlib import Path
+import random
+import shutil
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from episode import Episode
+
+
+def configure_optimizer(model, learning_rate, weight_decay, *blacklist_module_names):
+    """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)
+    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 = torch.optim.AdamW(optim_groups, lr=learning_rate)
+    return optimizer
+
+
+def init_weights(module):
+    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, module_name):
+    return OrderedDict({k.split('.', 1)[1]: v for k, v in state_dict.items() if k.startswith(module_name)})
+
+
+def set_seed(seed):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    random.seed(seed)
+
+
+def remove_dir(path, should_ask=False):
+    assert path.is_dir()
+    if (not should_ask) or input(f"Remove directory : {path} ? [Y/n] ").lower() != 'n':
+        shutil.rmtree(path)
+
+
+def compute_lambda_returns(rewards, values, ends, gamma, lambda_):
+    assert rewards.ndim == 2 or (rewards.ndim == 3 and rewards.size(2) == 1)
+    assert rewards.shape == ends.shape == values.shape, f"{rewards.shape}, {values.shape}, {ends.shape}"  # (B, T, 1)
+    t = rewards.size(1)
+    lambda_returns = torch.empty_like(values)
+    lambda_returns[:, -1] = values[:, -1]
+    lambda_returns[:, :-1] = rewards[:, :-1] + ends[:, :-1].logical_not() * gamma * (1 - lambda_) * values[:, 1:]
+
+    last = values[:, -1]
+    for i in list(range(t - 1))[::-1]:
+        lambda_returns[:, i] += ends[:, i].logical_not() * gamma * lambda_ * last
+        last = lambda_returns[:, i]
+
+    return lambda_returns
+
+
+class LossWithIntermediateLosses:
+    def __init__(self, **kwargs):
+        self.loss_total = sum(kwargs.values())
+        self.intermediate_losses = {k: v.item() for k, v in kwargs.items()}
+
+    def __truediv__(self, value):
+        for k, v in self.intermediate_losses.items():
+            self.intermediate_losses[k] = v / value
+        self.loss_total = self.loss_total / value
+        return self
+
+
+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()
+        episode.save(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()
+            episode.save(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,))