geometric_rgbd

geometric_rgbd

Lightweight geometry-aware RGBD encoder.

GeometricRGBDEncoder

GeometricRGBDEncoder(input_keys, embedding_dimension=512, number_of_heads=8, ffn_dimension=2048, decomposition_mode=value, initial_decay=2.0, decay_range=4.0, patch_size=16, pooling_method=value, pretrained=False, frozen=False, model_dtype=None)

Bases: RGBDEncoderMixin, Encoder

Single-layer geometry-aware RGBD encoder.

Initialize the geometric RGBD encoder.

Parameters:

Name	Type	Description	Default
input_keys	str | list[str]	Input keys for RGB and depth observations.	required
embedding_dimension	int	Dimension of patch embeddings and attention.	512
number_of_heads	int	Number of attention heads.	8
ffn_dimension	int	Hidden dimension of the feed-forward network.	2048
decomposition_mode	str	Attention computation strategy (full or separable).	value
initial_decay	float	Initial decay rate for spatial biases.	2.0
decay_range	float	Range of decay rates across heads.	4.0
patch_size	int	Size of image patches for the patch embedding.	16
pooling_method	str	Feature pooling method applied after attention.	value
pretrained	bool	Whether to use pretrained weights (not supported).	False
frozen	bool	Whether to freeze encoder weights (not supported).	False
model_dtype	str | None	Precision string from experiment config (e.g. "bf16-mixed").	None

Source code in src/versatil/models/encoding/encoders/cross_modal/rgbd/geometric_rgbd.py

def __init__(
    self,
    input_keys: str | list[str],
    embedding_dimension: int = 512,
    number_of_heads: int = 8,
    ffn_dimension: int = 2048,
    decomposition_mode: str = AttentionDecompositionMode.SEPARABLE.value,
    initial_decay: float = 2.0,
    decay_range: float = 4.0,
    patch_size: int = 16,
    pooling_method: str = PoolingMethod.AVERAGE.value,
    pretrained: bool = False,
    frozen: bool = False,
    model_dtype: str | None = None,
):
    """Initialize the geometric RGBD encoder.

    Args:
        input_keys: Input keys for RGB and depth observations.
        embedding_dimension: Dimension of patch embeddings and attention.
        number_of_heads: Number of attention heads.
        ffn_dimension: Hidden dimension of the feed-forward network.
        decomposition_mode: Attention computation strategy (full or separable).
        initial_decay: Initial decay rate for spatial biases.
        decay_range: Range of decay rates across heads.
        patch_size: Size of image patches for the patch embedding.
        pooling_method: Feature pooling method applied after attention.
        pretrained: Whether to use pretrained weights (not supported).
        frozen: Whether to freeze encoder weights (not supported).
        model_dtype: Precision string from experiment config (e.g. ``"bf16-mixed"``).
    """
    specification = EncoderInput(
        keys=input_keys,
        exactly_one_camera_modality=[CameraModality.RGB, CameraModality.DEPTH],
        required_camera_modalities=[CameraModality.RGB, CameraModality.DEPTH],
    )
    super().__init__(
        input_specification=specification,
        pretrained=pretrained,
        frozen=frozen,
        model_dtype=model_dtype,
    )
    if pretrained:
        logging.warning(
            "GeometricRGBDEncoder does not support pretrained weights. Continuing with random initialization."
        )
    if frozen:
        raise ValueError(
            "Freezing GeometricRGBDEncoder does not make sense as it has no pretrained weights. Set frozen=False."
        )
    self._setup_camera_keys(input_keys=self.input_specification.keys)
    self.embedding_dimension = embedding_dimension
    self.decomposition_mode = AttentionDecompositionMode(decomposition_mode)
    self.pooling_method = pooling_method
    self.patch_size = patch_size
    self.patch_embed = PatchEmbedding(
        patch_size=self.patch_size,
        in_chans=3,
        embed_dim=embedding_dimension,
        embed_type=PatchEmbedType.STANDARD.value,
        norm_layer=LayerNorm,
    )

    self.attention_block = GeometricAttentionEncoderBlock(
        decomposition_mode=AttentionDecompositionMode(decomposition_mode),
        embedding_dimension=embedding_dimension,
        number_of_heads=number_of_heads,
        ffn_dimension=ffn_dimension,
        initial_decay=initial_decay,
        decay_range=decay_range,
    )
    self.pre_attention_norm = nn.LayerNorm(embedding_dimension, eps=1e-6)
    self.post_attention_norm = nn.LayerNorm(embedding_dimension, eps=1e-6)
    self.pooling_head: PoolingHead | None = None
    self.output_dim: int | tuple[int, ...] = self.embedding_dimension
    if frozen:
        super()._freeze_weights()
    self._apply_model_dtype()

encode_features

encode_features(rgb_image, depth_map)

Encode RGB and depth into joint RGBD features using geometric attention.

Parameters:

Name	Type	Description	Default
rgb_image	Tensor	RGB image tensor of shape (B*T, C, H, W); forward() flattens the temporal axis into the batch before dispatching here.	required
depth_map	Tensor	Depth map tensor of shape (B*T, 1, H, W), temporally flattened like rgb_image.	required

