mode_act

mode_act

Mixture of Densities Action Transformer (MODE-ACT) for multi-modal action prediction.

This module implements a Mixture Density Network style transformer decoder that predicts multiple mixture components for each action, enabling multi-modal action distributions.

MixtureOfDensitiesActionTransformer

MixtureOfDensitiesActionTransformer(input_keys, action_space, action_heads, observation_space, observation_horizon, prediction_horizon, device, embedding_dimension=256, number_of_heads=8, number_of_key_value_heads=None, feedforward_dimension=None, number_of_layers=6, activation=value, normalization_type=value, attention_type=value, dropout_rate=0.1, attention_dropout=0.0, positional_encoding_type=value, num_mixture_components=8, gating_hidden_dims=None, gating_activation=value, gating_dropout=0.1, gating_normalization=True, temperature=1.0, learnable_temperature=False, gating_feature_key=None, gmm_init_strategy=value, inference_sampling_mode=value)

Bases: BaseParallelTransformerDecoder

Mixture Density Network Transformer for multi-modal action prediction.

Note

This architecture combines a transformer decoder with K expert action heads to predict mixture density parameters. For the mixture weight computation, either a mode learnable query token or an external feature token are used.

Initialize MODE-ACT decoder.

Parameters:

Name	Type	Description	Default
input_keys	list[str]	Feature keys from encoding pipeline.	required
action_space	ActionSpace	Action space configuration.	required
action_heads	dict[str, ActionHead]	Base action heads to clone K times.	required
observation_space	ObservationSpace	Observation space configuration.	required
observation_horizon	int	Number of observation timesteps.	required
prediction_horizon	int	Number of action timesteps to predict.	required
device	str	Device to place model on.	required
embedding_dimension	int	Transformer embedding dimension.	256
number_of_heads	int	Number of attention heads.	8
number_of_key_value_heads	int | None	Number of key-value heads for GQA.	None
feedforward_dimension	int | None	FFN hidden dimension.	None
number_of_layers	int	Number of decoder layers.	6
activation	str	Activation function.	value
normalization_type	str	Normalization type.	value
attention_type	str	Attention type.	value
dropout_rate	float	Dropout rate.	0.1
attention_dropout	float	Attention dropout rate.	0.0
positional_encoding_type	str | None	Positional encoding type.	value
num_mixture_components	int	Number of mixture components (K).	8
gating_hidden_dims	list[int] | None	Hidden dimensions for gating MLP.	None
gating_activation	str	Activation for gating MLP.	value
gating_dropout	float	Dropout rate in gating MLP.	0.1
gating_normalization	bool	Whether to normalize gating input.	True
temperature	float	Temperature for softmax scaling.	1.0
learnable_temperature	bool	Whether temperature is learnable.	False
gating_feature_key	str | None	If set, use this feature for gating instead of mode embedding.	None
gmm_init_strategy	str	Strategy for initializing GMM component means.	value
inference_sampling_mode	str	How to sample from the mixture at inference. DETERMINISTIC: argmax component, return mean. STOCHASTIC_MEAN: multinomial component, return mean (no noise). STOCHASTIC_SAMPLE: multinomial component, add Gaussian noise.	value

Source code in src/versatil/models/decoding/decoders/factory/mode_act.py

