episodic_dataset

episodic_dataset

EpisodicDataset

EpisodicDataset(zarr_path, action_space, observation_space, dataloader_config, pred_horizon, obs_horizon, train=True, seed=42, augment_images=None, replay_buffer=None)

Bases: Dataset

PyTorch Dataset for episodic robot demonstration data.

This class orchestrates modular components for: - Action processing - Data augmentation - Sample building - Episode splitting and management

Initialize episodic dataset.

Parameters:

Name	Type	Description	Default
zarr_path	str	Path to zarr replay buffer	required
action_space	ActionSpace	TaskSpace action space config (what to predict and how)	required
observation_space	ObservationSpace	TaskSpace observation space config (what to use as observation data)	required
dataloader_config	DataLoaderConfig	Data loading settings (splits, normalization types, augmentation pipelines, downsampling, padding).	required
pred_horizon	int	Prediction horizon, i.e. chunk size.	required
obs_horizon	int	Observation horizon, i.e. history size.	required
train	bool	Whether to use training mode.	True
seed	int	Random seed of the experiment.	42
augment_images	bool | None	Whether image augmentations are enabled. Defaults to train so existing training and validation behavior is unchanged.	None
replay_buffer	ReplayBuffer | None	Already-loaded replay buffer to reuse, avoiding a second in-memory copy when train and validation datasets share one store.	None

Source code in src/versatil/data/episodic_dataset.py

def __init__(
    self,
    zarr_path: str,
    action_space: ActionSpace,
    observation_space: ObservationSpace,
    dataloader_config: DataLoaderConfig,
    pred_horizon: int,
    obs_horizon: int,
    train: bool = True,
    seed: int = 42,
    augment_images: bool | None = None,
    replay_buffer: ReplayBuffer | None = None,
):
    """Initialize episodic dataset.

    Args:
        zarr_path: Path to zarr replay buffer
        action_space: TaskSpace action space config (what to predict and how)
        observation_space: TaskSpace observation space config (what to use as observation data)
        dataloader_config: Data loading settings (splits, normalization
            types, augmentation pipelines, downsampling, padding).
        pred_horizon: Prediction horizon, i.e. chunk size.
        obs_horizon: Observation horizon, i.e. history size.
        train: Whether to use training mode.
        seed: Random seed of the experiment.
        augment_images: Whether image augmentations are enabled. Defaults
            to ``train`` so existing training and validation behavior is
            unchanged.
        replay_buffer: Already-loaded replay buffer to reuse, avoiding a
            second in-memory copy when train and validation datasets share
            one store.
    """
    self.action_space = action_space
    self.observation_space = observation_space
    self.pred_horizon = pred_horizon
    self.obs_horizon = obs_horizon
    self.preload_data_in_memory = dataloader_config.preload_data_in_memory
    self.action_backward_shift = dataloader_config.action_backward_shift
    self.kinematics_norm_type = dataloader_config.kinematics_norm_type
    self.image_norm_type = dataloader_config.image_norm_type
    self.depth_norm_type = dataloader_config.depth_norm_type

    self.train = train
    self.augment_images = train if augment_images is None else augment_images
    self.seed = seed
    self.action_processor = ActionProcessor(action_space=action_space)
    self.image_processor = ImageProcessor(
        color_augmentation=dataloader_config.color_augmentation,
        spatial_augmentation=dataloader_config.spatial_augmentation,
        camera_metadata=observation_space.cameras,
        train=self.augment_images,
    )
    all_keys = list(
        set(
            observation_space.get_required_zarr_keys()
            + action_space.get_required_zarr_keys()
        )
    )  # Remove duplicates
    if replay_buffer is not None:
        self.replay_buffer = replay_buffer
    elif self.preload_data_in_memory:
        self.replay_buffer = ReplayBuffer.copy_from_path(
            zarr_path=zarr_path, keys=all_keys
        )
    else:
        self.replay_buffer = ReplayBuffer.create_from_path(zarr_path=zarr_path)
    missing_keys = set(all_keys) - set(self.replay_buffer.keys())
    if missing_keys:
        raise KeyError(f"Missing required keys in zarr: {missing_keys}")
    logging.info(f"Total episodes in buffer: {self.replay_buffer.n_episodes}")
    # Create episode mask (train/val split)
    episode_mask = self._create_episode_mask(
        val_ratio=dataloader_config.val_ratio,
        total_ratio=dataloader_config.total_ratio,
        train=train,
        seed=seed,
    )
    if dataloader_config.downsample_factor > 1:
        self._apply_downsampling(episode_mask, dataloader_config.downsample_factor)
        episode_mask = np.ones(self.replay_buffer.n_episodes, dtype=bool)
    self.episode_selection_mask = episode_mask
    self.episode_ends = self.replay_buffer.episode_ends[:]
    trailing_padded_actions = dataloader_config.trailing_padded_actions
    if trailing_padded_actions is None:
        trailing_padded_actions = self.pred_horizon - 1
    # Precomputed actions live in zarr at action_positions[t].
    # On-the-fly actions additionally require action_positions[t]+1 to exist.
    # Action backward shift is a "realtime hack" introduced by ACT to
    # compensate for latency in sensor data recording, but defaults to 0.
    next_position_slot = 0 if action_space.has_only_precomputed_actions else 1
    sequence_length = (
        self.obs_horizon
        + self.pred_horizon
        - 1
        + next_position_slot
        + self.action_backward_shift
    )
    self.sampler = SequenceSampler(
        replay_buffer=self.replay_buffer,
        sequence_length=sequence_length,
        pad_before=0,
        pad_after=trailing_padded_actions,
        episode_mask=episode_mask,
        key_first_k=dict.fromkeys(
            observation_space.cameras.keys(),
            self.obs_horizon + self.action_backward_shift,
        ),
        skip_initial=dataloader_config.skip_initial_episode_steps,
        pad_with_zeros=False,
    )
    self._setup_episode_indices()
    self.sample_builder = SampleBuilder(
        action_space=action_space,
        observation_space=observation_space,
        obs_horizon=obs_horizon,
        pred_horizon=pred_horizon,
        action_backward_shift=dataloader_config.action_backward_shift,
        image_processor=self.image_processor,
        action_processor=self.action_processor,
    )
    self.normalizer: LinearNormalizer | None = None

