Skip to content

base

base

Abstract dataset schema for defining dataset-specific configurations.

This module provides a framework for supporting multiple datasets with different formats and observation modalities. Schemas define the structure of stored data (what's in the zarr file). They do not define how data is used at runtime (that's determined by task space config).

DatasetSchema 

DatasetSchema(zarr_path, metadata, dataset_type)

Bases: ABC

Abstract base class for dataset schemas.

Defines common interface for zarr array creation, key management, and episode extraction. Subclasses (CsvDatasetSchema, Hdf5DatasetSchema) define format-specific data extraction.

Note

A schema defines the raw data structure for a specific dataset, including: - What observations are available (position, gripper state, etc.). - What action modalities exist (position, orientation, gripper). - How to extract episode data via extract_episode.

The schema is used when creating the zarr file from the raw episodic data.

Initialize the dataset schema.

Parameters:

Name Type Description Default
zarr_path str

Path to save/load the zarr file

 required
metadata DatasetMetadata

Metadata of the raw dataset

 required
dataset_type str

Type of dataset (e.g., 'libero', 'tso', 'metaworld')

 required
Source code in src/versatil/data/raw/schemas/base.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def __init__(
    self,
    zarr_path: str,
    metadata: DatasetMetadata,
    dataset_type: str,
):
    """Initialize the dataset schema.

    Args:
        zarr_path: Path to save/load the zarr file
        metadata: Metadata of the raw dataset
        dataset_type: Type of dataset (e.g., 'libero', 'tso', 'metaworld')
    """
    self.zarr_path = zarr_path
    self.metadata = metadata
    self.dataset_type = dataset_type

extract_episode abstractmethod 

extract_episode(episode_source)

Extract all data from an episode source.

Each schema defines its own extraction logic based on its metadata. The episode_source type depends on the schema format:

Parameters:

Name Type Description Default
episode_source Any

Format-specific episode data source

 required

Returns:

Type Description
dict[str, ndarray]

Dictionary mapping zarr keys to numpy arrays
Source code in src/versatil/data/raw/schemas/base.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
@abc.abstractmethod
def extract_episode(
    self,
    episode_source: Any,
) -> dict[str, np.ndarray]:
    """Extract all data from an episode source.

    Each schema defines its own extraction logic based on its metadata.
    The episode_source type depends on the schema format:

    Args:
        episode_source: Format-specific episode data source

    Returns:
        Dictionary mapping zarr keys to numpy arrays
    """
    raise NotImplementedError("Subclasses must implement extract_episode")

get_required_zarr_keys 

get_required_zarr_keys()

Get all required zarr keys based on the dataset metadata.

Returns:

Type Description
list[str]

List of all required zarr key names
Source code in src/versatil/data/raw/schemas/base.py 
70
71
72
73
74
75
76
def get_required_zarr_keys(self) -> list[str]:
    """Get all required zarr keys based on the dataset metadata.

    Returns:
        List of all required zarr key names
    """
    return self.metadata.get_all_keys()

get_zarr_array_specs 

get_zarr_array_specs()

Get specifications for all zarr arrays to create for the store.

Returns:

Type Description
dict

Dictionary mapping zarr key names to array specifications Each spec is a dict with: shape, chunks, dtype, needs_compressor
Source code in src/versatil/data/raw/schemas/base.py 
 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
def get_zarr_array_specs(self) -> dict:
    """Get specifications for all zarr arrays to create for the store.

    Returns:
        Dictionary mapping zarr key names to array specifications
            Each spec is a dict with: shape, chunks, dtype, needs_compressor
    """
    specs = {}
    for key, obs in self.metadata.observations.items():
        if isinstance(obs, CameraMetadata):
            specs[key] = {
                "shape": (0, obs.image_height, obs.image_width, obs.channels),
                "chunks": (16, obs.image_height, obs.image_width, obs.channels),
                "dtype": obs.dtype,
                "needs_compressor": True,
            }
        elif isinstance(obs, ObservationMetadata):
            needs_compression = obs.dtype != "str"
            specs[key] = {
                "shape": (0, obs.dimension),
                "chunks": (256, obs.dimension),
                "dtype": obs.dtype,
                "needs_compressor": needs_compression,
            }
    for key, action in self.metadata.precomputed_actions.items():
        if action.slice_start is not None and action.slice_end is not None:
            dim = action.slice_end - action.slice_start
        else:
            dim = action.storage_dimension
        specs[key] = {
            "shape": (0, dim),
            "chunks": (256, dim),
            "dtype": action.dtype,
            "needs_compressor": True,
        }
    return specs