Returns:

Type	Description
Tensor	Tuple of (features, H_patches, W_patches) where features has shape
int	(B*T, embedding_dimension, H_patches, W_patches).

Source code in src/versatil/models/encoding/encoders/cross_modal/rgbd/geometric_rgbd.py

def encode_features(
    self, rgb_image: torch.Tensor, depth_map: torch.Tensor
) -> tuple[torch.Tensor, int, int]:
    """Encode RGB and depth into joint RGBD features using geometric attention.

    Args:
        rgb_image: RGB image tensor of shape (B*T, C, H, W); ``forward()``
            flattens the temporal axis into the batch before
            dispatching here.
        depth_map: Depth map tensor of shape (B*T, 1, H, W), temporally
            flattened like ``rgb_image``.

    Returns:
        Tuple of (features, H_patches, W_patches) where features has shape
        (B*T, embedding_dimension, H_patches, W_patches).
    """
    features, H_patches, W_patches = self.patch_embed(
        rgb_image, return_patch_size=True
    )  # (B, N_patches, embedding_dimension)
    features = self.pre_attention_norm(features)
    features = features.reshape(
        rgb_image.shape[0], H_patches, W_patches, self.embedding_dimension
    )
    depth_map_resized = F.interpolate(
        depth_map, size=(H_patches, W_patches), mode="bilinear", align_corners=False
    )
    features = self.attention_block(
        features, depth_map_resized
    )  # (B, H_patches, W_patches, embedding_dimension)
    features = self.post_attention_norm(features)
    features = features.permute(
        0, 3, 1, 2
    ).contiguous()  # (B, embedding_dimension, H_patches, W_patches)
    return features, H_patches, W_patches

encode

encode(inputs)

Encode RGB + depth into fused features.

Parameters:

Name	Type	Description	Default
inputs	dict[str, Tensor]	Dict with RGB as (B, C, H, W) and depth as (B, 1, H, W).	required

Returns:

Type	Description
dict[str, Tensor]	Dict with RGBD features.

Source code in src/versatil/models/encoding/encoders/cross_modal/rgbd/geometric_rgbd.py

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

    Args:
        inputs: Dict with RGB as (B, C, H, W) and depth as (B, 1, H, W).

    Returns:
        Dict with RGBD features.
    """
    rgb_key = self._camera_key_for_modality(modality=CameraModality.RGB)
    depth_key = self._camera_key_for_modality(modality=CameraModality.DEPTH)
    rgb = inputs[rgb_key]
    depth = inputs[depth_key]

    if self.pooling_head is None:
        raise RuntimeError(
            "pooling_head is not initialized. Call set_image_size() before forward."
        )
    features, feature_map_height, feature_map_width = self.encode_features(
        rgb, depth
    )
    pooled_features = self.pooling_head(features)
    return {EncoderOutputKeys.RGBD.value: pooled_features}

set_image_size

set_image_size(image_height, image_width)

Compute feature map dimensions and create pooling head.

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/cross_modal/rgbd/geometric_rgbd.py

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

    Args:
        image_height: Target image height.
        image_width: Target image width.
    """
    mock_input_dtype = self._mock_forward_dtype()
    with torch.no_grad(), self._mock_forward_autocast():
        mock_rgb = torch.zeros(
            1, 3, image_height, image_width, dtype=mock_input_dtype
        )
        mock_depth = torch.zeros(
            1, 1, image_height, image_width, dtype=mock_input_dtype
        )
        _, spatial_height, spatial_width = self.encode_features(
            rgb_image=mock_rgb, depth_map=mock_depth
        )
    self._setup_pooling(spatial_height=spatial_height, spatial_width=spatial_width)
    self._apply_model_dtype()

get_explainability_targets

get_explainability_targets()

Return the geometric attention block for spatial attribution maps.

Returns:

Type	Description
list[VisionExplanationTarget]	One NHWC spatial feature-map target from the RGBD attention block.

Source code in src/versatil/models/encoding/encoders/cross_modal/rgbd/geometric_rgbd.py

def get_explainability_targets(self) -> list[VisionExplanationTarget]:
    """Return the geometric attention block for spatial attribution maps.

    Returns:
        One NHWC spatial feature-map target from the RGBD attention block.
    """
    return [
        VisionExplanationTarget(
            layer=self.attention_block,
            target_kind=ExplanationTargetKind.SPATIAL_FEATURE_MAP.value,
            activation_layout=ActivationLayout.NHWC.value,
        )
    ]

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 RGBD feature name and its pooled dimension.

Source code in src/versatil/models/encoding/encoders/cross_modal/rgbd/geometric_rgbd.py

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

    Returns:
        List of FeatureMetadata with RGBD feature name and its pooled dimension.
    """
    dimension = (
        (self.output_dim,) if isinstance(self.output_dim, int) else self.output_dim
    )
    return [
        FeatureMetadata(
            key=EncoderOutputKeys.RGBD.value,
            feature_type=infer_feature_type(dimension),
            dimension=dimension,
        )
    ]