openvla_oft

openvla_oft

OpenVLA-OFT-style continuous action-chunk decoder backed by a VLM.

OpenVLAOFTDecoder

OpenVLAOFTDecoder(action_heads, input_keys, action_space, observation_space, observation_horizon, prediction_horizon, device, vlm_backbone, slots_per_action_dimension=True, causal_action_slots=True, min_period=0.004, max_period=4.0)

Bases: LLMPrefixSuffixAttentionMixin, VLMBackboneDecoderMixin, ActionDecoder

Decode continuous action chunks from VLM prefix plus action slots.

Initialize a VLM-backed continuous action-chunk decoder.

Parameters:

Name	Type	Description	Default
action_heads	dict[str, ActionHead]	Exactly one joint action head that maps per-timestep decoder embeddings to the full continuous action vector. With slots_per_action_dimension=True, its input dimension must be action_dim * language_hidden_dimension. Otherwise, its input dimension must be language_hidden_dimension.	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 to predict.	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 language tower.	required
slots_per_action_dimension	bool	When True, each action scalar owns one VLM hidden-state slot before the joint action projection. When False, each timestep owns one slot.	True
causal_action_slots	bool	Whether action slots use causal self-attention.	True
min_period	float	Minimum period for sinusoidal timestep embeddings used by denoising algorithms.	0.004
max_period	float	Maximum period for sinusoidal timestep embeddings used by denoising algorithms.	4.0

Source code in src/versatil/models/decoding/decoders/factory/openvla_oft.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,
    slots_per_action_dimension: bool = True,
    causal_action_slots: bool = True,
    min_period: float = 4e-3,
    max_period: float = 4.0,
) -> None:
    """Initialize a VLM-backed continuous action-chunk decoder.

    Args:
        action_heads: Exactly one joint action head that maps per-timestep
            decoder embeddings to the full continuous action vector. With
            ``slots_per_action_dimension=True``, its input dimension must be
            ``action_dim * language_hidden_dimension``. Otherwise, its
            input dimension must be ``language_hidden_dimension``.
        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 to predict.
        device: Device used by decoder modules and generated tensors.
        vlm_backbone: Generative VLM that builds image-language prefix
            embeddings and exposes the language tower.
        slots_per_action_dimension: When ``True``, each action scalar owns
            one VLM hidden-state slot before the joint action projection.
            When ``False``, each timestep owns one slot.
        causal_action_slots: Whether action slots use causal self-attention.
        min_period: Minimum period for sinusoidal timestep embeddings used
            by denoising algorithms.
        max_period: Maximum period for sinusoidal timestep embeddings used
            by denoising algorithms.
    """
    self._validate_no_extra_input_keys(
        decoder_name=type(self).__name__,
        input_keys=input_keys,
    )
    self.slots_per_action_dimension = slots_per_action_dimension
    self.causal_action_slots = causal_action_slots
    self.min_period = min_period
    self.max_period = max_period
    self.causal_prefix = False
    self.language_hidden_dimension = int(vlm_backbone.hidden_dimension)
    action_projection_input_dimension = (
        self._resolve_action_projection_input_dimension(
            action_space=action_space,
            language_hidden_dimension=self.language_hidden_dimension,
            slots_per_action_dimension=slots_per_action_dimension,
        )
    )

    ActionDecoder.__init__(
        self,
        decoder_input=DecoderInput(
            keys=self._vlm_decoder_input_keys(
                input_keys=input_keys,
                vlm_backbone=vlm_backbone,
            ),
            requires_actions=False,
            needs_raw_observations=True,
        ),
        action_space=action_space,
        action_heads=action_heads,
        observation_space=observation_space,
        observation_horizon=observation_horizon,
        prediction_horizon=prediction_horizon,
        device=device,
    )
    self.vlm_backbone = vlm_backbone
    self.action_projection_input_dimension = action_projection_input_dimension
    self._validate_action_head_input_dimensions(
        expected_input_dimension=action_projection_input_dimension
    )
    self._build_action_slot_components()
    self._validate_context_capacity(
        vlm_backbone=vlm_backbone,
        includes_denoising_timestep=False,
    )
    self.to(self.device)

forward

forward(features, actions=None)

Predict a continuous action chunk from a VLM prefix.

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

def forward(
    self,
    features: dict[str, torch.Tensor],
    actions: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor]:
    """Predict a continuous action chunk from a VLM prefix."""
    if AlgorithmContextKey.TIMESTEP.value in features:
        self._validate_context_capacity(
            vlm_backbone=self.vlm_backbone,
            includes_denoising_timestep=True,
        )
        if actions is None:
            raise ValueError(
                "OpenVLAOFTDecoder with denoising algorithm requires "
                "ground truth actions during training."
            )
        batch_size, action_device = validate_action_tensors_against_dimensions(
            actions=actions,
            action_dimensions=self._predicted_action_dimensions(),
            prediction_horizon=self.prediction_horizon,
            decoder_name=type(self).__name__,
        )
        timestep = extract_timestep_conditioning(
            features=features,
            batch_size=batch_size,
            action_device=action_device,
        )
        prefix_tokens, prefix_mask = self._build_denoising_prefix(
            features=features,
            timestep=timestep,
        )
        action_slots = self._build_denoising_action_slots(
            actions=actions,
            dtype=prefix_tokens.dtype,
            device=prefix_tokens.device,
        )  # (B, action_slot_count, D)
    else:
        prefix_tokens, prefix_mask = self._build_prefix(features=features)
        action_slots = self._build_action_slots(
            batch_size=prefix_tokens.shape[0],
            dtype=prefix_tokens.dtype,
            device=prefix_tokens.device,
        )  # (B, action_slot_count, D)
    full_token_sequence, attention_mask = self._build_prefix_suffix_inputs(
        prefix_tokens=prefix_tokens,
        suffix_tokens=action_slots,
        prefix_mask=prefix_mask,
        causal_suffix=self.causal_action_slots,
    )
    sequence_output = self._run_language_model(
        tokens=full_token_sequence,
        attention_mask=attention_mask,
    )  # (B, P+action_slot_count, D)
    action_slot_states = sequence_output[
        :, -self.action_slot_count :, :
    ]  # (B, action_slot_count, D)
    action_embeddings = self._reshape_action_slot_states(action_slot_states)
    return self._project_action_output(action_embeddings)