base

base

Base class for action-decoding algorithms (BC, Diffusion, FlowMatching, etc.).

DecodingAlgorithm

DecodingAlgorithm()

Bases: Module, ABC

Base class for decoding algorithms.

Algorithms define the learning paradigm (behavioral cloning, diffusion, flow matching, etc.). Pure algorithms should not contain latent variable logic - use VariationalAlgorithm wrapper to add variational inference capabilities to any algorithm.

Examples:

Pure algorithm (deterministic)

algorithm = BehavioralCloning()

Algorithm with variational inference

algorithm = VariationalAlgorithm( base_algorithm=BehavioralCloning(), posterior_encoder=VAETransformerEncoder(...), prior=GaussianPrior(...) # or DiTPrior(...) for learned prior )

Initialize decoding algorithm.

Source code in src/versatil/models/decoding/algorithm/base.py

def __init__(self):
    """Initialize decoding algorithm."""
    super().__init__()

predicts_in_action_space property

predicts_in_action_space

Whether the network output lives in the action space.

When True (e.g. Behavioral Cloning), the network directly predicts actions and classification losses like BCE are valid. When False (e.g. Flow Matching, Diffusion with epsilon/velocity), the network predicts in a derived space (velocity field, noise) and classification losses on those outputs are meaningless.

forward abstractmethod

forward(network, features, actions=None)

Forward pass during training.

Parameters:

Name	Type	Description	Default
network	ActionDecoder	The action decoder network module	required
features	dict[str, Tensor]	Dictionary of encoded features from the encoding pipeline.	required
actions	dict[str, Tensor] | None	Optional dictionary of ground truth actions, if architecture requires it, e.g. ACT.	None

Returns:

Type	Description
dict[str, Tensor]	Decoder output dictionary containing predictions and any algorithm-specific outputs.

Source code in src/versatil/models/decoding/algorithm/base.py

@abstractmethod
def forward(
    self,
    network: ActionDecoder,
    features: dict[str, torch.Tensor],
    actions: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor]:
    """Forward pass during training.

    Args:
        network: The action decoder network module
        features: Dictionary of encoded features from the encoding pipeline.
        actions: Optional dictionary of ground truth actions, if architecture requires it, e.g. ACT.

    Returns:
        Decoder output dictionary containing predictions and any algorithm-specific outputs.
    """
    raise NotImplementedError

injected_feature_keys

injected_feature_keys()

Feature keys this algorithm adds to the decoder features itself.

Decoders may declare these in their input specification even though they never come from observations or the encoding pipeline.

Source code in src/versatil/models/decoding/algorithm/base.py

def injected_feature_keys(self) -> set[str]:
    """Feature keys this algorithm adds to the decoder features itself.

    Decoders may declare these in their input specification even though
    they never come from observations or the encoding pipeline.
    """
    return set()

get_auxiliary_output_keys

get_auxiliary_output_keys()

Get keys for auxiliary outputs this algorithm adds to the decoder output.

Override in subclasses that inject additional outputs (e.g. latent variables).

Returns:

Type	Description
set[str]	Set of auxiliary output key strings.

Source code in src/versatil/models/decoding/algorithm/base.py

def get_auxiliary_output_keys(self) -> set[str]:
    """Get keys for auxiliary outputs this algorithm adds to the decoder output.

    Override in subclasses that inject additional outputs (e.g. latent variables).

    Returns:
        Set of auxiliary output key strings.
    """
    return set()

get_targets

get_targets(algorithm_output, ground_truth_actions)

Return the correct regression targets for the loss.

The loss module is algorithm-agnostic: it computes error between predictions and targets. This method lets each algorithm specify what those targets are.

Default returns the ground-truth actions (correct for Behavioral Cloning). Generative algorithms override to return their algorithm-specific targets (velocity field, noise, etc.).

Parameters:

Name	Type	Description	Default
algorithm_output	dict[str, Tensor | dict[str, Tensor]]	Full output dict from forward().	required
ground_truth_actions	dict[str, Tensor]	Ground-truth actions from the data batch.	required

Returns:

Type	Description
dict[str, Tensor]	Per-key target tensors the loss should compare predictions against.

Source code in src/versatil/models/decoding/algorithm/base.py

def get_targets(
    self,
    algorithm_output: dict[str, torch.Tensor | dict[str, torch.Tensor]],
    ground_truth_actions: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Return the correct regression targets for the loss.

    The loss module is algorithm-agnostic: it computes error between
    predictions and targets. This method lets each algorithm specify
    what those targets are.

    Default returns the ground-truth actions (correct for Behavioral
    Cloning). Generative algorithms override to return their
    algorithm-specific targets (velocity field, noise, etc.).

    Args:
        algorithm_output: Full output dict from ``forward()``.
        ground_truth_actions: Ground-truth actions from the data batch.

    Returns:
        Per-key target tensors the loss should compare predictions against.
    """
    return ground_truth_actions

predict abstractmethod

predict(network, features)

Inference/prediction pass.

Parameters:

Name	Type	Description	Default
network	ActionDecoder	The action decoder network module	required
features	dict[str, Tensor]	Dict of encoded features from encoding pipeline	required

Returns:

Type	Description
dict[str, Tensor]	Decoder output dictionary containing predictions and any algorithm-specific outputs.

Source code in src/versatil/models/decoding/algorithm/base.py

@abstractmethod
def predict(
    self,
    network: ActionDecoder,
    features: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Inference/prediction pass.

    Args:
        network: The action decoder network module
        features: Dict of encoded features from encoding pipeline

    Returns:
        Decoder output dictionary containing predictions and any algorithm-specific outputs.
    """
    raise NotImplementedError

resolve_feature_reference

resolve_feature_reference(features)

Return batch size, device, and dtype for tensors created from features.

Feature dicts mix encoder outputs with integer token ids and boolean padding masks, and their ordering follows decoder configuration, so the reference must come from a floating-point feature rather than whichever value happens to be first.

Parameters:

Name	Type	Description	Default
features	dict[str, Tensor]	Encoded features keyed by name.	required

Raises:

Type	Description
ValueError	If the feature dict is empty.

Source code in src/versatil/models/decoding/algorithm/base.py

def resolve_feature_reference(
    features: dict[str, torch.Tensor],
) -> tuple[int, torch.device, torch.dtype]:
    """Return batch size, device, and dtype for tensors created from features.

    Feature dicts mix encoder outputs with integer token ids and boolean
    padding masks, and their ordering follows decoder configuration, so the
    reference must come from a floating-point feature rather than whichever
    value happens to be first.

    Args:
        features: Encoded features keyed by name.

    Raises:
        ValueError: If the feature dict is empty.
    """
    if len(features) == 0:
        raise ValueError("Cannot infer batch size from an empty feature dict.")
    for value in features.values():
        if isinstance(value, torch.Tensor) and value.is_floating_point():
            return value.shape[0], value.device, value.dtype
    first_value = next(iter(features.values()))
    return first_value.shape[0], first_value.device, torch.float32