interleaved_vlm

interleaved_vlm

Base decoder for interleaved (vision-language-model(VLM), action-expert) architectures, like pi0 and SmolVLA.

InterleavedLayerType

Bases: StrEnum

Layer routing types for interleaved VLM/action-expert decoders.

InterleavedVLMAttentionState dataclass

InterleavedVLMAttentionState(attention_mask, key_padding_mask, vlm_prefix_attention_mask, position_ids, expert_position_ids, expert_action_rope, prefix_token_count, action_token_count)

Masks, positions, and RoPE tensors shared by interleaved VLA decoders.

Attributes:

Name	Type	Description
attention_mask	Tensor	Expert-first attention mask with shape (batch_size, 1, action_token_count + prefix_token_count, action_token_count + prefix_token_count). True marks masked positions.
key_padding_mask	Tensor	Key padding mask in original prefix-first order with shape (batch_size, prefix_token_count + action_token_count).
vlm_prefix_attention_mask	Tensor | None	Additive (huggingface convention) VLM self-attention mask with shape (batch_size, 1, prefix_token_count, prefix_token_count) or None when no prefix positions are masked.
position_ids	Tensor	Prefix-first position ids with shape (batch_size, prefix_token_count + action_token_count).
expert_position_ids	Tensor	Action expert position ids with shape (batch_size, action_token_count).
expert_action_rope	tuple[Tensor, Tensor]	Rotary embeddings for action expert tokens.
prefix_token_count	int	Number of prefix tokens.
action_token_count	int	Number of action expert tokens.

BaseInterleavedVLMDecoder

BaseInterleavedVLMDecoder(input_keys, action_space, action_heads, observation_space, observation_horizon, prediction_horizon, device, vlm_backbone, requires_actions=True)

Bases: VLMBackboneDecoderMixin, ActionDecoder, ABC

Base class for VLA decoders with interleaved VLM and action expert layers.

Initialize decoder input wiring for a configured VLM backbone.

Parameters:

Name	Type	Description	Default
input_keys	list[str]	Encoded feature keys consumed by the action expert, excluding the raw observation keys consumed by vlm_backbone.	required
action_space	ActionSpace	Task action-space metadata.	required
action_heads	dict[str, ActionHead]	Exactly one joint action head mapping expert tokens to the full continuous action vector.	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 action timesteps predicted per sample.	required
device	str	Device used by decoder modules.	required
vlm_backbone	GenerativeVLM	VLM backbone that consumes normalized/tokenized observations and emits prefix embeddings with shape (B, P, D_vlm).	required
requires_actions	bool	Whether the decoder forward pass requires ground-truth actions.	True

Source code in src/versatil/models/decoding/decoders/interleaved_vlm.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,
    vlm_backbone: GenerativeVLM,
    requires_actions: bool = True,
) -> None:
    """Initialize decoder input wiring for a configured VLM backbone.

    Args:
        input_keys: Encoded feature keys consumed by the action expert,
            excluding the raw observation keys consumed by ``vlm_backbone``.
        action_space: Task action-space metadata.
        action_heads: Exactly one joint action head mapping expert tokens
            to the full continuous action vector.
        observation_space: Task observation-space metadata.
        observation_horizon: Number of observation timesteps in each sample.
        prediction_horizon: Number of action timesteps predicted per sample.
        device: Device used by decoder modules.
        vlm_backbone: VLM backbone that consumes normalized/tokenized
            observations and emits prefix embeddings with shape
            ``(B, P, D_vlm)``.
        requires_actions: Whether the decoder forward pass requires
            ground-truth actions.
    """
    if vlm_backbone is None:
        raise ValueError(f"{type(self).__name__} requires a vlm_backbone.")
    decoder_keys = self._vlm_decoder_input_keys(
        input_keys=input_keys,
        vlm_backbone=vlm_backbone,
    )
    decoder_input = DecoderInput(
        keys=decoder_keys,
        requires_actions=requires_actions,
        needs_raw_observations=True,
    )
    super().__init__(
        decoder_input=decoder_input,
        observation_space=observation_space,
        action_space=action_space,
        action_heads=action_heads,
        device=device,
        observation_horizon=observation_horizon,
        prediction_horizon=prediction_horizon,
    )
    self.vlm_backbone: GenerativeVLM = vlm_backbone
    self._encoder_cache_enabled = False
    self._encoder_cache_suppressed = False
    self._prefix_cache: ConditioningCache | None = None
    self._prefix_inputs_cache: tuple[torch.Tensor, torch.Tensor] | None = None

