transform

transform

normalize_sample

normalize_sample(sample, normalizer, observation_space, action_space)

Normalize a pre-built sample.

Tokenization is a separate step; see tokenize_observation.

Parameters:

Name	Type	Description	Default
sample	dict[str, dict[str, Tensor]]	Pre-built sample with observation and action dictionaries.	required
normalizer	LinearNormalizer	Normalizer to use for normalization.	required
observation_space	ObservationSpace	Observation space configuration.	required
action_space	ActionSpace	Action space configuration.	required

Returns:

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

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

def normalize_sample(
    sample: dict[str, dict[str, torch.Tensor]],
    normalizer: LinearNormalizer,
    observation_space: ObservationSpace,
    action_space: ActionSpace,
) -> dict[str, dict[str, torch.Tensor]]:
    """Normalize a pre-built sample.

    Tokenization is a separate step; see ``tokenize_observation``.

    Args:
        sample: Pre-built sample with observation and action dictionaries.
        normalizer: Normalizer to use for normalization.
        observation_space: Observation space configuration.
        action_space: Action space configuration.

    Returns:
        Normalized sample.
    """
    sample_copy = sample.copy()  # Avoid modifying in-place
    observation = sample_copy[SampleKey.OBSERVATION.value]
    sample_copy[SampleKey.OBSERVATION.value] = normalize_observation(
        observation=observation,
        normalizer=normalizer,
        observation_space=observation_space,
    )
    actions = sample_copy[SampleKey.ACTION.value]
    sample_copy[SampleKey.ACTION.value] = normalize_actions(
        actions=actions, normalizer=normalizer, action_space=action_space
    )
    return sample_copy

normalize_observation

normalize_observation(observation, normalizer, observation_space)

Normalize observations.

Parameters:

Name	Type	Description	Default
observation	dict[str, Tensor]	Observation dictionary.	required
normalizer	LinearNormalizer	Normalizer to use for normalization.	required
observation_space	ObservationSpace	Observation space configuration.	required

Returns:

Type	Description
dict[str, Tensor]	Normalized observation dictionary.

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

def normalize_observation(
    observation: dict[str, torch.Tensor],
    normalizer: LinearNormalizer,
    observation_space: ObservationSpace,
) -> dict[str, torch.Tensor]:
    """Normalize observations.

    Args:
        observation: Observation dictionary.
        normalizer: Normalizer to use for normalization.
        observation_space: Observation space configuration.

    Returns:
        Normalized observation dictionary.
    """
    normalized_observation = observation.copy()
    for key in observation_space.observations_metadata:
        if key in normalizer.params_dict.keys():
            normalized_observation[key] = normalizer[key].normalize(observation[key])
    return normalized_observation

normalize_actions

normalize_actions(actions, normalizer, action_space)

Normalize actions.

Parameters:

Name	Type	Description	Default
actions	dict[str, Tensor]	Action dictionary.	required
normalizer	LinearNormalizer	Normalizer to use for normalization.	required
action_space	ActionSpace	Action space configuration.	required

Returns:

Type	Description
dict[str, Tensor]	Normalized action dictionary.

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

def normalize_actions(
    actions: dict[str, torch.Tensor],
    normalizer: LinearNormalizer,
    action_space: ActionSpace,
) -> dict[str, torch.Tensor]:
    """Normalize actions.

    Args:
        actions: Action dictionary.
        normalizer: Normalizer to use for normalization.
        action_space: Action space configuration.

    Returns:
        Normalized action dictionary.
    """
    normalized_actions = actions.copy()
    for key in action_space.actions_metadata:
        if key in normalizer.params_dict.keys():
            normalized_actions[key] = normalizer[key].normalize(actions[key])
    return normalized_actions

unnormalize_actions

unnormalize_actions(normalized_actions, normalizer, action_space)

Unnormalize actions using the normalizer with per-key statistics.

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

def unnormalize_actions(
    normalized_actions: dict[str, torch.Tensor],
    normalizer: LinearNormalizer,
    action_space: ActionSpace,
) -> dict[str, torch.Tensor]:
    """Unnormalize actions using the normalizer with per-key statistics."""
    actions = normalized_actions.copy()
    for key in action_space.actions_metadata:
        if key in normalizer.params_dict.keys():
            actions[key] = normalizer[key].unnormalize(normalized_actions[key])
    return actions

tokenize_sample

tokenize_sample(sample, tokenizer, action_space)

Tokenize observations and actions according to the action space configuration.

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

def tokenize_sample(
    sample: dict[str, dict[str, torch.Tensor]],
    tokenizer: Tokenizer,
    action_space: ActionSpace,
) -> dict[str, dict[str, torch.Tensor]]:
    """Tokenize observations and actions according to the action space configuration."""
    if tokenizer.observation_tokenizer is not None:
        observation = sample[SampleKey.OBSERVATION.value]
        sample[SampleKey.OBSERVATION.value] = tokenize_observation(
            observation=observation, obs_tokenizer=tokenizer.observation_tokenizer
        )
    if tokenizer.action_tokenizer is not None:
        actions = sample[SampleKey.ACTION.value]
        sample[SampleKey.ACTION.value] = tokenize_actions(
            actions=actions,
            action_space=action_space,
            action_tokenizer=tokenizer.action_tokenizer,
        )
    return sample

tokenize_observation

tokenize_observation(observation, obs_tokenizer, batched=False)

Tokenize observations.

Parameters:

