task

task

Task space definitions for runtime data requirements.

This module defines what data an experiment uses at runtime:

ActionSpace

ActionSpace(actions_metadata, use_gripper_class_weights=False, denoise_actions=True, denoising_percentile=15.0)

Defines what actions the task will predict at runtime.

Attributes:

Name	Type	Description
actions_metadata		Dict of all action metadata, indexed by zarr store key. Values are OnTheFlyActionMetadata or PrecomputedActionMetadata subclasses.
use_gripper_class_weights		Whether to use class weights for binary gripper.
denoise_actions		Whether to apply denoising to actions.
denoising_percentile		Percentile for denoising threshold.

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

def __init__(
    self,
    actions_metadata: dict[str, ActionMetadata],
    use_gripper_class_weights: bool = False,
    denoise_actions: bool = True,
    denoising_percentile: float = 15.0,
):
    self.actions_metadata = resolve_dict_keys(actions_metadata)
    self.use_gripper_class_weights = use_gripper_class_weights
    self.denoise_actions = denoise_actions
    self.denoising_percentile = denoising_percentile

on_the_fly_actions property

on_the_fly_actions

Get actions computed on-the-fly from observations.

precomputed_actions property

precomputed_actions

Get precomputed actions loaded directly from zarr.

position_actions property

position_actions

Get all position actions (precomputed or on-the-fly).

orientation_actions property

orientation_actions

Get all orientation actions (precomputed or on-the-fly).

gripper_actions property

gripper_actions

Get all gripper actions (precomputed or on-the-fly).

predicted_action_keys property

predicted_action_keys

Return action keys predicted by the policy in metadata order.

predicted_action_dimensions property

predicted_action_dimensions

Return predicted action dimensions keyed by action name.

total_prediction_dimension property

total_prediction_dimension

Return total predicted continuous action dimension.

position_dim property

position_dim

Get total position action dimension.

orientation_dim property

orientation_dim

Get total orientation action dimension.

gripper_dim property

gripper_dim

Get total gripper action dimension.

has_on_the_fly_actions property

has_on_the_fly_actions

Check if there are any actions to compute on-the-fly.

has_precomputed_actions property

has_precomputed_actions

Check if there are any precomputed actions.

has_only_precomputed_actions property

has_only_precomputed_actions

Check if every action in the space is precomputed (no on-the-fly computation).

has_delta_actions property

has_delta_actions

Check if any actions are computed as deltas.

has_position_actions property

has_position_actions

Check if there are any position actions.

has_orientation_actions property

has_orientation_actions

Check if there are any orientation actions.

has_gripper_actions property

has_gripper_actions

Check if there are any gripper actions.

get_total_action_dim

get_total_action_dim()

Calculate total action dimension for predicted actions.

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

def get_total_action_dim(self) -> int:
    """Calculate total action dimension for predicted actions."""
    return self.total_prediction_dimension

validate_action_tensors

validate_action_tensors(actions, prediction_horizon, owner_name)

Validate action tensors against this action-space schema.

Parameters:

Name	Type	Description	Default
actions	dict[str, Tensor]	Action tensors keyed by action name.	required
prediction_horizon	int	Expected action chunk length.	required
owner_name	str	Name used in error messages.	required

Returns:

Type	Description
tuple[int, device]	Shared batch size and device.

Raises:

Type	Description
ValueError	If keys, ranks, dimensions, batch size, or devices are invalid.

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

