flat

flat

Flat RGB encoder producing (B, S, D) token sequences via timm forward_features.

FlatRGBEncoder

FlatRGBEncoder(input_keys, pretrained, frozen, pooling_method=value, backbone=value, image_size=None, intermediate_layer_index=None, model_dtype=None, lora_config=None)

Bases: RGBEncoderMixin, Encoder

RGB encoder for backbones that output flat token sequences.

Initialize flat RGB encoder with timm backbone.

Parameters:

Name	Type	Description	Default
input_keys	str | list[str]	Camera observation keys.	required
pretrained	bool	Whether to load pretrained weights.	required
frozen	bool	Whether to freeze all parameters.	required
pooling_method	str	Feature pooling strategy for patch tokens. Defaults to CLS token selection.	value
backbone	str	timm model name for the backbone.	value
image_size	int | tuple[int, int] | None	Optional image size passed to timm during backbone construction.	None
intermediate_layer_index	int | None	Optional intermediate layer index for feature extraction. Negative values index from the end.	None
model_dtype	str | None	Precision string from experiment config (e.g. "bf16-mixed").	None
lora_config	LoRAAdaptation | None	Optional PEFT LoRA adapter configuration.	None

Source code in src/versatil/models/encoding/encoders/rgb/flat.py

def __init__(
    self,
    input_keys: str | list[str],
    pretrained: bool,
    frozen: bool,
    pooling_method: str = PoolingMethod.DEFAULT.value,
    backbone: str = FlatBackboneType.DINOV2_VITB14.value,
    image_size: int | tuple[int, int] | None = None,
    intermediate_layer_index: int | None = None,
    model_dtype: str | None = None,
    lora_config: LoRAAdaptation | None = None,
) -> None:
    """Initialize flat RGB encoder with timm backbone.

    Args:
        input_keys: Camera observation keys.
        pretrained: Whether to load pretrained weights.
        frozen: Whether to freeze all parameters.
        pooling_method: Feature pooling strategy for patch tokens.
            Defaults to CLS token selection.
        backbone: timm model name for the backbone.
        image_size: Optional image size passed to timm during backbone
            construction.
        intermediate_layer_index: Optional intermediate layer index for
            feature extraction. Negative values index from the end.
        model_dtype: Precision string from experiment config (e.g. ``"bf16-mixed"``).
        lora_config: Optional PEFT LoRA adapter configuration.
    """
    specification = EncoderInput(
        keys=input_keys,
        required_camera_modalities=[CameraModality.RGB],
    )
    super().__init__(
        input_specification=specification,
        pretrained=pretrained,
        frozen=frozen,
        model_dtype=model_dtype,
    )
    valid_backbones = [e.value for e in FlatBackboneType]
    if backbone not in valid_backbones:
        raise ValueError(
            f"Invalid backbone '{backbone}'. Must be one of: {valid_backbones}"
        )

    pooling = PoolingMethod(pooling_method)
    if not pooling.supports_sequential:
        raise ValueError(
            f"Pooling method '{pooling_method}' is not compatible with "
            f"token sequences. Use one of: "
            f"{[p.value for p in PoolingMethod if p.supports_sequential]}"
        )
    self._setup_camera_keys(input_keys=self.input_specification.keys)
    self.pooling_method = pooling_method
    self.backbone_name = backbone
    self.image_size = image_size
    self.intermediate_layer_index = intermediate_layer_index
    self.lora_config = lora_config
    self._build_backbone()
    if (
        pooling_method == PoolingMethod.DEFAULT.value
        and self.backbone.num_prefix_tokens == 0
    ):
        raise ValueError(
            f"Backbone '{backbone}' has no class token, so DEFAULT pooling "
            "would silently return the first patch token. Use AVERAGE, "
            "LEARNED_AGGREGATION, or NONE pooling instead."
        )
    self.feature_dim: int = int(self.backbone.num_features)
    self.token_pooling_head = create_token_pooling_head(
        pooling_method=pooling_method,
        input_dimension=self.feature_dim,
        num_prefix_tokens=self.backbone.num_prefix_tokens,
    )
    self.output_dim = self.token_pooling_head.output_dim
    if frozen:
        super()._freeze_weights()
    self._apply_model_dtype()

encode

encode(inputs)

Encode images into features.

Parameters:

Name	Type	Description	Default
inputs	dict[str, Tensor]	Dict mapping camera keys to image tensors (B*T, C, H, W); forward() flattens the temporal axis into the batch before dispatching here.	required

Returns:

Type	Description
dict[str, Tensor]	Dict with RGB features. Single camera: key is rgb.
dict[str, Tensor]	Multiple cameras: keys are rgb:{camera_key} per camera.

