p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] input_pipeline/generate_atari_dataset.py (python)
"""
Generates a dataset of random-action Atari episodes.
Episodes are saved individually as memory-mapped files for efficient loading.
Replicates the behavior of generate_coinrun_dataset.py but for Atari.
"""
[content] input_pipeline/generate_atari_dataset.py (python)
    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)
[content] input_pipeline/generate_atari_dataset.py (python)
    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
[content] input_pipeline/generate_atari_dataset.py (python)
# 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)
[content] input_pipeline/generate_atari_dataset.py (python)
    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"]
)
[content] input_pipeline/generate_atari_dataset.py (python)
    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
[content] input_pipeline/generate_atari_dataset.py (python)
# 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
[content] input_pipeline/generate_atari_dataset.py (python)
    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)
[content] input_pipeline/generate_atari_dataset.py (python)
    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))
[content] input_pipeline/generate_atari_dataset.py (python)
        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()
[content] input_pipeline/generate_atari_dataset.py (python)
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
[content] input_pipeline/generate_atari_dataset.py (python)
        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)
[content] input_pipeline/generate_atari_dataset.py (python)
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)
[content] input_pipeline/generate_atari_dataset.py (python)
    """the id of the environment"""
    total_timesteps: int = 10000000
    """total timesteps of the experiments"""
    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"""
[content] input_pipeline/generate_atari_dataset.py (python)
    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: 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"""
[content] input_pipeline/generate_atari_dataset.py (python)
# docs and experiment results can be found at https://docs.cleanrl.dev/rl-algorithms/rainbow/#rainbow_ataripy
import collections
import math
import os
import random
import time
from collections import deque
[content] input_pipeline/generate_atari_dataset.py (python)
        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())
[content] input_pipeline/generate_atari_dataset.py (python)
        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)
[content] input_pipeline/generate_atari_dataset.py (python)
# 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] = min(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
[content] input_pipeline/generate_atari_dataset.py (python)
    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)
[content] input_pipeline/generate_atari_dataset.py (python)
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)
[content] input_pipeline/generate_atari_dataset.py (python)
        # optional early stop on dataset completion
        if args.capture_dataset and args.dataset_stop_on_complete and args.dataset_num_episodes > 0 and episodes_captured >= args.dataset_num_episodes:
            break
[content] input_pipeline/generate_atari_dataset.py (python)
                    # finalize episode capture
                    if args.capture_dataset and (args.dataset_num_episodes == 0 or episodes_captured < args.dataset_num_episodes):
                        current_len = len(observations_seq)
                        if current_len >= args.dataset_min_episode_length:
                            frames = np.concatenate(observations_seq, axis=0).astype(np.uint8)
                            acts = np.concatenate(actions_seq, axis=0).astype(np.int64)

                            episode_obs_chunks = []
                            episode_act_chunks = []
                            # chunk like coinrun (include remainder, warn if short)
                            start_idx = 0
                            while start_idx < current_len:
                                end_idx = min(start_idx + args.dataset_chunk_size, current_len)
                                if end_idx - start_idx < args.dataset_chunk_size:
                                    print(
                                        f"Warning: Inconsistent chunk_sizes. Episode has {current_len} frames, "
                                        f"which is smaller than the requested chunk_size: {args.dataset_chunk_size}. "
                                        "This might lead to performance degradation during training."
                                    )
                                episode_obs_chunks.append(frames[start_idx:end_idx][None, ...])
                                episode_act_chunks.append(acts[start_idx:end_idx][None, ...])
                                start_idx = end_idx

                            obs_chunks_data = [
                                np.concatenate(seq, axis=0).astype(np.uint8) for seq in episode_obs_chunks
                            ]
                            act_chunks_data = [np.concatenate(act, axis=0) for act in episode_act_chunks]
                            obs_chunks.extend(obs_chunks_data)
                            act_chunks.extend(act_chunks_data)

                            ep_metadata, obs_chunks, file_idx, act_chunks = save_chunks(
                                obs_chunks, file_idx, args.dataset_chunks_per_file, split_dir, act_chunks
                            )
                            episode_metadata.extend(ep_metadata)
                            episodes_captured += 1
                        else:
                            print(f"Episode too short ({current_len}), skipping capture...")

                        # reset buffers for next episode
                        observations_seq = []
                        actions_seq = []
[content] input_pipeline/generate_atari_dataset.py (python)
        # capture step
        if args.capture_dataset and (args.dataset_num_episodes == 0 or episodes_captured < args.dataset_num_episodes):
            observations_seq.append(next_obs.astype(np.uint8))
            actions_seq.append(actions.astype(np.int64))
[content] input_pipeline/generate_atari_dataset.py (python)
    # dataset capture state
    episodes_captured = 0
    obs_chunks: list[np.ndarray] = []
    act_chunks: list[np.ndarray] = []
    file_idx = 0
    episode_metadata: list[dict] = []
    split_dir = os.path.join(args.dataset_output_dir, args.dataset_split)
    if args.capture_dataset:
        os.makedirs(split_dir, exist_ok=True)
[content] input_pipeline/generate_atari_dataset.py (python)
    # Dataset capture (surgical additions)
    capture_dataset: bool = False
    dataset_output_dir: str = "data/atari_episodes"
    dataset_split: str = "train"  # one of {train,val,test}
    dataset_num_episodes: int = 0  # 0 = no limit
    dataset_chunk_size: int = 160
    dataset_chunks_per_file: int = 100
    dataset_min_episode_length: int = 1
    dataset_stop_on_complete: bool = True
[content] input_pipeline/generate_atari_dataset.py (python)
        env = gym.wrappers.ResizeObservation(env, (84, 84))  # type: ignore[attr-defined]
        env = gym.wrappers.GrayScaleObservation(env)  # type: ignore[attr-defined]
        env = gym.wrappers.FrameStack(env, 4)  # type: ignore[attr-defined]
[content] input_pipeline/generate_atari_dataset.py (python)
    # 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.dataset_output_dir, exist_ok=True)
        metadata_path = os.path.join(args.dataset_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_episodes_{split}", 0)
            metadata.setdefault(f"avg_episode_len_{split}", 0.0)
            metadata.setdefault(f"episode_metadata_{split}", [])

        split_key = args.dataset_split
        metadata[f"episode_metadata_{split_key}"].extend(episode_metadata)
        metadata[f"num_episodes_{split_key}"] = len(metadata[f"episode_metadata_{split_key}"])
        if metadata[f"num_episodes_{split_key}"] > 0:
            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)
"""
Generates a dataset of random-action Atari episodes.
Episodes are saved individually as memory-mapped files for efficient loading.
Replicates the behavior of generate_coinrun_dataset.py but for Atari.
"""

from dataclasses import dataclass

import gymnasium as gym
import numpy as np
import tyro
import json
import os
from cleanrl_utils.atari_wrappers import (
    ClipRewardEnv,
    EpisodicLifeEnv,
    FireResetEnv,
    MaxAndSkipEnv,
    NoopResetEnv,
)
from utils import save_chunks  # type: ignore


"""
Old dataset-only generator removed in favor of integrated Rainbow + capture mode.
"""