Name	Type	Description	Default
observation	dict[str, Tensor]	Observation dictionary.	required
obs_tokenizer	ObservationTokenizer	Fitted observation tokenizer.	required
batched	bool	Whether observations carry a leading batch dimension in front of the observation horizon.	False

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

def tokenize_observation(
    observation: dict[str, torch.Tensor],
    obs_tokenizer: ObservationTokenizer,
    batched: bool = False,
) -> dict[str, torch.Tensor]:
    """Tokenize observations.

    Args:
        observation: Observation dictionary.
        obs_tokenizer: Fitted observation tokenizer.
        batched: Whether observations carry a leading batch dimension in
            front of the observation horizon.
    """
    obs_copy = observation.copy()
    obs_to_tokenize = {}
    for key in obs_tokenizer.observation_keys:
        if key in obs_copy:
            obs_to_tokenize[key] = obs_copy[key]
        else:
            raise KeyError(
                f"Observation key '{key}' not found in sample for tokenization."
            )
    tokenized = obs_tokenizer.tokenize(obs_to_tokenize, batched=batched)
    obs_copy[SampleKey.TOKENIZED_OBSERVATIONS.value] = tokenized[
        SampleKey.TOKENIZED_OBSERVATIONS.value
    ]
    obs_copy[SampleKey.IS_PAD_OBSERVATION.value] = tokenized[
        SampleKey.IS_PAD_OBSERVATION.value
    ]
    return obs_copy

tokenize_actions

tokenize_actions(actions, action_tokenizer, action_space)

Tokenize actions.

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

def tokenize_actions(
    actions: dict[str, torch.Tensor],
    action_tokenizer: ActionTokenizer,
    action_space: ActionSpace,
) -> dict[str, torch.Tensor]:
    """Tokenize actions."""
    actions_to_tokenize = actions.copy()
    action_components = []
    # Iterate in metadata insertion order — the same canonical order
    # predicted_action_keys and the joint action layout use.
    for key, meta in action_space.actions_metadata.items():
        if meta.is_numerical and meta.requires_prediction_head:
            action_tensor = actions_to_tokenize[key]
            if action_tensor.ndim == 1:
                action_tensor = action_tensor.unsqueeze(-1)
            action_components.append(action_tensor)

    # Concatenate along last dimension: (pred_horizon, action_dim)
    action_tensor = torch.cat(action_components, dim=-1)
    is_pad_mask = actions_to_tokenize.get(SampleKey.IS_PAD_ACTION.value, None)
    tokenized = action_tokenizer.encode(action_tensor, is_pad_mask=is_pad_mask)
    actions_to_tokenize[SampleKey.TOKENIZED_ACTIONS.value] = tokenized[
        SampleKey.TOKENIZED_ACTIONS.value
    ]
    actions_to_tokenize[SampleKey.IS_PAD_ACTION.value] = tokenized[
        SampleKey.IS_PAD_ACTION.value
    ]  # Replace padding mask with tokenizer mask
    return actions_to_tokenize

detokenize_actions

detokenize_actions(action_tokens, action_tokenizer, action_space)

Detokenize actions.

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

def detokenize_actions(
    action_tokens: torch.Tensor,
    action_tokenizer: ActionTokenizer,
    action_space: ActionSpace,
) -> dict[str, torch.Tensor]:
    """Detokenize actions."""
    action_tokens_cpu = action_tokens.cpu()
    # Squeeze if needed: (B, pred_horizon, 1) -> (B, pred_horizon)
    if action_tokens_cpu.ndim == 3 and action_tokens_cpu.shape[-1] == 1:
        action_tokens_cpu = action_tokens_cpu.squeeze(-1)
    batch_size, pred_horizon = action_tokens_cpu.shape[:2]
    detokenized_actions = []
    for chunk in range(batch_size):
        actions_np = action_tokenizer.decode(action_tokens_cpu[chunk])
        detokenized_actions.append(torch.from_numpy(actions_np).float())
    actions = torch.stack(detokenized_actions, dim=0)  # (B, pred_horizon, action_dim)
    action_dict = {}
    current_idx = 0
    for key, meta in action_space.actions_metadata.items():
        if meta.is_numerical and meta.requires_prediction_head:
            dim = meta.prediction_dimension
            action_dict[key] = actions[..., current_idx : current_idx + dim]
            current_idx += dim
            if meta.action_type == ProprioceptiveType.GRIPPER.value:
                if isinstance(meta, GripperActionMetadata):
                    gripper_meta = meta
                elif isinstance(meta, OnTheFlyActionMetadata):
                    if not isinstance(meta.source_metadata, GripperObservationMetadata):
                        raise TypeError(
                            f"Expected GripperObservationMetadata, got {type(meta.source_metadata)}"
                        )
                    gripper_meta = meta.source_metadata
                else:
                    raise TypeError(
                        f"Unexpected metadata type for gripper action: {type(meta)}"
                    )
                if gripper_meta.gripper_type == GripperType.BINARY.value:
                    if (
                        gripper_meta.binary_gripper_range
                        == BinaryGripperRange.MINUS_ONE_ONE.value
                    ):
                        action_dict[key] = torch.where(
                            action_dict[key] > 0.0, 1, -1
                        ).long()
                    else:
                        if (
                            gripper_meta.binary_gripper_range
                            != BinaryGripperRange.ZERO_ONE.value
                        ):
                            logging.warning(
                                "Gripper type is binary but range is not set. "
                                "Assuming {0,1}."
                            )
                        action_dict[key] = (action_dict[key] > 0.5).long()

    return action_dict