transform_builder

transform_builder

TransformBuilder

TransformBuilder(replay_buffer, action_processor, prediction_horizon, observation_space, episode_ends, kinematics_norm_type, image_norm_type, depth_norm_type, depth_winsorize_quantiles=(0.01, 0.99), 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, episode_selection_mask=None, seed=42)

Builder for creating and configuring data normalizers and tokenizers.

Initialize transform builder.

Parameters:

Name	Type	Description	Default
replay_buffer	ReplayBuffer	Data source	required
action_processor	ActionProcessor	For computing actions and applying denoising	required
prediction_horizon	int	Prediction horizon for action chunking.	required
observation_space	ObservationSpace	The observation space configuration	required
episode_ends	ndarray	Episode boundaries	required
kinematics_norm_type	str	Normalization type for kinematics	required
image_norm_type	str	Normalization type for RGB images	required
depth_norm_type	str	Normalization type for depth images	required
depth_winsorize_quantiles	tuple[float, float] | None	Quantiles for depth winsorization (lower, upper).	(0.01, 0.99)
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 chunk windows (length prediction_horizon, aligned across keys) to stash on the normalizer per action key for downstream data-aware initialization (e.g. chunk-space mixture-density). Set to 0 to disable. Memory cost per action key is action_sample_size * prediction_horizon * action_dim * bytes_per_element.	2048
episode_selection_mask	ndarray | None	Optional boolean mask over episodes. Statistics, tokenizers, and denoising thresholds are fitted only on selected episodes (the training split), keeping validation data out of the fitted transforms.	None
seed	int	Seed for the subsampling used by depth-quantile estimation, so fitted statistics are reproducible.	42

Source code in src/versatil/data/processing/transform_builder.py

def __init__(
    self,
    replay_buffer: ReplayBuffer,
    action_processor: ActionProcessor,
    prediction_horizon: int,
    observation_space: ObservationSpace,
    episode_ends: np.ndarray,
    kinematics_norm_type: str,
    image_norm_type: str,
    depth_norm_type: str,
    depth_winsorize_quantiles: tuple[float, float] | None = (0.01, 0.99),
    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,
    episode_selection_mask: np.ndarray | None = None,
    seed: int = 42,
):
    """Initialize transform builder.

    Args:
        replay_buffer: Data source
        action_processor: For computing actions and applying denoising
        prediction_horizon: Prediction horizon for action chunking.
        observation_space: The observation space configuration
        episode_ends: Episode boundaries
        kinematics_norm_type: Normalization type for kinematics
        image_norm_type: Normalization type for RGB images
        depth_norm_type: Normalization type for depth images
        depth_winsorize_quantiles: Quantiles for depth winsorization (lower, upper).
        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 chunk windows (length
            ``prediction_horizon``, aligned across keys) to stash on the
            normalizer per action key for downstream data-aware
            initialization (e.g. chunk-space mixture-density).
            Set to 0 to disable. Memory cost per action key is
            ``action_sample_size * prediction_horizon * action_dim *
            bytes_per_element``.
        episode_selection_mask: Optional boolean mask over episodes.
            Statistics, tokenizers, and denoising thresholds are fitted
            only on selected episodes (the training split), keeping
            validation data out of the fitted transforms.
        seed: Seed for the subsampling used by depth-quantile estimation,
            so fitted statistics are reproducible.
    """
    self.replay_buffer = replay_buffer
    self.action_processor = action_processor
    self.observation_space = observation_space
    self.episode_ends = episode_ends
    self.episode_selection_mask = episode_selection_mask
    self._selected_step_mask = self._build_selected_step_mask()
    self.kinematics_norm_type = kinematics_norm_type
    self.image_norm_type = image_norm_type
    self.depth_norm_type = depth_norm_type
    self.depth_winsorize_quantiles = depth_winsorize_quantiles
    self.kinematics_winsorize_quantiles = kinematics_winsorize_quantiles
    self.tokenization_config = tokenization_config
    self.prediction_horizon = prediction_horizon
    self.clamp_kinematics_range = clamp_kinematics_range
    self.min_kinematics_std = min_kinematics_std
    self.min_kinematics_range = min_kinematics_range
    self.action_sample_size = action_sample_size
    self._random_generator = np.random.default_rng(seed)

