vq_encoder

vq_encoder

Posterior encoder with vector quantization for discrete latent variable models.

Uses the same transformer backbone as VAETransformerEncoder to produce a continuous embedding from actions and observations, then quantizes it via ResidualVQ to produce a discrete latent code. The quantized embedding is passed to the decoder via the VariationalAlgorithm. Commitment loss inputs (continuous z and quantized z) are stored in the output dict for external loss computation in the metrics module.

VQPosteriorEncoder

VQPosteriorEncoder(latent_dimension, num_codes, num_residual_layers, embedding_dimension, prediction_horizon, observation_horizon, device, ema_decay=0.99, dead_code_threshold=1.0, number_of_heads=4, feedforward_dimension=128, number_of_encoder_layers=1, activation=value, dropout_rate=0.0, attention_dropout=0.0, normalization_type=value, attention_type=value, positional_encoding_type=None, exclude_keys=None)

Bases: PosteriorLatentEncoder

Transformer posterior encoder with residual vector quantization.

Encodes actions and observations into a continuous embedding via a transformer encoder with a CLS token, then quantizes the embedding through a ResidualVQ bottleneck. The quantized output is the latent z passed to the decoder. The continuous pre-quantization embedding and codebook indices are stored in the output dict for commitment loss computation and prior training.

Parameters:

Name	Type	Description	Default
latent_dimension	int	Dimension of each codebook vector and the latent space passed to the decoder.	required
num_codes	int	Number of codebook entries per residual layer (K).	required
num_residual_layers	int	Number of cascading VQ layers.	required
ema_decay	float	EMA decay for codebook updates.	0.99
dead_code_threshold	float	Cluster size below which codes are replaced.	1.0
embedding_dimension	int	Transformer hidden dimension.	required
prediction_horizon	int	Number of action timesteps.	required
observation_horizon	int	Number of observation timesteps.	required
device	str	Device string.	required
number_of_heads	int	Number of attention heads.	4
feedforward_dimension	int	Feedforward network dimension.	128
number_of_encoder_layers	int	Number of transformer encoder layers.	1
activation	str	Activation function name.	value
dropout_rate	float	Dropout probability.	0.0
attention_type	str	Attention mechanism type (use AttentionType enum values).	value
exclude_keys	list[str] | None	Observation keys to exclude from encoding.	None

Source code in src/versatil/models/decoding/latent/posterior/vq_encoder.py

def __init__(
    self,
    latent_dimension: int,
    num_codes: int,
    num_residual_layers: int,
    embedding_dimension: int,
    prediction_horizon: int,
    observation_horizon: int,
    device: str,
    ema_decay: float = 0.99,
    dead_code_threshold: float = 1.0,
    number_of_heads: int = 4,
    feedforward_dimension: int = 128,
    number_of_encoder_layers: int = 1,
    activation: str = ActivationFunction.SWIGLU.value,
    dropout_rate: float = 0.0,
    attention_dropout: float = 0.0,
    normalization_type: str = NormalizationType.RMS_NORM.value,
    attention_type: str = AttentionType.MULTI_HEAD.value,
    positional_encoding_type: str | None = None,
    exclude_keys: list[str] | None = None,
):
    super().__init__(
        latent_dimension=latent_dimension,
        device=device,
    )
    self.code_dim = latent_dimension
    self.num_codes = num_codes
    self.num_residual_layers = num_residual_layers
    self.exclude_keys = exclude_keys if exclude_keys is not None else []
    self.embedding_dimension = embedding_dimension
    self.prediction_horizon = prediction_horizon
    self.observation_horizon = observation_horizon

    self.transformer_encoder = TransformerEncoder(
        number_of_layers=number_of_encoder_layers,
        embedding_dimension=embedding_dimension,
        number_of_heads=number_of_heads,
        feedforward_dimension=feedforward_dimension,
        activation=activation,
        dropout=dropout_rate,
        attention_dropout=attention_dropout,
        normalization_type=normalization_type,
        attention_type=attention_type,
        positional_encoding_type=positional_encoding_type,
    )

    temporal_positional_encoding = None
    if observation_horizon > 1:
        temporal_positional_encoding = LearnedPositionalEncoding1D(
            embedding_dimension=embedding_dimension
        )

    self.input_sequence_builder = TransformerInputBuilder(
        embedding_dimension=embedding_dimension,
        spatial_positional_encoding_layer=SinusoidalPositionalEncoding2D(
            embedding_dimension=embedding_dimension, normalize=True
        ),
        flat_positional_encoding_layer=SinusoidalPositionalEncoding1D(
            embedding_dimension=embedding_dimension,
            maximum_sequence_length=1000,
        ),
        temporal_positional_encoding_layer=temporal_positional_encoding,
    )

    self.cls_token = nn.Embedding(1, embedding_dimension)  # (1, emb_dim)
    self.latent_projection = nn.Linear(embedding_dimension, self.code_dim)

    self.residual_vq = ResidualVQ(
        input_dimension=self.code_dim,
        code_dim=self.code_dim,
        num_codes=num_codes,
        num_layers=num_residual_layers,
        ema_decay=ema_decay,
        dead_code_threshold=dead_code_threshold,
        kmeans_init=True,
    )

    self.to(device)

