base_prior

base_prior

Abstract base class for latent prior networks.

PriorLatentEncoder

PriorLatentEncoder(latent_dimension, device)

Bases: ModuleAttrMixin, ABC

Abstract base class for prior parametrizations over a latent space z, which can be either learned (through a NN) or fixed.

Latent priors model the conditional distribution p(z|s) where z is a latent variable and s is an optional conditioning of observations (if the prior distribution is learned).

Design

• forward() takes target latents and conditioning, returns predictions and targets
• sample_prior() generates latent samples for inference
• Loss computation happens in separate loss modules (not in forward())

Initialize latent prior.

Parameters:

Name	Type	Description	Default
latent_dimension	int	Dimension of latent variable z	required
device	str	Device to place prior on	required

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

def __init__(self, latent_dimension: int, device: str) -> None:
    """Initialize latent prior.

    Args:
        latent_dimension: Dimension of latent variable z
        device: Device to place prior on
    """
    super().__init__()
    self.latent_dimension = latent_dimension
    self.device = torch.device(device)

forward abstractmethod

forward(target_latents, observations)

Compute prior predictions for training.

This method takes clean latent samples (typically from a posterior encoder) and conditioning, and returns predictions and targets for loss computation. The actual loss is computed externally.

Parameters:

Name	Type	Description	Default
target_latents	Tensor | None	Clean latent samples, shape (B, latent_dim), or None when called from sample_prior() on learned priors. These are sampled from the approximate posterior q_phi(z|a,s) and should be detached to prevent gradients flowing twice to the approximate posterior.	required
observations	dict[str, Tensor]	Dictionary of conditioning features, typically the observations at current state.	required

Returns:

Type	Description
dict[str, Tensor]	Dictionary containing predictions and targets for loss computation.
dict[str, Tensor]	Keys depend on the specific prior type (e.g., "predicted_prior", "target_prior").

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

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

    This method takes clean latent samples (typically from a posterior
    encoder) and conditioning, and returns predictions and targets for
    loss computation. The actual loss is computed externally.

    Args:
        target_latents: Clean latent samples, shape (B, latent_dim), or
            None when called from sample_prior() on learned priors.
            These are sampled from the approximate posterior q_phi(z|a,s) and
            should be detached to prevent gradients flowing twice to the approximate posterior.
        observations: Dictionary of conditioning features, typically the observations at current state.

    Returns:
        Dictionary containing predictions and targets for loss computation.
        Keys depend on the specific prior type (e.g., "predicted_prior", "target_prior").
    """
    raise NotImplementedError

get_auxiliary_output_keys

get_auxiliary_output_keys()

Return the set of keys this prior adds to the predictions dict.

Base implementation returns keys common to Gaussian priors. Subclasses override to add or remove keys.

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

def get_auxiliary_output_keys(self) -> set[str]:
    """Return the set of keys this prior adds to the predictions dict.

    Base implementation returns keys common to Gaussian priors.
    Subclasses override to add or remove keys.
    """
    return {
        LatentKey.PRIOR_LATENT.value,
        LatentKey.PRIOR_MU.value,
        LatentKey.PRIOR_LOGVAR.value,
    }

build_training_target

build_training_target(posterior_output)

Select and detach the posterior latent used to train this prior.

Parameters:

Name	Type	Description	Default
posterior_output	dict[str, Tensor]	Posterior encoder output dictionary.	required

Returns:

Type	Description
Tensor	Detached posterior sample used as the default prior target.

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

def build_training_target(
    self, posterior_output: dict[str, torch.Tensor]
) -> torch.Tensor:
    """Select and detach the posterior latent used to train this prior.

    Args:
        posterior_output: Posterior encoder output dictionary.

    Returns:
        Detached posterior sample used as the default prior target.
    """
    return posterior_output[LatentKey.POSTERIOR_LATENT.value].detach()

sample_prior abstractmethod

sample_prior(batch_size, observations=None)

Sample latent variable from learned prior p(z|s).

During inference, we don't have ground-truth actions (and thus no posterior samples), so we sample directly from the (conditional) prior.

Parameters:

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

Returns:

Type	Description
Tensor	Sampled latent embeddings, shape (batch_size, embedding_dimension)

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

@abc.abstractmethod
def sample_prior(
    self,
    batch_size: int,
    observations: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
    """Sample latent variable from learned prior p(z|s).

    During inference, we don't have ground-truth actions (and thus no
    posterior samples), so we sample directly from the (conditional) prior.

    Args:
        batch_size: Number of samples to generate
        observations: Optional dictionary of conditioning features

    Returns:
        Sampled latent embeddings, shape (batch_size, embedding_dimension)
    """
    raise NotImplementedError