mixture

mixture

Negative log-likelihood losses for mixture distributions.

GaussianMixtureNLLoss

GaussianMixtureNLLoss(action_keys, weight=1.0, per_key_weights=None, learned_variance=True, sigmas=None, min_variance=0.0001)

Bases: ScalarWeightedLoss

Negative Log-Likelihood loss for Gaussian Mixture Model.

Supports both learned variance (from logvar predictions) and fixed variance (sigma parameter).

Two regimes are supported, dispatched on the shape of routing_weights:

Per-trajectory routing (B, K) — one mixture component is selected for the entire chunk (e.g., MoDEACT). The joint trajectory likelihood is: log p(a_{1:T}|s) = logsumexp_k [log π_k + Σ_t log N(a_t | μ_k(t), σ_k(t)²)] Components are forced to model coherent trajectories; without this formulation the gating collapses to one component that averages distinct trajectory modes.

Per-timestep routing (B, T, K) — a component is selected per timestep (e.g., PhaseACT). The per-step mixture likelihood applies: log p(a_t|s, t) = logsumexp_k [log π_kt + log N(a_t | μ_kt, σ_kt²)]

Initialize Gaussian mixture NLL loss.

Parameters:

Name	Type	Description	Default
action_keys	list[str]	List of continuous action keys.	required
weight	float	Overall loss weight.	1.0
per_key_weights	dict[str, float] | None	Optional per-key weights.	None
learned_variance	bool	If True, expects {action_key}_mean and {action_key}_logvar. If False, expects {action_key} (stacked means) and uses sigmas.	True
sigmas	dict[str, float] | None	Fixed stddev per action key (only used when learned_variance=False).	None
min_variance	float	Minimum variance for numerical stability (learned_variance=True).	0.0001

Source code in src/versatil/metrics/losses/mixture.py

def __init__(
    self,
    action_keys: list[str],
    weight: float = 1.0,
    per_key_weights: dict[str, float] | None = None,
    learned_variance: bool = True,
    sigmas: dict[str, float] | None = None,
    min_variance: float = 1e-4,
):
    """Initialize Gaussian mixture NLL loss.

    Args:
        action_keys: List of continuous action keys.
        weight: Overall loss weight.
        per_key_weights: Optional per-key weights.
        learned_variance: If True, expects {action_key}_mean and {action_key}_logvar.
            If False, expects {action_key} (stacked means) and uses sigmas.
        sigmas: Fixed stddev per action key (only used when learned_variance=False).
        min_variance: Minimum variance for numerical stability (learned_variance=True).
    """
    super().__init__()
    self.action_keys = action_keys
    self.weight = weight
    self.per_key_weights = per_key_weights or dict.fromkeys(action_keys, 1.0)
    self.learned_variance = learned_variance
    self.min_variance = min_variance
    if not learned_variance:
        self.sigmas = sigmas or dict.fromkeys(action_keys, 0.5)

get_required_keys

get_required_keys()

Get required target keys.

Source code in src/versatil/metrics/losses/mixture.py

def get_required_keys(self) -> set[str]:
    """Get required target keys."""
    return set(self.action_keys)

forward

forward(predictions, targets, is_pad=None)

Compute Gaussian mixture NLL loss.

Parameters:

Name	Type	Description	Default
predictions	dict[str, Tensor]	Dictionary containing: - routing_weights: (B, K) for per-trajectory or (B, T, K) for per-timestep. - If learned_variance: {action_key}_mean (B, T, K, D), {action_key}_logvar (B, T, K, D) - If fixed variance: {action_key} (B, T, K, D) stacked expert means	required
targets	dict[str, Tensor]	Dictionary with action_key targets (B, T, D).	required
is_pad	Tensor | None	Optional padding mask (B, T).	None

Returns:

Type	Description
LossOutput	LossOutput with Gaussian mixture NLL.

Source code in src/versatil/metrics/losses/mixture.py