get_auxiliary_output_keys

get_auxiliary_output_keys()

VQ posterior outputs quantized latent, codebook indices, and continuous z.

Source code in src/versatil/models/decoding/latent/posterior/vq_encoder.py

def get_auxiliary_output_keys(self) -> set[str]:
    """VQ posterior outputs quantized latent, codebook indices, and continuous z."""
    return {
        LatentKey.POSTERIOR_LATENT.value,
        LatentKey.VQ_INDICES.value,
        LatentKey.VQ_Z_CONTINUOUS.value,
        LatentKey.VQ_QUANTIZED.value,
    }

encode

encode(actions, observations=None)

Encode actions into a vector-quantized latent.

Parameters:

Name	Type	Description	Default
actions	dict[str, Tensor]	Dictionary of action tensors, shape (B, prediction_horizon, action_dim) per component.	required
observations	dict[str, Tensor] | None	Optional observation features for conditioning.	None

Returns:

Type

Description

dict[str, Tensor]

Dictionary containing: - LatentKey.POSTERIOR_LATENT: Quantized latent z (B, code_dim), with straight-through gradient for decoder training. - LatentKey.VQ_INDICES: Per-layer codebook indices, list of (B,) tensors, for prior training. - LatentKey.VQ_Z_CONTINUOUS: Per-layer pre-quantization encoder outputs in code space (L, B, code_dim). Carries gradient; paired with VQ_QUANTIZED for commitment loss. - LatentKey.VQ_QUANTIZED: Per-layer hard-quantized codebook vectors in code space (L, B, code_dim), detached. Used with VQ_Z_CONTINUOUS for per-layer commitment loss.

Source code in src/versatil/models/decoding/latent/posterior/vq_encoder.py

def encode(
    self,
    actions: dict[str, torch.Tensor],
    observations: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor]:
    """Encode actions into a vector-quantized latent.

    Args:
        actions: Dictionary of action tensors,
            shape (B, prediction_horizon, action_dim) per component.
        observations: Optional observation features for conditioning.

    Returns:
        Dictionary containing:
            - LatentKey.POSTERIOR_LATENT: Quantized latent z (B, code_dim),
                with straight-through gradient for decoder training.
            - LatentKey.VQ_INDICES: Per-layer codebook indices,
                list of (B,) tensors, for prior training.
            - LatentKey.VQ_Z_CONTINUOUS: Per-layer pre-quantization
                encoder outputs in code space (L, B, code_dim). Carries
                gradient; paired with VQ_QUANTIZED for commitment loss.
            - LatentKey.VQ_QUANTIZED: Per-layer hard-quantized codebook
                vectors in code space (L, B, code_dim), detached. Used
                with VQ_Z_CONTINUOUS for per-layer commitment loss.
    """
    if observations is not None:
        input_observations = {
            k: v for k, v in observations.items() if k not in self.exclude_keys
        }
    else:
        input_observations = {}

    action_feature_keys = []
    for action_key, action_tensor in actions.items():
        if action_key == SampleKey.IS_PAD_ACTION.value:
            continue
        input_observations[action_key] = action_tensor.to(
            self.cls_token.weight.device
        ).float()
        action_feature_keys.append(action_key)

    batch_size = list(input_observations.values())[0].size(0)
    is_pad = actions.get(SampleKey.IS_PAD_ACTION.value)
    if is_pad is None:
        logging.warning("No padding key found in actions; assuming no padding.")
        is_pad = torch.zeros(
            batch_size,
            self.prediction_horizon,
            dtype=torch.bool,
            device=self.cls_token.weight.device,
        )
        input_observations[SampleKey.IS_PAD_ACTION.value] = is_pad
    else:
        is_pad = is_pad.to(device=self.cls_token.weight.device, dtype=torch.bool)
        input_observations[SampleKey.IS_PAD_ACTION.value] = is_pad

    for action_key in action_feature_keys:
        input_observations[
            f"{action_key}_{EncoderOutputKeys.PADDING_MASK.value}"
        ] = is_pad

    cls_embedding = self.cls_token.weight.unsqueeze(0).repeat(
        batch_size, 1, 1
    )  # (B, 1, emb_dim)
    input_observations[AlgorithmContextKey.CLASS_TOKEN.value] = cls_embedding

    input_tokens, pos_encodings, padding_mask = self.input_sequence_builder(
        input_observations
    )  # (B, seq_len, emb_dim)

    hidden_states = input_tokens + pos_encodings
    encoder_output = self.transformer_encoder(
        hidden_states=hidden_states,
        padding_mask=padding_mask,
    )[:, -1, :]  # (B, emb_dim) — CLS token at last position

    z_continuous = self.latent_projection(encoder_output)  # (B, code_dim)

    z_q, all_indices, z_e_per_layer, z_q_per_layer = self.residual_vq(
        z_continuous
    )  # (B, code_dim), list[(B,)], (L, B, code_dim), (L, B, code_dim)

    return {
        LatentKey.POSTERIOR_LATENT.value: z_q,
        LatentKey.VQ_INDICES.value: all_indices,
        LatentKey.VQ_Z_CONTINUOUS.value: z_e_per_layer,
        LatentKey.VQ_QUANTIZED.value: z_q_per_layer,
    }