gpt_action_transformer

gpt_action_transformer

Action GPT Decoder for tokenized action prediction.

It uses a GPT-style autoregressive decoder (only self-attention) to generate sequences of tokenized actions.

GPTActionTransformer

GPTActionTransformer(action_heads, input_keys, action_space, observation_space, observation_horizon, prediction_horizon, device, max_seq_len=512, 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, temperature=1.0, learnable_temperature=False, deterministic=True)

Bases: AutoregressiveDecoderMixin, DiscreteDecoder

Autoregressive decoder for tokenized action prediction.

Uses pure GPT-style transformer with self-attention only (no cross-attention). Observation features are concatenated as prefix tokens, followed by action token embeddings for autoregressive generation. This is similar to Pi0 FAST but adapted to work with any feature encoder.

Initialize GPTActionTransformer decoder.

Parameters:

Name	Type	Description	Default
action_heads	dict[str, ActionHead]	Action heads for different action components (only DecoderOutputKey.ACTION_LOGITS.value used here).	required
input_keys	list[str]	Feature keys expected from encoder pipeline	required
action_space	ActionSpace	Action space configuration	required
observation_space	ObservationSpace	Observation space configuration	required
observation_horizon	int	Number of observation timesteps	required
prediction_horizon	int	Max action horizon for generation	required
device	str	Device to run model on	required
max_seq_len	int	Maximum sequence length for GPT (features + action tokens)	512
embedding_dimension	int	Common embedding dimension to bring input tokens to, also Transformer hidden size	256
number_of_heads	int	Number of query attention heads	8
number_of_key_value_heads	int | None	Number of K/V heads for GQA (None = same as heads = MHA)	None
feedforward_dimension	int | None	FFN hidden dimension (default: 4 * embedding_dimension)	None
number_of_layers	int	Number of transformer layers	6
activation	str	Activation function (swiglu, gelu, relu, silu)	value
normalization_type	str	Normalization type (rmsnorm, layernorm)	value
attention_type	str	Attention type (gqa, mha)	value
dropout_rate	float	Dropout probability	0.1
attention_dropout	float	Attention dropout probability	0.0
positional_encoding_type	str | None	Type of positional encoding (sinusoidal, rope, None)	value
temperature	float	Initial temperature for sampling (not used in greedy decoding)	1.0
learnable_temperature	bool	If True, make temperature a learnable parameter	False
deterministic	bool	If True, use greedy decoding during inference	True

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

def __init__(
    self,
    action_heads: dict[str, ActionHead],
    input_keys: list[str],
    action_space: ActionSpace,
    observation_space: ObservationSpace,
    observation_horizon: int,
    prediction_horizon: int,
    device: str,
    max_seq_len: int = 512,
    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,
    temperature: float = 1.0,
    learnable_temperature: bool = False,
    deterministic: bool = True,
) -> None:
    """Initialize GPTActionTransformer decoder.

    Args:
        action_heads: Action heads for different action components (only DecoderOutputKey.ACTION_LOGITS.value used here).
        input_keys: Feature keys expected from encoder pipeline
        action_space: Action space configuration
        observation_space: Observation space configuration
        observation_horizon: Number of observation timesteps
        prediction_horizon: Max action horizon for generation
        device: Device to run model on
        max_seq_len: Maximum sequence length for GPT (features + action tokens)
        embedding_dimension: Common embedding dimension to bring input tokens to, also Transformer hidden size
        number_of_heads: Number of query attention heads
        number_of_key_value_heads: Number of K/V heads for GQA (None = same as heads = MHA)
        feedforward_dimension: FFN hidden dimension (default: 4 * embedding_dimension)
        number_of_layers: Number of transformer layers
        activation: Activation function (swiglu, gelu, relu, silu)
        normalization_type: Normalization type (rmsnorm, layernorm)
        attention_type: Attention type (gqa, mha)
        dropout_rate: Dropout probability
        attention_dropout: Attention dropout probability
        positional_encoding_type: Type of positional encoding (sinusoidal, rope, None)
        temperature: Initial temperature for sampling (not used in greedy decoding)
        learnable_temperature: If True, make temperature a learnable parameter
        deterministic: If True, use greedy decoding during inference
    """
    self.action_space = action_space
    self.observation_space = observation_space
    self.observation_horizon = observation_horizon
    self.max_seq_len = max_seq_len
    self.embedding_dimension = embedding_dimension
    self.number_of_heads = number_of_heads
    self.number_of_key_value_heads = number_of_key_value_heads or number_of_heads
    self.feedforward_dimension = feedforward_dimension or (4 * embedding_dimension)
    self.number_of_layers = number_of_layers
    self.activation = activation
    self.normalization_type = normalization_type
    self.attention_type = attention_type
    self.dropout_rate = dropout_rate
    self.attention_dropout = attention_dropout
    self.positional_encoding_type = positional_encoding_type
    decoder_input = DecoderInput(
        keys=input_keys,
        requires_actions=True,
    )
    super().__init__(
        decoder_input=decoder_input,
        action_space=action_space,
        action_heads=action_heads,
        observation_space=observation_space,
        observation_horizon=observation_horizon,
        prediction_horizon=prediction_horizon,
        device=device,
        temperature=temperature,
        learnable_temperature=learnable_temperature,
        deterministic=deterministic,
    )
    self._build_transformer_components()
    self._init_action_bos_embedding(
        embedding_dimension=self.embedding_dimension,
        initializer_range=self.gpt_decoder.initializer_range,
    )
    self.to(self.device)

forward

forward(features, actions=None)

Forward pass.

Training: Teacher forcing with ground truth tokens Inference: Autoregressive generation with KV caching

Parameters:

Name	Type	Description	Default
features	dict[str, Tensor]	Encoded features from pipeline	required
actions	dict[str, Tensor] | None	Ground truth tokenized actions (training) or None (inference)	None

Returns:

Type	Description
dict[str, Tensor]	Dict with DecoderOutputKey.ACTION_LOGITS.value (training) or DecoderOutputKey.PREDICTED_ACTION_TOKENS.value (inference)

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

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

    Training: Teacher forcing with ground truth tokens
    Inference: Autoregressive generation with KV caching

    Args:
        features: Encoded features from pipeline
        actions: Ground truth tokenized actions (training) or None (inference)

    Returns:
        Dict with DecoderOutputKey.ACTION_LOGITS.value (training) or DecoderOutputKey.PREDICTED_ACTION_TOKENS.value (inference)
    """
    self._validate_action_tokenizer_is_set()
    feature_tokens, pos_encodings, feature_token_mask = self.input_sequence_builder(
        features
    )  # (B, token_len, embedding_dimension)
    feature_tokens = (
        feature_tokens + pos_encodings
        if pos_encodings is not None
        else feature_tokens
    )
    if actions is not None:
        predictions = self._forward_training(
            feature_tokens=feature_tokens,
            feature_token_mask=feature_token_mask,
            actions=actions,
        )
    else:
        predictions = self._forward_inference(
            feature_tokens=feature_tokens, feature_token_mask=feature_token_mask
        )

    return predictions