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)

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
Source code in src/versatil/data/preprocessing/create_zarr_from_synthetic.py 
26
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
def create_replay_buffer_from_synthetic(schema: SyntheticSchema) -> None:
    """Create a Zarr-based replay buffer from procedurally generated synthetic episodes.

    Args:
        schema: SyntheticSchema instance with generation parameters and zarr path.
    """
    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)
    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."
    )