def forward(
    self,
    predictions: dict[str, torch.Tensor],
    targets: dict[str, torch.Tensor],
    is_pad: torch.Tensor | None = None,
) -> LossOutput:
    """Compute Gaussian mixture NLL loss.

    Args:
        predictions: Dictionary containing:
            - routing_weights: (B, K) for per-trajectory or (B, T, K) for per-timestep.
            - If learned_variance: {action_key}_mean (B, T, K, D), {action_key}_logvar (B, T, K, D)
            - If fixed variance: {action_key} (B, T, K, D) stacked expert means
        targets: Dictionary with action_key targets (B, T, D).
        is_pad: Optional padding mask (B, T).

    Returns:
        LossOutput with Gaussian mixture NLL.
    """
    component_losses: dict[str, torch.Tensor] = {}
    log_components_by_key: dict[str, torch.Tensor] = {}
    mixing_probs = predictions[DecoderOutputKey.ROUTING_WEIGHTS.value]
    for action_key in self.action_keys:
        target = targets[action_key]  # (B, T, D)
        mean_key = f"{action_key}_{DecoderOutputKey.MEAN.value}"
        means = predictions.get(
            mean_key, predictions.get(action_key)
        )  # (B, T, K, D)
        if self.learned_variance:
            logvar_key = f"{action_key}_{DecoderOutputKey.LOGVAR.value}"
            logvars = predictions[logvar_key]  # (B, T, K, D)
            log_component = self._compute_learned_variance_log_pdf(
                target, means, logvars
            )
        else:
            sigma = self.sigmas.get(action_key, 0.5)
            log_component = self._compute_fixed_variance_log_pdf(
                target, means, sigma
            )
        log_components_by_key[action_key] = log_component
        per_key_nll = _aggregate_mixture_nll(
            log_component=log_component,
            mixing_probs=mixing_probs,
            is_pad=is_pad,
        )  # (B,)
        per_key_nll_reduced = per_key_nll.mean()
        component_losses[f"{action_key}_{MetricKey.GAUSSIAN_MIXTURE_NLL.value}"] = (
            per_key_nll_reduced
        )

    if len(self.action_keys) == 1:
        action_key = self.action_keys[0]
        joint_nll_reduced = component_losses[
            f"{action_key}_{MetricKey.GAUSSIAN_MIXTURE_NLL.value}"
        ]
        total_loss = self.per_key_weights.get(action_key, 1.0) * joint_nll_reduced
    else:
        weighted_log_components = [
            self.per_key_weights.get(action_key, 1.0)
            * log_components_by_key[action_key]
            for action_key in self.action_keys
        ]
        # A shared mixture component represents the full action tuple.
        joint_log_component = torch.stack(weighted_log_components, dim=0).sum(dim=0)
        joint_nll = _aggregate_mixture_nll(
            log_component=joint_log_component,
            mixing_probs=mixing_probs,
            is_pad=is_pad,
        )
        joint_nll_reduced = joint_nll.mean()
        total_loss = joint_nll_reduced

    component_losses[MetricKey.GAUSSIAN_MIXTURE_NLL.value] = joint_nll_reduced
    return LossOutput(
        total_loss=self.weight * total_loss, component_losses=component_losses
    )

GripperMixtureNLLoss

GripperMixtureNLLoss(key, actions_metadata, weight=1.0, learned_variance=False, sigma=0.5, min_variance=0.0001)

Bases: ScalarWeightedLoss

Negative Log-Likelihood loss for gripper with mixture distribution.

Binary gripper: p(a|z) = Σ_k π_k(z) · Bernoulli(a | p_k(z)) Continuous gripper: p(a|z) = Σ_k π_k(z) · N(a | μ_k(z), σ_k²)

Supports both fixed and learned variance for continuous gripper.

Initialize gripper mixture NLL loss.

Parameters:

