vamp_prior

vamp_prior

VampPrior: Variational Mixture of Posteriors Prior.

Implements the VampPrior from "VAE with a VampPrior" (Tomczak & Welling, 2018). The prior is a mixture of Gaussians where each component is defined by passing learnable pseudo-inputs through the posterior encoder.

This allows the prior to be more expressive than a standard Gaussian N(0,I) while maintaining the VAE framework.

VampPrior

VampPrior(latent_dimension, num_components, action_space, prediction_horizon, device, min_logvar=None)

Bases: PriorLatentEncoder

Variational Mixture of Posteriors Prior.

The posterior encoder in this policy family is conditional, q_phi(z | a, s). Therefore the VampPrior components are conditional pseudo-action posteriors, q_phi(z | u_k, s).

Parameters:

Name	Type	Description	Default
latent_dimension	int	Dimension of latent variable z	required
num_components	int	Number of mixture components K	required
action_space	ActionSpace	ActionSpace defining the action dimensions	required
prediction_horizon	int	Number of timesteps in action chunks	required
device	str	Device to place prior on	required
min_logvar	float | None	Optional minimum logvar clamp	None

Source code in src/versatil/models/decoding/latent/prior/vamp_prior.py

def __init__(
    self,
    latent_dimension: int,
    num_components: int,
    action_space: ActionSpace,
    prediction_horizon: int,
    device: str,
    min_logvar: float | None = None,
):
    super().__init__(latent_dimension=latent_dimension, device=device)
    self.num_components = num_components
    self.prediction_horizon = prediction_horizon
    self.action_dim = action_space.get_total_action_dim()
    self.min_logvar = min_logvar
    self.action_keys, self.action_dimensions = self._get_action_layout(
        action_space=action_space,
    )
    self.pseudo_inputs = nn.Parameter(
        torch.randn(num_components, prediction_horizon, self.action_dim)
    )
    self.log_weights = nn.Parameter(torch.zeros(num_components, 1, 1))
    self._encoder_reference: tuple[PosteriorLatentEncoder, ...] = ()
    self.to(torch.device(device))

encoder property

encoder

Get the posterior encoder.

get_auxiliary_output_keys

get_auxiliary_output_keys()

Return prediction keys produced by the mixture prior.

Source code in src/versatil/models/decoding/latent/prior/vamp_prior.py

def get_auxiliary_output_keys(self) -> set[str]:
    """Return prediction keys produced by the mixture prior."""
    return {
        LatentKey.PRIOR_LATENT.value,
        LatentKey.PRIOR_LOG_PROB.value,
    }

wire_posterior

wire_posterior(posterior)

Wire shared state from the posterior encoder.

Extracts the encoder reference needed to compute mixture components from learnable pseudo-inputs.

Parameters:

Name	Type	Description	Default
posterior	PosteriorLatentEncoder	Posterior encoder that maps inputs to (mu, logvar).	required

Source code in src/versatil/models/decoding/latent/prior/vamp_prior.py

def wire_posterior(self, posterior: PosteriorLatentEncoder) -> None:
    """Wire shared state from the posterior encoder.

    Extracts the encoder reference needed to compute mixture
    components from learnable pseudo-inputs.

    Args:
        posterior: Posterior encoder that maps inputs to (mu, logvar).
    """
    self._encoder_reference = (posterior,)

__getstate__

__getstate__()

Return copy-safe state without the posterior encoder reference.

Returns:

Type	Description
dict[str, Any]	Module state with runtime wiring cleared. The owning
dict[str, Any]	VariationalAlgorithm reconnects the copied prior to the copied
dict[str, Any]	posterior after deepcopy or unpickling.

Source code in src/versatil/models/decoding/latent/prior/vamp_prior.py

def __getstate__(self) -> dict[str, Any]:
    """Return copy-safe state without the posterior encoder reference.

    Returns:
        Module state with runtime wiring cleared. The owning
        VariationalAlgorithm reconnects the copied prior to the copied
        posterior after deepcopy or unpickling.
    """
    state = super().__getstate__()
    state["_encoder_reference"] = ()
    return state

get_mixture_params

get_mixture_params(observations=None)

Compute mixture component parameters by passing pseudo-inputs through encoder.

Returns:

Type	Description
tuple[Tensor, Tensor]	Tuple of conditional (means, logvars), each shaped (B, K, latent_dim).

Source code in src/versatil/models/decoding/latent/prior/vamp_prior.py

