base

base

Base contracts for action decoders.

DecoderInput dataclass

DecoderInput(keys, required_types=list(), raises_for_types=list(), requires_actions=False, needs_raw_observations=False, conditioning_key=None, conditioning_required=list(), conditioning_one_of_groups=list())

Structured input specification for decoder architectures.

__post_init__

__post_init__()

Post-initialization to ensure feature keys are consistent.

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

def __post_init__(self) -> None:
    """Post-initialization to ensure feature keys are consistent."""
    if self.conditioning_key:
        conditioning_set = {self.conditioning_key}
        missing_conditioning = set(self.conditioning_required) - conditioning_set
        if missing_conditioning:
            raise ValueError(
                f"Missing required conditioning for decoder input: {missing_conditioning}"
            )
        for group in self.conditioning_one_of_groups:
            matches = conditioning_set.intersection(group)
            if len(matches) != 1:
                raise ValueError(
                    f"Exactly one from {group} required for decoder input conditioning"
                )

validate_feature_types

validate_feature_types(available_features)

Validate that required feature types are available at instantiation time.

Parameters:

Name	Type	Description	Default
available_features	dict[str, FeatureMetadata]	Dict mapping feature names to their metadata.	required

Raises:

Type	Description
ValueError	If a required type is missing or a rejected type is present.

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

def validate_feature_types(
    self,
    available_features: dict[str, FeatureMetadata],
) -> None:
    """Validate that required feature types are available at instantiation time.

    Args:
        available_features: Dict mapping feature names to their metadata.

    Raises:
        ValueError: If a required type is missing or a rejected type is present.
    """
    for expected_type in self.required_types:
        matched = any(
            available_features[key].feature_type == expected_type
            for key in self.keys
            if key in available_features
        )
        if not matched:
            raise ValueError(
                f"Decoder requires at least one feature of type '{expected_type}', "
                f"but none found among: {self.keys}. "
                f"Available: {available_features}"
            )
    for rejected_type in self.raises_for_types:
        for key in self.keys:
            if (
                key in available_features
                and available_features[key].feature_type == rejected_type
            ):
                raise ValueError(
                    f"Decoder cannot accept {rejected_type} features, "
                    f"but '{key}' is {rejected_type}."
                )

ActionDecoder

ActionDecoder(decoder_input, observation_space, action_space, action_heads, device, observation_horizon, prediction_horizon)

Bases: ModuleAttrMixin, ABC

Abstract base class for neural network action decoders.

Initialize common action decoder state and action heads.

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

def __init__(
    self,
    decoder_input: DecoderInput,
    observation_space: ObservationSpace,
    action_space: ActionSpace,
    action_heads: dict[str, BaseActionHead],
    device: str,
    observation_horizon: int,
    prediction_horizon: int,
) -> None:
    """Initialize common action decoder state and action heads."""
    super().__init__()
    self.decoder_input = decoder_input
    resolved_heads = resolve_dict_keys(action_heads)
    self.action_heads = nn.ModuleDict(resolved_heads)
    self.observation_space = observation_space
    self.action_space = action_space
    self.observation_horizon = observation_horizon
    self.prediction_horizon = prediction_horizon
    self.device = torch.device(device)
    self._set_action_head_dimensions()
    self.validate_action_heads()
    self.normalizer: LinearNormalizer = LinearNormalizer()
    self.tokenizer: ActionTokenizer | None = None

encoder_cache_enabled property

encoder_cache_enabled

Whether an encoder cache is currently enabled.

has_history property

has_history

Whether the architecture processes temporal sequences.

action_dim property

action_dim

Get the total action dimension.

use_gripper_actions property

use_gripper_actions

Whether the architecture uses gripper actions.

gripper_dim property

gripper_dim

Get the gripper dimension if used.

use_orientation_actions property

use_orientation_actions

Whether the architecture uses orientation actions.

orientation_dim property

orientation_dim

Get the orientation dimension if used.

position_dim property

position_dim

Get the position dimension if used.

set_tokenizer

set_tokenizer(tokenizer=None)

Set tokenizer for decoders trained on tokenized actions.

This method is called by Policy.set_tokenizer() to pass the tokenizer to the decoder. Continuous decoders ignore it.

Parameters:

Name	Type	Description	Default
tokenizer	Tokenizer | None	Tokenizer instance from data pipeline (can be None)	None

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

def set_tokenizer(self, tokenizer: Tokenizer | None = None) -> None:
    """Set tokenizer for decoders trained on tokenized actions.

    This method is called by Policy.set_tokenizer() to pass the tokenizer
    to the decoder. Continuous decoders ignore it.

    Args:
        tokenizer: Tokenizer instance from data pipeline (can be None)
    """
    if not self.requires_tokenized_actions:
        self.tokenizer = None
        return
    if tokenizer is None or tokenizer.action_tokenizer is None:
        raise ValueError(
            "Tokenizer must be provided for tokenized action decoders."
        )
    self.tokenizer = tokenizer.action_tokenizer

