sample_builder

sample_builder

Sample construction for episodic dataset.

Builds the final training/validation samples by: - Processing images - Adding proprioceptive data - Adding actions - Computing padding masks - Normalizing (and optionally tokenizing) observations and actions

SampleBuilder

SampleBuilder(action_space, observation_space, obs_horizon, pred_horizon, action_backward_shift, image_processor, action_processor, tokenizer=None, normalizer=None)

Builds training samples from raw episode data.

Parameters:

Name	Type	Description	Default
action_space	ActionSpace	Action space of the experiment	required
observation_space	ObservationSpace	Observation space of the experiment	required
obs_horizon	int	Observation history length	required
pred_horizon	int	Action prediction horizon	required
action_backward_shift	int	Backward shift to give to action timesteps (if actions have latency)	required
image_processor	ImageProcessor	Handles augmentations	required
action_processor	ActionProcessor	Processes actions	required
tokenizer	Tokenizer | None	Unified tokenizer for observations and actions	None
normalizer	LinearNormalizer | None	Normalizer for observations and actions	None

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

def __init__(
    self,
    action_space: ActionSpace,
    observation_space: ObservationSpace,
    obs_horizon: int,
    pred_horizon: int,
    action_backward_shift: int,
    image_processor: ImageProcessor,
    action_processor: ActionProcessor,
    tokenizer: Tokenizer | None = None,
    normalizer: LinearNormalizer | None = None,
):
    """
    Args:
        action_space: Action space of the experiment
        observation_space: Observation space of the experiment
        obs_horizon: Observation history length
        pred_horizon: Action prediction horizon
        action_backward_shift: Backward shift to give to action timesteps (if actions have latency)
        image_processor: Handles augmentations
        action_processor: Processes actions
        tokenizer: Unified tokenizer for observations and actions
        normalizer: Normalizer for observations and actions
    """
    self.action_space = action_space
    self.observation_space = observation_space
    self.obs_horizon = obs_horizon
    self.pred_horizon = pred_horizon
    self.action_backward_shift = action_backward_shift
    self.image_processor = image_processor
    self.action_processor = action_processor
    self.tokenizer = tokenizer
    self.normalizer = normalizer

build_sample

build_sample(padded_data, action_data, action_meta, start_idx, sampler_indices)

Build a complete training sample.

Parameters:

Name	Type	Description	Default
padded_data	dict[str, ndarray]	Dictionary of padded episode data	required
action_data	dict[str, ndarray]	Dictionary of computed actions	required
action_meta	dict[str, ActionMetadata]	Dictionary of action metadata	required
start_idx	int	Starting index in sampler	required
sampler_indices	ndarray	Array of sampler indices	required

Note

Padded data layout: [historical buffer | observation window | future] [0:k] [k:k+H] [k+H:end] where k=action_backward_shift, H=obs_horizon, end= prediction_horizon+k+H

Returns:

Type	Description
dict[str, dict[str, Tensor]]	Dictionary containing observation and action dictionaries. Each sub-dictionary maps keys to tensors.

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

def build_sample(
    self,
    padded_data: dict[str, np.ndarray],
    action_data: dict[str, np.ndarray],
    action_meta: dict[str, ActionMetadata],
    start_idx: int,
    sampler_indices: np.ndarray,
) -> dict[str, dict[str, torch.Tensor]]:
    """Build a complete training sample.

    Args:
        padded_data: Dictionary of padded episode data
        action_data: Dictionary of computed actions
        action_meta: Dictionary of action metadata
        start_idx: Starting index in sampler
        sampler_indices: Array of sampler indices

    Note:
        Padded data layout: [historical buffer | observation window | future]
                                [0:k]             [k:k+H]            [k+H:end]
        where k=action_backward_shift, H=obs_horizon, end= prediction_horizon+k+H

    Returns:
        Dictionary containing observation and action dictionaries. Each sub-dictionary maps keys to tensors.
    """
    sample: dict[str, dict[str, torch.Tensor]] = {
        SampleKey.OBSERVATION.value: {},
        SampleKey.ACTION.value: {},
    }
    image_dict = self._get_sample_images(padded_data=padded_data)
    sample[SampleKey.OBSERVATION.value].update(image_dict)
    for key, metadata in self.observation_space.observations_metadata.items():
        if isinstance(metadata, ObservationMetadata):  # Excludes cameras
            sample[SampleKey.OBSERVATION.value].update(
                {
                    key: self._slice_observation_tensor(
                        key=key, metadata=metadata, padded_data=padded_data
                    )
                }
            )
    for key, _data in action_data.items():
        metadata = action_meta[key]
        sample[SampleKey.ACTION.value].update(
            {
                key: self._slice_action_data(
                    key=key, metadata=metadata, action_data=action_data
                )
            }
        )

    sample[SampleKey.ACTION.value][SampleKey.IS_PAD_ACTION.value] = (
        self._compute_action_padding_mask(
            start_idx=start_idx, sampler_indices=sampler_indices
        )
    )
    sample = self.normalize_and_tokenize_sample(sample=sample)
    return sample

normalize_and_tokenize_sample

normalize_and_tokenize_sample(sample)

Normalize and tokenize a pre-built sample.

Parameters:

Name	Type	Description	Default
sample	dict[str, dict[str, Tensor]]	Pre-built sample with observation and action dictionaries.	required

Returns:

Type	Description
dict[str, dict[str, Tensor]]	Normalized and tokenized sample.

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

def normalize_and_tokenize_sample(
    self,
    sample: dict[str, dict[str, torch.Tensor]],
) -> dict[str, dict[str, torch.Tensor]]:
    """Normalize and tokenize a pre-built sample.

    Args:
        sample: Pre-built sample with observation and action dictionaries.

    Returns:
        Normalized and tokenized sample.
    """
    if self.normalizer is not None:
        sample = normalize_sample(
            sample=sample,
            normalizer=self.normalizer,
            observation_space=self.observation_space,
            action_space=self.action_space,
        )
    if self.tokenizer is not None:
        sample = tokenize_sample(
            sample=sample, tokenizer=self.tokenizer, action_space=self.action_space
        )
    return sample