create_normalizer_and_tokenizer

create_normalizer_and_tokenizer(device=None)

Create and fit normalizer and optionally tokenizer to data. Pipeline: Raw data → Winsorize → Normalize → Tokenize

Parameters:

Name	Type	Description	Default
device	device | None	Target device for tensors	None

Returns:

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

Source code in src/versatil/data/processing/transform_builder.py

def create_normalizer_and_tokenizer(
    self,
    device: torch.device | None = None,
) -> tuple[LinearNormalizer, Tokenizer | None]:
    """Create and fit normalizer and optionally tokenizer to data.
    Pipeline: Raw data → Winsorize → Normalize → Tokenize

    Args:
        device: Target device for tensors

    Returns:
        Tuple of (normalizer, tokenizer) where tokenizer is None if not configured
    """
    self.compute_proprioceptive_denoising_thresholds()
    action_keys = self.action_processor.action_space.get_required_zarr_keys()
    action_source_data = {
        key: self.replay_buffer[key][:]
        for key in action_keys
        if key in self.replay_buffer
    }
    action_data, action_meta = self.action_processor.compute_sample_actions(
        padded_data=action_source_data,
        action_slice_start=0,
        action_slice_end=self.replay_buffer.n_steps - 1,
    )
    cross_indices = self.episode_ends[:-1] - 1
    valid_mask = np.ones(len(next(iter(action_data.values()))), dtype=bool)
    valid_mask[cross_indices] = False
    if self._selected_step_mask is not None:
        valid_mask &= self._selected_step_mask[: len(valid_mask)]
    valid_action_data = {key: data[valid_mask] for key, data in action_data.items()}
    # On-the-fly actions at episode-final rows straddle episode boundaries
    # and must stay out of the fitted statistics. Precomputed actions are
    # valid training targets there, so their stats cover every selected row.
    stats_action_data = {
        key: self._select_step_rows(array=action_source_data[key])
        if action_meta[key].is_precomputed
        else data
        for key, data in valid_action_data.items()
    }

    normalizer = self._create_normalizer(
        action_data=stats_action_data,
        action_meta=action_meta,
        device=device,
    )
    chunk_action_keys = [
        key for key, meta in action_meta.items() if meta.needs_normalization
    ]
    chunk_action_data = (
        action_source_data
        if self.action_processor.action_space.has_only_precomputed_actions
        else action_data
    )
    self._stash_action_chunk_samples(
        normalizer=normalizer,
        action_data=chunk_action_data,
        action_keys=chunk_action_keys,
        requires_next_observation=(
            not self.action_processor.action_space.has_only_precomputed_actions
        ),
    )
    tokenizer = None
    if self.tokenization_config and (
        self.tokenization_config.tokenize_observations
        or self.tokenization_config.tokenize_actions
    ):
        tokenizer_action_data = (
            {
                key: self._select_step_rows(array=action_source_data[key])
                for key in valid_action_data
            }
            if self.action_processor.action_space.has_only_precomputed_actions
            else valid_action_data
        )
        tokenizer = self._create_tokenizer(
            normalizer=normalizer,
            action_data=tokenizer_action_data,
            action_meta=action_meta,
            device=device,
        )
    return normalizer, tokenizer

compute_proprioceptive_denoising_thresholds

compute_proprioceptive_denoising_thresholds()

Compute denoising thresholds for proprioceptive observations.

Source code in src/versatil/data/processing/transform_builder.py

def compute_proprioceptive_denoising_thresholds(
    self,
) -> None:
    """Compute denoising thresholds for proprioceptive observations."""
    for key, meta in self.action_processor.action_space.actions_metadata.items():
        if isinstance(meta, OnTheFlyActionMetadata):
            source_meta = meta.source_metadata
            if isinstance(source_meta, PositionObservationMetadata):
                obs_data, selected_episode_ends = self._select_episodes_contiguous(
                    array=self.replay_buffer[key][:]
                )
                self.action_processor.compute_denoising_threshold(
                    obs_data=obs_data,
                    key=key,
                    meta=source_meta,
                    episode_ends=selected_episode_ends,
                )
    self.action_processor.log_movement_distribution()