__len__

__len__()

Dataset length depends on sampling mode.

Source code in src/versatil/data/episodic_dataset.py

def __len__(self) -> int:
    """Dataset length depends on sampling mode."""
    return len(self.sampler)

__getitem__

__getitem__(idx)

Get a training sample.

Source code in src/versatil/data/episodic_dataset.py

def __getitem__(
    self, idx: int
) -> dict[str, torch.Tensor] | dict[str, dict[str, torch.Tensor]]:
    """Get a training sample."""
    threadpool_limits(1)
    padded_data = self.sampler.sample_sequence(idx)
    action_slice_start = self.obs_horizon - 1
    action_slice_end = action_slice_start + self.pred_horizon
    action_data, action_meta = self.action_processor.compute_sample_actions(
        padded_data=padded_data,
        action_slice_start=action_slice_start,
        action_slice_end=action_slice_end,
    )
    sample = self.sample_builder.build_sample(
        padded_data=padded_data,
        action_data=action_data,
        action_meta=action_meta,
        start_idx=idx,
        sampler_indices=self.sampler.indices,
    )
    return sample

get_normalizer_and_tokenizer

get_normalizer_and_tokenizer(device=None, winsorize_depth=True, depth_winsorize_quantiles=(0.01, 0.99), winsorize_kinematics=False, kinematics_winsorize_quantiles=(0.01, 0.99), tokenization_config=None, clamp_kinematics_range=True, min_kinematics_std=0.02, min_kinematics_range=0.04, action_sample_size=2048)