def __init__(
    self,
    input_keys: list[str],
    action_space: ActionSpace,
    action_heads: dict[str, ActionHead],
    observation_space: ObservationSpace,
    observation_horizon: int,
    prediction_horizon: int,
    device: str,
    embedding_dimension: int = 256,
    number_of_heads: int = 8,
    number_of_key_value_heads: int | None = None,
    feedforward_dimension: int | None = None,
    number_of_layers: int = 6,
    activation: str = ActivationFunction.SWIGLU.value,
    normalization_type: str = NormalizationType.RMS_NORM.value,
    attention_type: str = AttentionType.MULTI_HEAD.value,
    dropout_rate: float = 0.1,
    attention_dropout: float = 0.0,
    positional_encoding_type: str | None = PositionalEncodingType.ROPE.value,
    num_mixture_components: int = 8,
    gating_hidden_dims: list[int] | None = None,
    gating_activation: str = ActivationFunction.SILU.value,
    gating_dropout: float = 0.1,
    gating_normalization: bool = True,
    temperature: float = 1.0,
    learnable_temperature: bool = False,
    gating_feature_key: str | None = None,
    gmm_init_strategy: str = GMMInitStrategy.KMEANS_PLUS_PLUS.value,
    inference_sampling_mode: str = MixtureSamplingMode.STOCHASTIC_MEAN.value,
) -> None:
    """Initialize MODE-ACT decoder.

    Args:
        input_keys: Feature keys from encoding pipeline.
        action_space: Action space configuration.
        action_heads: Base action heads to clone K times.
        observation_space: Observation space configuration.
        observation_horizon: Number of observation timesteps.
        prediction_horizon: Number of action timesteps to predict.
        device: Device to place model on.
        embedding_dimension: Transformer embedding dimension.
        number_of_heads: Number of attention heads.
        number_of_key_value_heads: Number of key-value heads for GQA.
        feedforward_dimension: FFN hidden dimension.
        number_of_layers: Number of decoder layers.
        activation: Activation function.
        normalization_type: Normalization type.
        attention_type: Attention type.
        dropout_rate: Dropout rate.
        attention_dropout: Attention dropout rate.
        positional_encoding_type: Positional encoding type.
        num_mixture_components: Number of mixture components (K).
        gating_hidden_dims: Hidden dimensions for gating MLP.
        gating_activation: Activation for gating MLP.
        gating_dropout: Dropout rate in gating MLP.
        gating_normalization: Whether to normalize gating input.
        temperature: Temperature for softmax scaling.
        learnable_temperature: Whether temperature is learnable.
        gating_feature_key: If set, use this feature for gating instead of mode embedding.
        gmm_init_strategy: Strategy for initializing GMM component means.
        inference_sampling_mode: How to sample from the mixture at inference.
            DETERMINISTIC: argmax component, return mean.
            STOCHASTIC_MEAN: multinomial component, return mean (no noise).
            STOCHASTIC_SAMPLE: multinomial component, add Gaussian noise.
    """
    decoder_input = DecoderInput(
        keys=input_keys,
        required_types=[FeatureType.SPATIAL.value],
        requires_actions=False,
    )
    super().__init__(
        decoder_input=decoder_input,
        action_space=action_space,
        action_heads=action_heads,
        observation_space=observation_space,
        prediction_horizon=prediction_horizon,
        observation_horizon=observation_horizon,
        device=device,
        embedding_dimension=embedding_dimension,
    )

    self.number_of_layers = number_of_layers
    self.activation = activation
    self.dropout_rate = dropout_rate
    self.feedforward_dimension = feedforward_dimension
    self.number_of_heads = number_of_heads
    self.number_of_key_value_heads = number_of_key_value_heads
    self.normalization_type = normalization_type
    self.attention_type = attention_type
    self.attention_dropout = attention_dropout
    self.positional_encoding_type = positional_encoding_type
    self.num_mixture_components = num_mixture_components
    self.gating_feature_key = gating_feature_key
    if gmm_init_strategy not in [s.value for s in GMMInitStrategy]:
        raise ValueError(
            f"Unknown gmm_init_strategy: {gmm_init_strategy}. "
            f"Expected one of {[s.value for s in GMMInitStrategy]}"
        )
    self.gmm_init_strategy = gmm_init_strategy
    if inference_sampling_mode not in [m.value for m in MixtureSamplingMode]:
        raise ValueError(
            f"Unknown inference_sampling_mode: {inference_sampling_mode}. "
            f"Expected one of {[m.value for m in MixtureSamplingMode]}"
        )
    self.inference_sampling_mode = inference_sampling_mode
    self.action_keys = list(self.action_heads.keys())
    self._build_transformer_components()
    self._build_mixture_heads()
    self._build_gating_network(
        input_dimension=embedding_dimension,
        hidden_dimensions=gating_hidden_dims,
        activation=gating_activation,
        dropout=gating_dropout,
        normalization=gating_normalization,
    )

    if learnable_temperature:
        self.temperature = nn.Parameter(
            torch.tensor(temperature, dtype=torch.float32), requires_grad=True
        )
    else:
        self.register_buffer(
            "temperature", torch.tensor(temperature, dtype=torch.float32)
        )

    self.to(self.device)

get_auxiliary_output_keys

get_auxiliary_output_keys()

MoDEACT produces routing weights for mixture density prediction.

Source code in src/versatil/models/decoding/decoders/factory/mode_act.py

def get_auxiliary_output_keys(self) -> set[str]:
    """MoDEACT produces routing weights for mixture density prediction."""
    keys = super().get_auxiliary_output_keys()
    keys.add(DecoderOutputKey.ROUTING_WEIGHTS.value)
    return keys

set_normalizer

set_normalizer(normalizer)

Set normalizer and initialize GMM components from output statistics.

Parameters:

Name	Type	Description	Default
normalizer	LinearNormalizer	Normalizer with fitted statistics.	required

Source code in src/versatil/models/decoding/decoders/factory/mode_act.py

def set_normalizer(self, normalizer: LinearNormalizer) -> None:
    """Set normalizer and initialize GMM components from output statistics.

    Args:
        normalizer: Normalizer with fitted statistics.
    """
    super().set_normalizer(normalizer)
    self._initialize_gating_network()
    for action_key in self.action_keys:
        if action_key not in normalizer.params_dict:
            continue
        field_normalizer = normalizer[action_key]
        output_stats = field_normalizer.get_output_stats()
        candidate_sample = field_normalizer.get_output_sample()
        base_head = self.action_heads[action_key]
        if isinstance(base_head, GaussianHead):
            self._initialize_gaussian_mixture(
                action_key=action_key,
                output_stats=output_stats,
                candidate_sample=candidate_sample,
            )

forward

forward(features, actions=None)

Forward pass dispatching based on training vs inference mode.

During training (actions provided): returns raw mixture outputs for loss computation. During inference (actions=None): returns sampled actions from mixture.

Parameters:

Name	Type	Description	Default
features	dict[str, Tensor]	Dictionary of encoded features.	required
actions	dict[str, Tensor] | None	Ground truth actions (used only to distinguish training from inference).	None

Returns:

Type	Description
dict[str, Tensor]	During training: Dict with mixture outputs (mean, logvar, routing_weights).
dict[str, Tensor]	During inference: Dict with sampled actions (B, T, D) for each action key.

Source code in src/versatil/models/decoding/decoders/factory/mode_act.py

def forward(
    self,
    features: dict[str, torch.Tensor],
    actions: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor]:
    """Forward pass dispatching based on training vs inference mode.

    During training (actions provided): returns raw mixture outputs for loss computation.
    During inference (actions=None): returns sampled actions from mixture.

    Args:
        features: Dictionary of encoded features.
        actions: Ground truth actions (used only to distinguish training from inference).

    Returns:
        During training: Dict with mixture outputs (mean, logvar, routing_weights).
        During inference: Dict with sampled actions (B, T, D) for each action key.
    """
    if actions is None:
        return self._sample_from_mixture_outputs(features)
    return self._forward_mixture(features)