replay_buffer

replay_buffer

ReplayBuffer

ReplayBuffer(root)

Manages a replay buffer dataset in Zarr or NumPy format for storing episodes of data.

The buffer organizes data into 'data' (arrays for observations, actions, etc.) and 'meta' (episode_ends array tracking cumulative step counts per episode). Supports creation from scratch, loading from paths/stores, copying with optional rechunking/recompression, saving, adding/dropping episodes, slicing, and backend-agnostic access (Zarr for disk, NumPy for in-memory). Ensures data consistency across keys. Uses cached properties for efficiency. Supports custom chunks and compressors per array key.

Attributes:

Name	Type	Description
root		Union[zarr.Group, Dict] holding 'data' and 'meta' groups/dicts.

Methods:

Name	Description
create_empty_zarr	Classmethod to create an empty Zarr-based buffer.
create_empty_numpy	Classmethod to create an empty NumPy-based buffer.
create_from_group	Classmethod to create from existing Zarr group.
create_from_path	Classmethod to load from Zarr file path.
copy_from_store	Classmethod to copy from source store with optional modifications.
copy_from_path	Classmethod to copy from Zarr path.
save_to_store	Save buffer to a store with optional rechunk/recompress.
save_to_path	Save to a file path.
resolve_compressor	Staticmethod to get Blosc compressor by name.
_resolve_array_compressor	Classmethod to resolve compressor for a key.
_resolve_array_chunks	Classmethod to resolve chunks for a key.
data	Cached property for data group/dict.
meta	Cached property for meta group/dict.
update_meta	Update meta with new key-value pairs as arrays.
episode_ends	Property for episode_ends array.
get_episode_idxs	Get array mapping steps to episode indices (Numba-optimized).
backend	Property detecting 'zarr' or 'numpy'.
__repr__	String representation (Zarr tree or default).
keys/values/items/__getitem__/__contains__	Dict-like access to data.
n_steps	Total number of steps.
n_episodes	Number of episodes.
chunk_size	First chunk size if Zarr.
episode_lengths	Array of episode lengths.
add_episode	Append a new episode (dict of arrays).
drop_episode	Remove the last episode.
pop_episode	Get and remove the last episode.
extend	Alias for add_episode.
get_episode	Get episode by index as dict of arrays.
get_episode_slice	Get slice for an episode.
get_steps_slice	Get sliced data dict for step range.
get_chunks	Get current chunks dict (Zarr only).
set_chunks	Set new chunks per key (Zarr only).
get_compressors	Get current compressors dict (Zarr only).
set_compressors	Set new compressors per key (Zarr only).

Initializes the ReplayBuffer with a root group or dict.

Validates presence of 'data', 'meta', 'episode_ends', and shape consistency across data arrays.

Parameters:

Name	Type	Description	Default
root	Group | dict[str, dict]	Zarr Group or dict with 'data' and 'meta' substructures.	required

Raises:

Type	Description
ValueError	If root is missing required keys or data shapes are inconsistent.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def __init__(self, root: zarr.Group | dict[str, dict]):
    """Initializes the ReplayBuffer with a root group or dict.

    Validates presence of 'data', 'meta', 'episode_ends', and shape consistency across data arrays.

    Args:
        root: Zarr Group or dict with 'data' and 'meta' substructures.

    Raises:
        ValueError: If root is missing required keys or data shapes are inconsistent.
    """
    for required_key in ("data", "meta"):
        if required_key not in root:
            raise ValueError(f"Root must contain a '{required_key}' group.")
    if "episode_ends" not in root["meta"]:
        raise ValueError("Meta group must contain 'episode_ends'.")
    for key in root["data"]:
        value = root["data"][key]
        expected_length = root["meta"]["episode_ends"][-1]
        if value.shape[0] != expected_length:
            raise ValueError(
                f"Data array '{key}' has length {value.shape[0]} but episode_ends indicates {expected_length}."
            )
    self.root = root

data cached property

data

Cached access to the 'data' group or dict.

meta cached property

meta

Cached access to the 'meta' group or dict.

episode_ends property

episode_ends

Property for the episode_ends array.

backend property

backend

Detects backend as 'zarr' or 'numpy' based on root type.

n_steps property

n_steps

Total steps: last episode_end or 0.

n_episodes property

n_episodes

Number of episodes: length of episode_ends.

chunk_size property

chunk_size

First dimension chunk size of first data array (Zarr only).

episode_lengths property

episode_lengths

Array of episode lengths from diffs of episode_ends.

create_empty_zarr classmethod

create_empty_zarr(storage=None, root=None)

Creates an empty Zarr-based ReplayBuffer.

Initializes 'data' and 'meta' groups, with 'episode_ends' as an empty int64 array.

Parameters:

Name	Type	Description	Default
storage	LocalStore | MemoryStore | None	Optional Zarr store (defaults to MemoryStore).	None
root	Group | None	Optional existing Zarr group to use.	None

Returns:

Type	Description
ReplayBuffer	ReplayBuffer instance.

Source code in src/versatil/data/preprocessing/replay_buffer.py