def validate_action_tensors(
    self,
    actions: dict[str, torch.Tensor],
    prediction_horizon: int,
    owner_name: str,
) -> tuple[int, torch.device]:
    """Validate action tensors against this action-space schema.

    Args:
        actions: Action tensors keyed by action name.
        prediction_horizon: Expected action chunk length.
        owner_name: Name used in error messages.

    Returns:
        Shared batch size and device.

    Raises:
        ValueError: If keys, ranks, dimensions, batch size, or devices are invalid.
    """
    expected_action_keys = self.predicted_action_keys
    if not expected_action_keys:
        raise ValueError(f"{owner_name} requires at least one predicted action.")

    actual_action_keys = set(actions.keys())
    expected_action_key_set = set(expected_action_keys)
    if actual_action_keys != expected_action_key_set:
        raise ValueError(
            f"{owner_name} expected action keys "
            f"{expected_action_keys}, got {sorted(actions.keys())}."
        )

    first_action_key = ""
    batch_size = 0
    action_device = torch.device("cpu")
    for action_key in expected_action_keys:
        action = actions[action_key]
        if action.ndim != 3:
            raise ValueError(
                f"Action '{action_key}' must have shape "
                f"(B, prediction_horizon, action_dim), got {action.shape}."
            )
        if action.shape[1] != prediction_horizon:
            raise ValueError(
                f"Action '{action_key}' must have prediction horizon "
                f"{prediction_horizon}, got {action.shape[1]}."
            )
        expected_dimension = self.predicted_action_dimensions[action_key]
        if action.shape[2] != expected_dimension:
            raise ValueError(
                f"Action '{action_key}' must have last dimension "
                f"{expected_dimension}, got {action.shape[2]}."
            )
        if first_action_key == "":
            first_action_key = action_key
            batch_size = action.shape[0]
            action_device = action.device
            continue
        if action.shape[0] != batch_size:
            raise ValueError(
                "All action tensors must have the same batch size, "
                f"got {action.shape[0]} for '{action_key}' and "
                f"{batch_size} for '{first_action_key}'."
            )
        if action.device != action_device:
            raise ValueError(
                "All action tensors must be on the same device, "
                f"got {action.device} for '{action_key}' and "
                f"{action_device} for '{first_action_key}'."
            )
    return batch_size, action_device

concatenate_action_tensors

concatenate_action_tensors(actions, prediction_horizon, owner_name, dtype=None, device=None)

Concatenate predicted action tensors in action-space order.

Parameters:

Name	Type	Description	Default
actions	dict[str, Tensor]	Action tensors keyed by action name.	required
prediction_horizon	int	Expected action chunk length.	required
owner_name	str	Name used in validation errors.	required
dtype	dtype | None	Optional dtype for the returned tensor.	None
device	device | None	Optional device for the returned tensor.	None

Returns:

Type	Description
Tensor	Concatenated action tensor with shape
Tensor	(B, prediction_horizon, total_prediction_dimension).

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

def concatenate_action_tensors(
    self,
    actions: dict[str, torch.Tensor],
    prediction_horizon: int,
    owner_name: str,
    dtype: torch.dtype | None = None,
    device: torch.device | None = None,
) -> torch.Tensor:
    """Concatenate predicted action tensors in action-space order.

    Args:
        actions: Action tensors keyed by action name.
        prediction_horizon: Expected action chunk length.
        owner_name: Name used in validation errors.
        dtype: Optional dtype for the returned tensor.
        device: Optional device for the returned tensor.

    Returns:
        Concatenated action tensor with shape
        ``(B, prediction_horizon, total_prediction_dimension)``.
    """
    self.validate_action_tensors(
        actions=actions,
        prediction_horizon=prediction_horizon,
        owner_name=owner_name,
    )
    action_tensors = []
    for action_key in self.predicted_action_keys:
        action = actions[action_key]
        if dtype is not None or device is not None:
            action = action.to(dtype=dtype, device=device)
        action_tensors.append(action)
    return torch.cat(action_tensors, dim=-1)

split_action_tensor

split_action_tensor(action_tensor, owner_name)

Split a joint action tensor into action-space components.

Parameters:

Name	Type	Description	Default
action_tensor	Tensor	Tensor with final dimension equal to the total predicted action dimension.	required
owner_name	str	Name used in validation errors.	required

Returns:

Type	Description
dict[str, Tensor]	Action tensors keyed by action name.

Raises:

Type	Description
ValueError	If the final dimension does not match the action space.

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

def split_action_tensor(
    self,
    action_tensor: torch.Tensor,
    owner_name: str,
) -> dict[str, torch.Tensor]:
    """Split a joint action tensor into action-space components.

    Args:
        action_tensor: Tensor with final dimension equal to the total
            predicted action dimension.
        owner_name: Name used in validation errors.

    Returns:
        Action tensors keyed by action name.

    Raises:
        ValueError: If the final dimension does not match the action space.
    """
    if action_tensor.shape[-1] != self.total_prediction_dimension:
        raise ValueError(
            f"{owner_name} expected joint action final dimension "
            f"{self.total_prediction_dimension}, got {action_tensor.shape[-1]}."
        )

    predictions = {}
    offset = 0
    for action_key, dimension in self.predicted_action_dimensions.items():
        predictions[action_key] = action_tensor[..., offset : offset + dimension]
        offset += dimension
    return predictions

