[tab] input_pipeline/generate_atari_dataset.py (python)
# adapted from https://github.com/vwxyzjn/cleanrl/blob/master/cleanrl/rainbow_atari.py
import collections
import math
import os
import random
import time
from collections import deque
from dataclasses import dataclass

import gymnasium as gym
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 cleanrl_utils.atari_wrappers import (
    ClipRewardEnv,
    EpisodicLifeEnv,
    FireResetEnv,
    MaxAndSkipEnv,
    NoopResetEnv,
)

# FIXME (f.srambical): remove the try-except block
try:
    from utils import save_chunks  # type: ignore
except Exception:  # pragma: no cover
    from input_pipeline.utils import save_chunks  # type: ignore
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"""
    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"""

    # Dataset capture
    capture_dataset: bool = True
    num_episodes_train: int = 10000
    num_episodes_val: int = 500
    num_episodes_test: int = 500
    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"
    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 = {
        "train": args.num_episodes_train,
        "val": args.num_episodes_val,
        "test": args.num_episodes_test,
    }
    # Determine splits to run (order: train -> val -> test)
    splits_in_order = [s for s in ["train", "val", "test"] if split_targets[s] > 0]

    episodes_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] = []
    for global_step in range(args.total_timesteps):
        # anneal PER beta to 1
        rb.beta = min(
            1.0,
            args.prioritized_replay_beta
            + global_step * (1.0 - args.prioritized_replay_beta) / args.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:
            observations_seq.append(next_obs.astype(np.uint8))
            actions_seq.append(actions.astype(np.int64))

        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(
                        episodes_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.concatenate(observations_seq, axis=0).astype(
                                np.uint8
                            )
                            acts = np.concatenate(actions_seq, axis=0).astype(np.int64)

                            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, obs_chunks, next_file_idx, act_chunks = (
                                save_chunks(
                                    obs_chunks,
                                    file_idx_by_split[current_split],
                                    args.chunks_per_file,
                                    split_dir,
                                    act_chunks,
                                )
                            )
                            file_idx_by_split[current_split] = next_file_idx
                            episode_metadata_by_split[current_split].extend(ep_metadata)

                            episodes_captured_per_split[current_split] += 1

                            if (
                                episodes_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(
                    episodes_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_episodes_{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_episodes_{split_key}"] = len(
                    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)