@classmethod
def create_empty_zarr(
    cls,
    storage: LocalStore | MemoryStore | None = None,
    root: zarr.Group | None = None,
) -> "ReplayBuffer":
    """Creates an empty Zarr-based ReplayBuffer.

    Initializes 'data' and 'meta' groups, with 'episode_ends' as an empty int64 array.

    Args:
        storage: Optional Zarr store (defaults to MemoryStore).
        root: Optional existing Zarr group to use.

    Returns:
        ReplayBuffer instance.
    """
    if root is None:
        if storage is None:
            storage = MemoryStore()
        root = zarr.open_group(store=storage, mode="w")
    root.create_group("data", overwrite=False)
    meta = root.create_group("meta", overwrite=False)
    if "episode_ends" not in meta:
        meta.create_array(
            "episode_ends", shape=(0,), dtype=np.int64, compressors=None
        )
    return cls(root=root)

create_empty_numpy classmethod

create_empty_numpy()

Creates an empty NumPy-based ReplayBuffer.

Initializes root as dict with empty 'data' dict and 'meta' with zero-length episode_ends array.

Returns:

Type	Description
ReplayBuffer	ReplayBuffer instance.

Source code in src/versatil/data/preprocessing/replay_buffer.py

@classmethod
def create_empty_numpy(cls) -> "ReplayBuffer":
    """Creates an empty NumPy-based ReplayBuffer.

    Initializes root as dict with empty 'data' dict and 'meta' with zero-length episode_ends array.

    Returns:
        ReplayBuffer instance.
    """
    root = {"data": {}, "meta": {"episode_ends": np.zeros((0,), dtype=np.int64)}}
    return cls(root=root)

create_from_group classmethod

create_from_group(group)

Creates ReplayBuffer from an existing Zarr group.

If 'data' missing, creates empty; else loads existing.

Parameters:

Name	Type	Description	Default
group	Group	Zarr group.	required

Returns:

Type	Description
ReplayBuffer	ReplayBuffer instance.

Source code in src/versatil/data/preprocessing/replay_buffer.py

@classmethod
def create_from_group(cls, group: zarr.Group) -> "ReplayBuffer":
    """Creates ReplayBuffer from an existing Zarr group.

    If 'data' missing, creates empty; else loads existing.

    Args:
        group: Zarr group.

    Returns:
        ReplayBuffer instance.
    """
    if "data" not in group:
        buffer = cls.create_empty_zarr(root=group)
    else:
        buffer = cls(root=group)
    return buffer

create_from_path classmethod

create_from_path(zarr_path)

Loads ReplayBuffer from a Zarr file path in read mode.

Parameters:

Name	Type	Description	Default
zarr_path	str	Path to Zarr directory.	required

Returns:

Type	Description
ReplayBuffer	ReplayBuffer instance.

Source code in src/versatil/data/preprocessing/replay_buffer.py

@classmethod
def create_from_path(cls, zarr_path: str) -> "ReplayBuffer":
    """Loads ReplayBuffer from a Zarr file path in read mode.

    Args:
        zarr_path: Path to Zarr directory.

    Returns:
        ReplayBuffer instance.
    """
    group = zarr.open_group(store=os.path.expanduser(zarr_path), mode="r")
    return cls.create_from_group(group)

copy_from_store classmethod

copy_from_store(src_store, store=None, keys=None, chunks=None, compressors=None, if_exists='replace', **kwargs)

Copies a ReplayBuffer from source store to new store or NumPy dict.

If store None, copies to NumPy dict; else to Zarr store. Selectively copies keys, applying custom chunks/compressors if specified, otherwise preserves or defaults.

Parameters:

Name	Type	Description	Default
src_store	LocalStore | MemoryStore	Source Zarr store.	required
store	LocalStore | MemoryStore | None	Optional destination store (None for NumPy).	None
keys	list[str] | None	Optional list of data keys to copy (all if None).	None
chunks	dict[str, tuple] | None	Dict of key to chunks tuple, or fallback to source/optimal.	None
compressors	dict | str | BloscCodec | None	Dict or single compressor, resolved per key.	None
if_exists	str	Zarr copy behavior ('replace', etc.).	'replace'
**kwargs	Any	Unused.	{}

Returns:

Type	Description
ReplayBuffer	ReplayBuffer instance from copied data.

Source code in src/versatil/data/preprocessing/replay_buffer.py