set_normalizer

set_normalizer(normalizer)

Set normalizer for data-dependent initialization.

Parameters:

Name	Type	Description	Default
normalizer	LinearNormalizer	LinearNormalizer instance with loaded state.	required

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

def set_normalizer(self, normalizer: "LinearNormalizer") -> None:
    """Set normalizer for data-dependent initialization.

    Args:
        normalizer: LinearNormalizer instance with loaded state.
    """
    self.normalizer = normalizer

enable_encoder_cache

enable_encoder_cache()

No-op. Override in decoders that support encoder caching.

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

def enable_encoder_cache(self) -> None:
    """No-op. Override in decoders that support encoder caching."""

disable_encoder_cache

disable_encoder_cache()

No-op. Override in decoders that support encoder caching.

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

def disable_encoder_cache(self) -> None:
    """No-op. Override in decoders that support encoder caching."""

set_encoder_cache_suppressed

set_encoder_cache_suppressed(suppressed)

No-op. Override in decoders that support encoder caching.

Parameters:

Name	Type	Description	Default
suppressed	bool	While True, cache toggles must not change cache state; attribution uses this to re-run forwards without caching.	required

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

def set_encoder_cache_suppressed(self, suppressed: bool) -> None:
    """No-op. Override in decoders that support encoder caching.

    Args:
        suppressed: While True, cache toggles must not change cache state;
            attribution uses this to re-run forwards without caching.
    """

forward abstractmethod

forward(features, actions=None)

Forward pass of the Neural Network architecture for action decoding.

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

@abstractmethod
def forward(
    self,
    features: dict[str, torch.Tensor],
    actions: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor]:
    """Forward pass of the Neural Network architecture for action decoding."""
    raise NotImplementedError(
        "Subclasses of ActionDecoder must implement the forward pass."
    )

get_auxiliary_output_keys

get_auxiliary_output_keys()

Get keys for auxiliary outputs this decoder produces beyond action predictions.

The base implementation includes tokenized action keys and MoE routing weights when applicable. Subclasses override to add decoder-specific auxiliary keys.

Returns:

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

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

def get_auxiliary_output_keys(self) -> set[str]:
    """Get keys for auxiliary outputs this decoder produces beyond action predictions.

    The base implementation includes tokenized action keys and MoE routing weights
    when applicable. Subclasses override to add decoder-specific auxiliary keys.

    Returns:
        Set of auxiliary output key strings.
    """
    auxiliary_keys: set[str] = set()
    if self.requires_tokenized_actions:
        auxiliary_keys.add(SampleKey.TOKENIZED_ACTIONS.value)
    if any(isinstance(head, MoEHead) for head in self.action_heads.values()):
        auxiliary_keys.add(DecoderOutputKey.ROUTING_WEIGHTS.value)
    return auxiliary_keys

get_loss_output_keys

get_loss_output_keys()

Return decoder output keys that can be supervised as actions.

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

def get_loss_output_keys(self) -> set[str]:
    """Return decoder output keys that can be supervised as actions."""
    match self.action_head_layout:
        case ActionHeadLayout.COMPONENT | ActionHeadLayout.VOCABULARY:
            return set(self.action_heads.keys())
        case ActionHeadLayout.JOINT:
            return set(self.action_space.predicted_action_keys)
        case ActionHeadLayout.NONE:
            if self.requires_tokenized_actions:
                return set()
            return set(self.action_space.predicted_action_keys)

get_prediction_output_keys

get_prediction_output_keys()

Return predicted action keys in action-space metadata order.

The order is the canonical action-key ordering: exported policies and compressed-artifact metadata index output tensors by it.

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

def get_prediction_output_keys(self) -> list[str]:
    """Return predicted action keys in action-space metadata order.

    The order is the canonical action-key ordering: exported policies and
    compressed-artifact metadata index output tensors by it.
    """
    return list(self.action_space.predicted_action_keys)

validate_action_heads

validate_action_heads()

Validate that configured heads match this decoder's head layout.

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

def validate_action_heads(self) -> None:
    """Validate that configured heads match this decoder's head layout."""
    match self.action_head_layout:
        case ActionHeadLayout.NONE:
            self._validate_no_action_heads()
        case ActionHeadLayout.COMPONENT:
            self._validate_component_action_heads()
        case ActionHeadLayout.JOINT:
            self._validate_joint_action_head()
        case ActionHeadLayout.VOCABULARY:
            self._validate_vocabulary_action_head()