autoregressive_vla

autoregressive_vla

Autoregressive VLA decoder for discrete action-token generation.

This module is the shared implementation behind the OpenVLA and pi0-FAST Hydra presets. It runs a generative VLM on raw image/text observations to build the conditioning prefix, then trains or samples discrete action tokens autoregressively through the VLM language model.

AutoregressiveVLADecoder

AutoregressiveVLADecoder(action_heads, input_keys, action_space, observation_space, observation_horizon, prediction_horizon, device, vlm_backbone, max_seq_len=512, temperature=1.0, learnable_temperature=False, deterministic=True, causal_prefix=False)

Bases: AutoregressiveDecoderMixin, LLMPrefixSuffixAttentionMixin, VLMBackboneDecoderMixin, DiscreteDecoder

Predict autoregressive action tokens from a VLM observation prefix.

Initialize a VLM-backed causal action-token decoder.

Parameters:

Name	Type	Description	Default
action_heads	dict[str, ActionHead]	Must be empty. This decoder predicts action tokens with the VLM language vocabulary head.	required
input_keys	list[str]	Must be empty. Raw observation keys are declared by vlm_backbone.input_specification.	required
action_space	ActionSpace	Task action-space metadata.	required
observation_space	ObservationSpace	Task observation-space metadata.	required
observation_horizon	int	Number of observation timesteps in each sample.	required
prediction_horizon	int	Number of future action timesteps represented by generated action tokens.	required
device	str	Device used by decoder modules and generated tensors.	required
vlm_backbone	GenerativeVLM	Generative VLM that builds image-language prefix embeddings and exposes the causal language model vocabulary.	required
max_seq_len	int	Maximum prefix plus generated action-token length.	512
temperature	float	Softmax temperature for stochastic inference.	1.0
learnable_temperature	bool	Whether temperature is optimized as a model parameter.	False
deterministic	bool	Whether inference uses greedy token selection.	True
causal_prefix	bool	Whether to use a standard causal padding mask (OpenVLA) for the whole sequence instead of bidirectional prefix attention (Pi0-FAST).	False

Source code in src/versatil/models/decoding/decoders/factory/autoregressive_vla.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,
    vlm_backbone: GenerativeVLM,
    max_seq_len: int = 512,
    temperature: float = 1.0,
    learnable_temperature: bool = False,
    deterministic: bool = True,
    causal_prefix: bool = False,
) -> None:
    """Initialize a VLM-backed causal action-token decoder.

    Args:
        action_heads: Must be empty. This decoder predicts action tokens
            with the VLM language vocabulary head.
        input_keys: Must be empty. Raw observation keys are declared by
            ``vlm_backbone.input_specification``.
        action_space: Task action-space metadata.
        observation_space: Task observation-space metadata.
        observation_horizon: Number of observation timesteps in each sample.
        prediction_horizon: Number of future action timesteps represented
            by generated action tokens.
        device: Device used by decoder modules and generated tensors.
        vlm_backbone: Generative VLM that builds image-language prefix
            embeddings and exposes the causal language model vocabulary.
        max_seq_len: Maximum prefix plus generated action-token length.
        temperature: Softmax temperature for stochastic inference.
        learnable_temperature: Whether ``temperature`` is optimized as a
            model parameter.
        deterministic: Whether inference uses greedy token selection.
        causal_prefix: Whether to use a standard causal padding mask (OpenVLA) for
            the whole sequence instead of bidirectional prefix attention (Pi0-FAST).
    """
    if action_heads:
        raise ValueError(
            "AutoregressiveVLADecoder predicts action tokens with the VLM language "
            "vocabulary head, so action_heads must be empty."
        )
    self._validate_no_extra_input_keys(
        decoder_name=type(self).__name__,
        input_keys=input_keys,
    )
    self.max_seq_len = max_seq_len
    self.causal_prefix = causal_prefix
    self.eos_token_id: int | None = None
    self.valid_generation_token_ids: torch.Tensor | None = None

    DiscreteDecoder.__init__(
        self,
        decoder_input=DecoderInput(
            keys=self._vlm_decoder_input_keys(
                input_keys=[],
                vlm_backbone=vlm_backbone,
            ),
            requires_actions=True,
            needs_raw_observations=True,
        ),
        action_space=action_space,
        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.vlm_backbone = vlm_backbone
    self.language_hidden_dimension = int(vlm_backbone.hidden_dimension)
    self.to(self.device)

set_tokenizer

set_tokenizer(tokenizer=None)

Set a language-vocabulary action tokenizer for autoregressive decoding.

Parameters:

Name	Type	Description	Default
tokenizer	Tokenizer | None	Tokenizer with action IDs mapped into the VLM language vocabulary.	None

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

def set_tokenizer(self, tokenizer: Tokenizer | None = None) -> None:
    """Set a language-vocabulary action tokenizer for autoregressive decoding.

    Args:
        tokenizer: Tokenizer with action IDs mapped into the VLM language
            vocabulary.
    """
    if tokenizer is None or tokenizer.action_tokenizer is None:
        raise ValueError(
            "AutoregressiveVLADecoder requires an action tokenizer with "
            "language-vocabulary token IDs."
        )
    action_tokenizer = tokenizer.action_tokenizer
    mapping_type = action_tokenizer.token_id_mapping.state_dict()["type"]
    if mapping_type != ActionTokenIdMappingType.LANGUAGE_VOCABULARY.value:
        raise ValueError(
            "AutoregressiveVLADecoder requires action_tokenizer.token_id_mapping.type="
            f"{ActionTokenIdMappingType.LANGUAGE_VOCABULARY.value}."
        )
    tokenizer_vocab_size = int(action_tokenizer.vocab_size)
    eos_token_id = int(action_tokenizer.eos_token_id)
    if eos_token_id < 0 or eos_token_id >= tokenizer_vocab_size:
        raise ValueError(
            "AutoregressiveVLADecoder received an action tokenizer with eos_token_id "
            "outside the model vocabulary: "
            f"eos_token_id={eos_token_id}, vocab_size={tokenizer_vocab_size}."
        )
    language_vocab_size = self._resize_vlm_to_action_vocabulary(
        tokenizer_vocab_size=tokenizer_vocab_size
    )
    self.tokenizer = action_tokenizer
    self.vocab_size = language_vocab_size
    self.eos_token_id = eos_token_id
    self.valid_generation_token_ids = self._build_valid_generation_token_ids(
        action_tokenizer=action_tokenizer,
        tokenizer_vocab_size=tokenizer_vocab_size,
        eos_token_id=eos_token_id,
    )

get_auxiliary_output_keys

get_auxiliary_output_keys()

Return token outputs produced without action heads.

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

def get_auxiliary_output_keys(self) -> set[str]:
    """Return token outputs produced without action heads."""
    return super().get_auxiliary_output_keys() | {
        DecoderOutputKey.ACTION_LOGITS.value,
        DecoderOutputKey.PREDICTED_ACTION_TOKENS.value,
    }

forward

forward(features, actions=None)

Run VLM-conditioned action-token prediction.

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

def forward(
    self,
    features: dict[str, torch.Tensor],
    actions: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor]:
    """Run VLM-conditioned action-token prediction."""
    self._validate_action_tokenizer_is_set()
    prefix_tokens, prefix_token_mask = self._build_projected_prefix(
        features=features
    )
    if actions is not None:
        return self._forward_action_token_training(
            actions=actions,
            prefix_tokens=prefix_tokens,
            prefix_token_mask=prefix_token_mask,
        )
    else:
        return self._forward_action_token_inference(
            prefix_tokens=prefix_tokens,
            prefix_token_mask=prefix_token_mask,
        )