[tab] ali-old-branch.diff (diff)
diff --git a/data/jasmine_data/atari/atari_utils.py b/data/jasmine_data/atari/atari_utils.py
deleted file mode 100644
index b057914..0000000
--- a/data/jasmine_data/atari/atari_utils.py
+++ /dev/null
@@ -1,337 +0,0 @@
-# Copyright notice
-#
-# This file contains code adapted from cleanrl
-# (https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl_utils/atari_wrappers.py)
-# licensed under the MIT License.
-#
-# Copyright (c) 2019-2023 Antonin Raffin, Ashley Hill, Anssi Kanervisto,
-# Maximilian Ernestus, Rinu Boney, Pavan Goli, and other contributors
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-from __future__ import annotations
-
-from typing import SupportsFloat
-
-import gymnasium as gym
-import numpy as np
-from gymnasium import spaces
-
-try:
-    import cv2
-
-    cv2.ocl.setUseOpenCL(False)
-except ImportError:
-    cv2 = None  # type: ignore[assignment]
-
-
-class StickyActionEnv(gym.Wrapper[np.ndarray, int, np.ndarray, int]):
-    """
-    Sticky action.
-
-    Paper: https://arxiv.org/abs/1709.06009
-    Official implementation: https://github.com/mgbellemare/Arcade-Learning-Environment
-
-    :param env: Environment to wrap
-    :param action_repeat_probability: Probability of repeating the last action
-    """
-
-    def __init__(self, env: gym.Env, action_repeat_probability: float) -> None:
-        super().__init__(env)
-        self.action_repeat_probability = action_repeat_probability
-        assert env.unwrapped.get_action_meanings()[0] == "NOOP"  # type: ignore[attr-defined]
-
-    def reset(self, **kwargs):
-        self._sticky_action = 0  # NOOP
-        return self.env.reset(**kwargs)
-
-    def step(self, action: int):
-        if self.np_random.random() >= self.action_repeat_probability:
-            self._sticky_action = action
-        return self.env.step(self._sticky_action)
-
-
-class NoopResetEnv(gym.Wrapper[np.ndarray, int, np.ndarray, int]):
-    """
-    Sample initial states by taking random number of no-ops on reset.
-    No-op is assumed to be action 0.
-
-    :param env: Environment to wrap
-    :param noop_max: Maximum value of no-ops to run
-    """
-
-    def __init__(self, env: gym.Env, noop_max: int = 30) -> None:
-        super().__init__(env)
-        self.noop_max = noop_max
-        self.override_num_noops = None
-        self.noop_action = 0
-        assert env.unwrapped.get_action_meanings()[0] == "NOOP"  # type: ignore[attr-defined]
-
-    def reset(self, **kwargs):
-        self.env.reset(**kwargs)
-        if self.override_num_noops is not None:
-            noops = self.override_num_noops
-        else:
-            noops = self.unwrapped.np_random.integers(1, self.noop_max + 1)
-        assert noops > 0
-        obs = np.zeros(0)
-        info: dict = {}
-        for _ in range(noops):
-            obs, _, terminated, truncated, info = self.env.step(self.noop_action)
-            if terminated or truncated:
-                obs, info = self.env.reset(**kwargs)
-        return obs, info
-
-
-class FireResetEnv(gym.Wrapper[np.ndarray, int, np.ndarray, int]):
-    """
-    Take action on reset for environments that are fixed until firing.
-
-    :param env: Environment to wrap
-    """
-
-    def __init__(self, env: gym.Env) -> None:
-        super().__init__(env)
-        assert env.unwrapped.get_action_meanings()[1] == "FIRE"  # type: ignore[attr-defined]
-        assert len(env.unwrapped.get_action_meanings()) >= 3  # type: ignore[attr-defined]
-
-    def reset(self, **kwargs):
-        self.env.reset(**kwargs)
-        obs, _, terminated, truncated, _ = self.env.step(1)
-        if terminated or truncated:
-            self.env.reset(**kwargs)
-        obs, _, terminated, truncated, _ = self.env.step(2)
-        if terminated or truncated:
-            self.env.reset(**kwargs)
-        return obs, {}
-
-
-class EpisodicLifeEnv(gym.Wrapper[np.ndarray, int, np.ndarray, int]):
-    """
-    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.
-
-    :param env: Environment to wrap
-    """
-
-    def __init__(self, env: gym.Env) -> None:
-        super().__init__(env)
-        self.lives = 0
-        self.was_real_done = True
-
-    def step(self, action: int):
-        obs, reward, terminated, truncated, info = self.env.step(action)
-        self.was_real_done = terminated or truncated
-        # check current lives, make loss of life terminal,
-        # then update lives to handle bonus lives
-        lives = self.env.unwrapped.ale.lives()  # type: ignore[attr-defined]
-        if 0 < lives < self.lives:
-            # for Qbert sometimes we stay in lives == 0 condition for a few frames
-            # so its important to keep lives > 0, so that we only reset once
-            # the environment advertises done.
-            terminated = True
-        self.lives = lives
-        return obs, reward, terminated, truncated, info
-
-    def reset(self, **kwargs):
-        """
-        Calls the Gym environment 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.
-
-        :param kwargs: Extra keywords passed to env.reset() call
-        :return: the first observation of the environment
-        """
-        if self.was_real_done:
-            obs, info = self.env.reset(**kwargs)
-        else:
-            # no-op step to advance from terminal/lost life state
-            obs, _, terminated, truncated, info = self.env.step(0)
-
-            # The no-op step can lead to a game over, so we need to check it again
-            # to see if we should reset the environment and avoid the
-            # monitor.py `RuntimeError: Tried to step environment that needs reset`
-            if terminated or truncated:
-                obs, info = self.env.reset(**kwargs)
-        self.lives = self.env.unwrapped.ale.lives()  # type: ignore[attr-defined]
-        return obs, info
-
-
-class MaxAndSkipEnv(gym.Wrapper[np.ndarray, int, np.ndarray, int]):
-    """
-    Return only every ``skip``-th frame (frameskipping)
-    and return the max between the two last frames.
-
-    :param env: Environment to wrap
-    :param skip: Number of ``skip``-th frame
-        The same action will be taken ``skip`` times.
-    """
-
-    def __init__(self, env: gym.Env, skip: int = 4) -> None:
-        super().__init__(env)
-        # most recent raw observations (for max pooling across time steps)
-        assert (
-            env.observation_space.dtype is not None
-        ), "No dtype specified for the observation space"
-        assert (
-            env.observation_space.shape is not None
-        ), "No shape defined for the observation space"
-        self._obs_buffer = np.zeros(
-            (2, *env.observation_space.shape), dtype=env.observation_space.dtype
-        )
-        self._skip = skip
-
-    def step(self, action: int):
-        """
-        Step the environment with the given action
-        Repeat action, sum reward, and max over last observations.
-
-        :param action: the action
-        :return: observation, reward, terminated, truncated, information
-        """
-        total_reward = 0.0
-        terminated = truncated = False
-        for i in range(self._skip):
-            obs, reward, terminated, truncated, info = self.env.step(action)
-            done = terminated or truncated
-            if i == self._skip - 2:
-                self._obs_buffer[0] = obs
-            if i == self._skip - 1:
-                self._obs_buffer[1] = obs
-            total_reward += float(reward)
-            if done:
-                break
-        # Note that the observation on the done=True frame
-        # doesn't matter
-        max_frame = self._obs_buffer.max(axis=0)
-
-        return max_frame, total_reward, terminated, truncated, info
-
-
-class ClipRewardEnv(gym.RewardWrapper):
-    """
-    Clip the reward to {+1, 0, -1} by its sign.
-
-    :param env: Environment to wrap
-    """
-
-    def __init__(self, env: gym.Env) -> None:
-        super().__init__(env)
-
-    def reward(self, reward: SupportsFloat) -> float:
-        """
-        Bin reward to {+1, 0, -1} by its sign.
-
-        :param reward:
-        :return:
-        """
-        return np.sign(float(reward))
-
-
-class WarpFrame(gym.ObservationWrapper[np.ndarray, int, np.ndarray]):
-    """
-    Convert to grayscale and warp frames to 84x84 (default)
-    as done in the Nature paper and later work.
-
-    :param env: Environment to wrap
-    :param width: New frame width
-    :param height: New frame height
-    """
-
-    def __init__(self, env: gym.Env, width: int = 84, height: int = 84) -> None:
-        super().__init__(env)
-        self.width = width
-        self.height = height
-        assert isinstance(
-            env.observation_space, spaces.Box
-        ), f"Expected Box space, got {env.observation_space}"
-
-        self.observation_space = spaces.Box(
-            low=0,
-            high=255,
-            shape=(self.height, self.width, 1),
-            dtype=env.observation_space.dtype,  # type: ignore[arg-type]
-        )
-
-    def observation(self, frame: np.ndarray) -> np.ndarray:
-        """
-        returns the current observation from a frame
-
-        :param frame: environment frame
-        :return: the observation
-        """
-        assert (
-            cv2 is not None
-        ), "OpenCV is not installed, you can do `pip install opencv-python`"
-        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
-        frame = cv2.resize(
-            frame, (self.width, self.height), interpolation=cv2.INTER_AREA
-        )
-        return frame[:, :, None]
-
-
-class AtariWrapper(gym.Wrapper[np.ndarray, int, np.ndarray, int]):
-    """
-    Atari 2600 preprocessings
-
-    Specifically:
-
-    * Noop reset: obtain initial state by taking random number of no-ops on reset.
-    * Frame skipping: 4 by default
-    * Max-pooling: most recent two observations
-    * Termination signal when a life is lost.
-    * Resize to a square image: 84x84 by default
-    * Grayscale observation
-    * Clip reward to {-1, 0, 1}
-    * Sticky actions: disabled by default
-
-    See https://danieltakeshi.github.io/2016/11/25/frame-skipping-and-preprocessing-for-deep-q-networks-on-atari-2600-games/
-    for a visual explanation.
-
-    .. warning::
-        Use this wrapper only with Atari v4 without frame skip: ``env_id = "*NoFrameskip-v4"``.
-
-    :param env: Environment to wrap
-    :param noop_max: Max number of no-ops
-    :param frame_skip: Frequency at which the agent experiences the game.
-        This correspond to repeating the action ``frame_skip`` times.
-    :param screen_size: Resize Atari frame
-    :param terminal_on_life_loss: If True, then step() returns done=True whenever a life is lost.
-    :param clip_reward: If True (default), the reward is clip to {-1, 0, 1} depending on its sign.
-    :param action_repeat_probability: Probability of repeating the last action
-    """
-
-    def __init__(
-        self,
-        env: gym.Env,
-        noop_max: int = 30,
-        frame_skip: int = 4,
-        screen_size: int = 84,
-        terminal_on_life_loss: bool = True,
-        clip_reward: bool = True,
-        action_repeat_probability: float = 0.0,
-    ) -> None:
-        if action_repeat_probability > 0.0:
-            env = StickyActionEnv(env, action_repeat_probability)
-        if noop_max > 0:
-            env = NoopResetEnv(env, noop_max=noop_max)
-        # frame_skip=1 is the same as no frame-skip (action repeat)
-        if frame_skip > 1:
-            env = MaxAndSkipEnv(env, skip=frame_skip)
-        if terminal_on_life_loss:
-            env = EpisodicLifeEnv(env)
-        if "FIRE" in env.unwrapped.get_action_meanings():  # type: ignore[attr-defined]
-            env = FireResetEnv(env)
-        env = WarpFrame(env, width=screen_size, height=screen_size)
-        if clip_reward:
-            env = ClipRewardEnv(env)
-
-        super().__init__(env)
diff --git a/data/jasmine_data/atari/generate_atari_dataset.py b/data/jasmine_data/atari/generate_atari_dataset.py
deleted file mode 100644
index fbefc1e..0000000
--- a/data/jasmine_data/atari/generate_atari_dataset.py
+++ /dev/null
@@ -1,808 +0,0 @@
-# adapted from https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/rainbow_atari.py
-import einops
-import collections
-import math
-import os
-import random
-import time
-from collections import deque
-from dataclasses import dataclass
-
-import gymnasium as gym
-import ale_py
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.optim as optim
-import tyro
-from typing import Optional, Any
-from torch.utils.tensorboard.writer import SummaryWriter
-
-from jasmine_data.atari.atari_utils import (
-    ClipRewardEnv,
-    EpisodicLifeEnv,
-    FireResetEnv,
-    MaxAndSkipEnv,
-    NoopResetEnv,
-)
-
-from jasmine_data.utils import save_chunks
-import json
-
-
-@dataclass
-class Args:
-    exp_name: str = os.path.basename(__file__)[: -len(".py")]
-    """the name of this experiment"""
-    seed: int = 1
-    """seed of the experiment"""
-    torch_deterministic: bool = True
-    """if toggled, `torch.backends.cudnn.deterministic=False`"""
-    cuda: bool = True
-    """if toggled, cuda will be enabled by default"""
-    track: bool = False
-    """if toggled, this experiment will be tracked with Weights and Biases"""
-    wandb_project_name: str = "cleanRL"
-    """the wandb's project name"""
-    wandb_entity: Optional[str] = None
-    """the entity (team) of wandb's project"""
-    capture_video: bool = False
-    """whether to capture videos of the agent performances (check out `videos` folder)"""
-    save_model: bool = False
-    """whether to save model into the `runs/{run_name}` folder"""
-    upload_model: bool = False
-    """whether to upload the saved model to huggingface"""
-    hf_entity: str = ""
-    """the user or org name of the model repository from the Hugging Face Hub"""
-
-    env_id: str = "ALE/Breakout-v5"
-    """the id of the environment"""
-    learning_rate: float = 0.0000625
-    """the learning rate of the optimizer"""
-    num_envs: int = 1
-    """the number of parallel game environments"""
-    buffer_size: int = 1000000
-    """the replay memory buffer size"""
-    gamma: float = 0.99
-    """the discount factor gamma"""
-    tau: float = 1.0
-    """the target network update rate"""
-    target_network_frequency: int = 8000
-    """the timesteps it takes to update the target network"""
-    batch_size: int = 32
-    """the batch size of sample from the reply memory"""
-    start_e: float = 1
-    """the starting epsilon for exploration"""
-    end_e: float = 0.01
-    """the ending epsilon for exploration"""
-    exploration_fraction: float = 0.10
-    """the fraction of `total-timesteps` it takes from start-e to go end-e"""
-    learning_starts: int = 80000
-    """timestep to start learning"""
-    train_frequency: int = 4
-    """the frequency of training"""
-    n_step: int = 3
-    """the number of steps to look ahead for n-step Q learning"""
-    prioritized_replay_alpha: float = 0.5
-    """alpha parameter for prioritized replay buffer"""
-    prioritized_replay_beta: float = 0.4
-    """beta parameter for prioritized replay buffer"""
-    prioritized_replay_eps: float = 1e-6
-    """epsilon parameter for prioritized replay buffer"""
-    n_atoms: int = 51
-    """the number of atoms"""
-    v_min: float = -10
-    """the return lower bound"""
-    v_max: float = 10
-    """the return upper bound"""
-
-    # Dataset capture
-    capture_dataset: bool = True
-    num_transitions_train: int = 10000000
-    num_transitions_val: int = 500000
-    num_transitions_test: int = 500000
-    output_dir: str = "data/atari_episodes"
-    min_episode_length: int = 1
-    chunk_size: int = 160
-    chunks_per_file: int = 100
-    stop_on_complete: bool = True
-
-
-def make_env(env_id, seed, idx, capture_video, run_name):
-    def thunk():
-        if capture_video and idx == 0:
-            env = gym.make(env_id, render_mode="rgb_array")
-            env = gym.wrappers.RecordVideo(env, f"videos/{run_name}")
-        else:
-            env = gym.make(env_id)
-        env = gym.wrappers.RecordEpisodeStatistics(env)
-
-        env = NoopResetEnv(env, noop_max=30)
-        env = MaxAndSkipEnv(env, skip=4)
-        env = EpisodicLifeEnv(env)
-        if "FIRE" in env.unwrapped.get_action_meanings():
-            env = FireResetEnv(env)
-        env = ClipRewardEnv(env)
-        env = gym.wrappers.ResizeObservation(env, (84, 84))
-        env = gym.wrappers.GrayScaleObservation(env)
-        env = gym.wrappers.FrameStack(env, 4)
-
-        env.action_space.seed(seed)
-        return env
-
-    return thunk
-
-
-class NoisyLinear(nn.Module):
-    def __init__(self, in_features, out_features, std_init=0.5):
-        super().__init__()
-        self.in_features = in_features
-        self.out_features = out_features
-        self.std_init = std_init
-
-        self.weight_mu = nn.Parameter(torch.FloatTensor(out_features, in_features))
-        self.weight_sigma = nn.Parameter(torch.FloatTensor(out_features, in_features))
-        self.register_buffer(
-            "weight_epsilon", torch.FloatTensor(out_features, in_features)
-        )
-        self.bias_mu = nn.Parameter(torch.FloatTensor(out_features))
-        self.bias_sigma = nn.Parameter(torch.FloatTensor(out_features))
-        self.register_buffer("bias_epsilon", torch.FloatTensor(out_features))
-        # factorized gaussian noise
-        self.reset_parameters()
-        self.reset_noise()
-
-    def reset_parameters(self):
-        mu_range = 1 / math.sqrt(self.in_features)
-        self.weight_mu.data.uniform_(-mu_range, mu_range)
-        self.weight_sigma.data.fill_(self.std_init / math.sqrt(self.in_features))
-        self.bias_mu.data.uniform_(-mu_range, mu_range)
-        self.bias_sigma.data.fill_(self.std_init / math.sqrt(self.out_features))
-
-    def reset_noise(self):
-        self.weight_epsilon.normal_()
-        self.bias_epsilon.normal_()
-
-    def forward(self, input):
-        if self.training:
-            weight = self.weight_mu + self.weight_sigma * self.weight_epsilon
-            bias = self.bias_mu + self.bias_sigma * self.bias_epsilon
-        else:
-            weight = self.weight_mu
-            bias = self.bias_mu
-        return F.linear(input, weight, bias)
-
-
-# ALGO LOGIC: initialize agent here:
-class NoisyDuelingDistributionalNetwork(nn.Module):
-    def __init__(self, env, n_atoms, v_min, v_max):
-        super().__init__()
-        self.n_atoms = n_atoms
-        self.v_min = v_min
-        self.v_max = v_max
-        self.delta_z = (v_max - v_min) / (n_atoms - 1)
-        self.n_actions = env.single_action_space.n
-        self.register_buffer("support", torch.linspace(v_min, v_max, n_atoms))
-
-        self.network = nn.Sequential(
-            nn.Conv2d(4, 32, 8, stride=4),
-            nn.ReLU(),
-            nn.Conv2d(32, 64, 4, stride=2),
-            nn.ReLU(),
-            nn.Conv2d(64, 64, 3, stride=1),
-            nn.ReLU(),
-            nn.Flatten(),
-        )
-        conv_output_size = 3136
-
-        self.value_head = nn.Sequential(
-            NoisyLinear(conv_output_size, 512), nn.ReLU(), NoisyLinear(512, n_atoms)
-        )
-
-        self.advantage_head = nn.Sequential(
-            NoisyLinear(conv_output_size, 512),
-            nn.ReLU(),
-            NoisyLinear(512, n_atoms * self.n_actions),
-        )
-
-    def forward(self, x):
-        h = self.network(x / 255.0)
-        value = self.value_head(h).view(-1, 1, self.n_atoms)
-        advantage = self.advantage_head(h).view(-1, self.n_actions, self.n_atoms)
-        q_atoms = value + advantage - advantage.mean(dim=1, keepdim=True)
-        q_dist = F.softmax(q_atoms, dim=2)
-        return q_dist
-
-    def reset_noise(self):
-        for layer in self.value_head:
-            if isinstance(layer, NoisyLinear):
-                layer.reset_noise()
-        for layer in self.advantage_head:
-            if isinstance(layer, NoisyLinear):
-                layer.reset_noise()
-
-
-PrioritizedBatch = collections.namedtuple(
-    "PrioritizedBatch",
-    [
-        "observations",
-        "actions",
-        "rewards",
-        "next_observations",
-        "dones",
-        "indices",
-        "weights",
-    ],
-)
-
-
-# adapted from: https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py
-class SumSegmentTree:
-    def __init__(self, capacity):
-        self.capacity = capacity
-        self.tree_size = 2 * capacity - 1
-        self.tree = np.zeros(self.tree_size, dtype=np.float32)
-
-    def _propagate(self, idx):
-        parent = (idx - 1) // 2
-        while parent >= 0:
-            self.tree[parent] = self.tree[parent * 2 + 1] + self.tree[parent * 2 + 2]
-            parent = (parent - 1) // 2
-
-    def update(self, idx, value):
-        tree_idx = idx + self.capacity - 1
-        self.tree[tree_idx] = value
-        self._propagate(tree_idx)
-
-    def total(self):
-        return self.tree[0]
-
-    def retrieve(self, value):
-        idx = 0
-        while idx * 2 + 1 < self.tree_size:
-            left = idx * 2 + 1
-            right = left + 1
-            if value <= self.tree[left]:
-                idx = left
-            else:
-                value -= self.tree[left]
-                idx = right
-        return idx - (self.capacity - 1)
-
-
-# adapted from: https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py
-class MinSegmentTree:
-    def __init__(self, capacity):
-        self.capacity = capacity
-        self.tree_size = 2 * capacity - 1
-        self.tree = np.full(self.tree_size, float("inf"), dtype=np.float32)
-
-    def _propagate(self, idx):
-        parent = (idx - 1) // 2
-        while parent >= 0:
-            self.tree[parent] = np.minimum(
-                self.tree[parent * 2 + 1], self.tree[parent * 2 + 2]
-            )
-            parent = (parent - 1) // 2
-
-    def update(self, idx, value):
-        tree_idx = idx + self.capacity - 1
-        self.tree[tree_idx] = value
-        self._propagate(tree_idx)
-
-    def min(self):
-        return self.tree[0]
-
-
-class PrioritizedReplayBuffer:
-    def __init__(
-        self, capacity, obs_shape, device, n_step, gamma, alpha=0.6, beta=0.4, eps=1e-6
-    ):
-        self.capacity = capacity
-        self.device = device
-        self.n_step = n_step
-        self.gamma = gamma
-        self.alpha = alpha
-        self.beta = beta
-        self.eps = eps
-
-        self.buffer_obs = np.zeros((capacity,) + obs_shape, dtype=np.uint8)
-        self.buffer_next_obs = np.zeros((capacity,) + obs_shape, dtype=np.uint8)
-        self.buffer_actions = np.zeros(capacity, dtype=np.int64)
-        self.buffer_rewards = np.zeros(capacity, dtype=np.float32)
-        self.buffer_dones = np.zeros(capacity, dtype=np.bool_)
-
-        self.pos = 0
-        self.size = 0
-        self.max_priority = 1.0
-
-        self.sum_tree = SumSegmentTree(capacity)
-        self.min_tree = MinSegmentTree(capacity)
-
-        # For n-step returns
-        self.n_step_buffer = deque(maxlen=n_step)
-
-    def _get_n_step_info(self):
-        reward = 0.0
-        next_obs = self.n_step_buffer[-1][3]
-        done = self.n_step_buffer[-1][4]
-
-        for i in range(len(self.n_step_buffer)):
-            reward += self.gamma**i * self.n_step_buffer[i][2]
-            if self.n_step_buffer[i][4]:
-                next_obs = self.n_step_buffer[i][3]
-                done = True
-                break
-        return reward, next_obs, done
-
-    def add(self, obs, action, reward, next_obs, done):
-        self.n_step_buffer.append((obs, action, reward, next_obs, done))
-
-        if len(self.n_step_buffer) < self.n_step:
-            return
-
-        reward, next_obs, done = self._get_n_step_info()
-        obs = self.n_step_buffer[0][0]
-        action = self.n_step_buffer[0][1]
-
-        idx = self.pos
-        self.buffer_obs[idx] = obs
-        self.buffer_next_obs[idx] = next_obs
-        self.buffer_actions[idx] = action
-        self.buffer_rewards[idx] = reward
-        self.buffer_dones[idx] = done
-
-        priority = self.max_priority**self.alpha
-        self.sum_tree.update(idx, priority)
-        self.min_tree.update(idx, priority)
-
-        self.pos = (self.pos + 1) % self.capacity
-        self.size = min(self.size + 1, self.capacity)
-
-        if done:
-            self.n_step_buffer.clear()
-
-    def sample(self, batch_size):
-        indices = []
-        p_total = self.sum_tree.total()
-        segment = p_total / batch_size
-
-        for i in range(batch_size):
-            a = segment * i
-            b = segment * (i + 1)
-            upperbound = np.random.uniform(a, b)
-            idx = self.sum_tree.retrieve(upperbound)
-            indices.append(idx)
-
-        samples = {
-            "observations": torch.from_numpy(self.buffer_obs[indices]).to(self.device),
-            "actions": torch.from_numpy(self.buffer_actions[indices])
-            .to(self.device)
-            .unsqueeze(1),
-            "rewards": torch.from_numpy(self.buffer_rewards[indices])
-            .to(self.device)
-            .unsqueeze(1),
-            "next_observations": torch.from_numpy(self.buffer_next_obs[indices]).to(
-                self.device
-            ),
-            "dones": torch.from_numpy(self.buffer_dones[indices])
-            .to(self.device)
-            .unsqueeze(1),
-        }
-
-        probs = np.array(
-            [self.sum_tree.tree[idx + self.capacity - 1] for idx in indices]
-        )
-        weights = (self.size * probs / p_total) ** -self.beta
-        weights = weights / weights.max()
-        samples["weights"] = torch.from_numpy(weights).to(self.device).unsqueeze(1)
-        samples["indices"] = indices
-
-        return PrioritizedBatch(**samples)
-
-    def update_priorities(self, indices, priorities):
-        priorities = np.abs(priorities) + self.eps
-        self.max_priority = max(self.max_priority, priorities.max())
-
-        for idx, priority in zip(indices, priorities):
-            priority = priority**self.alpha
-            self.sum_tree.update(idx, priority)
-            self.min_tree.update(idx, priority)
-
-
-if __name__ == "__main__":
-    args = tyro.cli(Args)
-    assert args.num_envs == 1, "vectorized envs are not supported at the moment"
-    gym.register_envs(ale_py)
-    run_name = f"{args.env_id}__{args.exp_name}__{args.seed}__{int(time.time())}"
-    if args.track:
-        import wandb
-
-        wandb.init(
-            project=args.wandb_project_name,
-            entity=args.wandb_entity,
-            sync_tensorboard=True,
-            config=vars(args),
-            name=run_name,
-            monitor_gym=True,
-            save_code=True,
-        )
-    writer = SummaryWriter(f"runs/{run_name}")
-    writer.add_text(
-        "hyperparameters",
-        "|param|value|
|-|-|
%s"
-        % ("
".join([f"|{key}|{value}|" for key, value in vars(args).items()])),
-    )
-
-    # TRY NOT TO MODIFY: seeding
-    random.seed(args.seed)
-    np.random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    torch.backends.cudnn.deterministic = args.torch_deterministic
-
-    device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")
-
-    # env setup
-    envs = gym.vector.SyncVectorEnv(
-        [
-            make_env(args.env_id, args.seed + i, i, args.capture_video, run_name)
-            for i in range(args.num_envs)
-        ]
-    )
-    assert isinstance(
-        envs.single_action_space, gym.spaces.Discrete
-    ), "only discrete action space is supported"
-
-    q_network = NoisyDuelingDistributionalNetwork(
-        envs, args.n_atoms, args.v_min, args.v_max
-    ).to(device)
-    optimizer = optim.Adam(q_network.parameters(), lr=args.learning_rate, eps=1.5e-4)
-    target_network = NoisyDuelingDistributionalNetwork(
-        envs, args.n_atoms, args.v_min, args.v_max
-    ).to(device)
-    target_network.load_state_dict(q_network.state_dict())
-
-    rb = PrioritizedReplayBuffer(
-        args.buffer_size,
-        envs.single_observation_space.shape,
-        device,
-        args.n_step,
-        args.gamma,
-        args.prioritized_replay_alpha,
-        args.prioritized_replay_beta,
-        args.prioritized_replay_eps,
-    )
-
-    # dataset capture state
-    split_targets: dict[str, int] = {
-        "train": args.num_transitions_train,
-        "val": args.num_transitions_val,
-        "test": args.num_transitions_test,
-    }
-    # Determine splits to run (order: train -> val -> test)
-    splits_in_order = [s for s in ["train", "val", "test"] if split_targets[s] > 0]
-
-    transitions_captured_per_split: dict[str, int] = {
-        s: 0 for s in ["train", "val", "test"]
-    }
-    file_idx_by_split: dict[str, int] = {s: 0 for s in ["train", "val", "test"]}
-    episode_metadata_by_split: dict[str, list[dict]] = {
-        s: [] for s in ["train", "val", "test"]
-    }
-
-    obs_chunks: list[np.ndarray] = []
-    act_chunks: list[np.ndarray] = []
-
-    current_split_idx = 0
-    current_split = splits_in_order[0]
-    split_dir = os.path.join(args.output_dir, current_split)
-    if args.capture_dataset:
-        os.makedirs(split_dir, exist_ok=True)
-
-    start_time = time.time()
-
-    # TRY NOT TO MODIFY: start the game
-    obs, _ = envs.reset(seed=args.seed)
-    observations_seq: list[np.ndarray] = []
-    actions_seq: list[np.ndarray] = []
-    total_timesteps = (
-        args.num_transitions_train
-        + args.num_transitions_val
-        + args.num_transitions_test
-    )
-    for global_step in range(total_timesteps):
-        # anneal PER beta to 1
-        rb.beta = min(
-            1.0,
-            args.prioritized_replay_beta
-            + global_step * (1.0 - args.prioritized_replay_beta) / total_timesteps,
-        )
-
-        # ALGO LOGIC: put action logic here
-        with torch.no_grad():
-            q_dist = q_network(torch.Tensor(obs).to(device))
-            q_values = torch.sum(q_dist * q_network.support, dim=2)
-            actions = torch.argmax(q_values, dim=1).cpu().numpy()
-
-        # TRY NOT TO MODIFY: execute the game and log data.
-        next_obs, rewards, terminations, truncations, infos = envs.step(actions)
-
-        if args.capture_dataset:
-            assert (
-                obs.shape[0] == 1 and actions.shape[0] == 1
-            ), "Vectorized envs are currently not supported during data capture."
-            obs_to_save = obs  # (1, F, H, W)
-            actions_to_save = actions  # (1,)
-
-            # remove frame-stacking
-            obs_to_save = obs_to_save[:, -1, :, :]  # (1, H, W)
-
-            # remove asserted singleton vectorization dimension
-            obs_to_save = einops.rearrange(obs_to_save, "1 H W -> H W")  # (H, W)
-            actions_to_save = einops.rearrange(actions_to_save, "1 ->")  # ()
-            observations_seq.append(obs_to_save.astype(np.uint8))
-            actions_seq.append(actions_to_save.astype(np.int8))
-
-        if "final_info" in infos:
-            for info in infos["final_info"]:
-                if info and "episode" in info:
-                    print(
-                        f"global_step={global_step}, episodic_return={info['episode']['r']}"
-                    )
-                    writer.add_scalar(
-                        "charts/episodic_return", info["episode"]["r"], global_step
-                    )
-                    writer.add_scalar(
-                        "charts/episodic_length", info["episode"]["l"], global_step
-                    )
-
-                    continue_capturing_multi = any(
-                        transitions_captured_per_split[s] < split_targets[s]
-                        for s in splits_in_order
-                    )
-                    if args.capture_dataset and continue_capturing_multi:
-                        current_len = len(observations_seq)
-                        if current_len >= args.min_episode_length:
-                            frames = np.stack(observations_seq, axis=0).astype(
-                                np.uint8
-                            )  # (T, H, W)
-                            acts = np.stack(actions_seq, axis=0).astype(np.int8)
-                            # Convert frames to three-channel RGB by repeating grayscale values
-                            frames = np.stack(
-                                [frames, frames, frames], axis=-1
-                            )  # (T, H, W, 3)
-
-                            episode_obs_chunks = []
-                            episode_act_chunks = []
-                            start_idx = 0
-                            while start_idx < current_len:
-                                end_idx = min(start_idx + args.chunk_size, current_len)
-                                if end_idx - start_idx < args.chunk_size:
-                                    print(
-                                        f"Warning: Inconsistent chunk_sizes. Episode has {current_len} frames, "
-                                        f"which is smaller than the requested chunk_size: {args.chunk_size}. "
-                                        "This might lead to performance degradation during training."
-                                    )
-                                episode_obs_chunks.append(frames[start_idx:end_idx])
-                                episode_act_chunks.append(acts[start_idx:end_idx])
-                                start_idx = end_idx
-
-                            obs_chunks_data = [
-                                seq.astype(np.uint8) for seq in episode_obs_chunks
-                            ]
-                            act_chunks_data = [act for act in episode_act_chunks]
-                            obs_chunks.extend(obs_chunks_data)
-                            act_chunks.extend(act_chunks_data)
-
-                            # Save to the active split
-                            (
-                                ep_metadata,
-                                file_idx_by_split[current_split],
-                                obs_chunks,
-                                act_chunks,
-                            ) = save_chunks(
-                                file_idx_by_split[current_split],
-                                args.chunks_per_file,
-                                split_dir,
-                                obs_chunks,
-                                act_chunks,
-                            )
-                            episode_metadata_by_split[current_split].extend(ep_metadata)
-
-                            transitions_captured_per_split[current_split] += current_len
-
-                            if (
-                                transitions_captured_per_split[current_split]
-                                >= split_targets[current_split]
-                            ):
-                                if len(obs_chunks) > 0:
-                                    print(
-                                        f"Warning: Dropping {len(obs_chunks)} chunks before switching split '",
-                                        {current_split},
-                                        "' for consistent number of chunks per file.",
-                                        "Consider changing the chunk_size and chunks_per_file parameters to prevent data-loss.",
-                                    )
-                                    obs_chunks = []
-                                    act_chunks = []
-                                if current_split_idx + 1 < len(splits_in_order):
-                                    current_split_idx += 1
-                                    current_split = splits_in_order[current_split_idx]
-                                    split_dir = os.path.join(
-                                        args.output_dir, current_split
-                                    )
-                                    os.makedirs(split_dir, exist_ok=True)
-                        else:
-                            print(
-                                f"Episode too short ({current_len}), skipping capture..."
-                            )
-
-                        observations_seq = []
-                        actions_seq = []
-
-        # TRY NOT TO MODIFY: save data to reply buffer; handle `final_observation`
-        real_next_obs = next_obs.copy()
-        for idx, trunc in enumerate(truncations):
-            if trunc:
-                real_next_obs[idx] = infos["final_observation"][idx]
-        rb.add(obs, actions, rewards, real_next_obs, terminations)
-
-        # TRY NOT TO MODIFY: CRUCIAL step easy to overlook
-        obs = next_obs
-
-        # ALGO LOGIC: training.
-        if global_step > args.learning_starts:
-            if global_step % args.train_frequency == 0:
-                # reset the noise for both networks
-                q_network.reset_noise()
-                target_network.reset_noise()
-                data = rb.sample(args.batch_size)
-
-                with torch.no_grad():
-                    next_dist = target_network(
-                        data.next_observations
-                    )  # [B, num_actions, n_atoms]
-                    support = target_network.support  # [n_atoms]
-                    next_q_values = torch.sum(
-                        next_dist * support, dim=2
-                    )  # [B, num_actions]
-
-                    # double q-learning
-                    next_dist_online = q_network(
-                        data.next_observations
-                    )  # [B, num_actions, n_atoms]
-                    next_q_online = torch.sum(
-                        next_dist_online * support, dim=2
-                    )  # [B, num_actions]
-                    best_actions = torch.argmax(next_q_online, dim=1)  # [B]
-                    next_pmfs = next_dist[
-                        torch.arange(args.batch_size), best_actions
-                    ]  # [B, n_atoms]
-
-                    # compute the n-step Bellman update.
-                    gamma_n = args.gamma**args.n_step
-                    next_atoms = data.rewards + gamma_n * support * (
-                        1 - data.dones.float()
-                    )
-                    tz = next_atoms.clamp(q_network.v_min, q_network.v_max)
-
-                    # projection
-                    delta_z = q_network.delta_z
-                    b = (tz - q_network.v_min) / delta_z  # shape: [B, n_atoms]
-                    l = b.floor().clamp(0, args.n_atoms - 1)
-                    u = b.ceil().clamp(0, args.n_atoms - 1)
-
-                    # (l == u).float() handles the case where bj is exactly an integer
-                    # example bj = 1, then the upper ceiling should be uj= 2, and lj= 1
-                    d_m_l = (
-                        u.float() + (l == b).float() - b
-                    ) * next_pmfs  # [B, n_atoms]
-                    d_m_u = (b - l) * next_pmfs  # [B, n_atoms]
-
-                    target_pmfs = torch.zeros_like(next_pmfs)
-                    for i in range(target_pmfs.size(0)):
-                        target_pmfs[i].index_add_(0, l[i].long(), d_m_l[i])
-                        target_pmfs[i].index_add_(0, u[i].long(), d_m_u[i])
-
-                dist = q_network(data.observations)  # [B, num_actions, n_atoms]
-                pred_dist = dist.gather(
-                    1, data.actions.unsqueeze(-1).expand(-1, -1, args.n_atoms)
-                ).squeeze(1)
-                log_pred = torch.log(pred_dist.clamp(min=1e-5, max=1 - 1e-5))
-
-                loss_per_sample = -(target_pmfs * log_pred).sum(dim=1)
-                loss = (loss_per_sample * data.weights.squeeze()).mean()
-
-                # update priorities
-                new_priorities = loss_per_sample.detach().cpu().numpy()
-                rb.update_priorities(data.indices, new_priorities)
-
-                if global_step % 100 == 0:
-                    writer.add_scalar("losses/td_loss", loss.item(), global_step)
-                    q_values = (pred_dist * q_network.support).sum(dim=1)  # [B]
-                    writer.add_scalar(
-                        "losses/q_values", q_values.mean().item(), global_step
-                    )
-                    sps = int(global_step / (time.time() - start_time))
-                    print("SPS:", sps)
-                    writer.add_scalar("charts/SPS", sps, global_step)
-                    writer.add_scalar("charts/beta", rb.beta, global_step)
-
-                # optimize the model
-                optimizer.zero_grad()
-                loss.backward()
-                optimizer.step()
-
-            # update target network
-            if global_step % args.target_network_frequency == 0:
-                for target_param, param in zip(
-                    target_network.parameters(), q_network.parameters()
-                ):
-                    target_param.data.copy_(
-                        args.tau * param.data + (1.0 - args.tau) * target_param.data
-                    )
-
-        # optional early stop on dataset completion
-        if args.capture_dataset and args.stop_on_complete:
-            all_done = (
-                all(
-                    transitions_captured_per_split[s] >= split_targets[s]
-                    for s in splits_in_order
-                )
-                and len(splits_in_order) > 0
-            )
-            if all_done:
-                break
-
-    envs.close()
-    writer.close()
-
-    # write metadata for dataset
-    if args.capture_dataset:
-        if len(obs_chunks) > 0:
-            print(
-                f"Warning: Dropping {len(obs_chunks)} chunks for consistent number of chunks per file.",
-                "Consider changing the chunk_size and chunks_per_file parameters to prevent data-loss.",
-            )
-
-        os.makedirs(args.output_dir, exist_ok=True)
-        metadata_path = os.path.join(args.output_dir, "metadata.json")
-        if os.path.exists(metadata_path):
-            try:
-                with open(metadata_path, "r") as f:
-                    metadata = json.load(f)
-            except Exception:
-                metadata = {}
-        else:
-            metadata = {}
-
-        metadata.setdefault("env", args.env_id)
-        metadata.setdefault("num_actions", int(envs.single_action_space.n))
-        for split in ["train", "val", "test"]:
-            metadata.setdefault(f"num_transitions_{split}", 0)
-            metadata.setdefault(f"avg_episode_len_{split}", 0.0)
-            metadata.setdefault(f"episode_metadata_{split}", [])
-
-        for split_key in splits_in_order:
-            ep_meta_list = episode_metadata_by_split[split_key]
-            if ep_meta_list:
-                metadata[f"episode_metadata_{split_key}"].extend(ep_meta_list)
-                metadata[f"num_transitions_{split_key}"] = int(
-                    np.sum(
-                        [
-                            ep["avg_seq_len"]
-                            for ep in metadata[f"episode_metadata_{split_key}"]
-                        ]
-                    )
-                )
-                metadata[f"avg_episode_len_{split_key}"] = float(
-                    np.mean(
-                        [
-                            ep["avg_seq_len"]
-                            for ep in metadata[f"episode_metadata_{split_key}"]
-                        ]
-                    )
-                )
-
-        with open(metadata_path, "w") as f:
-            json.dump(metadata, f)
diff --git a/data/jasmine_data/atari/visualize_array_record.py b/data/jasmine_data/atari/visualize_array_record.py
deleted file mode 100644
index fca5c13..0000000
--- a/data/jasmine_data/atari/visualize_array_record.py
+++ /dev/null
@@ -1,210 +0,0 @@
-import os
-import math
-import argparse
-import pickle
-from typing import Optional, List
-
-import numpy as np
-
-from PIL import Image, ImageDraw
-
-from array_record.python.array_record_module import ArrayRecordReader
-
-
-def infer_hw_from_bytes(
-    total_elements: int, seq_len: int, channels: int
-) -> tuple[int, int]:
-    assert seq_len > 0 and channels > 0, "sequence_length and channels must be positive"
-    base = total_elements // (seq_len * channels)
-    side = int(math.isqrt(base))
-    if side * side != base:
-        raise ValueError(
-            f"Could not infer square HxW from buffer. elements={total_elements}, seq_len={seq_len}, channels={channels}"
-        )
-    return side, side
-
-
-def load_one_record(input_path: str) -> dict:
-    assert (
-        ArrayRecordReader is not None
-    ), "array_record is not installed. pip install array-record"
-    reader = ArrayRecordReader(str(input_path))
-
-    raw = reader.read()
-    if raw is None:
-        raise RuntimeError(f"No record could be read from {input_path}")
-
-    return pickle.loads(raw)
-
-
-def get_action_meanings(env_id: Optional[str]) -> Optional[List[str]]:
-    if env_id is None:
-        return None
-    try:
-        import gymnasium as gym
-        import ale_py
-
-        gym.register_envs(ale_py)
-        env = gym.make(env_id)
-        try:
-            meanings = list(env.unwrapped.get_action_meanings())  # type: ignore[attr-defined]
-        finally:
-            env.close()
-        return meanings
-    except Exception:
-        print(f"Error getting action meanings for {env_id}")
-        return None
-
-
-def save_frames_with_actions(
-    record: dict,
-    output_dir: str,
-    channels: int,
-    height: Optional[int],
-    width: Optional[int],
-    action_meanings: Optional[List[str]],
-    fps: int,
-) -> None:
-    os.makedirs(output_dir, exist_ok=True)
-
-    seq_len: int = int(record["sequence_length"])
-    raw_video: bytes = record["raw_video"]
-    actions: Optional[np.ndarray] = record.get("actions")
-    if actions is not None:
-        actions = np.asarray(actions).reshape(-1)
-
-    # Decode frames
-    arr = np.frombuffer(raw_video, dtype=np.uint8)
-    total_elements = arr.size
-
-    if height is not None and width is not None:
-        h, w = int(height), int(width)
-        expected = seq_len * h * w * channels
-        assert (
-            expected == total_elements
-        ), f"Expected {expected} elements, got {total_elements}"
-    else:
-        h, w = infer_hw_from_bytes(total_elements, seq_len, channels)
-
-    frames = arr.reshape(seq_len, h, w, channels)
-
-    if actions is not None:
-        assert (
-            actions.shape[0] == seq_len
-        ), f"Expected {seq_len} actions, got {actions.shape[0]}"
-
-    # Save an actions.txt index for quick inspection
-    if actions is not None:
-        with open(os.path.join(output_dir, "actions.txt"), "w") as f:
-            for t in range(seq_len):
-                a = int(actions[t])
-                name = (
-                    action_meanings[a]
-                    if action_meanings is not None and 0 <= a < len(action_meanings)
-                    else None
-                )
-                f.write(f"{t}	{a}" + (f"	{name}" if name is not None else "") + "
")
-
-    # Build frames with overlays and save as GIF
-    duration_ms = max(1, int(1000 / max(1, fps)))
-    imgs: List[Image.Image] = []
-    for t in range(seq_len):
-        img_np_rgb = frames[t]
-        img = Image.fromarray(img_np_rgb, mode="RGB")
-        draw = ImageDraw.Draw(img)
-        if actions is not None:
-            a = int(actions[t])
-            name = (
-                action_meanings[a]
-                if action_meanings is not None and 0 <= a < len(action_meanings)
-                else None
-            )
-            text = f"{a}" if name is None else f"{a} {name}"
-        else:
-            text = "?"
-        draw.text((2, 2), text, fill=(255, 255, 255))
-        imgs.append(img)
-
-    if not imgs:
-        raise RuntimeError("No frames to render.")
-
-    os.makedirs(output_dir, exist_ok=True)
-    gif_path = os.path.join(output_dir, "sequence.gif")
-
-    imgs[0].save(
-        gif_path,
-        save_all=True,
-        append_images=imgs[1:],
-        duration=duration_ms,
-        loop=0,
-    )
-
-    print(
-        f"Saved GIF to {gif_path} with {seq_len} frames (H={h}, W={w}, C={channels}, fps={fps})."
-    )
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Visualize a random sequence from an ArrayRecord by saving frames with actions."
-    )
-    parser.add_argument("--input", required=True, help="Path to .array_record file")
-    parser.add_argument(
-        "--output_dir",
-        required=True,
-        help="Directory to save output images and actions.txt",
-    )
-    parser.add_argument(
-        "--channels",
-        type=int,
-        default=3,
-        help="Number of channels in frames (default 3 for RGB)",
-    )
-    parser.add_argument(
-        "--height",
-        type=int,
-        default=84,
-        help="Frame height (if mismatch, will be inferred)",
-    )
-    parser.add_argument(
-        "--width",
-        type=int,
-        default=84,
-        help="Frame width (if mismatch, will be inferred)",
-    )
-    parser.add_argument(
-        "--env_id",
-        type=str,
-        default=None,
-        help="Gymnasium env id to map action indices to names",
-    )
-    parser.add_argument("--fps", type=int, default=10, help="GIF frames per second")
-
-    args = parser.parse_args()
-
-    assert args.channels == 3, "Only 3 channels are currently supported"
-
-    record = load_one_record(args.input)
-
-    # Print quick summary
-    seq_len = int(record.get("sequence_length", -1))
-    has_actions = "actions" in record and record["actions"] is not None
-    print(
-        f"Loaded record with sequence_length={seq_len}, actions_present={has_actions}"
-    )
-
-    action_meanings = get_action_meanings(args.env_id)
-
-    save_frames_with_actions(
-        record=record,
-        output_dir=args.output_dir,
-        channels=int(args.channels),
-        height=int(args.height) if args.height else None,
-        width=int(args.width) if args.width else None,
-        action_meanings=action_meanings,
-        fps=int(args.fps),
-    )
-
-
-if __name__ == "__main__":
-    main()
diff --git a/data/jasmine_data/coinrun/generate_coinrun_dataset.py b/data/jasmine_data/coinrun/generate_coinrun_dataset.py
index d933dc1..7c34a91 100644
--- a/data/jasmine_data/coinrun/generate_coinrun_dataset.py
+++ b/data/jasmine_data/coinrun/generate_coinrun_dataset.py
@@ -11,11 +11,12 @@ from procgen import ProcgenGym3Env
 import tyro
 import json
 import os
-from jasmine_data.utils import save_chunks
+from data.jasmine_data.utils import save_chunks
 
 
 @dataclass
 class Args:
+    env_name: str = "coinrun"
     num_episodes_train: int = 10000
     num_episodes_val: int = 500
     num_episodes_test: int = 500
@@ -39,7 +40,7 @@ if args.min_episode_length < args.chunk_size:
 
 
 # --- Generate episodes ---
-def generate_episodes(num_episodes, split):
+def generate_episodes(num_episodes, split, start_seed, env_name):
     episode_idx = 0
     episode_metadata = []
     obs_chunks = []
@@ -47,8 +48,9 @@ def generate_episodes(num_episodes, split):
     file_idx = 0
     output_dir_split = os.path.join(args.output_dir, split)
     while episode_idx < num_episodes:
-        seed = np.random.randint(0, 10000)
-        env = ProcgenGym3Env(num=1, env_name="coinrun", start_level=seed)
+        env = ProcgenGym3Env(
+            num=1, env_name=env_name, start_level=start_seed + episode_idx
+        )
 
         observations_seq = []
         actions_seq = []
@@ -116,21 +118,31 @@ def generate_episodes(num_episodes, split):
 
 
 def get_action_space():
-    env = ProcgenGym3Env(num=1, env_name="coinrun", start_level=0)
+    env = ProcgenGym3Env(num=1, env_name=args.env_name, start_level=0)
     return env.ac_space.eltype.n
 
 
 def main():
     # Set random seed and create dataset directories
     np.random.seed(args.seed)
+    train_start_seed = np.random.randint(0, 1000)
+    val_start_seed = train_start_seed + args.num_episodes_train
+    test_start_seed = val_start_seed + args.num_episodes_val
+
     # --- Generate episodes ---
-    train_episode_metadata = generate_episodes(args.num_episodes_train, "train")
-    val_episode_metadata = generate_episodes(args.num_episodes_val, "val")
-    test_episode_metadata = generate_episodes(args.num_episodes_test, "test")
+    train_episode_metadata = generate_episodes(
+        args.num_episodes_train, "train", train_start_seed, args.env_name
+    )
+    val_episode_metadata = generate_episodes(
+        args.num_episodes_val, "val", val_start_seed, args.env_name
+    )
+    test_episode_metadata = generate_episodes(
+        args.num_episodes_test, "test", test_start_seed, args.env_name
+    )
 
     # --- Save metadata ---
     metadata = {
-        "env": "coinrun",
+        "env": args.env_name,
         "num_actions": get_action_space(),
         "num_episodes_train": args.num_episodes_train,
         "num_episodes_val": args.num_episodes_val,
diff --git a/data/jasmine_data/minatar/generate_minatar_breakout_dataset.py b/data/jasmine_data/minatar/generate_minatar_breakout_dataset.py
deleted file mode 100644
index 7d8599b..0000000
--- a/data/jasmine_data/minatar/generate_minatar_breakout_dataset.py
+++ /dev/null
@@ -1,175 +0,0 @@
-"""
-Generates a dataset of random-action Breakout episodes using MinAtar.
-Episodes are saved individually as memory-mapped files for efficient loading.
-"""
-
-from dataclasses import dataclass
-import numpy as np
-import os
-import json
-import tyro
-
-from minatar import Environment
-from jasmine_data.utils import save_chunks
-
-
-@dataclass
-class Args:
-    num_episodes_train: int = 5000000
-    num_episodes_val: int = 25000
-    num_episodes_test: int = 25000
-    output_dir: str = "data/breakout_episodes"
-    min_episode_length: int = 20
-    max_episode_length: int = 500
-    chunk_size: int = 50
-    chunks_per_file: int = 100
-    seed: int = 0
-
-
-args = tyro.cli(Args)
-assert (
-    args.max_episode_length >= args.min_episode_length
-), "Maximum episode length must be greater than or equal to minimum episode length."
-
-if args.min_episode_length < args.chunk_size:
-    print(
-        "Warning: Minimum episode length is smaller than chunk size. "
-        "Episodes shorter than the chunk size will be discarded."
-    )
-
-
-def _obs_to_rgb(obs):
-    # Define a color matrix for each boolean combination
-    color_matrix = np.array(
-        [
-            [0, 0, 0],  # Black for all False
-            [128, 0, 0],  # Maroon
-            [0, 128, 0],  # Dark Green
-            [0, 0, 128],  # Navy
-            [128, 128, 0],  # Olive
-            [128, 0, 128],  # Purple
-            [0, 128, 128],  # Teal
-            [192, 192, 192],  # Silver
-            [128, 128, 128],  # Gray
-            [255, 0, 0],  # Red
-            [0, 255, 0],  # Green
-            [0, 0, 255],  # Blue
-            [255, 255, 0],  # Yellow
-            [255, 0, 255],  # Magenta
-            [0, 255, 255],  # Cyan
-            [255, 255, 255],  # White
-        ],
-        dtype=np.uint8,
-    )
-
-    # Convert boolean array to integer indices
-    indices = obs.dot(1 << np.arange(obs.shape[-1] - 1, -1, -1))
-
-    # Map indices to colors using matrix multiplication
-    obs = color_matrix[indices]
-    return obs
-
-
-def generate_episodes(num_episodes: int, split: str, seed: int):
-    episode_idx = 0
-    episode_metadata = []
-    obs_chunks = []
-    act_chunks = []
-    file_idx = 0
-    output_dir_split = os.path.join(args.output_dir, split)
-    os.makedirs(output_dir_split, exist_ok=True)
-
-    while episode_idx < num_episodes:
-        env = Environment("breakout", sticky_action_prob=0.0)
-        env.seed(seed)
-        env.reset()
-        obs_seq, act_seq = [], []
-        episode_obs_chunks, episode_act_chunks = [], []
-
-        step_t = 0
-        for step_t in range(args.max_episode_length):
-            obs = env.state()  # shape: (10,10,num_channels)
-            obs = _obs_to_rgb(obs)
-            action = np.random.randint(env.num_actions())
-            _, done = env.act(action)
-            obs_seq.append(obs.astype(np.uint8))
-            act_seq.append(action)
-
-            if len(obs_seq) == args.chunk_size:
-                episode_obs_chunks.append(np.stack(obs_seq))
-                episode_act_chunks.append(np.array(act_seq))
-                obs_seq, act_seq = [], []
-
-            if done:
-                break
-
-        if step_t + 1 >= args.min_episode_length:
-            if obs_seq:
-                if len(obs_seq) < args.chunk_size:
-                    print(
-                        f"Warning: Inconsistent chunk sizes. Episode has {len(obs_seq)} frames "
-                        f"(less than chunk_size {args.chunk_size})."
-                    )
-                episode_obs_chunks.append(np.stack(obs_seq))
-                episode_act_chunks.append(np.array(act_seq))
-
-            obs_chunks_data = episode_obs_chunks
-            act_chunks_data = episode_act_chunks
-            obs_chunks.extend(obs_chunks_data)
-            act_chunks.extend(act_chunks_data)
-
-            ep_metadata, file_idx, obs_chunks, act_chunks = save_chunks(
-                file_idx, args.chunks_per_file, output_dir_split, obs_chunks, act_chunks
-            )
-            episode_metadata.extend(ep_metadata)
-
-            print(f"[{split}] Episode {episode_idx} completed, length: {step_t + 1}")
-            episode_idx += 1
-        else:
-            print(f"Episode too short ({step_t + 1}), resampling...")
-        seed += 1
-
-    if len(obs_chunks) > 0:
-        print(
-            f"Warning: Dropping {len(obs_chunks)} leftover chunks. "
-            "Consider adjusting chunk_size or chunks_per_file."
-        )
-
-    print(f"Done generating {split} split.")
-    return episode_metadata, seed
-
-
-def get_action_space() -> int:
-    env = Environment("breakout")
-    return env.num_actions()
-
-
-def main():
-    os.makedirs(args.output_dir, exist_ok=True)
-    seed = args.seed
-    train_meta, seed = generate_episodes(args.num_episodes_train, "train", seed)
-    val_meta, seed = generate_episodes(args.num_episodes_val, "val", seed)
-    test_meta, seed = generate_episodes(args.num_episodes_test, "test", seed)
-
-    metadata = {
-        "env": "MinAtar-Breakout",
-        "num_actions": get_action_space(),
-        "num_episodes_train": args.num_episodes_train,
-        "num_episodes_val": args.num_episodes_val,
-        "num_episodes_test": args.num_episodes_test,
-        "avg_episode_len_train": np.mean([ep["avg_seq_len"] for ep in train_meta]),
-        "avg_episode_len_val": np.mean([ep["avg_seq_len"] for ep in val_meta]),
-        "avg_episode_len_test": np.mean([ep["avg_seq_len"] for ep in test_meta]),
-        "episode_metadata_train": train_meta,
-        "episode_metadata_val": val_meta,
-        "episode_metadata_test": test_meta,
-    }
-
-    with open(os.path.join(args.output_dir, "metadata.json"), "w") as f:
-        json.dump(metadata, f)
-
-    print("Done generating dataset.")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/data/pyproject.toml b/data/pyproject.toml
index c389ab4..3c6136b 100644
--- a/data/pyproject.toml
+++ b/data/pyproject.toml
@@ -18,12 +18,6 @@ dependencies = [
   "pillow>=11.3.0",
   "tqdm>=4.67.1",
   "tyro>=0.8.5",
-  "gymnasium[atari, accept-rom-license]==0.29.1",
-  "minatar>=1.0.15",
-  "torch>=2.8.0",
-  "tensorboard>=2.20.0",
-  "opencv-python>=4.6.0.66,<5",
-  "einops>=0.8.0",
 ]
 
 [build-system]
diff --git a/jasmine/train_dynamics.py b/jasmine/train_dynamics.py
index 06cd966..ba6f5b9 100644
--- a/jasmine/train_dynamics.py
+++ b/jasmine/train_dynamics.py
@@ -39,8 +39,8 @@ class Args:
     seed: int = 0
     seq_len: int = 16
     image_channels: int = 3
-    image_height: int = 64
-    image_width: int = 64
+    image_height: int = 90
+    image_width: int = 160
     data_dir: str = ""
     save_ckpt: bool = False
     restore_ckpt: bool = False
@@ -50,7 +50,7 @@ class Args:
     max_lr: float = 3e-5
     decay_end: float = 0.0
     wsd_decay_steps: int = (
-        20_000  # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
+        10000  # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
     )
     warmup_steps: int = 5000
     lr_schedule: str = "wsd"  # supported options: wsd, cos
@@ -59,7 +59,7 @@ class Args:
     tokenizer_ffn_dim: int = 2048
     latent_patch_dim: int = 32
     num_patch_latents: int = 1024
-    patch_size: int = 16
+    patch_size: int = 4
     tokenizer_num_blocks: int = 4
     tokenizer_num_heads: int = 8
     tokenizer_checkpoint: str = ""
@@ -80,27 +80,26 @@ class Args:
     dyna_num_heads: int = 8
     dropout: float = 0.0
     mask_limit: float = 0.5
-    z_loss_weight: float = 0.0
     param_dtype = jnp.float32
     dtype = jnp.bfloat16
     use_flash_attention: bool = True
     use_gt_actions: bool = False
     # Logging
-    log: bool = True
+    log: bool = False
     entity: str = ""
     project: str = ""
     name: str = "train_dynamics"
     tags: list[str] = field(default_factory=lambda: ["dynamics"])
-    log_interval: int = 50
-    log_image_interval: int = 1000
+    log_interval: int = 5
+    log_image_interval: int = 250
     ckpt_dir: str = ""
-    log_checkpoint_interval: int = 5000
-    log_checkpoint_keep_period: int = 20_000
+    log_checkpoint_interval: int = 25000
+    log_checkpoint_keep_period: int = 20000
     log_gradients: bool = False
     val_data_dir: str = ""
     val_interval: int = 20_000
     val_steps: int = 50
-    eval_full_frame: bool = True
+    eval_full_frame: bool = False
     val_maskgit_steps: int = 25
     val_temperature: float = 1
     val_sample_argmax: bool = False
@@ -326,48 +325,20 @@ def restore_or_initialize_components(
     return step, optimizer, train_iterator, val_iterator, rng
 
 
-def _calculate_top_k_accuracy(
-    token_logits_BTNV: jax.Array,
-    video_tokens_BTN: jax.Array,
-    mask_BTN: jax.Array,
-    k: int,
-) -> jax.Array:
-    _, topk_indices_BTNK = jax.lax.top_k(token_logits_BTNV, k)
-    topk_correct = jnp.any(
-        topk_indices_BTNK == video_tokens_BTN[..., jnp.newaxis], axis=-1
-    )
-    topk_acc = (mask_BTN * topk_correct).sum() / mask_BTN.sum()
-    return topk_acc
-
-
 def _calculate_step_metrics(
     outputs: dict[str, jax.Array],
     gt: jax.Array,
     num_actions: int,
     num_patch_latents: int,
 ) -> tuple[jax.Array, dict]:
-    mask_BTN = outputs["mask"]
+    mask = outputs["mask"]
     outputs["token_logits"] = outputs["token_logits"].astype(jnp.float32)
     ce_loss = optax.softmax_cross_entropy_with_integer_labels(
         outputs["token_logits"], outputs["video_tokens"]
     )
-    ce_loss = (mask_BTN * ce_loss).sum() / mask_BTN.sum()
-    z_val = jax.nn.logsumexp(outputs["token_logits"], axis=-1)
-    z_loss_metric = (mask_BTN * (z_val**2)).sum() / mask_BTN.sum()
-
-    masked_token_top_1_acc = _calculate_top_k_accuracy(
-        outputs["token_logits"], outputs["video_tokens"], mask_BTN, 1
-    )
-    masked_token_top_2_acc = _calculate_top_k_accuracy(
-        outputs["token_logits"], outputs["video_tokens"], mask_BTN, 2
-    )
-    masked_token_top_5_acc = _calculate_top_k_accuracy(
-        outputs["token_logits"], outputs["video_tokens"], mask_BTN, 5
-    )
-    masked_token_top_16_acc = _calculate_top_k_accuracy(
-        outputs["token_logits"], outputs["video_tokens"], mask_BTN, 16
-    )
-
+    ce_loss = (mask * ce_loss).sum() / mask.sum()
+    acc = outputs["token_logits"].argmax(-1) == outputs["video_tokens"]
+    acc = (mask * acc).sum() / mask.sum()
     select_probs = jax.nn.softmax(outputs["token_logits"])
     gt_val = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
     recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
@@ -381,14 +352,10 @@ def _calculate_step_metrics(
     codebook_usage_tokenizer = (index_counts_tokenizer != 0).mean()
     metrics = dict(
         cross_entropy_loss=ce_loss,
-        masked_token_top1_accuracy=masked_token_top_1_acc,
-        masked_token_top2_accuracy=masked_token_top_2_acc,
-        masked_token_top5_accuracy=masked_token_top_5_acc,
-        masked_token_top16_accuracy=masked_token_top_16_acc,
+        masked_token_accuracy=acc,
         select_logit=outputs["token_logits"].max(-1).mean(),
         select_p=select_probs.max(-1).mean(),
         entropy=jax.scipy.special.entr(select_probs).sum(-1).mean(),
-        z_loss=z_loss_metric,
         psnr=psnr,
         ssim=ssim,
         codebook_usage_tokenizer=codebook_usage_tokenizer,
@@ -492,10 +459,7 @@ def main(args: Args) -> None:
         ce_loss, metrics = _calculate_step_metrics(
             outputs, gt, args.num_actions, args.num_patch_latents
         )
-        z_loss = metrics["z_loss"]
-        total_loss = ce_loss + args.z_loss_weight * z_loss
-        metrics["total_loss"] = total_loss
-        return total_loss, (outputs["recon"], metrics)
+        return ce_loss, (outputs["recon"], metrics)
 
     @nnx.jit(donate_argnums=0)
     def train_step(
diff --git a/jasmine/train_lam.py b/jasmine/train_lam.py
index 95fcb67..9d9874a 100644
--- a/jasmine/train_lam.py
+++ b/jasmine/train_lam.py
@@ -38,8 +38,8 @@ class Args:
     seed: int = 0
     seq_len: int = 16
     image_channels: int = 3
-    image_height: int = 64
-    image_width: int = 64
+    image_height: int = 90
+    image_width: int = 160
     data_dir: str = ""
     save_ckpt: bool = False
     restore_ckpt: bool = False
@@ -50,7 +50,7 @@ class Args:
     max_lr: float = 3e-5
     decay_end: float = 0.0
     wsd_decay_steps: int = (
-        20_000  # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
+        10000  # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
     )
     warmup_steps: int = 5000
     lr_schedule: str = "wsd"  # supported options: wsd, cos
@@ -67,22 +67,22 @@ class Args:
     codebook_dropout: float = 0.0
     param_dtype = jnp.float32
     dtype = jnp.bfloat16
-    use_flash_attention: bool = True
     # Logging
-    log: bool = True
+    log: bool = False
     entity: str = ""
     project: str = ""
     name: str = "train_lam"
     tags: list[str] = field(default_factory=lambda: ["lam"])
-    log_interval: int = 50
-    log_image_interval: int = 1000
+    log_interval: int = 5
+    log_image_interval: int = 250
     ckpt_dir: str = ""
-    log_checkpoint_interval: int = 5000
-    log_checkpoint_keep_period: int = 20_000
+    log_checkpoint_interval: int = 10000
+    log_checkpoint_keep_period: int = 20000
     val_data_dir: str = ""
     val_interval: int = 20_000
     val_steps: int = 50
     wandb_id: str = ""
+    use_flash_attention: bool = True
 
 
 def build_model(args: Args, rng: jax.Array) -> tuple[LatentActionModel, jax.Array]:
diff --git a/jasmine/train_tokenizer.py b/jasmine/train_tokenizer.py
index 0bdcd7b..33206b4 100644
--- a/jasmine/train_tokenizer.py
+++ b/jasmine/train_tokenizer.py
@@ -38,8 +38,8 @@ class Args:
     seed: int = 0
     seq_len: int = 16
     image_channels: int = 3
-    image_height: int = 64
-    image_width: int = 64
+    image_height: int = 90
+    image_width: int = 160
     data_dir: str = ""
     save_ckpt: bool = False
     restore_ckpt: bool = False
@@ -50,7 +50,7 @@ class Args:
     max_lr: float = 3e-4
     decay_end: float = 0.0
     wsd_decay_steps: int = (
-        30_000  # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
+        20000  # NOTE: wsd_decay_steps will only be used when using a wsd-schedule
     )
     lr_schedule: str = "wsd"  # supported options: wsd, cos
     warmup_steps: int = 10000
@@ -59,30 +59,30 @@ class Args:
     ffn_dim: int = 2048
     latent_dim: int = 32
     num_latents: int = 1024
-    patch_size: int = 16
+    patch_size: int = 4
     num_blocks: int = 4
     num_heads: int = 8
     dropout: float = 0.0
     codebook_dropout: float = 0.01
     param_dtype = jnp.float32
     dtype = jnp.bfloat16
-    use_flash_attention: bool = True
     # Logging
-    log: bool = True
+    log: bool = False
     entity: str = ""
     project: str = ""
     name: str = "train_tokenizer"
     tags: list[str] = field(default_factory=lambda: ["tokenizer"])
-    log_interval: int = 50
-    log_image_interval: int = 1000
+    log_interval: int = 5
+    log_image_interval: int = 250
     ckpt_dir: str = ""
-    log_checkpoint_interval: int = 1000
-    log_checkpoint_keep_period: int = 20_000
+    log_checkpoint_interval: int = 10000
+    log_checkpoint_keep_period: int = 20000
     log_gradients: bool = False
     val_data_dir: str = ""
     val_interval: int = 20_000
     val_steps: int = 50
     wandb_id: str = ""
+    use_flash_attention: bool = True
 
 
 def build_model(args: Args, rng: jax.Array) -> tuple[TokenizerVQVAE, jax.Array]: