Skip to content

create_zarr_arrays

create_zarr_arrays

Shared zarr replay buffer creation utilities.

Provides a unified pipeline for creating zarr stores from any dataset schema, with WebP compression for uint8 images and lz4 for numerical arrays.

is_uint8_image_spec 

is_uint8_image_spec(spec)

Check if a zarr array spec corresponds to a uint8 image array.

Source code in src/versatil/data/preprocessing/create_zarr_arrays.py 
22
23
24
def is_uint8_image_spec(spec: dict) -> bool:
    """Check if a zarr array spec corresponds to a uint8 image array."""
    return len(spec["shape"]) == 4 and spec["dtype"] == "uint8"

create_zarr_arrays 

create_zarr_arrays(data_group, schema, image_codec, numeric_compressor)

Create zarr arrays with codec selection based on data type.

Uses WebP compression for uint8 image arrays and Blosc/lz4 for numerical arrays.

Parameters:

Name Type Description Default
data_group Group

Zarr group to create arrays in.

 required
schema DatasetSchema

Dataset schema providing array specifications.

 required
image_codec WebPCodec

WebP codec for uint8 image arrays.

 required
numeric_compressor BloscCodec

Blosc codec for numerical arrays.

 required
Source code in src/versatil/data/preprocessing/create_zarr_arrays.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def create_zarr_arrays(
    data_group: zarr.Group,
    schema: DatasetSchema,
    image_codec: WebPCodec,
    numeric_compressor: BloscCodec,
) -> None:
    """Create zarr arrays with codec selection based on data type.

    Uses WebP compression for uint8 image arrays and Blosc/lz4 for
    numerical arrays.

    Args:
        data_group: Zarr group to create arrays in.
        schema: Dataset schema providing array specifications.
        image_codec: WebP codec for uint8 image arrays.
        numeric_compressor: Blosc codec for numerical arrays.
    """
    specs = schema.get_zarr_array_specs()
    for key, spec in specs.items():
        dtype = str if spec["dtype"] == "str" else getattr(np, spec["dtype"])
        if is_uint8_image_spec(spec):
            # One image per chunk for WebP per-frame compression
            chunks = (1, *spec["shape"][1:])
            data_group.create_array(
                name=key,
                shape=spec["shape"],
                chunks=chunks,
                dtype=dtype,
                serializer=image_codec,
                compressors=None,
            )
        else:
            data_group.create_array(
                name=key,
                shape=spec["shape"],
                chunks=spec["chunks"],
                dtype=dtype,
                compressors=[numeric_compressor] if spec["needs_compressor"] else None,
            )

create_zarr_replay_buffer 

create_zarr_replay_buffer(schema, episodes, total_episodes=None)

Create a zarr replay buffer from an iterable of episodes.

Handles store creation, array setup with proper codecs, episode insertion, and metadata. Uses WebP for uint8 images and lz4 for numerical data.

Parameters:

Name Type Description Default
schema DatasetSchema

Dataset schema with zarr_path and array specs.

 required
episodes Iterable[dict[str, ndarray]]

Iterable yielding episode dicts mapping keys to arrays.

 required
total_episodes int | None

Total number of episodes for progress reporting.

 None
Source code in src/versatil/data/preprocessing/create_zarr_arrays.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def create_zarr_replay_buffer(
    schema: DatasetSchema,
    episodes: Iterable[dict[str, np.ndarray]],
    total_episodes: int | None = None,
) -> None:
    """Create a zarr replay buffer from an iterable of episodes.

    Handles store creation, array setup with proper codecs, episode insertion,
    and metadata. Uses WebP for uint8 images and lz4 for numerical data.

    Args:
        schema: Dataset schema with zarr_path and array specs.
        episodes: Iterable yielding episode dicts mapping keys to arrays.
        total_episodes: Total number of episodes for progress reporting.
    """
    logging.info(
        msg=f"Creating Zarr dataset at {schema.zarr_path} "
        f"from schema {schema.__class__.__name__}"
    )

    store = zarr.storage.LocalStore(schema.zarr_path)
    root = zarr.open_group(store=store, mode="w")
    data_group = root.create_group("data")
    meta_group = root.create_group("meta")

    image_codec = WebPCodec(level=WEBP_QUALITY)
    numeric_compressor = BloscCodec(
        cname="lz4", clevel=5, shuffle=BloscShuffle.noshuffle
    )
    create_zarr_arrays(
        data_group=data_group,
        schema=schema,
        image_codec=image_codec,
        numeric_compressor=numeric_compressor,
    )

    episode_ends: list[int] = []
    cumulative_len = 0

    with threadpool_limits(limits=1):
        for i, episode_data in enumerate(episodes):
            if total_episodes is not None and i % 50 == 0:
                logging.info(msg=f"  Processing episode {i + 1}/{total_episodes}...")
            for key, array in episode_data.items():
                data_group[key].append(array)
            cumulative_len += len(next(iter(episode_data.values())))
            episode_ends.append(cumulative_len)

    meta_group.create_array(
        "episode_ends",
        data=np.array(episode_ends),
        chunks=(len(episode_ends),),
        compressors=None,
    )

    logging.info(
        msg=f"Created Zarr dataset with {len(episode_ends)} episodes, "
        f"{cumulative_len} total steps."
    )