p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] input_pipeline/utils.py (python)
import os
import pickle
import numpy as np
from array_record.python.array_record_module import ArrayRecordWriter


def save_chunks(obs_chunks, file_idx, chunks_per_file, output_dir, act_chunks=None):
    os.makedirs(output_dir, exist_ok=True)

    metadata = []
    while len(obs_chunks) >= chunks_per_file:
        chunk_batch = obs_chunks[:chunks_per_file]
        obs_chunks = obs_chunks[chunks_per_file:]
        act_chunk_batch = None
        if act_chunks:
            act_chunk_batch = act_chunks[:chunks_per_file]
            act_chunks = act_chunks[chunks_per_file:]
        episode_path = os.path.join(output_dir, f"data_{file_idx:04d}.array_record")
        writer = ArrayRecordWriter(str(episode_path), "group_size:1")
        seq_lens = []
        for idx, chunk in enumerate(chunk_batch):
            seq_len = chunk.shape[0]
            seq_lens.append(seq_len)
            chunk_record = {
                "raw_video": chunk.tobytes(),
                "sequence_length": seq_len,
            }
            if act_chunk_batch:
                assert len(chunk) == len(
                    act_chunk_batch[idx]
                ), f"Observation data length and action sequence length do not match: {len(chunk)} != {len(act_chunk_batch[idx])}"
                chunk_record["actions"] = act_chunk_batch[idx]
            writer.write(pickle.dumps(chunk_record))
        writer.close()
        file_idx += 1
        metadata.append(
            {
                "path": episode_path,
                "num_chunks": len(chunk_batch),
                "avg_seq_len": np.mean(seq_lens),
            }
        )
        print(f"Created {episode_path} with {len(chunk_batch)} video chunks")

    return metadata, obs_chunks, file_idx, act_chunks
[content] input_pipeline/generate_atari_dataset.py (python)
        # Update metadata for the splits we generated this run
        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)
                # Maintain existing convention: num_episodes_* equals number of files (metadata entries)
                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}"]])
                )
[content] input_pipeline/generate_atari_dataset.py (python)
        if args.capture_dataset and args.dataset_stop_on_complete:
            if use_multi_split:
                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
            else:
                if args.dataset_num_episodes > 0 and episodes_captured >= args.dataset_num_episodes:
                    break
[content] input_pipeline/generate_atari_dataset.py (python)
                            file_idx_by_split[current_split] = next_file_idx
                            episode_metadata_by_split[current_split].extend(ep_metadata)

                            # Increment counters
                            if use_multi_split:
                                episodes_captured_per_split[current_split] += 1
                            else:
                                episodes_captured += 1

                            # If current split target reached, advance to next split
                            if use_multi_split and episodes_captured_per_split[current_split] >= split_targets[current_split]:
                                # Drop any leftover chunks for clean split boundary
                                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 = []
                                # Advance to next active split if any
                                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.dataset_output_dir, current_split)
                                    os.makedirs(split_dir, exist_ok=True)
[content] input_pipeline/generate_atari_dataset.py (python)
                            # 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.dataset_chunks_per_file,
                                split_dir,
                                act_chunks,
[content] input_pipeline/generate_atari_dataset.py (python)
                    # In multi-split mode, keep capturing until all requested splits reach their targets.
                    continue_capturing_single = (args.dataset_num_episodes == 0 or episodes_captured < args.dataset_num_episodes)
                    continue_capturing_multi = (
                        not use_multi_split
                        or any(
                            episodes_captured_per_split[s] < split_targets[s]
                            for s in splits_in_order
                        )
                    )
                    if args.capture_dataset and (continue_capturing_single and continue_capturing_multi):
[content] input_pipeline/generate_atari_dataset.py (python)

    # Single-split fallback episodes counter
    episodes_captured = 0

    # Active split
    current_split_idx = 0
    current_split = splits_in_order[0]
    split_dir = os.path.join(args.dataset_output_dir, current_split)
[content] input_pipeline/generate_atari_dataset.py (python)
    # Multi-split support: if any of the per-split targets > 0, route episodes accordingly.
    split_targets = {
        "train": args.dataset_num_episodes_train,
        "val": args.dataset_num_episodes_val,
        "test": args.dataset_num_episodes_test,
    }
    use_multi_split = any(t > 0 for t in split_targets.values())

    # Determine splits to run (order: train -> val -> test when requested)
    splits_in_order = [s for s in ["train", "val", "test"] if split_targets[s] > 0] if use_multi_split else [args.dataset_split]

    # Per-split tracking
    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"]}

    # Shared chunk accumulators (flushed when switching splits to avoid mixing)
[content] input_pipeline/generate_atari_dataset.py (python)
    # Multi-split episode targets (prefer these over single-split args if any > 0)
    dataset_num_episodes_train: int = 10000
    dataset_num_episodes_val: int = 500
    dataset_num_episodes_test: int = 500
[content] input_pipeline/generate_atari_dataset.py (python)
        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}"]])
            )
[content] input_pipeline/generate_atari_dataset.py (python)
        samples: dict[str, Any] = {}
        samples["observations"] = torch.from_numpy(self.buffer_obs[indices]).to(self.device)
        samples["actions"] = torch.from_numpy(self.buffer_actions[indices]).to(self.device).unsqueeze(1)
        samples["rewards"] = torch.from_numpy(self.buffer_rewards[indices]).to(self.device).unsqueeze(1)
        samples["next_observations"] = torch.from_numpy(self.buffer_next_obs[indices]).to(self.device)
        samples["dones"] = torch.from_numpy(self.buffer_dones[indices]).to(self.device).unsqueeze(1)
[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),
        }
[content] input_pipeline/generate_atari_dataset.py (python)
        if args.capture_dataset and args.dataset_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
[content] input_pipeline/generate_atari_dataset.py (python)
                            # 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.dataset_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)

                            # Increment counters
                            episodes_captured_per_split[current_split] += 1

                            # If current split target reached, advance to next split
                            if episodes_captured_per_split[current_split] >= split_targets[current_split]:
                                # Drop any leftover chunks for clean split boundary
                                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 = []
                                # Advance to next active split if any
                                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.dataset_output_dir, current_split)
                                    os.makedirs(split_dir, exist_ok=True)
[content] input_pipeline/generate_atari_dataset.py (python)
                    # Keep capturing until all requested splits reach their targets.
                    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:
[content] input_pipeline/generate_atari_dataset.py (python)
    # Multi-split support: if any of the per-split targets > 0, route episodes accordingly.
    split_targets = {
        "train": args.dataset_num_episodes_train,
        "val": args.dataset_num_episodes_val,
        "test": args.dataset_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]

    # Per-split tracking
    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"]}

    # Shared chunk accumulators (flushed when switching splits to avoid mixing)