variational

variational

Variational inference wrapper for decoding algorithms.

This module provides a compositional wrapper that adds variational inference capabilities to any base decoding algorithm. Variational Inference enables modeling of multi-modal action distributions by introducing a latent variable z that is conditioned on both observations and actions: p(a|s) = p(a|z,s) p(z|s) Where: - p(a|z,s) is the base algorithm's decoder (BC, FlowMatching, etc.) - p(z|s) is the learned conditional prior (Gaussian or learned through a NN)

VariationalAlgorithm

VariationalAlgorithm(base_algorithm, posterior_encoder, prior=None, sampling_from_prior_probability=0.0, posterior_decoder_noise_std=0.0)

Bases: DecodingAlgorithm

Compositional wrapper adding variational inference to any decoding algorithm.

Wraps a base algorithm with variational inference, enabling multi-modal action prediction through a learned posterior q_phi(z|a,s) and prior p(z|s).

Training

1. Encode actions via approximated posterior: z ~ q_phi(z|a,s)
2. Train prior to match posterior (if learned prior)
3. Decode actions via posterior + decoding algorithm: p(a|z,s)q_phi(z|a,s)

Inference

1. Sample latent from prior: z ~ p(z|s)
2. Decode actions via base algorithm: p(a|z,s)

Parameters:

Name	Type	Description	Default
base_algorithm	DecodingAlgorithm	The underlying decoding algorithm (BC, FlowMatching, Diffusion, etc.)	required
posterior_encoder	PosteriorLatentEncoder	Latent action encoder for posterior q_phi(z|a,s) (e.g., a Transformer Encoder)	required
prior	PriorLatentEncoder | None	Latent prior for p(z|s). If None, auto-creates GaussianPrior.	None
sampling_from_prior_probability	float	Probability of sampling from prior during training.	0.0
posterior_decoder_noise_std	float	Fixed Gaussian jitter added to posterior latents before decoder training. This is applied only when training decodes from the posterior, not when decoding from prior samples.	0.0

Initialize variational algorithm wrapper.

Source code in src/versatil/models/decoding/algorithm/variational.py

def __init__(
    self,
    base_algorithm: DecodingAlgorithm,
    posterior_encoder: PosteriorLatentEncoder,
    prior: PriorLatentEncoder | None = None,
    sampling_from_prior_probability: float = 0.0,
    posterior_decoder_noise_std: float = 0.0,
):
    """Initialize variational algorithm wrapper."""
    super().__init__()
    if posterior_decoder_noise_std < 0.0:
        raise ValueError(
            "posterior_decoder_noise_std must be non-negative, "
            f"got {posterior_decoder_noise_std}."
        )
    self.base_algorithm = base_algorithm
    self.posterior_encoder = posterior_encoder
    self.p_prior = sampling_from_prior_probability
    self.posterior_decoder_noise_std = posterior_decoder_noise_std
    if prior is None:
        device = str(posterior_encoder.device)
        self.prior = GaussianPrior(
            latent_dimension=self.posterior_encoder.latent_dimension,
            device=device,
        )
        logging.info(
            f"Auto-created GaussianPrior with latent_dim={self.posterior_encoder.latent_dimension}. "
        )
    else:
        self.prior = prior

    if self.prior.latent_dimension != self.posterior_encoder.latent_dimension:
        raise ValueError(
            f"Latent dimension mismatch: prior.latent_dim={self.prior.latent_dimension} "
            f"!= posterior_encoder.latent_dim={self.posterior_encoder.latent_dimension}"
        )
    self._wire_prior_to_posterior()

predicts_in_action_space property

predicts_in_action_space

Delegates to the wrapped base algorithm.

__setstate__

__setstate__(state)

Restore module state and reconnect posterior-owned prior references.

Parameters:

Name	Type	Description	Default
state	dict[str, Any]	Pickled or deep-copied module state produced by PyTorch.	required

Source code in src/versatil/models/decoding/algorithm/variational.py

def __setstate__(self, state: dict[str, Any]) -> None:
    """Restore module state and reconnect posterior-owned prior references.

    Args:
        state: Pickled or deep-copied module state produced by PyTorch.
    """
    super().__setstate__(state)
    self._wire_prior_to_posterior()

get_callbacks

get_callbacks(experiment_config)

Provide callbacks for variational latent monitoring and prior fitting.

Parameters:

Name	Type	Description	Default
experiment_config	ExperimentConfig	Experiment-level callback configuration.	required

Returns:

Type	Description
list	Callbacks required by the variational algorithm.

Source code in src/versatil/models/decoding/algorithm/variational.py

def get_callbacks(self, experiment_config: ExperimentConfig) -> list:
    """Provide callbacks for variational latent monitoring and prior fitting.

    Args:
        experiment_config: Experiment-level callback configuration.

    Returns:
        Callbacks required by the variational algorithm.
    """
    callbacks = [
        LatentVisualizationCallback(
            log_every_n_epochs=experiment_config.val_every,
        )
    ]
    if isinstance(self.prior, CallbackProvider):
        callbacks.extend(
            self.prior.get_callbacks(experiment_config=experiment_config)
        )
    return callbacks

get_auxiliary_output_keys

get_auxiliary_output_keys()

Variational algorithm adds latent variable keys to the output.

Source code in src/versatil/models/decoding/algorithm/variational.py

