p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] cleanrl/ppo_atari_envpool.py (python)
            if traj_saver is not None:
                # Use true termination flags when available
                terminated = np.asarray(info.get("terminated", next_done.detach().cpu().numpy()), dtype=bool)
                traj_saver.add_step(obs_to_save, actions_to_save, terminated)
[content] cleanrl/ppo_atari_envpool.py (python)
            # Capture current observation and action prior to stepping the envs
            if traj_saver is not None:
                obs_to_save = next_obs.detach().cpu().numpy().astype(np.uint8)
                actions_to_save = action.detach().cpu().numpy()
[content] cleanrl/ppo_atari_envpool.py (python)
    # Optional: trajectory saver
    traj_saver = None
    if args.capture_trajectories:
        traj_saver = TrajectorySaver(
            output_dir=args.trajectories_output_dir,
            num_envs=args.num_envs,
            chunk_size=args.trajectories_chunk_size,
            chunks_per_file=args.trajectories_chunks_per_file,
        )
[content] cleanrl/ppo_atari_envpool.py (python)
    # Trajectory capture
    capture_trajectories: bool = False
    """if toggled, save observation frames and actions during training"""
    trajectories_output_dir: str = "data/atari_trajectories"
    """directory to store trajectory chunks (.npz)"""
    trajectories_chunk_size: int = 160
    """number of time steps per chunk before writing"""
    trajectories_chunks_per_file: int = 100
    """number of chunks grouped per npz file"""
[content] cleanrl/ppo_atari_envpool.py (python)
        # Initialize for type-checkers
        approx_kl = torch.tensor(0.0)
        old_approx_kl = torch.tensor(0.0)
        v_loss = torch.tensor(0.0)
        pg_loss = torch.tensor(0.0)
        entropy_loss = torch.tensor(0.0)
[content] cleanrl/ppo_atari_envpool.py (python)
            if traj_saver is not None and obs_to_save is not None and actions_to_save is not None:
                # Use true termination flags when available
                terminated = np.asarray(info.get("terminated", next_done.detach().cpu().numpy()), dtype=bool)
                traj_saver.add_step(obs_to_save, actions_to_save, terminated)
[content] cleanrl/ppo_atari_envpool.py (python)
            # Capture current observation and action prior to stepping the envs
            obs_to_save = None
            actions_to_save = None
            if traj_saver is not None:
                obs_to_save = next_obs.detach().cpu().numpy().astype(np.uint8)
                actions_to_save = action.detach().cpu().numpy()
[content] cleanrl/ppo_atari_envpool.py (python)
    assert isinstance(envs.action_space, spaces.Discrete), "only discrete action space is supported"

    # Optional: trajectory saver
    traj_saver = None
    if args.capture_trajectories:
        traj_saver = TrajectorySaver(
            output_dir=args.trajectories_output_dir,
            num_envs=args.num_envs,
            chunk_size=args.trajectories_chunk_size,
            chunks_per_file=args.trajectories_chunks_per_file,
        )
[tab] /home/franz.srambical/jafar/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