Skip to content

action_transformer

action_transformer

Decoder-only transformer for parallel action chunk prediction, with cross-attention to encoding features.

ActionTransformer 

ActionTransformer(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)

Bases: BaseParallelTransformerDecoder

Bidirectional Transformer decoder which decodes action chunks with cross-attention to observation tokens.

Source code in src/versatil/models/decoding/decoders/factory/action_transformer.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
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,
) -> None:
    decoder_input = DecoderInput(
        keys=input_keys,
        required_types=[],
        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._build_transformer_components()
    self.to(self.device)

forward 

forward(features, actions=None)

Forward pass of the transformer decoder architecture.

Parameters:

Name Type Description Default
features dict[str, Tensor]

Dictionary of encoded features from EncodingPipeline Expected to contain flat features (B, embedding_dimension) or (B, T, embedding_dimension)

 required
actions dict[str, Tensor] | None

Not used here

 None

Returns:

Type Description
dict[str, Tensor]

Dictionary containing: - Action head predictions (e.g. position, orientation, gripper)
Source code in src/versatil/models/decoding/decoders/factory/action_transformer.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
def forward(
    self,
    features: dict[str, torch.Tensor],
    actions: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor]:
    """Forward pass of the transformer decoder architecture.

    Args:
        features: Dictionary of encoded features from EncodingPipeline
            Expected to contain flat features (B, embedding_dimension) or (B, T, embedding_dimension)
        actions: Not used here

    Returns:
        Dictionary containing:
            - Action head predictions (e.g. position, orientation, gripper)
    """
    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]
    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,
        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