Get normalizer and optionally tokenizer for this dataset.

Parameters:

Name	Type	Description	Default
device	device | None	Target device for tensors	None
winsorize_depth	bool	Apply winsorization to depth values	True
depth_winsorize_quantiles	tuple[float, float] | None	Quantiles for depth winsorization	(0.01, 0.99)
winsorize_kinematics	bool	Apply winsorization to kinematics	False
kinematics_winsorize_quantiles	tuple[float, float] | None	Quantiles for kinematics winsorization	(0.01, 0.99)
tokenization_config	TokenizationConfig | None	Tokenization configuration. If None, no tokenizer created.	None
clamp_kinematics_range	bool	Whether to clamp std/range to minimum values.	True
min_kinematics_std	float	Minimum std for Gaussian mode when clamp_kinematics_range=True.	0.02
min_kinematics_range	float	Minimum range for MinMax mode when clamp_kinematics_range=True.	0.04
action_sample_size	int	Number of action rows to stash on the normalizer per action key for downstream data-aware initialization. Set to 0 to disable.	2048

Returns:

Type	Description
tuple[LinearNormalizer, Tokenizer | None]	Tuple of (normalizer, tokenizer) where tokenizer is None if not configured

Source code in src/versatil/data/episodic_dataset.py

def get_normalizer_and_tokenizer(
    self,
    device: torch.device | None = None,
    winsorize_depth: bool = True,
    depth_winsorize_quantiles: tuple[float, float] | None = (0.01, 0.99),
    winsorize_kinematics: bool = False,
    kinematics_winsorize_quantiles: tuple[float, float] | None = (0.01, 0.99),
    tokenization_config: TokenizationConfig | None = None,
    clamp_kinematics_range: bool = True,
    min_kinematics_std: float = 2e-2,
    min_kinematics_range: float = 4e-2,
    action_sample_size: int = 2048,
) -> tuple[LinearNormalizer, Tokenizer | None]:
    """Get normalizer and optionally tokenizer for this dataset.

    Args:
        device: Target device for tensors
        winsorize_depth: Apply winsorization to depth values
        depth_winsorize_quantiles: Quantiles for depth winsorization
        winsorize_kinematics: Apply winsorization to kinematics
        kinematics_winsorize_quantiles: Quantiles for kinematics winsorization
        tokenization_config: Tokenization configuration. If None, no tokenizer created.
        clamp_kinematics_range: Whether to clamp std/range to minimum values.
        min_kinematics_std: Minimum std for Gaussian mode when clamp_kinematics_range=True.
        min_kinematics_range: Minimum range for MinMax mode when clamp_kinematics_range=True.
        action_sample_size: Number of action rows to stash on the normalizer per
            action key for downstream data-aware initialization. Set to 0 to disable.

    Returns:
        Tuple of (normalizer, tokenizer) where tokenizer is None if not configured
    """
    normalizer_builder = TransformBuilder(
        replay_buffer=self.replay_buffer,
        action_processor=self.action_processor,
        observation_space=self.observation_space,
        episode_ends=self.episode_ends,
        kinematics_norm_type=self.kinematics_norm_type,
        image_norm_type=self.image_norm_type,
        depth_norm_type=self.depth_norm_type,
        depth_winsorize_quantiles=depth_winsorize_quantiles
        if winsorize_depth
        else None,
        kinematics_winsorize_quantiles=kinematics_winsorize_quantiles
        if winsorize_kinematics
        else None,
        tokenization_config=tokenization_config,
        prediction_horizon=self.pred_horizon,
        clamp_kinematics_range=clamp_kinematics_range,
        min_kinematics_std=min_kinematics_std,
        min_kinematics_range=min_kinematics_range,
        action_sample_size=action_sample_size,
        episode_selection_mask=self.episode_selection_mask,
        seed=self.seed,
    )

    return normalizer_builder.create_normalizer_and_tokenizer(
        device=device,
    )

