Skip to content

create_zarr_from_csv

create_zarr_from_csv

Creates a Zarr-based replay buffer dataset from robot demonstration CSV files and associated images.

create_replay_buffer

create_replay_buffer(schema, datasets_paths, image_frames_per_shard=DEFAULT_IMAGE_FRAMES_PER_SHARD)

Creates a Zarr-based replay buffer using a Hydra-instantiated dataset schema.

Parameters:

Name Type Description Default
schema CsvDatasetSchema

CsvDatasetSchema instance (instantiated by Hydra)

required
datasets_paths list[str]

List of paths to episode CSV files

required
image_frames_per_shard int | None

Number of image frames stored in each shard, or None to disable sharding.

DEFAULT_IMAGE_FRAMES_PER_SHARD

Raises:

Type Description
TypeError

If image_frames_per_shard is not an integer.

ValueError

If image_frames_per_shard is not a positive multiple of the image chunk length.

Source code in src/versatil/data/preprocessing/create_zarr_from_csv.py
def create_replay_buffer(
    schema: CsvDatasetSchema,
    datasets_paths: list[str],
    image_frames_per_shard: int | None = DEFAULT_IMAGE_FRAMES_PER_SHARD,
) -> None:
    """Creates a Zarr-based replay buffer using a Hydra-instantiated dataset schema.

    Args:
        schema: CsvDatasetSchema instance (instantiated by Hydra)
        datasets_paths: List of paths to episode CSV files
        image_frames_per_shard: Number of image frames stored in each shard, or
            None to disable sharding.

    Raises:
        TypeError: If image_frames_per_shard is not an integer.
        ValueError: If image_frames_per_shard is not a positive multiple of the
            image chunk length.
    """
    sorted_paths = sorted(datasets_paths, key=lambda x: int(Path(x).parent.name))
    episodes = _iter_csv_episodes(
        schema=schema,
        sorted_paths=sorted_paths,
    )

    create_zarr_replay_buffer(
        schema=schema,
        episodes=episodes,
        total_episodes=len(datasets_paths),
        image_frames_per_shard=image_frames_per_shard,
    )