Skip to content

create_zarr_from_synthetic

create_zarr_from_synthetic

Creates a Zarr-based replay buffer dataset from synthetic episode generators.

create_replay_buffer_from_synthetic 

create_replay_buffer_from_synthetic(schema, image_frames_per_shard=DEFAULT_IMAGE_FRAMES_PER_SHARD)

Create a Zarr-based replay buffer from procedurally generated synthetic episodes.

Parameters:

Name Type Description Default
schema SyntheticSchema

SyntheticSchema instance with generation parameters and zarr path.

 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_synthetic.py 
 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
 66
 67
 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
def create_replay_buffer_from_synthetic(
    schema: SyntheticSchema,
    image_frames_per_shard: int | None = DEFAULT_IMAGE_FRAMES_PER_SHARD,
) -> None:
    """Create a Zarr-based replay buffer from procedurally generated synthetic episodes.

    Args:
        schema: SyntheticSchema instance with generation parameters and zarr path.
        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.
    """
    logging.info(
        f"Creating synthetic Zarr at {schema.zarr_path} "
        f"(task={schema.task_name}, episodes={schema.num_episodes})"
    )

    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")
    compressor = BloscCodec(cname="lz4", clevel=5, shuffle=BloscShuffle.noshuffle)
    _create_zarr_arrays(
        data_group=data_group,
        schema=schema,
        compressor=compressor,
        image_frames_per_shard=image_frames_per_shard,
    )
    episodes = generate_task_episodes(
        task_name=schema.task_name,
        num_episodes=schema.num_episodes,
        seed=schema.seed,
        image_size=schema.image_size,
        num_modes=schema.num_modes,
        trajectory_length=schema.trajectory_length,
        noise_std=schema.noise_std,
        num_styles=schema.num_styles,
        mode_weights=schema.mode_weights,
    )
    episode_ends = []
    cumulative_length = 0
    for episode in episodes:
        for generator_key, zarr_key in GENERATOR_KEY_TO_ZARR_KEY.items():
            if zarr_key in data_group:
                data_group[zarr_key].append(episode[generator_key])
        episode_length = len(episode["position"])
        cumulative_length += episode_length
        episode_ends.append(cumulative_length)
    meta_group.create_array(
        "episode_ends",
        data=np.array(episode_ends),
        chunks=(len(episode_ends),),
        compressors=None,
    )
    _save_training_visualization(
        episodes=episodes,
        task_name=schema.task_name,
        zarr_path=schema.zarr_path,
        num_modes=schema.num_modes,
        num_styles=schema.num_styles,
        noise_std=schema.noise_std,
    )
    logging.info(
        f"Created Zarr dataset with {len(episode_ends)} episodes, "
        f"{cumulative_length} total steps."
    )