Skip to content

parallel_transformer

parallel_transformer

Shared helpers for parallel transformer action decoders.

BaseParallelTransformerDecoder 

BaseParallelTransformerDecoder(*, decoder_input, action_space, action_heads, observation_space, prediction_horizon, observation_horizon, device, embedding_dimension)

Bases: ActionDecoder

Base class for transformer decoders that predict action chunks.

Initialize shared parallel transformer decoder state.

Source code in src/versatil/models/decoding/decoders/parallel_transformer.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def __init__(
    self,
    *,
    decoder_input: DecoderInput,
    action_space: ActionSpace,
    action_heads: dict[str, BaseActionHead],
    observation_space: ObservationSpace,
    prediction_horizon: int,
    observation_horizon: int,
    device: str,
    embedding_dimension: int,
) -> None:
    """Initialize shared parallel transformer decoder state."""
    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,
    )
    self.embedding_dimension = embedding_dimension