Source code in src/versatil/models/encoding/encoders/rgb/flat.py

def encode(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Encode images into features.

    Args:
        inputs: Dict mapping camera keys to image tensors (B*T, C, H, W);
            ``forward()`` flattens the temporal axis into the batch
            before dispatching here.

    Returns:
        Dict with RGB features. Single camera: key is ``rgb``.
        Multiple cameras: keys are ``rgb:{camera_key}`` per camera.
    """
    return self._encode_vision(inputs)

set_image_size

set_image_size(image_height, image_width)

Rebuild the backbone with the target image size.

Parameters:

Name	Type	Description	Default
image_height	int	Target image height.	required
image_width	int	Target image width.	required

Source code in src/versatil/models/encoding/encoders/rgb/flat.py

def set_image_size(self, image_height: int, image_width: int) -> None:
    """Rebuild the backbone with the target image size.

    Args:
        image_height: Target image height.
        image_width: Target image width.
    """
    self.image_size = (image_height, image_width)
    self._build_backbone()
    self.feature_dim: int = int(self.backbone.num_features)
    self.token_pooling_head = create_token_pooling_head(
        pooling_method=self.pooling_method,
        input_dimension=self.feature_dim,
        num_prefix_tokens=self.backbone.num_prefix_tokens,
    )
    self.output_dim = self.token_pooling_head.output_dim
    if self.frozen:
        self._freeze_weights()
    self._apply_model_dtype()

validate_input_metadata

validate_input_metadata(key, metadata)

Validate that input metadata is camera metadata.

Parameters:

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

Returns:

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

Source code in src/versatil/models/encoding/encoders/rgb/flat.py

def validate_input_metadata(self, key: str, metadata: BaseMetadata) -> str | None:
    """Validate that input metadata is camera metadata.

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

    Returns:
        Error message if incompatible, None if valid.
    """
    if not isinstance(metadata, CameraMetadata):
        return f"Expected CameraMetadata for '{key}', got {type(metadata).__name__}"
    return None

get_explainability_targets

get_explainability_targets()

Return a transformer block for patch-token attribution maps.

Returns:

Type	Description
list[VisionExplanationTarget]	One token-sequence target with NLC layout, prefix-token count, and
list[VisionExplanationTarget]	patch-grid metadata when available. Returns an empty list for
list[VisionExplanationTarget]	backbones that do not expose a ViT blocks sequence.

Note

Standard CLS-token ViTs often read only the CLS token in the final head, so final patch-token outputs can be uninformative. When the backbone exposes at least two blocks, this selects the block before the last block.

Source code in src/versatil/models/encoding/encoders/rgb/flat.py

def get_explainability_targets(self) -> list[VisionExplanationTarget]:
    """Return a transformer block for patch-token attribution maps.

    Returns:
        One token-sequence target with NLC layout, prefix-token count, and
        patch-grid metadata when available. Returns an empty list for
        backbones that do not expose a ViT ``blocks`` sequence.

    Note:
        Standard CLS-token ViTs often read only the CLS token in the final
        head, so final patch-token outputs can be uninformative. When the
        backbone exposes at least two blocks, this selects the block before
        the last block.
    """
    blocks = getattr(self.backbone, "blocks", None)
    if blocks is None or len(blocks) == 0:
        return []
    target_block = blocks[-2] if len(blocks) > 1 else blocks[-1]
    return [
        VisionExplanationTarget(
            layer=target_block,
            target_kind=ExplanationTargetKind.TOKEN_SEQUENCE.value,
            activation_layout=ActivationLayout.NLC.value,
            prefix_token_count=int(getattr(self.backbone, "num_prefix_tokens", 0)),
            patch_grid=self._get_patch_grid(),
        )
    ]

get_output_specification

get_output_specification()

Get structured output specification with feature names and dimensions.

Returns:

Type	Description
list[FeatureMetadata]	List of FeatureMetadata with per-camera feature names and pooled dimensions.

Source code in src/versatil/models/encoding/encoders/rgb/flat.py

def get_output_specification(self) -> list[FeatureMetadata]:
    """Get structured output specification with feature names and dimensions.

    Returns:
        List of FeatureMetadata with per-camera feature names and pooled dimensions.
    """
    feature_names = self._get_vision_feature_names()
    dimension = (
        (self.output_dim,) if isinstance(self.output_dim, int) else self.output_dim
    )
    return [
        FeatureMetadata(
            key=name,
            feature_type=infer_feature_type(dimension),
            dimension=dimension,
        )
        for name in feature_names
    ]