@classmethod
def copy_from_store(
    cls,
    src_store: LocalStore | MemoryStore,
    store: LocalStore | MemoryStore | None = None,
    keys: list[str] | None = None,
    chunks: dict[str, tuple] | None = None,
    compressors: dict | str | BloscCodec | None = None,
    if_exists: str = "replace",
    **kwargs: Any,
) -> "ReplayBuffer":
    """Copies a ReplayBuffer from source store to new store or NumPy dict.

    If store None, copies to NumPy dict; else to Zarr store. Selectively copies keys,
    applying custom chunks/compressors if specified, otherwise preserves or defaults.

    Args:
        src_store: Source Zarr store.
        store: Optional destination store (None for NumPy).
        keys: Optional list of data keys to copy (all if None).
        chunks: Dict of key to chunks tuple, or fallback to source/optimal.
        compressors: Dict or single compressor, resolved per key.
        if_exists: Zarr copy behavior ('replace', etc.).
        **kwargs: Unused.

    Returns:
        ReplayBuffer instance from copied data.
    """
    if chunks is None:
        chunks = {}
    if compressors is None:
        compressors = {}
    src_root = zarr.open_group(store=src_store, mode="r")
    root = None
    if store is None:
        meta = {}
        for key in src_root["meta"]:
            value = src_root["meta"][key]
            if isinstance(value, zarr.Group):
                continue
            if value.shape == ():
                meta[key] = np.array(value)
            else:
                meta[key] = value[:]
        if keys is None:
            keys = src_root["data"].keys()
        data = {}
        for key in keys:
            arr = src_root["data"][key]
            data[key] = arr[:]
        root = {"meta": meta, "data": data}
    else:
        root = zarr.open_group(store=store, mode="w")
        meta_group = root.create_group("meta", overwrite=True)
        for key in src_root["meta"]:
            value = src_root["meta"][key]
            if isinstance(value, zarr.Group):
                continue
            _ = meta_group.create_array(
                name=key,
                data=value[:] if value.shape != () else np.array(value),
                chunks=value.shape,
                compressors=None,
            )
        # Manually copy data
        # TODO: update when issue is closed on https://github.com/zarr-developers/zarr-python/issues/2407
        data_group = root.create_group("data", overwrite=True)
        if keys is None:
            keys = src_root["data"].keys()
        for key in keys:
            value = src_root["data"][key]
            cks = cls._resolve_array_chunks(chunks=chunks, key=key, array=value)
            cpr = cls._resolve_array_compressor(
                compressors=compressors, key=key, array=value
            )
            _ = _create_zarr_data_array(
                group=data_group,
                name=key,
                data=value[:],
                chunks=cks,
                shards=value.shards,
                codec=cpr,
            )
    buffer = cls(root=root)
    return buffer

copy_from_path classmethod

copy_from_path(zarr_path, backend=None, store=None, keys=None, chunks=None, compressors=None, if_exists='replace', **kwargs)

Copies ReplayBuffer from Zarr path, optionally to store or NumPy.

Warns if backend specified (deprecated). Expands user path.

Parameters:

Name	Type	Description	Default
zarr_path	str	Source Zarr directory path.	required
backend	str | None	Deprecated; use store=None for NumPy.	None
store	LocalStore | MemoryStore | None	Destination store (None for NumPy).	None
keys	list[str] | None	Optional list of data keys to copy (all if None).	None
chunks	dict[str, tuple] | None	Dict of key to chunks tuple, or fallback to source/optimal.	None
compressors	dict | str | BloscCodec | None	Dict or single compressor, resolved per key.	None
if_exists	str	Zarr copy behavior ('replace', etc.).	'replace'
**kwargs	Any	Passed to copy_from_store.	{}

Returns:

Type	Description
ReplayBuffer	ReplayBuffer instance.

Source code in src/versatil/data/preprocessing/replay_buffer.py

@classmethod
def copy_from_path(
    cls,
    zarr_path: str,
    backend: str | None = None,
    store: LocalStore | MemoryStore | None = None,
    keys: list[str] | None = None,
    chunks: dict[str, tuple] | None = None,
    compressors: dict | str | BloscCodec | None = None,
    if_exists: str = "replace",
    **kwargs: Any,
) -> "ReplayBuffer":
    """Copies ReplayBuffer from Zarr path, optionally to store or NumPy.

    Warns if backend specified (deprecated). Expands user path.

    Args:
        zarr_path: Source Zarr directory path.
        backend: Deprecated; use store=None for NumPy.
        store: Destination store (None for NumPy).
        keys: Optional list of data keys to copy (all if None).
        chunks: Dict of key to chunks tuple, or fallback to source/optimal.
        compressors: Dict or single compressor, resolved per key.
        if_exists: Zarr copy behavior ('replace', etc.).
        **kwargs: Passed to copy_from_store.

    Returns:
        ReplayBuffer instance.
    """
    if backend == "numpy":
        logging.warning(msg="backend argument is deprecated!")
        store = None
    group = zarr.open_group(store=os.path.expanduser(zarr_path), mode="r")
    return cls.copy_from_store(
        src_store=group.store,
        store=store,
        keys=keys,
        chunks=chunks,
        compressors=compressors,
        if_exists=if_exists,
        **kwargs,
    )

save_to_store

save_to_store(store, chunks=None, compressors=None, if_exists='replace', **kwargs)

Write the replay buffer into a Zarr store with optional chunking and codecs.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def save_to_store(
    self,
    store: LocalStore | MemoryStore,
    chunks: dict[str, tuple] | None = None,
    compressors: str | BloscCodec | dict | None = None,
    if_exists: str = "replace",
    **kwargs,
) -> LocalStore | MemoryStore:
    """Write the replay buffer into a Zarr store with optional chunking and codecs."""
    if chunks is None:
        chunks = {}
    if compressors is None:
        compressors = {}
    root = zarr.open_group(store, mode="w" if if_exists == "replace" else "a")

    # Manually copy meta
    # TODO: update when issue is closed on https://github.com/zarr-developers/zarr-python/issues/2407
    meta_group = root.create_group("meta", overwrite=True)
    for key in self.meta:
        value = self.meta[key]
        data_to_save = value[:] if isinstance(value, zarr.Array) else value
        _ = meta_group.create_array(
            name=key,
            data=data_to_save,
            chunks=value.shape,
            compressors=None,
        )
    # Manually copy data
    # TODO: update when issue is closed on https://github.com/zarr-developers/zarr-python/issues/2407
    data_group = root.create_group("data", overwrite=True)
    for key in self.root["data"]:
        value = self.root["data"][key]
        cks = self._resolve_array_chunks(chunks=chunks, key=key, array=value)
        cpr = self._resolve_array_compressor(
            compressors=compressors, key=key, array=value
        )
        data_to_save = value[:] if isinstance(value, zarr.Array) else value
        _ = _create_zarr_data_array(
            group=data_group,
            name=key,
            data=data_to_save,
            chunks=cks,
            shards=value.shards if isinstance(value, zarr.Array) else None,
            codec=cpr,
        )
    return store

