vlm_encoder

vlm_encoder

VLM encoder with separate vision and language blocks (CLIP-style).

VLMEncoder

VLMEncoder(input_keys, pretrained, frozen, pooling_method, model_name=value, attention_type=value, model_dtype=None, lora_config=None)

Bases: LanguageEncoderMixin, RGBEncoderMixin, Encoder

VLM encoder with separate vision and language outputs.

Initialize the VLM encoder.

Parameters:

Name	Type	Description	Default
input_keys	str | list[str]	Input keys for cameras and tokenized text.	required
pretrained	bool	Whether to load pretrained HuggingFace weights.	required
frozen	bool	Whether to freeze all encoder weights.	required
pooling_method	str	Feature pooling strategy for vision and language outputs.	required
model_name	str	HuggingFace model identifier for the VLM.	value
attention_type	str	Attention implementation (e.g. SDPA, eager).	value
model_dtype	str | None	Precision string from experiment config (e.g. "bf16-mixed").	None
lora_config	LoRAAdaptation | None	Optional LoRA adapter configuration.	None

Source code in src/versatil/models/encoding/encoders/cross_modal/vision_language/vlm_encoder.py

def __init__(
    self,
    input_keys: str | list[str],
    pretrained: bool,
    frozen: bool,
    pooling_method: str,
    model_name: str = ImageTextModelType.CLIP_VITB32.value,
    attention_type: str = AttentionImplementation.SDPA.value,
    model_dtype: str | None = None,
    lora_config: LoRAAdaptation | None = None,
):
    """Initialize the VLM encoder.

    Args:
        input_keys: Input keys for cameras and tokenized text.
        pretrained: Whether to load pretrained HuggingFace weights.
        frozen: Whether to freeze all encoder weights.
        pooling_method: Feature pooling strategy for vision and language outputs.
        model_name: HuggingFace model identifier for the VLM.
        attention_type: Attention implementation (e.g. SDPA, eager).
        model_dtype: Precision string from experiment config (e.g. ``"bf16-mixed"``).
        lora_config: Optional LoRA adapter configuration.
    """
    if isinstance(input_keys, str):
        input_keys = [input_keys]
    all_keys = list(input_keys) + [
        SampleKey.TOKENIZED_OBSERVATIONS.value,
        SampleKey.IS_PAD_OBSERVATION.value,
    ]
    specification = EncoderInput(
        keys=all_keys,
        required_camera_modalities=[CameraModality.RGB],
        required=[SampleKey.TOKENIZED_OBSERVATIONS.value],
        requires_tokenized=True,
    )
    super().__init__(
        input_specification=specification,
        pretrained=pretrained,
        frozen=frozen,
        model_dtype=model_dtype,
    )
    self._setup_camera_keys(input_keys=input_keys)
    self._setup_language_keys(output_modality=EncoderOutputKeys.LANGUAGE.value)
    self.pooling_method = pooling_method
    self.model_name = model_name
    self.lora_config = lora_config
    config = AutoConfig.from_pretrained(model_name)
    if pretrained:
        self.encoder = AutoModel.from_pretrained(
            model_name, attn_implementation=attention_type
        )
    else:
        self.encoder = AutoModel.from_config(
            config, attn_implementation=attention_type
        )
    self.encoder = apply_lora_config(
        model=self.encoder,
        lora_config=lora_config,
        frozen=frozen,
    )
    self.image_processor = AutoImageProcessor.from_pretrained(
        model_name,
        do_rescale=False,
        do_normalize=False,
        do_convert_rgb=False,
        do_resize=False,
    )
    vision_config = self.encoder.vision_model.config
    self.image_size: int | None = getattr(vision_config, "image_size", None)
    self.max_text_length: int = (
        self.encoder.text_model.config.max_position_embeddings
    )
    self.hidden_vision_dim: int = vision_config.hidden_size
    self.hidden_language_dim: int = self.encoder.text_model.config.hidden_size
    vision_has_cls = (
        hasattr(self.encoder.vision_model.embeddings, "class_embedding")
        and self.encoder.vision_model.embeddings.class_embedding is not None
    )
    self.vision_pooling_head = create_token_pooling_head(
        pooling_method=pooling_method,
        input_dimension=self.hidden_vision_dim,
        num_prefix_tokens=1 if vision_has_cls else 0,
    )
    self.language_pooling_head = create_token_pooling_head(
        pooling_method=pooling_method,
        input_dimension=self.hidden_language_dim,
        sequence_length=self.max_text_length,
        num_prefix_tokens=0,
    )
    self.output_vision_dim = self.vision_pooling_head.output_dim
    self.output_language_dim = self.language_pooling_head.output_dim
    self.output_padding_mask_dim = (
        (self.max_text_length,) if pooling_method == PoolingMethod.NONE.value else 1
    )
    if frozen:
        super()._freeze_weights()
    self._apply_model_dtype()

encode

encode(inputs)

Encode images and text through the VLM.

Parameters:

