spatial_backbone

spatial_backbone

Shared spatial image encoder producing (B, C, H, W) feature maps via timm.

SpatialBackboneEncoder

SpatialBackboneEncoder(input_keys, backbone=value, pooling_method=value, batch_norm_handling=value, intermediate_layer_index=None, pretrained=False, frozen=False, model_dtype=None, lora_config=None)

Bases: ImageEncoderMixin, Encoder

Shared implementation for spatial image encoders backed by timm.

Supports any timm backbone compatible with features_only=True, regardless of whether the architecture is convolutional (ResNet, EfficientNet, ConvNeXt) or attention-based (Swin, TinyViT). Handles both NCHW and NHWC output layouts transparently. Subclasses pair this with a camera-modality mixin and set _input_channels and _camera_modality.

Initialize the spatial encoder with a timm backbone.

Parameters:

Name	Type	Description	Default
input_keys	str | list[str]	Camera observation keys.	required
backbone	str	timm model name from SpatialBackboneType.	value
pooling_method	str	Feature pooling strategy.	value
batch_norm_handling	str	How to handle batch normalization layers.	value
intermediate_layer_index	int | None	Optional timm intermediate layer index to pool. Negative values index from the end; None uses the last layer.	None
pretrained	bool	Whether to load pretrained weights.	False
frozen	bool	Whether to freeze all parameters.	False
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/spatial_backbone.py

def __init__(
    self,
    input_keys: str | list[str],
    backbone: str = SpatialBackboneType.RESNET18.value,
    pooling_method: str = PoolingMethod.AVERAGE.value,
    batch_norm_handling: str = BatchNormHandling.FROZEN.value,
    intermediate_layer_index: int | None = None,
    pretrained: bool = False,
    frozen: bool = False,
    model_dtype: str | None = None,
    lora_config: LoRAAdaptation | None = None,
) -> None:
    """Initialize the spatial encoder with a timm backbone.

    Args:
        input_keys: Camera observation keys.
        backbone: timm model name from SpatialBackboneType.
        pooling_method: Feature pooling strategy.
        batch_norm_handling: How to handle batch normalization layers.
        intermediate_layer_index: Optional timm intermediate layer index
            to pool. Negative values index from the end; ``None`` uses
            the last layer.
        pretrained: Whether to load pretrained weights.
        frozen: Whether to freeze all parameters.
        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=[self._camera_modality],
    )
    super().__init__(
        input_specification=specification,
        pretrained=pretrained,
        frozen=frozen,
        model_dtype=model_dtype,
    )
    valid_backbones = [e.value for e in SpatialBackboneType]
    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_spatial:
        raise ValueError(
            f"Pooling method '{pooling_method}' is not compatible with "
            f"spatial feature maps. Use one of: "
            f"{[p.value for p in PoolingMethod if p.supports_spatial]}"
        )
    self._setup_camera_keys(input_keys=self.input_specification.keys)
    self.batch_norm_handling = batch_norm_handling
    self.pooling_method = pooling_method
    self.intermediate_layer_index = intermediate_layer_index
    self.backbone_name = backbone
    self.lora_config = lora_config
    self._channels_last = False
    self._build_backbone()
    self.feature_dim = self._get_intermediate_layer_channels()
    self.pooling_head: PoolingHead | None = None
    self.output_dim: int | tuple[int, ...] = self.feature_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 image features. Single camera: key is the modality name.
dict[str, Tensor]	Multiple cameras: keys are {modality}:{camera_key} per camera.

Source code in src/versatil/models/encoding/encoders/spatial_backbone.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 image features. Single camera: key is the modality name.
        Multiple cameras: keys are ``{modality}:{camera_key}`` per camera.
    """
    return self._encode_vision(inputs)

set_image_size

set_image_size(image_height, image_width)

Compute feature map dimensions and create pooling head.

For backbones with strict input size requirements (e.g. Swin), this rebuilds the backbone with the target dimensions.

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/spatial_backbone.py

def set_image_size(self, image_height: int, image_width: int) -> None:
    """Compute feature map dimensions and create pooling head.

    For backbones with strict input size requirements (e.g. Swin), this
    rebuilds the backbone with the target dimensions.

    Args:
        image_height: Target image height.
        image_width: Target image width.
    """
    if self._has_strict_image_size():
        self._build_backbone(img_size=(image_height, image_width))
        self.feature_dim = self._get_intermediate_layer_channels()
        if self.frozen:
            self._freeze_weights()
        # The rebuilt backbone is float32; recast before the mock
        # forward below runs an input through it.
        self._apply_model_dtype()

    mock_input_dtype = self._mock_forward_dtype()
    with torch.no_grad(), self._mock_forward_autocast():
        mock_input = torch.zeros(
            1,
            self._input_channels,
            image_height,
            image_width,
            dtype=mock_input_dtype,
        )
        intermediate_outputs = self.backbone(mock_input)
        layer_index = self._resolve_intermediate_layer_index(
            intermediate_layer_index=self.intermediate_layer_index,
            output_count=len(intermediate_outputs),
        )
        mock_features = intermediate_outputs[layer_index]

    expected_channels = self.feature_dim
    if mock_features.shape[1] == expected_channels:
        self._channels_last = False
        _, _, spatial_height, spatial_width = mock_features.shape
    elif mock_features.shape[-1] == expected_channels:
        self._channels_last = True
        _, spatial_height, spatial_width, _ = mock_features.shape
    else:
        raise RuntimeError(
            f"Backbone '{self.backbone_name}' output shape {mock_features.shape} "
            f"does not match expected channels {expected_channels} in "
            f"either NCHW or NHWC layout."
        )

    self._setup_pooling(spatial_height=spatial_height, spatial_width=spatial_width)
    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/spatial_backbone.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 spatial feature-map targets for visual explanations.

The target layout follows the backbone output layout detected in set_image_size(). CNN backbones usually emit NCHW; Swin/TinyViT style feature backbones may emit NHWC.

Returns:

Type	Description
list[VisionExplanationTarget]	One spatial feature-map target for visual attribution methods.

Raises:

Type	Description
RuntimeError	If no supported backbone target layer can be selected.

Source code in src/versatil/models/encoding/encoders/spatial_backbone.py

def get_explainability_targets(self) -> list[VisionExplanationTarget]:
    """Return spatial feature-map targets for visual explanations.

    The target layout follows the backbone output layout detected in
    ``set_image_size()``. CNN backbones usually emit NCHW; Swin/TinyViT
    style feature backbones may emit NHWC.

    Returns:
        One spatial feature-map target for visual attribution methods.

    Raises:
        RuntimeError: If no supported backbone target layer can be selected.
    """
    layout = (
        ActivationLayout.NHWC.value
        if self._channels_last
        else ActivationLayout.NCHW.value
    )
    return [
        VisionExplanationTarget(
            layer=self._get_backbone_explanation_layer(),
            target_kind=ExplanationTargetKind.SPATIAL_FEATURE_MAP.value,
            activation_layout=layout,
        )
    ]

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/spatial_backbone.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
    ]