save_to_path

save_to_path(zarr_path, chunks=None, compressors=None, if_exists='replace', **kwargs)

Saves to a local Zarr path using LocalStore.

Parameters:

Name	Type	Description	Default
zarr_path	str	Destination path.	required
chunks	dict[str, tuple] | None	Per-key chunks for arrays.	None
compressors	str | BloscCodec | dict | None	Per-key or global compressor.	None
if_exists	str	Zarr write mode ('replace', etc.).	'replace'
**kwargs	Any	Passed to save_to_store.	{}

Returns:

Type	Description
LocalStore	The store.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def save_to_path(
    self,
    zarr_path: str,
    chunks: dict[str, tuple] | None = None,
    compressors: str | BloscCodec | dict | None = None,
    if_exists: str = "replace",
    **kwargs: Any,
) -> LocalStore:
    """Saves to a local Zarr path using LocalStore.

    Args:
        zarr_path: Destination path.
        chunks: Per-key chunks for arrays.
        compressors: Per-key or global compressor.
        if_exists: Zarr write mode ('replace', etc.).
        **kwargs: Passed to save_to_store.

    Returns:
        The store.
    """
    store = LocalStore(os.path.expanduser(zarr_path))
    return self.save_to_store(
        store, chunks=chunks, compressors=compressors, if_exists=if_exists, **kwargs
    )

resolve_compressor staticmethod

resolve_compressor(compressor='default')

Resolves compressor string to BloscCodec instance.

'default': lz4 level 5, no shuffle. 'disk': zstd level 5, bitshuffle.

Parameters:

Name	Type	Description	Default
compressor	str | BloscCodec	'default', 'disk', or BloscCodec instance.	'default'

Returns:

Type	Description
BloscCodec	BloscCodec.

Source code in src/versatil/data/preprocessing/replay_buffer.py

@staticmethod
def resolve_compressor(
    compressor: str | BloscCodec = "default",
) -> BloscCodec:
    """Resolves compressor string to BloscCodec instance.

    'default': lz4 level 5, no shuffle.
    'disk': zstd level 5, bitshuffle.

    Args:
        compressor: 'default', 'disk', or BloscCodec instance.

    Returns:
        BloscCodec.
    """
    if compressor == "default":
        compressor = BloscCodec(
            cname="lz4", clevel=5, shuffle=BloscShuffle.noshuffle
        )
    elif compressor == "disk":
        compressor = BloscCodec(
            cname="zstd", clevel=5, shuffle=BloscShuffle.bitshuffle
        )
    return compressor

update_meta

update_meta(data)

Updates meta with new key-value pairs as NumPy arrays.

Converts values to arrays if needed, overwrites existing keys. For Zarr, creates arrays with no compression; for NumPy, dict update.

Parameters:

Name	Type	Description	Default
data	dict[str, ndarray | list | int | float]	Dict of key to value (scalar/list/array).	required

Returns:

Type	Description
Group | dict[str, ndarray]	Updated meta group/dict.

Raises:

Type	Description
TypeError	If value can't be converted to non-object array.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def update_meta(
    self, data: dict[str, np.ndarray | list | int | float]
) -> zarr.Group | dict[str, np.ndarray]:
    """Updates meta with new key-value pairs as NumPy arrays.

    Converts values to arrays if needed, overwrites existing keys.
    For Zarr, creates arrays with no compression; for NumPy, dict update.

    Args:
        data: Dict of key to value (scalar/list/array).

    Returns:
        Updated meta group/dict.

    Raises:
        TypeError: If value can't be converted to non-object array.
    """
    np_data = {}
    for key, value in data.items():
        if isinstance(value, np.ndarray):
            np_data[key] = value
        else:
            arr = np.array(value)
            if arr.dtype == object:
                raise TypeError(f"Invalid value type {type(value)}")
            np_data[key] = arr
    meta_group = self.meta
    if self.backend == "zarr":
        for key, value in np_data.items():
            _ = meta_group.create_array(
                name=key,
                data=value,
                chunks=value.shape,
                overwrite=True,
                compressors=None,
            )
    else:
        meta_group.update(np_data)

    return meta_group

get_episode_idxs

get_episode_idxs()

Computes array mapping each step to its episode index.

Uses Numba-jitted function for efficiency.

Returns:

Type	Description
ndarray	NumPy array of episode indices per step.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def get_episode_idxs(self) -> np.ndarray:
    """Computes array mapping each step to its episode index.

    Uses Numba-jitted function for efficiency.

    Returns:
        NumPy array of episode indices per step.
    """

    @numba.jit(nopython=True)
    def _get_episode_idxs(episode_ends):
        result = np.zeros((episode_ends[-1],), dtype=np.int64)
        for i in range(len(episode_ends)):  # This will work with numpy array
            start = 0
            if i > 0:
                start = episode_ends[i - 1]
            end = episode_ends[i]
            for idx in range(start, end):
                result[idx] = i
        return result

    # Convert zarr array to numpy first
    episode_ends_np = self.episode_ends[:]
    return _get_episode_idxs(episode_ends_np)