Name	Type	Description	Default
inputs	dict[str, Tensor]	Dict with camera images as (B, C, H, W) per camera key, tokenized text as (B, S), and optional padding mask.	required

Returns:

Type	Description
dict[str, Tensor]	Dict with per-camera RGB features, language features, and padding mask.

Source code in src/versatil/models/encoding/encoders/cross_modal/vision_language/vlm_encoder.py

def encode(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Encode images and text through the VLM.

    Args:
        inputs: Dict with camera images as (B, C, H, W) per camera key,
            tokenized text as (B, S), and optional padding mask.

    Returns:
        Dict with per-camera RGB features, language features, and padding mask.
    """
    text_input_ids, language_mask = self._extract_text_inputs(inputs=inputs)
    batch_size = text_input_ids.shape[0]
    text_input_ids, language_mask = self._pad_text_inputs(
        text_input_ids=text_input_ids,
        language_mask=language_mask,
        max_length=self.max_text_length,
    )
    attention_mask = self._build_attention_mask(
        language_mask=language_mask, text_input_ids=text_input_ids
    )
    text_output = self.encoder.text_model(
        input_ids=text_input_ids,
        attention_mask=attention_mask,
    )
    language_features = self._pool_features(
        text_output,
        modality=EncoderOutputKeys.LANGUAGE.value,
        padding_mask=~attention_mask.bool(),
    )
    token_padding_mask = self._build_output_padding_mask(
        attention_mask=attention_mask,
        pooling_method=self.pooling_method,
        batch_size=batch_size,
        device=language_features.device,
        num_prefix_tokens=self.language_pooling_head.num_prefix_tokens,
    )
    result = {
        EncoderOutputKeys.LANGUAGE.value: language_features,
        self.padding_mask_name: token_padding_mask,
    }
    result.update(self._encode_vision(inputs))
    return result

validate_input_metadata

validate_input_metadata(key, metadata)

Validate that camera keys have camera metadata and token keys are not images.

Parameters:

Name	Type	Description	Default
key	str	Observation key being validated.	required
metadata	BaseMetadata	Metadata from the observation space for this key.	required

Returns:

Type	Description
str | None	Error message if incompatible, None if valid.

Source code in src/versatil/models/encoding/encoders/cross_modal/vision_language/vlm_encoder.py

def validate_input_metadata(self, key: str, metadata: BaseMetadata) -> str | None:
    """Validate that camera keys have camera metadata and token keys are not images.

    Args:
        key: Observation key being validated.
        metadata: Metadata from the observation space for this key.

    Returns:
        Error message if incompatible, None if valid.
    """
    if key in self.camera_keys:
        if not isinstance(metadata, CameraMetadata):
            return f"Expected CameraMetadata for '{key}', got {type(metadata).__name__}"
    else:
        if isinstance(metadata, CameraMetadata):
            return (
                f"VLMEncoder cannot process image data for '{key}'. "
                f"Got CameraMetadata, expected tokenized text input."
            )
    return None

get_output_specification

get_output_specification()

Return the output feature names and dimensions for this encoder.

Returns:

Type	Description
list[FeatureMetadata]	List of FeatureMetadata with per-camera vision features, language features,
list[FeatureMetadata]	and the language padding mask.

Source code in src/versatil/models/encoding/encoders/cross_modal/vision_language/vlm_encoder.py

def get_output_specification(self) -> list[FeatureMetadata]:
    """Return the output feature names and dimensions for this encoder.

    Returns:
        List of FeatureMetadata with per-camera vision features, language features,
        and the language padding mask.
    """
    vision_names = self._get_vision_feature_names()
    vision_dim = (
        (self.output_vision_dim,)
        if isinstance(self.output_vision_dim, int)
        else self.output_vision_dim
    )
    language_dim = (
        (self.output_language_dim,)
        if isinstance(self.output_language_dim, int)
        else self.output_language_dim
    )
    padding_dim = (
        (self.output_padding_mask_dim,)
        if isinstance(self.output_padding_mask_dim, int)
        else self.output_padding_mask_dim
    )
    result = [
        FeatureMetadata(
            key=name,
            feature_type=infer_feature_type(vision_dim),
            dimension=vision_dim,
        )
        for name in vision_names
    ]
    result.append(
        FeatureMetadata(
            key=EncoderOutputKeys.LANGUAGE.value,
            feature_type=infer_feature_type(language_dim),
            dimension=language_dim,
        )
    )
    result.append(
        FeatureMetadata(
            key=self.padding_mask_name,
            feature_type=FeatureType.FLAT.value,
            dimension=padding_dim,
        )
    )
    return result

get_vocab_size

get_vocab_size()

Get the vocabulary size of the text encoder.

Returns:

Type	Description
int	Vocabulary size of the language model component

Source code in src/versatil/models/encoding/encoders/cross_modal/vision_language/vlm_encoder.py

def get_vocab_size(self) -> int:
    """Get the vocabulary size of the text encoder.

    Returns:
        Vocabulary size of the language model component
    """
    return self.encoder.text_model.config.vocab_size