base

base

Base classes for encoder input/output specifications.

EncodingMixin

EncodingMixin(input_specification, pretrained=False, frozen=False, device=None, model_dtype=None)

Bases: ModuleAttrMixin, ABC

Base interface for all encoders, conditional and non-conditional.

Initialize base encoder.

Parameters:

Name	Type	Description	Default
input_specification	EncoderInput	Structured input specification for this encoder	required
pretrained	bool	Whether to use pretrained weights	False
frozen	bool	Whether to freeze encoder weights	False
device	str | None	Device to place the encoder on (will use "cuda" if available, else "cpu" if None)	None
model_dtype	str | None	Precision string from experiment config (e.g. "bf16-mixed"). Resolved to torch.dtype via PrecisionType.	None

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

def __init__(
    self,
    input_specification: EncoderInput,
    pretrained: bool = False,
    frozen: bool = False,
    device: str | None = None,
    model_dtype: str | None = None,
) -> None:
    """Initialize base encoder.

    Args:
        input_specification: Structured input specification for this encoder
        pretrained: Whether to use pretrained weights
        frozen: Whether to freeze encoder weights
        device: Device to place the encoder on (will use "cuda" if available, else "cpu" if None)
        model_dtype: Precision string from experiment config (e.g. ``"bf16-mixed"``).
            Resolved to ``torch.dtype`` via ``PrecisionType``.
    """
    super().__init__()
    input_specification.validate()
    self.input_specification = input_specification
    self.pretrained = pretrained
    self.frozen = frozen
    if model_dtype is not None:
        valid_values = [p.value for p in PrecisionType]
        if model_dtype not in valid_values:
            raise ValueError(
                f"Invalid model_dtype '{model_dtype}'. "
                f"Must be one of: {valid_values}"
            )
        self.precision_type: PrecisionType | None = PrecisionType(model_dtype)
        self.model_dtype: torch.dtype | None = self.precision_type.get_model_dtype()
    else:
        self.precision_type = None
        self.model_dtype = None
    if device is None:
        device = "cuda" if _cuda_runtime_available() else "cpu"
    self.device = torch.device(device)

train

train(mode=True)

Set train/eval mode while keeping fully frozen encoders eval-locked.

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

def train(self, mode: bool = True) -> "EncodingMixin":
    """Set train/eval mode while keeping fully frozen encoders eval-locked."""
    super().train(mode)
    parameters = list(self.parameters())
    if (
        mode
        and self.frozen
        and parameters
        and all(not parameter.requires_grad for parameter in parameters)
    ):
        super().train(False)
    return self

get_output_specification abstractmethod

get_output_specification()

Get encoder structured output specification.

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

@abstractmethod
def get_output_specification(self) -> list[FeatureMetadata]:
    """Get encoder structured output specification."""
    raise NotImplementedError

get_vocab_size

get_vocab_size()

Get vocabulary size if applicable, else None.

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

def get_vocab_size(self) -> int | None:
    """Get vocabulary size if applicable, else None."""
    return None

validate_input_metadata

validate_input_metadata(key, metadata)

Check that observation metadata is compatible with this encoder.

Note

Called by the encoding pipeline per input key during setup. Subclasses override to check type, channels, dimensions, etc.

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

def validate_input_metadata(self, key: str, metadata: BaseMetadata) -> str | None:
    """Check that observation metadata is compatible with this encoder.

    Note:
        Called by the encoding pipeline per input key during setup.
        Subclasses override to check type, channels, dimensions, etc.

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

    Returns:
        Error message if incompatible, None if valid.
    """
    return None

set_image_size

set_image_size(image_height, image_width)

Set the target image size for this encoder.

Called by the encoding pipeline after Hydra instantiation, with per-camera image dimensions from the observation space. Image encoders override this to configure their backbone (ViT img_size) or pooling head (CNN spatial dims).

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

def set_image_size(self, image_height: int, image_width: int) -> None:
    """Set the target image size for this encoder.

    Called by the encoding pipeline after Hydra instantiation, with per-camera
    image dimensions from the observation space. Image encoders override this
    to configure their backbone (ViT img_size) or pooling head (CNN spatial dims).

    Args:
        image_height: Target image height.
        image_width: Target image width.
    """
    pass

get_explainability_targets

get_explainability_targets()

Return target layers and layout metadata for visual explanations.

Encoders that do not process images return an empty list. Vision encoders override this to expose the exact module whose activations should be captured, plus enough layout metadata for the explainability package to reshape gradients back to image space.

Returns:

Type	Description
list[VisionExplanationTarget]	Explainability targets supported by this encoder.

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

def get_explainability_targets(self) -> list[VisionExplanationTarget]:
    """Return target layers and layout metadata for visual explanations.

    Encoders that do not process images return an empty list. Vision
    encoders override this to expose the exact module whose activations
    should be captured, plus enough layout metadata for the explainability
    package to reshape gradients back to image space.

    Returns:
        Explainability targets supported by this encoder.
    """
    return []

is_vision_encoder

is_vision_encoder()

Return whether this encoder exposes visual explanation targets.

Returns:

Type	Description
bool	True when at least one explainability target is available.

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

def is_vision_encoder(self) -> bool:
    """Return whether this encoder exposes visual explanation targets.

    Returns:
        True when at least one explainability target is available.
    """
    return bool(self.get_explainability_targets())