__repr__

__repr__()

String representation: Zarr tree or default repr.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def __repr__(self) -> str:
    """String representation: Zarr tree or default repr."""
    if self.backend == "zarr":
        try:
            return str(self.root.tree())
        except ImportError:
            # 'rich' library not installed, fall back to simple repr
            return f"<ReplayBuffer zarr backend: {self.n_episodes} episodes, {self.n_steps} steps>"
    else:
        return super().__repr__()

keys

keys()

Keys of the data dict/group.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def keys(self) -> list[str]:
    """Keys of the data dict/group."""
    return self.data.keys()

values

values()

Values of the data dict/group.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def values(self) -> list[zarr.Array | np.ndarray]:
    """Values of the data dict/group."""
    if self.backend == "zarr":
        return self.data.array_values()
    else:
        return self.data.values()

items

items()

Items of the data dict/group.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def items(self) -> list[tuple[str, zarr.Array | np.ndarray]]:
    """Items of the data dict/group."""
    if self.backend == "zarr":
        return self.data.members()
    else:
        return self.data.items()

__getitem__

__getitem__(key)

Getitem for data[key].

Source code in src/versatil/data/preprocessing/replay_buffer.py

def __getitem__(self, key: str) -> zarr.Array | np.ndarray:
    """Getitem for data[key]."""
    return self.data[key]

__contains__

__contains__(key)

Contains check for data.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def __contains__(self, key: str) -> bool:
    """Contains check for data."""
    return key in self.data

add_episode

add_episode(data, chunks=None, compressors=None)

Adds an episode as dict of arrays, resizing all data arrays.

Creates new keys if needed with resolved chunks/compressors. Appends to episode_ends, rechunks if grown significantly (Zarr).

Parameters:

Name	Type	Description	Default
data	dict[str, ndarray]	Dict of key to NumPy array (consistent lengths).	required
chunks	dict[str, tuple] | None	Per-key chunks for new arrays.	None
compressors	str | BloscCodec | dict | None	Per-key or global for new arrays.	None

Raises:

Type	Description
ValueError	If empty data, inconsistent lengths, or mismatched shapes.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def add_episode(
    self,
    data: dict[str, np.ndarray],
    chunks: dict[str, tuple] | None = None,
    compressors: str | BloscCodec | dict | None = None,
) -> None:
    """Adds an episode as dict of arrays, resizing all data arrays.

    Creates new keys if needed with resolved chunks/compressors.
    Appends to episode_ends, rechunks if grown significantly (Zarr).

    Args:
        data: Dict of key to NumPy array (consistent lengths).
        chunks: Per-key chunks for new arrays.
        compressors: Per-key or global for new arrays.

    Raises:
        ValueError: If empty data, inconsistent lengths, or mismatched shapes.
    """
    if chunks is None:
        chunks = {}
    if compressors is None:
        compressors = {}
    if len(data) == 0:
        raise ValueError("Episode data must not be empty.")
    is_zarr = self.backend == "zarr"
    curr_len = self.n_steps
    episode_length = None
    for key, value in data.items():
        if len(value.shape) < 1:
            raise ValueError(f"Array '{key}' must be at least 1-dimensional.")
        if episode_length is None:
            episode_length = len(value)
        elif episode_length != len(value):
            raise ValueError(
                f"Inconsistent episode lengths: expected {episode_length}, got {len(value)} for key '{key}'."
            )
    new_len = curr_len + episode_length
    for key in data:
        value = data[key]
        new_shape = (new_len,) + value.shape[1:]
        if key not in self.data:
            if is_zarr:
                cks = self._resolve_array_chunks(
                    chunks=chunks, key=key, array=value
                )
                cpr = self._resolve_array_compressor(
                    compressors=compressors, key=key, array=value
                )
                zarr_dtype = str if value.dtype == np.dtype("O") else value.dtype
                arr = _create_zarr_data_array(
                    group=self.data,
                    name=key,
                    chunks=cks,
                    codec=cpr,
                    shape=new_shape,
                    dtype=zarr_dtype,
                    fill_value=0,
                )
            else:
                arr = np.zeros(shape=new_shape, dtype=value.dtype)
                self.data[key] = arr
        else:
            arr = self.data[key]
            if value.shape[1:] != arr.shape[1:]:
                raise ValueError(
                    f"Shape mismatch for '{key}': existing {arr.shape[1:]}, got {value.shape[1:]}."
                )
            if is_zarr:
                # Convert to plain Python ints for zarr v3
                arr.resize(tuple(int(x) for x in new_shape))
            else:
                arr.resize(new_shape, refcheck=False)
        arr[-value.shape[0] :] = value

    episode_ends = self.episode_ends
    if is_zarr:
        episode_ends.resize(episode_ends.shape[0] + 1)
    else:
        episode_ends.resize(episode_ends.shape[0] + 1, refcheck=False)
    episode_ends[-1] = new_len
    if is_zarr and episode_ends.chunks[0] < episode_ends.shape[0]:
        rechunk_recompress_array(
            self.meta, "episode_ends", chunk_length=int(episode_ends.shape[0] * 1.5)
        )

