base

base

Base classes for loss computation and metrics tracking.

LossOutput dataclass

LossOutput(total_loss, component_losses=dict(), metadata=dict())

Output from loss computation containing total loss and component losses.

Attributes:

Name	Type	Description
total_loss	Tensor	Scalar tensor representing the total weighted loss
component_losses	dict[str, Tensor]	Dictionary mapping loss component names to their values
metadata	dict[str, Any]	Optional dictionary for additional information (e.g., predictions for metrics)

to_dict

to_dict()

Convert loss output to dictionary of scalar values.

Returns:

Type	Description
dict[str, float]	Dictionary with MetricKey.TOTAL_LOSS and all component losses as floats

Source code in src/versatil/metrics/base.py

def to_dict(self) -> dict[str, float]:
    """Convert loss output to dictionary of scalar values.

    Returns:
        Dictionary with MetricKey.TOTAL_LOSS and all component losses as floats
    """
    result = {MetricKey.TOTAL_LOSS.value: self.total_loss.item()}
    for key, value in self.component_losses.items():
        result[key] = value.item() if isinstance(value, torch.Tensor) else value
    return result

BaseLoss

Bases: Module, ABC

Abstract base class for loss computation modules.

Subclasses override weights and set_weights to expose their tunable scalars.

weights property

weights

Current public weight structure. Default: no tunables.

requires_action_space_targets property

requires_action_space_targets

Whether this loss expects targets in the action space.

When True, the loss uses classification-style objectives (e.g. BCE) that only make sense when targets are actual action labels. Algorithms that predict outside the action space (velocity, noise) are incompatible with such losses.

Default is False (regression losses like MSE/L1 work in any space).

set_weights

set_weights(new_weights)

Replace the full weight structure for this node.

Source code in src/versatil/metrics/base.py

def set_weights(self, new_weights: WeightsDictionary) -> None:
    """Replace the full weight structure for this node."""
    self._validate_weights(new_weights)

update_weights

update_weights(override_weights)

Apply a partial override by merging onto the current structure.

Source code in src/versatil/metrics/base.py

def update_weights(self, override_weights: WeightsDictionary) -> None:
    """Apply a partial override by merging onto the current structure."""
    self.set_weights(
        _merge_weights(
            existing_weights=self.weights, override_weights=override_weights
        )
    )

forward abstractmethod

forward(predictions, targets, is_pad=None)

Compute loss given predictions and targets.

Parameters:

Name	Type	Description	Default
predictions	dict[str, Tensor]	Dictionary of model predictions	required
targets	dict[str, Tensor]	Dictionary of ground truth values	required
is_pad	Tensor | None	Optional boolean tensor indicating padded positions (B, horizon)	None

Returns:

Type	Description
LossOutput	LossOutput containing total loss and component losses

Source code in src/versatil/metrics/base.py

@abstractmethod
def forward(
    self,
    predictions: dict[str, torch.Tensor],
    targets: dict[str, torch.Tensor],
    is_pad: torch.Tensor | None = None,
) -> LossOutput:
    """Compute loss given predictions and targets.

    Args:
        predictions: Dictionary of model predictions
        targets: Dictionary of ground truth values
        is_pad: Optional boolean tensor indicating padded positions (B, horizon)

    Returns:
        LossOutput containing total loss and component losses
    """
    raise NotImplementedError

get_required_keys abstractmethod

get_required_keys()

Get the set of output keys this loss consumes.

Used at experiment validation to check every key the loss reads is produced by the decoder, the algorithm, or a non-predicted action-space entry (see experiment.validate_loss_keys). Losses that read their targets from the predictions dictionary report those keys here. For composite losses, this should recursively collect keys from all sub-losses.

Returns:

Type	Description
set[str]	Set of output keys required by this loss

Source code in src/versatil/metrics/base.py

@abstractmethod
def get_required_keys(self) -> set[str]:
    """Get the set of output keys this loss consumes.

    Used at experiment validation to check every key the loss reads is
    produced by the decoder, the algorithm, or a non-predicted action-space
    entry (see ``experiment.validate_loss_keys``). Losses that read their
    targets from the predictions dictionary report those keys here. For
    composite losses, this should recursively collect keys from all
    sub-losses.

    Returns:
        Set of output keys required by this loss
    """
    raise NotImplementedError

ScalarWeightedLoss

Bases: BaseLoss, ABC

A BaseLoss with exactly one tunable scalar weight stored as self.weight.

weights property

weights

Getter that returns dictionary with weight keys and scalar coefficients.

set_weights

set_weights(new_weights)

Setter that updates the weight scalar coefficients.

Source code in src/versatil/metrics/base.py

def set_weights(self, new_weights: WeightsDictionary) -> None:
    """Setter that updates the weight scalar coefficients."""
    self._validate_weights(new_weights)
    self.weight = new_weights["weight"]

reduce_loss_with_padding

reduce_loss_with_padding(loss_tensor, is_pad, reduction='mean')

Apply padding-aware reduction to a loss tensor.

Parameters:

Name	Type	Description	Default
loss_tensor	Tensor	Loss values of shape (B, horizon, ...)	required
is_pad	Tensor | None	Boolean mask of shape (B, horizon) where True indicates padding	required
reduction	str	Reduction mode ('mean', 'sum', or 'none')	'mean'

Returns:

Type	Description
Tensor	Reduced loss tensor

Source code in src/versatil/metrics/base.py

def reduce_loss_with_padding(
    loss_tensor: torch.Tensor,
    is_pad: torch.Tensor | None,
    reduction: str = "mean",
) -> torch.Tensor:
    """Apply padding-aware reduction to a loss tensor.

    Args:
        loss_tensor: Loss values of shape (B, horizon, ...)
        is_pad: Boolean mask of shape (B, horizon) where True indicates padding
        reduction: Reduction mode ('mean', 'sum', or 'none')

    Returns:
        Reduced loss tensor
    """
    if is_pad is None:
        if reduction == "mean":
            return loss_tensor.mean()
        elif reduction == "sum":
            return loss_tensor.sum()
        else:
            return loss_tensor
    pad_mask = (~is_pad).float()
    while pad_mask.dim() < loss_tensor.dim():
        pad_mask = pad_mask.unsqueeze(-1)
    masked_loss = loss_tensor * pad_mask
    if reduction == "mean":
        # Divide by valid elements, not valid timesteps, so the masked mean
        # matches the scale of loss_tensor.mean() on unpadded inputs.
        valid_elements = pad_mask.expand_as(loss_tensor).sum()
        return masked_loss.sum() / (valid_elements + 1e-8)
    elif reduction == "sum":
        return masked_loss.sum()
    else:
        return masked_loss