get_required_zarr_keys

get_required_zarr_keys()

Get zarr keys needed for this action space.

Returns:

Type	Description
list[str]	List of keys to load from replay buffer

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

def get_required_zarr_keys(self) -> list[str]:
    """Get zarr keys needed for this action space.

    Returns:
        List of keys to load from replay buffer
    """
    return list(self.actions_metadata.keys())

ObservationSpace

ObservationSpace(observations_metadata)

Defines what observations the task will load at runtime.

Attributes:

Name	Type	Description
observations_metadata		Dict of all observation metadata, indexed by zarr store key. Values are ObservationMetadata subclasses or CameraMetadata.

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

def __init__(
    self,
    observations_metadata: dict[str, ObservationMetadata | CameraMetadata],
):
    self.observations_metadata = resolve_dict_keys(observations_metadata)
    validate_camera_metadata_keys(self.cameras)

cameras property

cameras

Get all camera observations.

depth_cameras property

depth_cameras

Get all depth camera observations.

rgb_cameras property

rgb_cameras

Get all RGB camera observations.

position_observations property

position_observations

Get all position observations.

orientation_observations property

orientation_observations

Get all orientation observations.

gripper_observations property

gripper_observations

Get all gripper state observations.

proprioceptive_observations property

proprioceptive_observations

Get robot proprioceptive observations (position, orientation, gripper).

numerical_observations property

numerical_observations

Get all numerical non-image observations.

custom_observations property

custom_observations

Get custom observations (not position, orientation, gripper, or camera).

has_cameras property

has_cameras

Check if any camera observations are included.

has_gripper_state property

has_gripper_state

Check if gripper state observation is included.

has_proprioceptive_state property

has_proprioceptive_state

Check if any proprioceptive observations are included.

has_proprioceptive_position property

has_proprioceptive_position

Check if any position observations are included.

has_proprioceptive_orientation property

has_proprioceptive_orientation

Check if any orientation observations are included.

get_required_zarr_keys

get_required_zarr_keys()

Get all zarr keys needed for this observation space.

Returns:

Type	Description
list[str]	List of keys to load from replay buffer

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

def get_required_zarr_keys(self) -> list[str]:
    """Get all zarr keys needed for this observation space.

    Returns:
        List of keys to load from replay buffer
    """
    return list(self.observations_metadata.keys())

TaskSpace

TaskSpace(dataset_schema, dataloader, action_space, observation_space, prediction_horizon=16, observation_horizon=1)

Combines action/observation space with dataset schema for runtime.

The task space validates that requested keys exist in the dataset schema and that metadata is consistent between the schema and task requirements.

Initialize task space.

Parameters:

Name	Type	Description	Default
dataset_schema	DatasetSchema	Schema defining what's in the zarr store.	required
dataloader	DataLoaderConfig	Data loading configuration.	required
action_space	ActionSpace	Actions to predict.	required
observation_space	ObservationSpace	Observations to load.	required
prediction_horizon	int	Number of timesteps to predict.	16
observation_horizon	int	Number of history timesteps to load.	1

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

def __init__(
    self,
    dataset_schema: DatasetSchema,
    dataloader: DataLoaderConfig,
    action_space: ActionSpace,
    observation_space: ObservationSpace,
    prediction_horizon: int = 16,
    observation_horizon: int = 1,
):
    """Initialize task space.

    Args:
        dataset_schema: Schema defining what's in the zarr store.
        dataloader: Data loading configuration.
        action_space: Actions to predict.
        observation_space: Observations to load.
        prediction_horizon: Number of timesteps to predict.
        observation_horizon: Number of history timesteps to load.
    """
    self.dataset_schema = dataset_schema
    self.dataloader = dataloader
    self.action_space = action_space
    self.observation_space = observation_space
    self.prediction_horizon = prediction_horizon
    self.observation_horizon = observation_horizon
    self._validate()