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.

Note

Initialization always runs in chunk space; the strategy only selects how the trajectory centers are spread (k-means++ over the candidate chunks or uniform envelope interpolation). Candidates are the demonstrated chunk subsamples carried by the normalizer, or constant trajectories sampled uniformly from the normalized action range when no subsample is available.

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.

    Note:
        Initialization always runs in chunk space; the strategy only
        selects how the trajectory centers are spread (k-means++ over the
        candidate chunks or uniform envelope interpolation). Candidates
        are the demonstrated chunk subsamples carried by the normalizer,
        or constant trajectories sampled uniformly from the normalized
        action range when no subsample is available.

    Args:
        normalizer: Normalizer with fitted statistics.
    """
    super().set_normalizer(normalizer)
    self._initialize_gating_network()
    gaussian_keys = [
        action_key
        for action_key in self.action_keys
        if action_key in normalizer.params_dict
        and isinstance(self.action_heads[action_key], GaussianHead)
    ]
    if not gaussian_keys:
        return
    chunk_samples = self._collect_chunk_samples(
        normalizer=normalizer, action_keys=gaussian_keys
    )
    if chunk_samples is None:
        chunk_samples = self._constant_chunk_candidates(
            normalizer=normalizer, action_keys=gaussian_keys
        )
    self._initialize_gaussian_mixture(
        normalizer=normalizer,
        action_keys=gaussian_keys,
        chunk_samples=chunk_samples,
    )

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)