gaussian

gaussian

GaussianHead predicts Gaussian distribution parameters (mean, logvar).

GaussianHead

GaussianHead(input_dimension, blocks=None, min_logvar=-10.0, max_logvar=4.0)

Bases: BaseActionHead

Action head that outputs Gaussian distribution parameters (mean, logvar).

Initialize Gaussian head.

Parameters:

Name	Type	Description	Default
input_dimension	int	Input embedding dimension from decoder.	required
blocks	list[ActionHeadBlock] | None	Blocks to apply before output projection.	None
min_logvar	float	Minimum value for logvar clamping.	-10.0
max_logvar	float	Maximum value for logvar clamping.	4.0

Source code in src/versatil/models/decoding/action_heads/gaussian.py

def __init__(
    self,
    input_dimension: int,
    blocks: list[ActionHeadBlock] | None = None,
    min_logvar: float = -10.0,
    max_logvar: float = 4.0,
) -> None:
    """Initialize Gaussian head.

    Args:
        input_dimension: Input embedding dimension from decoder.
        blocks: Blocks to apply before output projection.
        min_logvar: Minimum value for logvar clamping.
        max_logvar: Maximum value for logvar clamping.
    """
    super().__init__(input_dimension=input_dimension, blocks=blocks)
    self.min_logvar = min_logvar
    self.max_logvar = max_logvar
    self._logvar_proj: nn.Linear | None = None

set_output_dim

set_output_dim(dim)

Create both mean and logvar projections.

Parameters:

Name	Type	Description	Default
dim	int	Output action dimension.	required

Source code in src/versatil/models/decoding/action_heads/gaussian.py

def set_output_dim(self, dim: int) -> None:
    """Create both mean and logvar projections.

    Args:
        dim: Output action dimension.
    """
    super().set_output_dim(dim)
    hidden_dimension = self._get_hidden_dim()
    self._logvar_proj = nn.Linear(hidden_dimension, dim)

forward

forward(action_embedding)

Forward pass returning mean and clamped logvar.

Parameters:

Name	Type	Description	Default
action_embedding	Tensor	(B, T, embedding_dimension)	required

Returns:

Type	Description
dict[str, Tensor]	Dict with "mean" and "logvar" keys.

Source code in src/versatil/models/decoding/action_heads/gaussian.py

def forward(self, action_embedding: torch.Tensor) -> dict[str, torch.Tensor]:
    """Forward pass returning mean and clamped logvar.

    Args:
        action_embedding: (B, T, embedding_dimension)

    Returns:
        Dict with "mean" and "logvar" keys.
    """
    if self.output_proj is None or self._logvar_proj is None:
        raise RuntimeError("output_dim not set. Call set_output_dim() first.")
    action_embedding = self._apply_blocks(action_embedding)
    mean = self.output_proj(action_embedding)
    logvar = self._logvar_proj(action_embedding)
    logvar = logvar.clamp(min=self.min_logvar, max=self.max_logvar)
    return {
        DecoderOutputKey.MEAN.value: mean,
        DecoderOutputKey.LOGVAR.value: logvar,
    }