feature_projection

feature_projection

Feature projection torch.nn.Module.

Across the module, the following abbreviations are used: - B: Batch size - T: Temporal length (if applicable) - C: Channels/ Original Feature dimension - H: Height (for spatial features) - W: Width (for spatial features) - Emb: Embedding dimension

FeatureProjection

FeatureProjection(embedding_dimension)

Bases: ModuleAttrMixin

Projects features to a common embedding dimension.

It supports algorithm-context features (B, C), temporal vectors (B, T, C), token sequences (B, T, S, C), and spatial features (B, T, C, H, W). Only 5D tensors are treated as spatial; every other rank is projected with a linear layer over the final dimension.

This module uses lazy initialization - projection layers are created on first forward pass. To support loading checkpoints, it overrides _load_from_state_dict to create layers dynamically from the state dict.

Example

feature_projection = FeatureProjection(embedding_dimension=256)
>>> flat_features = {
...     "language": torch.randn(8, 64),
...     "proprio": torch.randn(8, 128),
... }
>>> projected = feature_projection(flat_features)
>>> projected["language"].shape  # (8, 256)
>>> projected["proprio"].shape   # (8, 256)

Initialize feature projection module.

Parameters:

Name	Type	Description	Default
embedding_dimension	int	Target embedding dimension for all features	required

Source code in src/versatil/models/layers/feature_projection.py

def __init__(
    self,
    embedding_dimension: int,
):
    """Initialize feature projection module.

    Args:
        embedding_dimension: Target embedding dimension for all features
    """
    super().__init__()
    self.embedding_dimension = embedding_dimension
    # These two dicts are doing exactly the same mathematical operation. But using linear projections
    # for spatial features would require transposing the tensors,  so we use conv2d instead.
    self.linear_projections = nn.ModuleDict()
    self.spatial_projections = nn.ModuleDict()

forward

forward(features)

Project features to common embedding dimension.

The caller controls which features to project by filtering the dictionary before passing it to this method.

Parameters:

Name	Type	Description	Default
features	dict[str, Tensor]	Dictionary of features to project	required

Returns:

Type	Description
dict[str, Tensor]	Dictionary of projected features. Spatial inputs come back as
dict[str, Tensor]	(B, T, Emb, H, W); all other ranks keep their shape with the final
dict[str, Tensor]	dimension projected to Emb.

Source code in src/versatil/models/layers/feature_projection.py

def forward(
    self,
    features: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Project features to common embedding dimension.

    The caller controls which features to project by filtering the
    dictionary before passing it to this method.

    Args:
        features: Dictionary of features to project

    Returns:
        Dictionary of projected features. Spatial inputs come back as
        (B, T, Emb, H, W); all other ranks keep their shape with the final
        dimension projected to Emb.
    """
    projected = {}
    for feature_name, raw_feature in features.items():
        B, T = None, None
        # Pipeline features always carry (B, T, ...): 5D spatial maps,
        # 4D token sequences, 3D vectors. 2D tensors are algorithm
        # context (e.g. latents) without a time axis.
        is_spatial = raw_feature.ndim == 5
        feature = raw_feature
        if is_spatial:
            B, T = feature.shape[0], feature.shape[1]
            feature = feature.reshape(B * T, *feature.shape[2:])  # (B*T, C, H, W)
        projection_dict = (
            self.spatial_projections if is_spatial else self.linear_projections
        )
        if feature_name not in projection_dict:
            projection_dict[feature_name] = self._create_projection_layer(
                feature=feature, is_spatial=is_spatial
            )
        feature_projection = projection_dict[feature_name](feature)
        if B is not None and T is not None:
            feature_projection = feature_projection.reshape(
                B, T, *feature_projection.shape[1:]
            )
        projected[feature_name] = feature_projection
    return projected

project_and_concatenate

project_and_concatenate(features, concatenation_dimension=-1)

Project features and concatenate them.

The caller controls which features to include by filtering the dictionary before passing it to this method.

Parameters:

Name	Type	Description	Default
features	dict[str, Tensor]	Dictionary of features to project and concatenate	required
concatenation_dimension	int	Dimension to concatenate along (default: -1)	-1

Returns:

Type	Description
Tensor	Concatenated projected features

Source code in src/versatil/models/layers/feature_projection.py

def project_and_concatenate(
    self,
    features: dict[str, torch.Tensor],
    concatenation_dimension: int = -1,
) -> torch.Tensor:
    """Project features and concatenate them.

    The caller controls which features to include by filtering the
    dictionary before passing it to this method.

    Args:
        features: Dictionary of features to project and concatenate
        concatenation_dimension: Dimension to concatenate along (default: -1)

    Returns:
        Concatenated projected features
    """
    projected = self.forward(features)
    sorted_tensors = [projected[key] for key in sorted(projected.keys())]
    if len(sorted_tensors) == 0:
        raise ValueError("No features to concatenate")
    elif len(sorted_tensors) == 1:
        return sorted_tensors[0]
    else:
        # Ensure tensors have the same shape except for concatenation dimension
        base_shape = list(sorted_tensors[0].shape)
        for tensor in sorted_tensors[1:]:
            for dim in range(len(base_shape)):
                if dim == concatenation_dimension:
                    continue
                if tensor.shape[dim] != base_shape[dim]:
                    raise ValueError(
                        f"Feature shapes do not match for concatenation: "
                        f"{base_shape} vs {tensor.shape} at dim {dim}"
                    )
        return torch.cat(sorted_tensors, dim=concatenation_dimension)