Name	Type	Description	Default
key	str	Key for gripper actions.	required
actions_metadata	dict[str, ActionMetadata]	Dict of metadata of the action space.	required
weight	float	Loss weight.	1.0
learned_variance	bool	If True, expects {key}_mean and {key}_logvar for continuous. If False, expects {key} (stacked means) and uses sigma.	False
sigma	float	Fixed std for continuous gripper (only used when learned_variance=False).	0.5
min_variance	float	Minimum variance for numerical stability (learned_variance=True).	0.0001

Source code in src/versatil/metrics/losses/mixture.py

def __init__(
    self,
    key: str,
    actions_metadata: dict[str, ActionMetadata],
    weight: float = 1.0,
    learned_variance: bool = False,
    sigma: float = 0.5,
    min_variance: float = 1e-4,
):
    """Initialize gripper mixture NLL loss.

    Args:
        key: Key for gripper actions.
        actions_metadata: Dict of metadata of the action space.
        weight: Loss weight.
        learned_variance: If True, expects {key}_mean and {key}_logvar for continuous.
            If False, expects {key} (stacked means) and uses sigma.
        sigma: Fixed std for continuous gripper (only used when learned_variance=False).
        min_variance: Minimum variance for numerical stability (learned_variance=True).
    """
    super().__init__()
    self.key = key
    self.weight = weight
    self.learned_variance = learned_variance
    self.sigma = sigma
    self.min_variance = min_variance
    self.gripper_type, self.binary_gripper_range = resolve_gripper_metadata(
        key=key, actions_metadata=actions_metadata
    )

get_required_keys

get_required_keys()

Return the prediction key this loss consumes.

Source code in src/versatil/metrics/losses/mixture.py

def get_required_keys(self) -> set[str]:
    """Return the prediction key this loss consumes."""
    return {self.key}

forward

forward(predictions, targets, is_pad=None)

Compute gripper mixture NLL.

Parameters:

Name	Type	Description	Default
predictions	dict[str, Tensor]	Dictionary containing: - routing_weights: (B, K) or (B, T, K) mixing probabilities - For binary: {key} stacked logits (B, T, K) or (B, T, K, 1) - For continuous with learned_variance: {key}_mean (B, T, K, D), {key}_logvar (B, T, K, D) - For continuous with fixed variance: {key} stacked means (B, T, K, D)	required
targets	dict[str, Tensor]	Dictionary with gripper targets.	required
is_pad	Tensor | None	Optional padding mask (B, T).	None

Returns:

Type	Description
LossOutput	LossOutput with gripper NLL.

Source code in src/versatil/metrics/losses/mixture.py

def forward(
    self,
    predictions: dict[str, torch.Tensor],
    targets: dict[str, torch.Tensor],
    is_pad: torch.Tensor | None = None,
) -> LossOutput:
    """Compute gripper mixture NLL.

    Args:
        predictions: Dictionary containing:
            - routing_weights: (B, K) or (B, T, K) mixing probabilities
            - For binary: {key} stacked logits (B, T, K) or (B, T, K, 1)
            - For continuous with learned_variance: {key}_mean (B, T, K, D), {key}_logvar (B, T, K, D)
            - For continuous with fixed variance: {key} stacked means (B, T, K, D)
        targets: Dictionary with gripper targets.
        is_pad: Optional padding mask (B, T).

    Returns:
        LossOutput with gripper NLL.
    """
    if self.key not in targets:
        raise ValueError(
            f"Targets must contain '{self.key}' for GripperMixtureNLLoss."
        )
    target = targets[self.key]
    mixing_probs = predictions[DecoderOutputKey.ROUTING_WEIGHTS.value]
    if self.gripper_type == GripperType.BINARY.value:
        log_component = self._compute_binary_log_component(predictions, target)
    else:
        log_component = self._compute_continuous_log_component(predictions, target)
    nll_per_batch = _aggregate_mixture_nll(
        log_component=log_component,
        mixing_probs=mixing_probs,
        is_pad=is_pad,
    )
    nll_reduced = nll_per_batch.mean()
    return LossOutput(
        total_loss=self.weight * nll_reduced,
        component_losses={MetricKey.GRIPPER_NLL.value: nll_reduced},
    )