drop_episode

drop_episode()

Drops the last episode by resizing arrays backward.

Raises:

Type	Description
ValueError	If no episodes to drop.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def drop_episode(self) -> None:
    """Drops the last episode by resizing arrays backward.

    Raises:
        ValueError: If no episodes to drop.
    """
    is_zarr = self.backend == "zarr"
    episode_ends = self.episode_ends[:].copy()
    if len(episode_ends) == 0:
        raise ValueError("Cannot drop episode from an empty buffer.")
    start_idx = 0
    if len(episode_ends) > 1:
        start_idx = episode_ends[-2]
    for key in self.data:
        value = self.data[key]
        new_shape = (start_idx,) + value.shape[1:]
        if is_zarr:
            # Convert to plain Python ints for zarr v3
            value.resize(tuple(int(x) for x in new_shape))
        else:
            value.resize(new_shape, refcheck=False)
    if is_zarr:
        self.episode_ends.resize(len(episode_ends) - 1)
    else:
        self.episode_ends.resize(len(episode_ends) - 1, refcheck=False)

pop_episode

pop_episode()

Gets the last episode and drops it.

Returns:

Type	Description
dict[str, ndarray]	Dict of arrays for the episode.

Raises:

Type	Description
ValueError	If no episodes to pop.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def pop_episode(self) -> dict[str, np.ndarray]:
    """Gets the last episode and drops it.

    Returns:
        Dict of arrays for the episode.

    Raises:
        ValueError: If no episodes to pop.
    """
    if self.n_episodes == 0:
        raise ValueError("Cannot pop episode from an empty buffer.")
    episode = self.get_episode(self.n_episodes - 1, copy=True)
    self.drop_episode()
    return episode

extend

extend(data)

Alias for add_episode.

Parameters:

Name	Type	Description	Default
data	dict[str, ndarray]	As in add_episode.	required

Source code in src/versatil/data/preprocessing/replay_buffer.py

def extend(self, data: dict[str, np.ndarray]) -> None:
    """Alias for add_episode.

    Args:
        data: As in add_episode.
    """
    self.add_episode(data)

get_episode

get_episode(idx, copy=False)

Gets an episode by index as dict of sliced arrays.

Parameters:

Name	Type	Description	Default
idx	int	Episode index (handles negative via list).	required
copy	bool	If True, copy NumPy arrays.	False

Returns:

Type	Description
dict[str, ndarray]	Dict of key to array slice.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def get_episode(self, idx: int, copy: bool = False) -> dict[str, np.ndarray]:
    """Gets an episode by index as dict of sliced arrays.

    Args:
        idx: Episode index (handles negative via list).
        copy: If True, copy NumPy arrays.

    Returns:
        Dict of key to array slice.
    """
    idx = list(range(self.episode_ends.shape[0]))[idx]
    start_idx = 0
    if idx > 0:
        start_idx = self.episode_ends[idx - 1]
    end_idx = self.episode_ends[idx]
    result = self.get_steps_slice(start_idx, end_idx, copy=copy)
    return result

get_episode_slice

get_episode_slice(idx)

Gets the slice object for an episode's steps.

Parameters:

Name	Type	Description	Default
idx	int	Episode index.	required

Returns:

Type	Description
slice	slice(start, end)

Source code in src/versatil/data/preprocessing/replay_buffer.py

def get_episode_slice(self, idx: int) -> slice:
    """Gets the slice object for an episode's steps.

    Args:
        idx: Episode index.

    Returns:
        slice(start, end)
    """
    start_idx = 0
    if idx > 0:
        start_idx = self.episode_ends[idx - 1]
    end_idx = self.episode_ends[idx]
    return slice(start_idx, end_idx)

get_steps_slice

get_steps_slice(start, stop, step=None, copy=False)

Gets dict of sliced arrays for a step range.

Parameters:

Name	Type	Description	Default
start	int	Start step index.	required
stop	int	Stop step index.	required
step	int | None	Step size for slicing.	None
copy	bool	Copy if NumPy backend.	False

Returns:

Type	Description
dict[str, ndarray]	Dict of key to sliced array.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def get_steps_slice(
    self,
    start: int,
    stop: int,
    step: int | None = None,
    copy: bool = False,
) -> dict[str, np.ndarray]:
    """Gets dict of sliced arrays for a step range.

    Args:
        start: Start step index.
        stop: Stop step index.
        step: Step size for slicing.
        copy: Copy if NumPy backend.

    Returns:
        Dict of key to sliced array.
    """
    start = int(start) if isinstance(start, np.integer) else start
    stop = int(stop) if isinstance(stop, np.integer) else stop
    _slice = slice(start, stop, step)
    result = {}
    for key in self.data:
        value = self.data[key]
        x = value[_slice]
        if copy and isinstance(value, np.ndarray):
            x = x.copy()
        result[key] = x
    return result

get_chunks

get_chunks()

Gets current chunks per data key (Zarr only).

Returns:

Type	Description
dict[str, tuple[int, ...]]	Dict of key to chunks tuple.

Raises:

Type	Description
RuntimeError	If not Zarr backend.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def get_chunks(self) -> dict[str, tuple[int, ...]]:
    """Gets current chunks per data key (Zarr only).

    Returns:
        Dict of key to chunks tuple.

    Raises:
        RuntimeError: If not Zarr backend.
    """
    if self.backend != "zarr":
        raise RuntimeError("get_chunks is only supported on Zarr backend.")
    chunks = {}
    for key in self.data:
        value = self.data[key]
        chunks[key] = value.chunks
    return chunks