def get_auxiliary_output_keys(self) -> set[str]:
    """Variational algorithm adds latent variable keys to the output."""
    keys = self.base_algorithm.get_auxiliary_output_keys()
    keys.update(self.posterior_encoder.get_auxiliary_output_keys())
    keys.update(self.prior.get_auxiliary_output_keys())
    return keys

injected_feature_keys

injected_feature_keys()

Latent keys added on top of the base algorithm's injections.

Source code in src/versatil/models/decoding/algorithm/variational.py

def injected_feature_keys(self) -> set[str]:
    """Latent keys added on top of the base algorithm's injections."""
    return self.base_algorithm.injected_feature_keys() | {
        LatentKey.POSTERIOR_LATENT.value,
        LatentKey.PRIOR_LATENT.value,
    }

forward

forward(network, features, actions=None)

Forward pass with variational latent encoding.

Training mode (self.training=True): - Encodes actions via approximated posterior q_phi(z|a,s) - Trains learned prior to match approximated posterior (if learned) - Stochastic mixing: samples from prior with probability p_prior - Delegates to base algorithm

Validation/eval mode (self.training=False): - Still encodes posterior for loss computation (KL term) - But uses ONLY prior samples for action decoding (like inference) - Better evaluates actual inference performance

Parameters:

Name	Type	Description	Default
network	ActionDecoder	Action decoder network	required
features	dict[str, Tensor]	Dictionary of input features from encoding pipeline	required
actions	dict[str, Tensor] | None	Optional ground-truth actions (for training)	None

Returns:

Type	Description
dict[str, Tensor]	Dictionary containing: - Action predictions from base algorithm - Latent variables (mu, logvar) if training - Prior outputs (prior_prediction, prior_target) if learned prior

Source code in src/versatil/models/decoding/algorithm/variational.py

def forward(
    self,
    network: ActionDecoder,
    features: dict[str, torch.Tensor],
    actions: dict[str, torch.Tensor] | None = None,
) -> dict[str, torch.Tensor]:
    """Forward pass with variational latent encoding.

    Training mode (self.training=True):
        - Encodes actions via approximated posterior q_phi(z|a,s)
        - Trains learned prior to match approximated posterior (if learned)
        - Stochastic mixing: samples from prior with probability p_prior
        - Delegates to base algorithm

    Validation/eval mode (self.training=False):
        - Still encodes posterior for loss computation (KL term)
        - But uses ONLY prior samples for action decoding (like inference)
        - Better evaluates actual inference performance

    Args:
        network: Action decoder network
        features: Dictionary of input features from encoding pipeline
        actions: Optional ground-truth actions (for training)

    Returns:
        Dictionary containing:
            - Action predictions from base algorithm
            - Latent variables (mu, logvar) if training
            - Prior outputs (prior_prediction, prior_target) if learned prior
    """
    if actions is None:
        raise ValueError(
            "Actions must be provided during training for variational algorithm."
        )
    posterior_output, prior_output = self._variational_step(
        features=features,
        actions=actions,
    )
    if self.training:
        sample_from_prior = torch.rand(1).item() < self.p_prior
    else:
        sample_from_prior = True
    if sample_from_prior:
        if LatentKey.PRIOR_LATENT.value not in prior_output:
            # Sample from prior only when needed (avoids costly denoising inference during training)
            batch_size, _, _ = resolve_feature_reference(features=features)
            prior_output[LatentKey.PRIOR_LATENT.value] = self.prior.sample_prior(
                batch_size=batch_size, observations=features
            )
        latent = prior_output[LatentKey.PRIOR_LATENT.value]
    else:
        latent = self._posterior_latent_for_decoder(posterior_output)
    features_with_latent = {
        **features,
        LatentKey.POSTERIOR_LATENT.value: latent,
    }  # (B, latent_dimension)
    predictions = self.base_algorithm.forward(
        network=network,
        features=features_with_latent,
        actions=actions,
    )

    predictions.update(posterior_output)
    predictions.update(prior_output)
    return predictions

get_targets

get_targets(algorithm_output, ground_truth_actions)

Delegate to the base algorithm's target selection.

Source code in src/versatil/models/decoding/algorithm/variational.py

def get_targets(
    self,
    algorithm_output: dict[str, torch.Tensor],
    ground_truth_actions: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Delegate to the base algorithm's target selection."""
    return self.base_algorithm.get_targets(algorithm_output, ground_truth_actions)

predict

predict(network, features)

Inference mode: sample latent z from prior and decode.

Parameters:

Name	Type	Description	Default
network	ActionDecoder	Action decoder network	required
features	dict[str, Tensor]	Dictionary of input features	required

Returns:

Type	Description
dict[str, Tensor]	Dictionary containing action predictions

Source code in src/versatil/models/decoding/algorithm/variational.py

def predict(
    self,
    network: ActionDecoder,
    features: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Inference mode: sample latent z from prior and decode.

    Args:
        network: Action decoder network
        features: Dictionary of input features

    Returns:
        Dictionary containing action predictions
    """
    batch_size, _, _ = resolve_feature_reference(features=features)
    latent_embedding = self._sample_prior(features=features, batch_size=batch_size)
    features_with_latent = {
        **features,
        LatentKey.POSTERIOR_LATENT.value: latent_embedding,
    }  # (B, latent_dimension)
    return self.base_algorithm.predict(
        network=network,
        features=features_with_latent,
    )