lact

lact

Latent Action Transformer (LACT) architecture for action decoding.

LACT is an Action Transformer with latent-conditioned decoding through adaptive normalization.

LACT

LACT(input_keys, action_space, action_heads, observation_space, observation_horizon, prediction_horizon, device, latent_dimension, 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, positional_encoding_type=value, dropout_rate=0.1, attention_dropout=0.0, use_gating=True)

Bases: BaseParallelTransformerDecoder

Latent Action Transformer for generative action decoding.

Forward pass steps

Build observation tokens from spatial/flat features Condition learnable queries with latent via AdaLN-Zero Decode actions cross-attention to observation tokens with latent modulation at each layer Apply action heads to produce predictions

Initialize LACT decoder.

Parameters:

Name	Type	Description	Default
input_keys	list[str]	List of feature keys expected from encoder pipeline	required
action_space	ActionSpace	Action space configuration	required
action_heads	dict[str, ActionHead]	Dictionary of action prediction heads	required
observation_space	ObservationSpace	Observation space configuration	required
observation_horizon	int	Number of observation timesteps	required
prediction_horizon	int	Number of actions to predict	required
device	str	Device to run the model on	required
latent_dimension	int	Dimension of latent conditioning vector	required
embedding_dimension	int	Transformer hidden dimension	256
number_of_heads	int	Number of attention heads	8
number_of_key_value_heads	int | None	Number of K/V heads for GQA (None for MHA)	None
feedforward_dimension	int | None	FFN hidden dimension (default: 4 * embedding_dimension)	None
number_of_layers	int	Number of conditional transformer decoder layers	6
activation	str	Activation function name	value
normalization_type	str	Type of adaptive normalization layer	value
attention_type	str	Type of attention mechanism (multi-head, grouped query, etc.)	value
positional_encoding_type	str | None	Type of positional encoding.	value
dropout_rate	float	Dropout probability for residual connections	0.1
attention_dropout	float	Dropout probability for attention weights	0.0
use_gating	bool	Whether to use AdaLN-Zero gating on residual connections	True

Source code in src/versatil/models/decoding/decoders/factory/lact.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,
    latent_dimension: int,
    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,
    positional_encoding_type: str | None = PositionalEncodingType.ROPE.value,
    dropout_rate: float = 0.1,
    attention_dropout: float = 0.0,
    use_gating: bool = True,
) -> None:
    """Initialize LACT decoder.

    Args:
        input_keys: List of feature keys expected from encoder pipeline
        action_space: Action space configuration
        action_heads: Dictionary of action prediction heads
        observation_space: Observation space configuration
        observation_horizon: Number of observation timesteps
        prediction_horizon: Number of actions to predict
        device: Device to run the model on
        latent_dimension: Dimension of latent conditioning vector
        embedding_dimension: Transformer hidden dimension
        number_of_heads: Number of attention heads
        number_of_key_value_heads: Number of K/V heads for GQA (None for MHA)
        feedforward_dimension: FFN hidden dimension (default: 4 * embedding_dimension)
        number_of_layers: Number of conditional transformer decoder layers
        activation: Activation function name
        normalization_type: Type of adaptive normalization layer
        attention_type: Type of attention mechanism (multi-head, grouped query, etc.)
        positional_encoding_type: Type of positional encoding.
        dropout_rate: Dropout probability for residual connections
        attention_dropout: Dropout probability for attention weights
        use_gating: Whether to use AdaLN-Zero gating on residual connections
    """
    decoder_input = DecoderInput(
        keys=input_keys,
        required_types=[],
        requires_actions=False,
        conditioning_key=LatentKey.POSTERIOR_LATENT.value,
        conditioning_required=[LatentKey.POSTERIOR_LATENT.value],
    )
    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.latent_dimension = latent_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.positional_encoding_type = positional_encoding_type
    self.attention_dropout = attention_dropout
    self.use_gating = use_gating
    self._build_components()
    self.to(self.device)

forward

forward(features, actions=None)

Forward pass of LACT architecture.

Parameters:

Name	Type	Description	Default
features	dict[str, Tensor]	Dictionary of encoded features from EncodingPipeline. Must contain LatentKey.POSTERIOR_LATENT.value with shape (B, latent_dimension).	required
actions	dict[str, Tensor] | None	Not used, present for API compatibility.	None

Returns:

Type	Description
dict[str, Tensor]	Dictionary containing action head predictions (e.g. position, orientation, gripper)

Raises:

Type	Description
ValueError	If LatentKey.POSTERIOR_LATENT.value is not present in features

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

def forward(
    self,
    features: dict[str, torch.Tensor],
    actions: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor]:
    """Forward pass of LACT architecture.

    Args:
        features: Dictionary of encoded features from EncodingPipeline.
            Must contain LatentKey.POSTERIOR_LATENT.value with shape (B, latent_dimension).
        actions: Not used, present for API compatibility.

    Returns:
        Dictionary containing action head predictions (e.g. position, orientation, gripper)

    Raises:
        ValueError: If LatentKey.POSTERIOR_LATENT.value is not present in features
    """
    if LatentKey.POSTERIOR_LATENT.value not in features:
        raise ValueError(
            f"LACT requires '{LatentKey.POSTERIOR_LATENT.value}' in features. "
            f"Make sure to use a variational algorithm that provides latent embeddings. "
            f"Available features: {list(features.keys())}"
        )
    latent = features[LatentKey.POSTERIOR_LATENT.value]  # (B, latent_dim)
    obs_tokens, obs_padding_mask = self._build_parallel_observation_tokens(
        input_sequence_builder=self.input_sequence_builder,
        features=features,
        add_positional_encodings=True,
    )  # (B, observation_token_count, embedding_dimension), (B, observation_token_count)
    batch_size = obs_tokens.shape[0]
    latent = self._validate_latent(
        latent=latent,
        batch_size=batch_size,
        observation_device=obs_tokens.device,
    )
    query = self._expand_parallel_query_embedding(
        query_embedding=self.learnable_query,
        batch_size=batch_size,
    )  # (B, prediction_horizon, embedding_dimension)
    action_embeddings = self.action_decoder(
        hidden_states=query,
        condition=latent,
        encoded_features=obs_tokens,
        query_padding_mask=None,
        memory_padding_mask=obs_padding_mask,
    )  # (B, prediction_horizon, embedding_dimension)
    predictions = self._apply_action_heads(action_embeddings)
    return predictions