set_chunks

set_chunks(chunks)

Sets new chunks per data key if changed (Zarr only).

Uses rechunk_recompress_array.

Parameters:

Name	Type	Description	Default
chunks	dict[str, tuple[int, ...]]	Dict of key to new chunks tuple.	required

Raises:

Type	Description
RuntimeError	If not Zarr backend.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def set_chunks(self, chunks: dict[str, tuple[int, ...]]) -> None:
    """Sets new chunks per data key if changed (Zarr only).

    Uses rechunk_recompress_array.

    Args:
        chunks: Dict of key to new chunks tuple.

    Raises:
        RuntimeError: If not Zarr backend.
    """
    if self.backend != "zarr":
        raise RuntimeError("set_chunks is only supported on Zarr backend.")
    for key in chunks:
        value = chunks[key]
        if key in self.data:
            arr = self.data[key]
            if value != arr.chunks:
                check_chunks_compatible(chunks=value, shape=arr.shape)
                rechunk_recompress_array(self.data, key, chunks=value)

get_compressors

get_compressors()

Gets current codec per data key (Zarr only).

Returns BloscCodec for numerically compressed arrays, WebPCodec for image arrays using WebP serialization, or None if uncompressed.

Returns:

Type	Description
dict[str, BloscCodec | WebPCodec | None]	Dict mapping key to BloscCodec, WebPCodec, or None.

Raises:

Type	Description
RuntimeError	If not Zarr backend.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def get_compressors(self) -> dict[str, BloscCodec | WebPCodec | None]:
    """Gets current codec per data key (Zarr only).

    Returns BloscCodec for numerically compressed arrays, WebPCodec for
    image arrays using WebP serialization, or None if uncompressed.

    Returns:
        Dict mapping key to BloscCodec, WebPCodec, or None.

    Raises:
        RuntimeError: If not Zarr backend.
    """
    if self.backend != "zarr":
        raise RuntimeError("get_compressors is only supported on Zarr backend.")
    compressors = {}
    for key in self.data:
        array = self.data[key]
        if array.compressors:
            compressors[key] = array.compressors[-1]
        else:
            compressors[key] = _get_serializer_codec(array)
    return compressors

set_compressors

set_compressors(compressors)

Sets new compressor per data key if changed (Zarr only).

Resolves strings, uses rechunk_recompress_array.

Parameters:

Name	Type	Description	Default
compressors	dict[str, str | BloscCodec | WebPCodec]	Dict of key to compressor/str.	required

Raises:

Type	Description
RuntimeError	If not Zarr backend.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def set_compressors(
    self, compressors: dict[str, str | BloscCodec | WebPCodec]
) -> None:
    """Sets new compressor per data key if changed (Zarr only).

    Resolves strings, uses rechunk_recompress_array.

    Args:
        compressors: Dict of key to compressor/str.

    Raises:
        RuntimeError: If not Zarr backend.
    """
    if self.backend != "zarr":
        raise RuntimeError("set_compressors is only supported on Zarr backend.")
    for key, value in compressors.items():
        if key in self.data:
            arr = self.data[key]
            compressor = self.resolve_compressor(value)
            arr_cpr = (
                arr.compressors[-1]
                if arr.compressors
                else _get_serializer_codec(arr)
            )
            if compressor != arr_cpr:
                rechunk_recompress_array(self.data, key, compressor=compressor)

check_chunks_compatible

check_chunks_compatible(chunks, shape)

Checks if given chunks are compatible with the array shape.

Ensures that chunks and shape have the same dimensionality, each chunk size is a positive integer, and implicitly that chunks do not exceed shape dimensions (though not explicitly checked here).

Parameters:

Name	Type	Description	Default
chunks	tuple[int, ...]	Tuple of chunk sizes for each dimension.	required
shape	tuple[int, ...]	Tuple of array shape dimensions.	required

Raises:

Type	Description
ValueError	If lengths differ or chunks are not positive.
TypeError	If chunks are not integral.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def check_chunks_compatible(chunks: tuple[int, ...], shape: tuple[int, ...]) -> None:
    """Checks if given chunks are compatible with the array shape.

    Ensures that chunks and shape have the same dimensionality, each chunk size is a positive integer,
    and implicitly that chunks do not exceed shape dimensions (though not explicitly checked here).

    Args:
        chunks: Tuple of chunk sizes for each dimension.
        shape: Tuple of array shape dimensions.

    Raises:
        ValueError: If lengths differ or chunks are not positive.
        TypeError: If chunks are not integral.
    """
    if len(shape) != len(chunks):
        raise ValueError(
            f"Chunks dimensionality {len(chunks)} does not match shape dimensionality {len(shape)}."
        )
    for c in chunks:
        if not isinstance(c, numbers.Integral):
            raise TypeError(f"Chunk size must be an integer, got {type(c)}.")
        if c <= 0:
            raise ValueError(f"Chunk size must be positive, got {c}.")

rechunk_recompress_array

rechunk_recompress_array(group, name, chunks=None, chunk_length=None, compressor=None, tmp_key='_temp')

Rechunk and/or recompress an existing zarr array.