encoder_cache_enabled property

encoder_cache_enabled

Whether VLM prefix caching is currently enabled.

enable_encoder_cache

enable_encoder_cache()

Enable reusable VLM prefix caching for inference.

Source code in src/versatil/models/decoding/decoders/interleaved_vlm.py

def enable_encoder_cache(self) -> None:
    """Enable reusable VLM prefix caching for inference."""
    if self._encoder_cache_suppressed:
        return
    self._encoder_cache_enabled = True
    self._prefix_cache = None
    self._prefix_inputs_cache = None

disable_encoder_cache

disable_encoder_cache()

Disable reusable VLM prefix caching and clear stored cache.

Source code in src/versatil/models/decoding/decoders/interleaved_vlm.py

def disable_encoder_cache(self) -> None:
    """Disable reusable VLM prefix caching and clear stored cache."""
    if self._encoder_cache_suppressed:
        return
    self._encoder_cache_enabled = False
    self._prefix_cache = None
    self._prefix_inputs_cache = None

set_encoder_cache_suppressed

set_encoder_cache_suppressed(suppressed)

Freeze or unfreeze the cache toggles during attribution.

Parameters:

Name	Type	Description	Default
suppressed	bool	While True, enable/disable calls are ignored.	required

Source code in src/versatil/models/decoding/decoders/interleaved_vlm.py

def set_encoder_cache_suppressed(self, suppressed: bool) -> None:
    """Freeze or unfreeze the cache toggles during attribution.

    Args:
        suppressed: While True, enable/disable calls are ignored.
    """
    self._encoder_cache_suppressed = suppressed

set_vlm_backbone

set_vlm_backbone(vlm_backbone)

Attach a VLM backbone and initialize architecture-specific layers.

Parameters:

Name	Type	Description	Default
vlm_backbone	GenerativeVLM	VLM backbone exposing transformer layers, RoPE, text config, hidden size, and a prefix builder returning embeddings with shape (B, P, D_vlm).	required

Source code in src/versatil/models/decoding/decoders/interleaved_vlm.py

def set_vlm_backbone(self, vlm_backbone: GenerativeVLM) -> None:
    """Attach a VLM backbone and initialize architecture-specific layers.

    Args:
        vlm_backbone: VLM backbone exposing transformer layers, RoPE, text
            config, hidden size, and a prefix builder returning embeddings
            with shape ``(B, P, D_vlm)``.
    """
    self.vlm_backbone = vlm_backbone
    self.build_action_expert(
        vlm_layers=vlm_backbone.layers,
        rotary_emb=vlm_backbone.rotary_embedding,
        vlm_hidden_dimension=vlm_backbone.hidden_dimension,
        vlm_text_config=vlm_backbone.text_config,
    )

build_action_expert abstractmethod

build_action_expert(vlm_layers, rotary_emb, vlm_hidden_dimension, vlm_text_config)

Build decoder-specific action expert layers from VLM internals.

Note

This is called by BaseInterleavedVLMDecoder.set_vlm_backbone() after the full VLM backbone is attached.

Parameters:

Name	Type	Description	Default
vlm_layers	ModuleList	Transformer layers copied or referenced by the decoder.	required
rotary_emb	Module	Rotary embedding module used by the VLM layers.	required
vlm_hidden_dimension	int	Hidden dimension of VLM prefix tokens.	required
vlm_text_config	PretrainedConfig	HuggingFace text config for the VLM language tower.	required

Source code in src/versatil/models/decoding/decoders/interleaved_vlm.py

@abc.abstractmethod
def build_action_expert(
    self,
    vlm_layers: torch.nn.ModuleList,
    rotary_emb: torch.nn.Module,
    vlm_hidden_dimension: int,
    vlm_text_config: PretrainedConfig,
) -> None:
    """Build decoder-specific action expert layers from VLM internals.

    Note:
        This is called by ``BaseInterleavedVLMDecoder.set_vlm_backbone()``
        after the full VLM backbone is attached.

    Args:
        vlm_layers: Transformer layers copied or referenced by the decoder.
        rotary_emb: Rotary embedding module used by the VLM layers.
        vlm_hidden_dimension: Hidden dimension of VLM prefix tokens.
        vlm_text_config: HuggingFace text config for the VLM language tower.
    """
    raise NotImplementedError