set_tokenizer

set_tokenizer(tokenizer)

Set tokenizer for the sample builder.

Parameters:

Name	Type	Description	Default
tokenizer	Tokenizer | None	Unified tokenizer containing observation and action tokenizers	required

Source code in src/versatil/data/episodic_dataset.py

def set_tokenizer(self, tokenizer: Tokenizer | None) -> None:
    """Set tokenizer for the sample builder.

    Args:
        tokenizer: Unified tokenizer containing observation and action tokenizers
    """
    self.sample_builder.tokenizer = tokenizer

set_normalizer

set_normalizer(normalizer)

Set normalizer for the dataset.

Parameters:

Name	Type	Description	Default
normalizer	LinearNormalizer	Normalizer for observations and actions	required

Source code in src/versatil/data/episodic_dataset.py

def set_normalizer(self, normalizer: LinearNormalizer) -> None:
    """Set normalizer for the dataset.

    Args:
        normalizer: Normalizer for observations and actions
    """
    self.sample_builder.normalizer = normalizer

get_gripper_positive_class_imbalance_weight

get_gripper_positive_class_imbalance_weight()

Get class imbalance weight for binary gripper actions.

This is only meaningful for binary grippers where we want to compute the ratio of negative to positive samples for class-weighted BCE loss.

Returns:

Type	Description
float	Weight for positive class (ratio of negative to positive samples)

Raises:

Type	Description
ValueError	If gripper is not configured or is not binary type

Source code in src/versatil/data/episodic_dataset.py

def get_gripper_positive_class_imbalance_weight(self) -> float:
    """Get class imbalance weight for binary gripper actions.

    This is only meaningful for binary grippers where we want to compute
    the ratio of negative to positive samples for class-weighted BCE loss.

    Returns:
        Weight for positive class (ratio of negative to positive samples)

    Raises:
        ValueError: If gripper is not configured or is not binary type
    """
    if not self.action_space.has_gripper_actions:
        raise ValueError("Gripper actions are not being predicted")
    if len(self.action_space.gripper_actions) != 1:
        raise ValueError(
            "Class imbalance weights only supported for single gripper action"
        )
    key, meta = next(iter(self.action_space.gripper_actions.items()))
    if isinstance(meta, GripperActionMetadata):
        gripper_type = meta.gripper_type
    elif isinstance(meta, OnTheFlyActionMetadata):
        if not isinstance(meta.source_metadata, GripperObservationMetadata):
            raise TypeError(
                f"Expected GripperObservationMetadata, got {type(meta.source_metadata)}"
            )
        gripper_type = meta.source_metadata.gripper_type
    else:
        raise ValueError(f"Unsupported gripper action metadata type: {type(meta)}")
    if gripper_type != GripperType.BINARY.value:
        raise ValueError(
            f"Class imbalance weights only supported for binary grippers, "
            f"got gripper_type={gripper_type} for key={key}"
        )
    gripper_actions = self.replay_buffer[key][:]
    if not bool(np.all(self.episode_selection_mask)):
        step_mask = np.zeros(len(gripper_actions), dtype=bool)
        episode_start = 0
        for episode_index, episode_end in enumerate(self.episode_ends):
            if self.episode_selection_mask[episode_index]:
                step_mask[episode_start:episode_end] = True
            episode_start = int(episode_end)
        gripper_actions = gripper_actions[step_mask]
    gripper_actions = gripper_actions.reshape(-1)
    number_of_positive_actions = int((gripper_actions == 1).sum())
    number_of_negative_actions = len(gripper_actions) - number_of_positive_actions
    if number_of_positive_actions == 0 or number_of_negative_actions == 0:
        raise ValueError(
            "Class-imbalance weighting needs both gripper classes in the "
            f"training data; got {number_of_positive_actions} positive and "
            f"{number_of_negative_actions} negative samples for key={key}."
        )
    return number_of_negative_actions / number_of_positive_actions