def get_mixture_params(
    self,
    observations: dict[str, torch.Tensor] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute mixture component parameters by passing pseudo-inputs through encoder.

    Returns:
        Tuple of conditional (means, logvars), each shaped (B, K, latent_dim).
    """
    observations = self._require_observations(observations=observations)
    batch_size = self._get_observation_batch_size(observations=observations)
    encoder_observations = self._repeat_observations_for_components(
        observations=observations,
    )

    pseudo_actions = self._build_pseudo_actions(batch_size=batch_size)
    encoder_output = self.encoder.encode(
        actions=pseudo_actions,
        observations=encoder_observations,
    )
    mu = encoder_output[LatentKey.POSTERIOR_MU.value]  # (K, latent_dim)
    logvar = encoder_output[LatentKey.POSTERIOR_LOGVAR.value]  # (K, latent_dim)
    mu = mu.reshape(batch_size, self.num_components, self.latent_dimension)
    logvar = logvar.reshape(
        batch_size,
        self.num_components,
        self.latent_dimension,
    )
    if self.min_logvar is not None:
        logvar = torch.clamp(logvar, min=self.min_logvar)
    return mu, logvar

build_training_target

build_training_target(posterior_output)

Use a differentiable posterior sample for the sample-based KL.

Source code in src/versatil/models/decoding/latent/prior/vamp_prior.py

def build_training_target(
    self,
    posterior_output: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Use a differentiable posterior sample for the sample-based KL."""
    return posterior_output[LatentKey.POSTERIOR_LATENT.value]

sample_prior

sample_prior(batch_size, observations=None)

Sample from the VampPrior mixture.

Parameters:

Name	Type	Description	Default
batch_size	int	Number of samples to generate	required
observations	dict[str, Tensor] | None	Required conditioning features	None

Returns:

Type	Description
Tensor	Sampled latents of shape (batch_size, latent_dim)

Source code in src/versatil/models/decoding/latent/prior/vamp_prior.py

def sample_prior(
    self,
    batch_size: int,
    observations: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
    """Sample from the VampPrior mixture.

    Args:
        batch_size: Number of samples to generate
        observations: Required conditioning features

    Returns:
        Sampled latents of shape (batch_size, latent_dim)
    """
    mu, logvar = self.get_mixture_params(observations=observations)
    weights = F.softmax(self.log_weights.view(-1), dim=0)  # (K,)
    component_indices = torch.multinomial(
        weights, batch_size, replacement=True
    )  # (batch_size,)
    if mu.size(0) != batch_size:
        raise ValueError(
            "VampPrior conditional mixture batch size mismatch: "
            f"got {mu.size(0)} component batches for requested "
            f"batch_size {batch_size}."
        )
    batch_indices = torch.arange(batch_size, device=mu.device)
    selected_mu = mu[batch_indices, component_indices]
    selected_logvar = logvar[batch_indices, component_indices]
    std = (selected_logvar / 2).exp()
    eps = torch.randn_like(std)
    z = selected_mu + std * eps
    return z

forward

forward(target_latents, observations)

Compute prior log probability for training.

VampPrior is a mixture of Gaussians. Unlike single-Gaussian priors, we return log_prob directly since mu/logvar don't represent a mixture.

Parameters:

Name	Type	Description	Default
target_latents	Tensor | None	Target latents from posterior (B, latent_dim)	required
observations	dict[str, Tensor]	Conditioning features	required

Returns:

Type	Description
dict[str, Tensor]	Dictionary with LatentKey.PRIOR_LOG_PROB

Source code in src/versatil/models/decoding/latent/prior/vamp_prior.py

def forward(
    self,
    target_latents: torch.Tensor | None,
    observations: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Compute prior log probability for training.

    VampPrior is a mixture of Gaussians. Unlike single-Gaussian priors,
    we return log_prob directly since mu/logvar don't represent a mixture.

    Args:
        target_latents: Target latents from posterior (B, latent_dim)
        observations: Conditioning features

    Returns:
        Dictionary with LatentKey.PRIOR_LOG_PROB
    """
    if target_latents is None:
        raise ValueError(
            "VampPrior.forward() requires target_latents for log-prob "
            "computation. Use sample_prior() for inference."
        )
    log_prob = self.log_prob(target_latents, observations=observations)
    return {
        LatentKey.PRIOR_LOG_PROB.value: log_prob,
    }

log_prob

log_prob(z, observations=None)

Compute log probability of z under the VampPrior mixture.

Parameters:

Name	Type	Description	Default
z	Tensor	Latent samples of shape (B, latent_dim)	required
observations	dict[str, Tensor] | None	Required conditioning features	None

Returns:

Type	Description
Tensor	Log probability of shape (B,) - summed over latent dimensions

Source code in src/versatil/models/decoding/latent/prior/vamp_prior.py

def log_prob(
    self,
    z: torch.Tensor,
    observations: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
    """Compute log probability of z under the VampPrior mixture.

    Args:
        z: Latent samples of shape (B, latent_dim)
        observations: Required conditioning features

    Returns:
        Log probability of shape (B,) - summed over latent dimensions
    """
    mu, logvar = self.get_mixture_params(observations=observations)
    weights = F.softmax(self.log_weights.view(-1), dim=0)  # (K,)
    log_p_components = log_normal_diag(
        z.unsqueeze(1),
        mu,
        logvar,
    ).sum(dim=-1)  # (B, K)
    log_p_weighted = log_p_components + torch.log(weights).unsqueeze(0)
    log_prob = torch.logsumexp(log_p_weighted, dim=1)  # (B,)
    return log_prob

log_normal_diag

log_normal_diag(z, mu, logvar)

Compute log probability of z under diagonal Gaussian N(mu, exp(logvar)).

Parameters:

Name	Type	Description	Default
z	Tensor	Samples, shape (..., latent_dim)	required
mu	Tensor	Mean, shape (..., latent_dim)	required
logvar	Tensor	Log variance, shape (..., latent_dim)	required

Returns:

Type	Description
Tensor	Log probability, shape (..., latent_dim) - per-dimension log probs

Source code in src/versatil/models/decoding/latent/prior/vamp_prior.py

def log_normal_diag(
    z: torch.Tensor, mu: torch.Tensor, logvar: torch.Tensor
) -> torch.Tensor:
    """Compute log probability of z under diagonal Gaussian N(mu, exp(logvar)).

    Args:
        z: Samples, shape (..., latent_dim)
        mu: Mean, shape (..., latent_dim)
        logvar: Log variance, shape (..., latent_dim)

    Returns:
        Log probability, shape (..., latent_dim) - per-dimension log probs
    """
    log_p = -0.5 * (
        torch.log(2 * torch.pi * torch.ones(1, device=z.device))
        + logvar
        + (z - mu) ** 2 / torch.exp(logvar)
    )
    return log_p