[tab] cleanrl/trajectory_saver.py (python)
import os
import json
from typing import List, Dict, Any

import numpy as np


class TrajectorySaver:
    """
    Collects per-environment observation frames and actions, chunks them,
    and writes chunked files to disk. Designed to be lightweight and only
    enabled when trajectory capture is requested.

    Behavior mirrors the high level flow in `generate_coinrun_dataset.py`:
    - Build sequences of length `chunk_size` from steps
    - At episode end, if a partial sequence exists (< chunk_size), include it
      with a warning about inconsistent chunk sizes
    - Group `chunks_per_file` chunks per output file. Any remainder at the
      very end is dropped with a warning
    - Store both observation chunks and action chunks side-by-side
    """

    def __init__(
        self,
        output_dir: str,
        num_envs: int,
        chunk_size: int = 160,
        chunks_per_file: int = 100,
    ) -> None:
        self.output_dir = output_dir
        self.chunks_dir = os.path.join(output_dir, "chunks")
        os.makedirs(self.chunks_dir, exist_ok=True)
        self.num_envs = num_envs
        self.chunk_size = int(chunk_size)
        self.chunks_per_file = int(chunks_per_file)

        # Per-env rolling buffers for the current episode
        self.env_buffers: List[Dict[str, Any]] = [
            {"obs_seq": [], "act_seq": [], "episode_obs_chunks": [], "episode_act_chunks": []}
            for _ in range(self.num_envs)
        ]

        # Global chunk buffers waiting to be flushed to disk
        self.obs_chunks: List[np.ndarray] = []
        self.act_chunks: List[np.ndarray] = []
        self.file_idx: int = 0

        # Metadata (optional). We keep a minimal JSONL with per-file info
        self.metadata_path = os.path.join(self.output_dir, "metadata.jsonl")

    def _finalize_episode_for_env(self, env_idx: int) -> None:
        buf = self.env_buffers[env_idx]
        obs_seq: List[np.ndarray] = buf["obs_seq"]
        act_seq: List[np.ndarray] = buf["act_seq"]

        if len(obs_seq) > 0:
            # Partial chunk at episode end
            if len(obs_seq) < self.chunk_size:
                print(
                    f"Warning: Inconsistent chunk sizes. Env {env_idx} episode ended with "
                    f"{len(obs_seq)} frames (< chunk_size={self.chunk_size}). "
                    "Including the partial chunk may impact training I/O performance."
                )
            self._append_chunk_from_sequences(env_idx)

        # Move episode chunks to global buffers
        if len(buf["episode_obs_chunks"]) > 0:
            self.obs_chunks.extend(buf["episode_obs_chunks"])
            self.act_chunks.extend(buf["episode_act_chunks"])
            buf["episode_obs_chunks"].clear()
            buf["episode_act_chunks"].clear()

        self._save_if_ready()

    def _append_chunk_from_sequences(self, env_idx: int) -> None:
        buf = self.env_buffers[env_idx]
        obs_seq: List[np.ndarray] = buf["obs_seq"]
        act_seq: List[np.ndarray] = buf["act_seq"]

        if len(obs_seq) == 0:
            return

        # Stack along time dimension. Shapes become:
        #   obs_chunk: (T, C, H, W)  dtype=uint8
        #   act_chunk: (T,)          dtype=int64 (or int)
        obs_chunk = np.stack(obs_seq, axis=0).astype(np.uint8)
        act_chunk = np.asarray(act_seq, dtype=np.int64)

        buf["episode_obs_chunks"].append(obs_chunk)
        buf["episode_act_chunks"].append(act_chunk)

        # Reset rolling sequences
        buf["obs_seq"] = []
        buf["act_seq"] = []

    def _save_if_ready(self) -> None:
        while len(self.obs_chunks) >= self.chunks_per_file:
            obs_to_write = self.obs_chunks[: self.chunks_per_file]
            act_to_write = self.act_chunks[: self.chunks_per_file]

            # Remove from buffers
            self.obs_chunks = self.obs_chunks[self.chunks_per_file :]
            self.act_chunks = self.act_chunks[self.chunks_per_file :]

            file_stem = f"chunks_{self.file_idx:06d}"
            out_path = os.path.join(self.chunks_dir, f"{file_stem}.npz")

            # Write as a .npz with numbered arrays to preserve variable chunk lengths
            save_dict: Dict[str, Any] = {}
            for idx, (o, a) in enumerate(zip(obs_to_write, act_to_write)):
                save_dict[f"obs_{idx:03d}"] = o
                save_dict[f"acts_{idx:03d}"] = a

            np.savez_compressed(out_path, **save_dict)

            # Minimal metadata per file for bookkeeping
            meta = {
                "file": os.path.basename(out_path),
                "num_chunks": len(obs_to_write),
                "chunk_size": self.chunk_size,
                "chunk_indices": [i for i in range(len(obs_to_write))],
            }
            with open(self.metadata_path, "a") as f:
                f.write(json.dumps(meta) + "
")

            self.file_idx += 1

    def add_step(
        self,
        obs_batch: np.ndarray,
        actions_batch: np.ndarray,
        episode_done_batch: np.ndarray,
    ) -> None:
        """
        Add a single environment step across all envs.

        Args:
            obs_batch: np.ndarray with shape (num_envs, C, H, W). Values should be [0, 255].
            actions_batch: np.ndarray with shape (num_envs,).
            episode_done_batch: np.ndarray with shape (num_envs,), boolean flags indicating end-of-episode.
        """
        assert obs_batch.shape[0] == self.num_envs
        assert actions_batch.shape[0] == self.num_envs
        assert episode_done_batch.shape[0] == self.num_envs

        for env_idx in range(self.num_envs):
            buf = self.env_buffers[env_idx]
            buf["obs_seq"].append(obs_batch[env_idx])
            buf["act_seq"].append(int(actions_batch[env_idx]))

            if len(buf["obs_seq"]) == self.chunk_size:
                self._append_chunk_from_sequences(env_idx)

            # Episode ended: finalize the episode
            if bool(episode_done_batch[env_idx]):
                self._finalize_episode_for_env(env_idx)

    def close(self) -> None:
        # Drop remainder to keep consistent chunks-per-file contract
        remainder = len(self.obs_chunks)
        if remainder > 0:
            print(
                f"Warning: Dropping {remainder} chunks to keep a consistent number of chunks per file. "
                "Consider adjusting chunk_size and chunks_per_file to reduce data loss."
            )
        # Nothing else to do; per-episode buffers do not carry data unless an episode just ended