Preserves the original codec and shard shape unless a new compressor is provided. WebPCodec is obtained from the array serializer.

Parameters:

Name	Type	Description	Default
group	Group	Zarr group containing the array.	required
name	str	Name of the array to rechunk/recompress.	required
chunks	tuple | None	New chunk sizes. Defaults to existing chunks.	None
chunk_length	int | None	Shorthand to override only the first chunk dimension.	None
compressor	BloscCodec | WebPCodec | None	New codec. Defaults to the array's existing codec.	None
tmp_key	str	Temporary key used during rechunking (unused, kept for API).	'_temp'

Source code in src/versatil/data/preprocessing/replay_buffer.py

def rechunk_recompress_array(
    group: zarr.Group,
    name: str,
    chunks: tuple | None = None,
    chunk_length: int | None = None,
    compressor: BloscCodec | WebPCodec | None = None,
    tmp_key: str = "_temp",
) -> zarr.Array:
    """Rechunk and/or recompress an existing zarr array.

    Preserves the original codec and shard shape unless a new compressor is
    provided. WebPCodec is obtained from the array serializer.

    Args:
        group: Zarr group containing the array.
        name: Name of the array to rechunk/recompress.
        chunks: New chunk sizes. Defaults to existing chunks.
        chunk_length: Shorthand to override only the first chunk dimension.
        compressor: New codec. Defaults to the array's existing codec.
        tmp_key: Temporary key used during rechunking (unused, kept for API).
    """
    old_arr = group[name]
    if chunks is None:
        if chunk_length is not None:
            chunks = (chunk_length,) + old_arr.chunks[1:]
        else:
            chunks = old_arr.chunks
    check_chunks_compatible(chunks=chunks, shape=old_arr.shape)

    if compressor is None:
        if old_arr.compressors:
            compressor = old_arr.compressors[-1]
        else:
            compressor = _get_serializer_codec(old_arr)

    old_codec = (
        old_arr.compressors[-1]
        if old_arr.compressors
        else _get_serializer_codec(old_arr)
    )
    if (chunks == old_arr.chunks) and (compressor == old_codec):
        return old_arr

    # Manually rechunk/recompress instead of using zarr.copy
    # TODO: update when issue is closed on https://github.com/zarr-developers/zarr-python/issues/2407
    data = old_arr[:]
    shards = old_arr.shards
    del group[name]
    arr = _create_zarr_data_array(
        group=group,
        name=name,
        chunks=chunks,
        shards=shards,
        codec=compressor,
        data=data,
    )
    return arr

get_optimal_chunks

get_optimal_chunks(shape, dtype, target_chunk_bytes=2000000.0, max_chunk_length=None)

Computes optimal chunk sizes for an array to target a specific chunk byte size.

Reverses shape to prioritize inner dimensions, caps outer dimension if max_chunk_length given. Finds the split point where adding another dimension exceeds target bytes, then adjusts the last chunk to meet the target. Pads with 1s if needed for dimensionality.

Parameters:

Name	Type	Description	Default
shape	tuple[int, ...]	Tuple of array dimensions.	required
dtype	dtype | type	NumPy dtype of the array.	required
target_chunk_bytes	float	Desired approximate bytes per chunk (default 2MB).	2000000.0
max_chunk_length	int | None	Optional cap on the outermost chunk size.	None

Returns:

Type	Description
tuple[int, ...]	Tuple of optimal chunk sizes.

Source code in src/versatil/data/preprocessing/replay_buffer.py

def get_optimal_chunks(
    shape: tuple[int, ...],
    dtype: np.dtype | type,
    target_chunk_bytes: float = 2e6,
    max_chunk_length: int | None = None,
) -> tuple[int, ...]:
    """Computes optimal chunk sizes for an array to target a specific chunk byte size.

    Reverses shape to prioritize inner dimensions, caps outer dimension if max_chunk_length given.
    Finds the split point where adding another dimension exceeds target bytes, then adjusts the
    last chunk to meet the target. Pads with 1s if needed for dimensionality.

    Args:
        shape: Tuple of array dimensions.
        dtype: NumPy dtype of the array.
        target_chunk_bytes: Desired approximate bytes per chunk (default 2MB).
        max_chunk_length: Optional cap on the outermost chunk size.

    Returns:
        Tuple of optimal chunk sizes.
    """
    itemsize = np.dtype(dtype).itemsize
    rshape = list(shape[::-1])
    if max_chunk_length is not None:
        rshape[-1] = int(max_chunk_length)
    split_idx = len(shape) - 1
    for i in range(len(shape) - 1):
        this_chunk_bytes = itemsize * np.prod(rshape[:i])
        next_chunk_bytes = itemsize * np.prod(rshape[: i + 1])
        if (
            this_chunk_bytes <= target_chunk_bytes
            and next_chunk_bytes > target_chunk_bytes
        ):
            split_idx = i
    rchunks = rshape[:split_idx]
    item_chunk_bytes = itemsize * np.prod(rshape[:split_idx])
    this_max_chunk_length = rshape[split_idx]
    next_chunk_length = min(
        this_max_chunk_length, math.ceil(target_chunk_bytes / item_chunk_bytes)
    )
    rchunks.append(next_chunk_length)
    len_diff = len(shape) - len(rchunks)
    rchunks.extend([1] * len_diff)
    chunks = tuple(rchunks[::-1])
    return chunks